ESQL: Fix variable shadowing when pushing down past Project

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java

@@ -31,6 +31,7 @@
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.FieldAttribute;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
+import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
 import org.elasticsearch.xpack.esql.core.expression.predicate.Range;
 import org.elasticsearch.xpack.esql.core.index.EsIndex;
 import org.elasticsearch.xpack.esql.core.session.Configuration;
@@ -155,10 +156,6 @@ public static Literal of(Object value) {
         return of(Source.EMPTY, value);
     }
 
-    public static Configuration randomConfiguration() {
-        return new Configuration(randomZone(), randomAlphaOfLength(10), randomAlphaOfLength(10));
-    }
-
     /**
      * Utility method for creating 'in-line' Literals (out of values instead of expressions).
      */
@@ -169,10 +166,18 @@ public static Literal of(Source source, Object value) {
         return new Literal(source, value, DataType.fromJava(value));
     }
 
+    public static ReferenceAttribute referenceAttribute(String name, DataType type) {
+        return new ReferenceAttribute(EMPTY, name, type);
+    }
+
     public static Range rangeOf(Expression value, Expression lower, boolean includeLower, Expression upper, boolean includeUpper) {
         return new Range(EMPTY, value, lower, includeLower, upper, includeUpper, randomZone());
     }
 
+    public static Configuration randomConfiguration() {
+        return new Configuration(randomZone(), randomAlphaOfLength(10), randomAlphaOfLength(10));
+    }
+
     public static EsRelation relation() {
         return new EsRelation(EMPTY, new EsIndex(randomAlphaOfLength(8), emptyMap()), IndexMode.STANDARD, randomBoolean());
     }

x-pack/plugin/esql/qa/testFixtures/src/main/resources/stats.csv-spec

@@ -663,6 +663,61 @@ ca:l | cx:l | l:i
 1    | 1    | null
 ;
 
+///////////////////////////////////////////////////////////////
+// Test edge case interaction with push down optimization rules
+// https://github.com/elastic/elasticsearch/issues/108008
+///////////////////////////////////////////////////////////////
+
+// TODO: we update capabilities
+countSameFieldWithEval
+from employees | stats  b = count(gender), c = count(gender) by gender | eval b = gender | sort c asc
+;
+
+c:l | gender:s | b:s
+0   | null     | null
+33  | F        | F
+57  | M        | M
+;
+
+countSameFieldWithDissect
+from employees | stats b = count(gender), c = count(gender) by gender | dissect gender "%{b}" | sort c asc
+;
+
+c:l | gender:s | b:s
+0   | null     | null
+33  | F        | F
+57  | M        | M
+;
+
+countSameFieldWithGrok
+from employees | stats  b = count(gender), c = count(gender) by gender | grok gender "%{USERNAME:b}" | sort c asc
+;
+
+c:l | gender:s | b:s
+0   | null     | null
+33  | F        | F
+57  | M        | M
+;
+
+countSameFieldWithEnrich
+required_capability: enrich_load
+from employees | stats  b = count(gender), c = count(gender) by gender | enrich languages_policy on gender with b = language_name | sort c asc
+;
+
+c:l | gender:s | b:s
+0   | null     | null
+33  | F        | null
+57  | M        | null
+;
+
+countSameFieldWithEnrichLimit0
+from employees | stats  b = count(gender), c = count(gender) by gender | enrich languages_policy on gender with b = language_name | sort c asc | limit 0
+;
+
+c:l | gender:s | b:s
+;
+///////////////////////////////////////////////////////////////
+
 aggsWithoutStats
 from employees | stats by gender | sort gender;
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/NamedExpressions.java

@@ -33,7 +33,8 @@ public static List<Attribute> mergeOutputAttributes(
     /**
      * Merges output expressions of a command given the new attributes plus the existing inputs that are emitted as outputs.
      * As a general rule, child output will come first in the list, followed by the new fields.
-     * In case of name collisions, only last entry is preserved (previous expressions with the same name are discarded)
+     * In case of name collisions, only the last entry is preserved (previous expressions with the same name are discarded)
+     * and the new attributes have precedence over the child output.
      * @param fields the fields added by the command
      * @param childOutput the command input that has to be propagated as output
      * @return

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizer.java

@@ -15,6 +15,7 @@
 import org.elasticsearch.xpack.esql.core.expression.AttributeMap;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Expressions;
+import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
 import org.elasticsearch.xpack.esql.core.expression.Order;
 import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
 import org.elasticsearch.xpack.esql.core.rule.ParameterizedRule;
@@ -64,6 +65,7 @@
 import org.elasticsearch.xpack.esql.optimizer.rules.SubstituteSpatialSurrogates;
 import org.elasticsearch.xpack.esql.optimizer.rules.SubstituteSurrogates;
 import org.elasticsearch.xpack.esql.optimizer.rules.TranslateMetricsAggregate;
+import org.elasticsearch.xpack.esql.plan.GeneratingPlan;
 import org.elasticsearch.xpack.esql.plan.logical.Eval;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.plan.logical.OrderBy;
@@ -74,8 +76,10 @@
 import org.elasticsearch.xpack.esql.type.EsqlDataTypes;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 import static java.util.Arrays.asList;
@@ -209,6 +213,7 @@ public static LogicalPlan skipPlan(UnaryPlan plan, LocalSupplier supplier) {
         return new LocalRelation(plan.source(), plan.output(), supplier);
     }
 
+    // TODO: update
     /**
      * Pushes LogicalPlans which generate new attributes (Eval, Grok/Dissect, Enrich), past OrderBys and Projections.
      * Although it seems arbitrary whether the OrderBy or the Eval is executed first, this transformation ensures that OrderBys only
@@ -235,11 +240,10 @@ public static LogicalPlan skipPlan(UnaryPlan plan, LocalSupplier supplier) {
      *
      * ... | eval $$a = a | eval a = b + 1 | sort $$a | drop $$a
      */
-    public static LogicalPlan pushGeneratingPlanPastProjectAndOrderBy(UnaryPlan generatingPlan, List<Attribute> generatedAttributes) {
+    public static <Plan extends UnaryPlan & GeneratingPlan<Plan>> LogicalPlan pushGeneratingPlanPastProjectAndOrderBy(Plan generatingPlan) {
         LogicalPlan child = generatingPlan.child();
-
         if (child instanceof OrderBy orderBy) {
-            Set<String> evalFieldNames = new LinkedHashSet<>(Expressions.names(generatedAttributes));
+            Set<String> evalFieldNames = new LinkedHashSet<>(Expressions.names(generatingPlan.generatedAttributes()));
 
             // Look for attributes in the OrderBy's expressions and create aliases with temporary names for them.
             AttributeReplacement nonShadowedOrders = renameAttributesInExpressions(evalFieldNames, orderBy.order());
@@ -260,9 +264,77 @@ public static LogicalPlan pushGeneratingPlanPastProjectAndOrderBy(UnaryPlan gene
             }
 
             return orderBy.replaceChild(generatingPlan.replaceChild(orderBy.child()));
-        } else if (child instanceof Project) {
-            var projectWithEvalChild = pushDownPastProject(generatingPlan);
-            return projectWithEvalChild.withProjections(mergeOutputExpressions(generatedAttributes, projectWithEvalChild.projections()));
+        } else if (child instanceof Project project) {
+            // We need to account for attribute shadowing. E.g.
+            // Eval[[2 * x{f}#1 AS y]]
+            // \_Project[[x{f}#1, y{f}#2, y{f}#2 AS z]]
+            // ..\_....
+            // Just moving the Eval down breaks z because we shadow y{f}#2.
+            // Instead, we use a different alias in the Eval, eventually renaming back to y:
+            // Project[[x{f}#1, y{f}#2 as z, $$y{r}#3 as y]]
+            // \_Eval[[2 * x{f}#1 as $$y]]
+            // ..\_....
+
+            AttributeMap.Builder<Expression> aliasBuilder = AttributeMap.builder();
+            project.forEachExpression(Alias.class, a -> aliasBuilder.put(a.toAttribute(), a.child()));
+            var aliases = aliasBuilder.build();
+
+            // Resolve Project's renames in the generating plan.
+            @SuppressWarnings("unchecked")
+            Plan generatingPlanWithResolvedExpressions = (Plan) generatingPlan.transformExpressionsOnly(
+                ReferenceAttribute.class,
+                r -> aliases.resolve(r, r)
+            );
+
+            // Look for generated Attributes that currently shadow any of the childs's references.
+            // We need to generate them using a different, non-shadowing name to avoid inconsistencies.
+            List<Attribute> generatedAttributes = generatingPlan.generatedAttributes();
+            Set<String> projectReferencedNames = project.references().names();
+            Map<String, String> renameGeneratedAttributeTo = new HashMap<>();
+            for (Attribute attr : generatedAttributes) {
+                String name = attr.name();
+                if (projectReferencedNames.contains(name)) {
+                    renameGeneratedAttributeTo.putIfAbsent(
+                        name,
+                        // TODO: Use e.g. AtomicLong to make sure generated temp names can not clash.
+                        // Do not use the attribute's id, as multiple attributes with the same name can occur.
+                        SubstituteSurrogates.rawTemporaryName(name, "temp_name", "")
+                    );
+                }
+            }
+
+            List<String> newNames = generatedAttributes.stream()
+                .map(attr -> renameGeneratedAttributeTo.getOrDefault(attr.name(), attr.name()))
+                .toList();
+            Plan generatingPlanWithRenamedAttributes = generatingPlanWithResolvedExpressions.withGeneratedNames(newNames);
+
+            // Put the project at the top, but include the generated attributes.
+            // Any generated attributes that had to be renamed need to be re-renamed to their original names.
+            List<NamedExpression> generatedAttributesRenamedToOriginal = new ArrayList<>(generatedAttributes.size());
+            List<Attribute> renamedGeneratedAttributes = generatingPlanWithRenamedAttributes.generatedAttributes();
+            for (int i = 0; i < generatedAttributes.size(); i++) {
+                Attribute originalAttribute = generatedAttributes.get(i);
+                Attribute renamedAttribute = renamedGeneratedAttributes.get(i);
+                if (originalAttribute.name().equals(renamedAttribute.name())) {
+                    generatedAttributesRenamedToOriginal.add(renamedAttribute);
+                } else {
+                    generatedAttributesRenamedToOriginal.add(
+                        new Alias(
+                            originalAttribute.source(),
+                            originalAttribute.name(),
+                            originalAttribute.qualifier(),
+                            renamedAttribute,
+                            originalAttribute.id(),
+                            originalAttribute.synthetic()
+                        )
+                    );
+                }
+            }
+
+            Project projectWithGeneratingChild = project.replaceChild(generatingPlanWithRenamedAttributes.replaceChild(project.child()));
+            return projectWithGeneratingChild.withProjections(
+                mergeOutputExpressions(generatedAttributesRenamedToOriginal, projectWithGeneratingChild.projections())
+            );
         }
 
         return generatingPlan;
@@ -286,6 +358,7 @@ private static AttributeReplacement renameAttributesInExpressions(
             rewrittenExpressions.add(expr.transformUp(Attribute.class, attr -> {
                 if (attributeNamesToRename.contains(attr.name())) {
                     Alias renamedAttribute = aliasesForReplacedAttributes.computeIfAbsent(attr, a -> {
+                        // TODO: Use e.g. AtomicLong to make sure generated temp names can not clash.
                         String tempName = SubstituteSurrogates.rawTemporaryName(a.name(), "temp_name", a.id().toString());
                         // TODO: this should be synthetic
                         return new Alias(a.source(), tempName, null, a, null, false);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/OptimizerRules.java

@@ -91,6 +91,11 @@ protected AttributeSet references(LogicalPlan plan) {
                 // But they are not actually referring to attributes from the input plan - only the match field does.
                 return enrich.matchField().references();
             }
+            if (plan instanceof RegexExtract re) {
+                // Similarly as for Enrich: the extractedFields are Aliases for ReferenceAttributes, which are in turn created by the
+                // RegexExtract node.
+                return re.input().references();
+            }
             return super.references(plan);
         }
 
@@ -107,7 +112,7 @@ protected AttributeSet generates(LogicalPlan logicalPlan) {
                 return new AttributeSet(Expressions.asAttributes(eval.fields()));
             }
             if (logicalPlan instanceof RegexExtract extract) {
-                return new AttributeSet(extract.extractedFields());
+                return new AttributeSet(Expressions.asAttributes(extract.extractedFields()));
             }
             if (logicalPlan instanceof MvExpand mvExpand) {
                 return new AttributeSet(mvExpand.expanded());
@@ -143,7 +148,7 @@ protected AttributeSet generates(PhysicalPlan physicalPlan) {
                 return new AttributeSet(Expressions.asAttributes(eval.fields()));
             }
             if (physicalPlan instanceof RegexExtractExec extract) {
-                return new AttributeSet(extract.extractedFields());
+                return new AttributeSet(Expressions.asAttributes(extract.extractedFields()));
             }
             if (physicalPlan instanceof MvExpandExec mvExpand) {
                 return new AttributeSet(mvExpand.expanded());

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/PushDownEnrich.java

@@ -11,11 +11,9 @@
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 
-import static org.elasticsearch.xpack.esql.core.expression.Expressions.asAttributes;
-
 public final class PushDownEnrich extends OptimizerRules.OptimizerRule<Enrich> {
     @Override
     protected LogicalPlan rule(Enrich en) {
-        return LogicalPlanOptimizer.pushGeneratingPlanPastProjectAndOrderBy(en, asAttributes(en.enrichFields()));
+        return LogicalPlanOptimizer.pushGeneratingPlanPastProjectAndOrderBy(en);
     }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/PushDownEval.java

@@ -11,11 +11,9 @@
 import org.elasticsearch.xpack.esql.plan.logical.Eval;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 
-import static org.elasticsearch.xpack.esql.core.expression.Expressions.asAttributes;
-
 public final class PushDownEval extends OptimizerRules.OptimizerRule<Eval> {
     @Override
     protected LogicalPlan rule(Eval eval) {
-        return LogicalPlanOptimizer.pushGeneratingPlanPastProjectAndOrderBy(eval, asAttributes(eval.fields()));
+        return LogicalPlanOptimizer.pushGeneratingPlanPastProjectAndOrderBy(eval);
     }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/PushDownRegexExtract.java

@@ -14,6 +14,6 @@
 public final class PushDownRegexExtract extends OptimizerRules.OptimizerRule<RegexExtract> {
     @Override
     protected LogicalPlan rule(RegexExtract re) {
-        return LogicalPlanOptimizer.pushGeneratingPlanPastProjectAndOrderBy(re, re.extractedFields());
+        return LogicalPlanOptimizer.pushGeneratingPlanPastProjectAndOrderBy(re);
     }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/LogicalPlanBuilder.java

@@ -26,7 +26,6 @@
 import org.elasticsearch.xpack.esql.core.expression.MetadataAttribute;
 import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
 import org.elasticsearch.xpack.esql.core.expression.Order;
-import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
 import org.elasticsearch.xpack.esql.core.expression.UnresolvedAttribute;
 import org.elasticsearch.xpack.esql.core.expression.UnresolvedStar;
 import org.elasticsearch.xpack.esql.core.parser.ParserUtils;
@@ -66,7 +65,6 @@
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.function.Function;
 
 import static org.elasticsearch.common.logging.HeaderWarning.addWarning;
@@ -198,20 +196,10 @@ public PlanFactory visitDissectCommand(EsqlBaseParser.DissectCommandContext ctx)
 
             try {
                 DissectParser parser = new DissectParser(pattern, appendSeparator);
-                Set<String> referenceKeys = parser.referenceKeys();
-                if (referenceKeys.size() > 0) {
-                    throw new ParsingException(
-                        src,
-                        "Reference keys not supported in dissect patterns: [%{*{}}]",
-                        referenceKeys.iterator().next()
-                    );
-                }
-                List<Attribute> keys = new ArrayList<>();
-                for (var x : parser.outputKeys()) {
-                    if (x.isEmpty() == false) {
-                        keys.add(new ReferenceAttribute(src, x, DataType.KEYWORD));
-                    }
-                }
+
+                Dissect.Parser esqlDissectParser = new Dissect.Parser(pattern, appendSeparator, parser);
+                List<Attribute> keys = esqlDissectParser.keyAttributes(src);
+
                 return new Dissect(src, p, expression(ctx.primaryExpression()), new Dissect.Parser(pattern, appendSeparator, parser), keys);
             } catch (DissectException e) {
                 throw new ParsingException(src, "Invalid pattern for dissect: [{}]", pattern);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/GeneratingPlan.java

@@ -0,0 +1,18 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.plan;
+
+import org.elasticsearch.xpack.esql.core.expression.Attribute;
+
+import java.util.List;
+
+public interface GeneratingPlan<PlanType extends GeneratingPlan<PlanType>> {
+    List<Attribute> generatedAttributes();
+
+    PlanType withGeneratedNames(List<String> newNames);
+}