ESQL: Fix bug when combining projections (elastic#107131)

Recursive aliases (eval x = 1, x1 = x) were not taken into account when
 combining projections causing the target field to be lost (and only the
 immediate intermediate named expression to be used instead which became
 invalid).

Fix elastic#107083

(cherry picked from commit a9388e1)

docs/changelog/107131.yaml

@@ -0,0 +1,6 @@
+pr: 107131
+summary: "ESQL: Fix bug when combining projections"
+area: ES|QL
+type: bug
+issues:
+ - 107083

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

@@ -1315,3 +1315,37 @@ FROM employees
 rows:l
 6
 ;
+
+evalOverridingKey#[skip:-8.13.1,reason:supported in 8.13.2]
+FROM employees
+| EVAL k = languages
+| STATS c = COUNT() BY languages, k
+| DROP k
+| SORT languages
+;
+
+c:l| languages:i
+15 | 1
+19 | 2
+17 | 3
+18 | 4
+21 | 5
+10 | null
+;
+
+evalMultipleOverridingKeys#[skip:-8.13.1,reason:supported in 8.13.2]
+FROM employees
+| EVAL k = languages, k1 = k
+| STATS c = COUNT() BY languages, k, k1, languages
+| DROP k
+| SORT languages
+;
+
+c:l | k1:i | languages:i
+15  | 1    | 1
+19  | 2    | 2
+17  | 3    | 3
+18  | 4    | 4
+21  | 5    | 5
+10  | null | null
+;

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

@@ -363,11 +363,6 @@ private static List<? extends NamedExpression> projectAggregations(
             List<? extends NamedExpression> upperProjection,
             List<? extends NamedExpression> lowerAggregations
         ) {
-            AttributeMap<Expression> lowerAliases = new AttributeMap<>();
-            for (NamedExpression ne : lowerAggregations) {
-                lowerAliases.put(ne.toAttribute(), Alias.unwrap(ne));
-            }
-
             AttributeSet seen = new AttributeSet();
             for (NamedExpression upper : upperProjection) {
                 Expression unwrapped = Alias.unwrap(upper);
@@ -391,19 +386,26 @@ private static List<NamedExpression> combineProjections(
             List<? extends NamedExpression> lower
         ) {
 
-            // collect aliases in the lower list
-            AttributeMap<NamedExpression> aliases = new AttributeMap<>();
+            // collect named expressions declaration in the lower list
+            AttributeMap<NamedExpression> namedExpressions = new AttributeMap<>();
+            // while also collecting the alias map for resolving the source (f1 = 1, f2 = f1, etc..)
+            AttributeMap<Expression> aliases = new AttributeMap<>();
             for (NamedExpression ne : lower) {
-                if ((ne instanceof Attribute) == false) {
-                    aliases.put(ne.toAttribute(), ne);
+                // record the alias
+                aliases.put(ne.toAttribute(), Alias.unwrap(ne));
+
+                // record named expression as is
+                if (ne instanceof Alias as) {
+                    Expression child = as.child();
+                    namedExpressions.put(ne.toAttribute(), as.replaceChild(aliases.resolve(child, child)));
                 }
             }
             List<NamedExpression> replaced = new ArrayList<>();
 
             // replace any matching attribute with a lower alias (if there's a match)
             // but clean-up non-top aliases at the end
             for (NamedExpression ne : upper) {
-                NamedExpression replacedExp = (NamedExpression) ne.transformUp(Attribute.class, a -> aliases.resolve(a, a));
+                NamedExpression replacedExp = (NamedExpression) ne.transformUp(Attribute.class, a -> namedExpressions.resolve(a, a));
                 replaced.add((NamedExpression) trimNonTopLevelAliases(replacedExp));
             }
             return replaced;
@@ -436,7 +438,10 @@ private List<Expression> replacePrunedAliasesUsedInGroupBy(
 
             var newGroupings = new ArrayList<Expression>(groupings.size());
             for (Expression group : groupings) {
-                newGroupings.add(group.transformUp(Attribute.class, a -> removedAliases.resolve(a, a)));
+                var transformed = group.transformUp(Attribute.class, a -> removedAliases.resolve(a, a));
+                if (Expressions.anyMatch(newGroupings, g -> Expressions.equalsAsAttribute(g, transformed)) == false) {
+                    newGroupings.add(transformed);
+                }
             }
 
             return newGroupings;

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java

@@ -306,6 +306,52 @@ public void testCombineProjections() {
         var relation = as(limit.child(), EsRelation.class);
     }
 
+    /**
+     * Expects
+     * Project[[languages{f}#12 AS f2]]
+     * \_Limit[1000[INTEGER]]
+     *   \_EsRelation[test][_meta_field{f}#15, emp_no{f}#9, first_name{f}#10, g..]
+     */
+    public void testCombineProjectionsWithEvalAndDrop() {
+        var plan = plan("""
+            from test
+            | eval f1 = languages, f2 = f1
+            | keep f2
+            """);
+
+        var keep = as(plan, Project.class);
+        assertThat(Expressions.names(keep.projections()), contains("f2"));
+        assertThat(Expressions.name(Alias.unwrap(keep.projections().get(0))), is("languages"));
+        var limit = as(keep.child(), Limit.class);
+        var relation = as(limit.child(), EsRelation.class);
+
+    }
+
+    /**
+     * Expects
+     * Project[[last_name{f}#26, languages{f}#25 AS f2, f4{r}#13]]
+     * \_Eval[[languages{f}#25 + 3[INTEGER] AS f4]]
+     *   \_Limit[1000[INTEGER]]
+     *     \_EsRelation[test][_meta_field{f}#28, emp_no{f}#22, first_name{f}#23, ..]
+     */
+    public void testCombineProjectionsWithEval() {
+        var plan = plan("""
+            from test
+            | eval f1 = languages, f2 = f1, f3 = 1 + 2, f4 = f3 + languages
+            | keep emp_no, *name, salary, f*
+            | drop f3
+            | keep last_name, f2, f4
+            """);
+
+        var keep = as(plan, Project.class);
+        assertThat(Expressions.names(keep.projections()), contains("last_name", "f2", "f4"));
+        var eval = as(keep.child(), Eval.class);
+        assertThat(Expressions.names(eval.fields()), contains("f4"));
+        var add = as(Alias.unwrap(eval.fields().get(0)), Add.class);
+        var limit = as(eval.child(), Limit.class);
+        var relation = as(limit.child(), EsRelation.class);
+    }
+
     public void testCombineProjectionWithFilterInBetween() {
         var plan = plan("""
             from test
@@ -348,6 +394,27 @@ public void testCombineProjectionWithAggregation() {
         assertThat(Expressions.names(agg.groupings()), contains("last_name", "first_name"));
     }
 
+    /**
+     * Expects
+     * Limit[1000[INTEGER]]
+     * \_Aggregate[[last_name{f}#23, first_name{f}#20, k{r}#4],[SUM(salary{f}#24) AS s, last_name{f}#23, first_name{f}#20, first_n
+     * ame{f}#20 AS k]]
+     *   \_EsRelation[test][_meta_field{f}#25, emp_no{f}#19, first_name{f}#20, ..]
+     */
+    public void testCombineProjectionWithAggregationAndEval() {
+        var plan = plan("""
+            from test
+            | eval k = first_name, k1 = k
+            | stats s = sum(salary) by last_name, first_name, k, k1
+            | keep s, last_name, first_name, k
+            """);
+
+        var limit = as(plan, Limit.class);
+        var agg = as(limit.child(), Aggregate.class);
+        assertThat(Expressions.names(agg.aggregates()), contains("s", "last_name", "first_name", "k"));
+        assertThat(Expressions.names(agg.groupings()), contains("last_name", "first_name", "k"));
+    }
+
     /**
      * Project[[s{r}#4 AS d, s{r}#4, last_name{f}#21, first_name{f}#18]]
      * \_Limit[1000[INTEGER]]