ESQL: Fix missing refs due to pruning renamed grouping columns (#107328) (#107599)

alex-spies · web-flow · commit ecb06ef904e4 · 2024-04-18T05:03:51.000-04:00
Sometimes, CombineProjections does not correctly update an aggregation's groupings when combining with a preceding projection. Fix this by resolving any aliases used in the groupings and de-duplicating them. --------- Co-authored-by: Andrei Stefan <astefan@users.noreply.github.com> (cherry picked from commit adaa476) # Conflicts: # x-pack/plugin/esql/qa/testFixtures/src/main/resources/stats.csv-spec # x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java
diff --git a/docs/changelog/107328.yaml b/docs/changelog/107328.yaml
@@ -0,0 +1,7 @@
+pr: 107328
+summary: "ESQL: Fix missing refs due to pruning renamed grouping columns"
+area: ES|QL
+type: bug
+issues:
+ - 107083
+ - 107166
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/stats.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/stats.csv-spec
@@ -1349,3 +1349,20 @@ c:l | k1:i | languages:i
 21  | 5    | 5
 10  | null | null
 ;
+
+evalMultipleOverridingKeysWithAggregateExpr#[skip:-8.13.2,reason:supported in 8.13.3]
+FROM employees
+| EVAL k = languages, k1 = k
+| STATS c = 3*COUNT() BY languages, k, k1, languages
+| DROP k
+| SORT languages
+;
+
+c:l | k1:i | languages:i
+45  | 1    | 1
+57  | 2    | 2
+51  | 3    | 3
+54  | 4    | 4
+63  | 5    | 5
+30  | null | null
+;
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizer.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizer.java
@@ -318,6 +318,7 @@ static class CombineProjections extends OptimizerRules.OptimizerRule<UnaryPlan>
         }
 
         @Override
+        @SuppressWarnings("unchecked")
         protected LogicalPlan rule(UnaryPlan plan) {
             LogicalPlan child = plan.child();
 
@@ -348,7 +349,22 @@ protected LogicalPlan rule(UnaryPlan plan) {
             // Agg with underlying Project (group by on sub-queries)
             if (plan instanceof Aggregate a) {
                 if (child instanceof Project p) {
-                    plan = new Aggregate(a.source(), p.child(), a.groupings(), combineProjections(a.aggregates(), p.projections()));
+                    var groupings = a.groupings();
+                    List<Attribute> groupingAttrs = new ArrayList<>(a.groupings().size());
+                    for (Expression grouping : groupings) {
+                        if (grouping instanceof Attribute attribute) {
+                            groupingAttrs.add(attribute);
+                        } else {
+                            // After applying ReplaceStatsNestedExpressionWithEval, groupings can only contain attributes.
+                            throw new EsqlIllegalArgumentException("Expected an Attribute, got {}", grouping);
+                        }
+                    }
+                    plan = new Aggregate(
+                        a.source(),
+                        p.child(),
+                        combineUpperGroupingsAndLowerProjections(groupingAttrs, p.projections()),
+                        combineProjections(a.aggregates(), p.projections())
+                    );
                 }
             }
 
@@ -411,6 +427,26 @@ private static List<NamedExpression> combineProjections(
             return replaced;
         }
 
+        private static List<Expression> combineUpperGroupingsAndLowerProjections(
+            List<? extends Attribute> upperGroupings,
+            List<? extends NamedExpression> lowerProjections
+        ) {
+            // Collect the alias map for resolving the source (f1 = 1, f2 = f1, etc..)
+            AttributeMap<Attribute> aliases = new AttributeMap<>();
+            for (NamedExpression ne : lowerProjections) {
+                // Projections are just aliases for attributes, so casting is safe.
+                aliases.put(ne.toAttribute(), (Attribute) Alias.unwrap(ne));
+            }
+
+            // Replace any matching attribute directly with the aliased attribute from the projection.
+            AttributeSet replaced = new AttributeSet();
+            for (Attribute attr : upperGroupings) {
+                // All substitutions happen before; groupings must be attributes at this point.
+                replaced.add(aliases.resolve(attr, attr));
+            }
+            return new ArrayList<>(replaced);
+        }
+
         /**
          * Replace grouping alias previously contained in the aggregations that might have been projected away.
          */
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java
@@ -397,8 +397,8 @@ public void testCombineProjectionWithAggregation() {
     /**
      * 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]]
+     * \_Aggregate[[last_name{f}#23, first_name{f}#20],[SUM(salary{f}#24) AS s, last_name{f}#23, first_name{f}#20, first_name{f}#2
+     * 0 AS k]]
      *   \_EsRelation[test][_meta_field{f}#25, emp_no{f}#19, first_name{f}#20, ..]
      */
     public void testCombineProjectionWithAggregationAndEval() {
@@ -412,7 +412,7 @@ public void testCombineProjectionWithAggregationAndEval() {
         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"));
+        assertThat(Expressions.names(agg.groupings()), contains("last_name", "first_name"));
     }
 
     /**
@@ -454,6 +454,12 @@ public void testQlComparisonOptimizationsApply() {
         assertThat(con.value(), equalTo(5));
     }
 
+    /**
+     * Expects
+     * Limit[1000[INTEGER]]
+     * \_Aggregate[[first_name{f}#12],[COUNT(salary{f}#16) AS count(salary), first_name{f}#12 AS x]]
+     *   \_EsRelation[test][_meta_field{f}#17, emp_no{f}#11, first_name{f}#12, ..]
+     */
     public void testCombineProjectionWithPruning() {
         var plan = plan("""
             from test
@@ -465,19 +471,17 @@ public void testCombineProjectionWithPruning() {
         var limit = as(plan, Limit.class);
         var agg = as(limit.child(), Aggregate.class);
         assertThat(Expressions.names(agg.aggregates()), contains("count(salary)", "x"));
-        assertThat(Expressions.names(agg.groupings()), contains("x"));
+        assertThat(Expressions.names(agg.groupings()), contains("first_name"));
         var alias = as(agg.aggregates().get(1), Alias.class);
         var field = as(alias.child(), FieldAttribute.class);
         assertThat(field.name(), is("first_name"));
-        var group = as(agg.groupings().get(0), Attribute.class);
-        assertThat(group, is(alias.toAttribute()));
         var from = as(agg.child(), EsRelation.class);
     }
 
     /**
      * Expects
      * Limit[1000[INTEGER]]
-     * \_Aggregate[[f{r}#7],[SUM(emp_no{f}#15) AS s, COUNT(first_name{f}#16) AS c, first_name{f}#16 AS f]]
+     * \_Aggregate[[first_name{f}#16],[SUM(emp_no{f}#15) AS s, COUNT(first_name{f}#16) AS c, first_name{f}#16 AS f]]
      *   \_EsRelation[test][_meta_field{f}#21, emp_no{f}#15, first_name{f}#16, ..]
      */
     public void testCombineProjectionWithAggregationFirstAndAliasedGroupingUsedInAgg() {
@@ -501,13 +505,13 @@ public void testCombineProjectionWithAggregationFirstAndAliasedGroupingUsedInAgg
         as = as(aggs.get(2), Alias.class);
         assertThat(Expressions.name(as.child()), is("first_name"));
 
-        assertThat(Expressions.names(agg.groupings()), contains("f"));
+        assertThat(Expressions.names(agg.groupings()), contains("first_name"));
     }
 
     /**
      * Expects
      * Limit[1000[INTEGER]]
-     * \_Aggregate[[f{r}#7],[SUM(emp_no{f}#15) AS s, first_name{f}#16 AS f]]
+     * \_Aggregate[[first_name{f}#16],[SUM(emp_no{f}#15) AS s, first_name{f}#16 AS f]]
      *   \_EsRelation[test][_meta_field{f}#21, emp_no{f}#15, first_name{f}#16, ..]
      */
     public void testCombineProjectionWithAggregationFirstAndAliasedGroupingUnused() {
@@ -527,7 +531,7 @@ public void testCombineProjectionWithAggregationFirstAndAliasedGroupingUnused()
         as = as(aggs.get(1), Alias.class);
         assertThat(Expressions.name(as.child()), is("first_name"));
 
-        assertThat(Expressions.names(agg.groupings()), contains("f"));
+        assertThat(Expressions.names(agg.groupings()), contains("first_name"));
     }
 
     /**
@@ -2688,6 +2692,27 @@ public void testEliminateDuplicateAggsNonCount() {
         var source = as(agg.child(), EsRelation.class);
     }
 
+    /**
+     * Expects
+     * Limit[1000[INTEGER]]
+     * \_Aggregate[[salary{f}#12],[salary{f}#12, salary{f}#12 AS x]]
+     *   \_EsRelation[test][_meta_field{f}#13, emp_no{f}#7, first_name{f}#8, ge..]
+     */
+    public void testEliminateDuplicateRenamedGroupings() {
+        var plan = plan("""
+            from test
+            | eval x = salary
+            | stats by salary, x
+            """);
+
+        var limit = as(plan, Limit.class);
+        var agg = as(limit.child(), Aggregate.class);
+        var relation = as(agg.child(), EsRelation.class);
+
+        assertThat(Expressions.names(agg.groupings()), contains("salary"));
+        assertThat(Expressions.names(agg.aggregates()), contains("salary", "x"));
+    }
+
     /**
      * Expected
      * Limit[2[INTEGER]]
@@ -2734,7 +2759,7 @@ public void testRenameStatsDropGroup() {
     /**
      * Expected
      * Limit[1000[INTEGER]]
-     * \_Aggregate[[a{r}#2, bar{r}#8],[COUNT([2a][KEYWORD]) AS baz, b{r}#4 AS bar]]
+     * \_Aggregate[[a{r}#3, b{r}#5],[COUNT([2a][KEYWORD]) AS baz, b{r}#5 AS bar]]
      *   \_Row[[1[INTEGER] AS a, 2[INTEGER] AS b]]
      */
     public void testMultipleRenameStatsDropGroup() {
@@ -2746,15 +2771,15 @@ public void testMultipleRenameStatsDropGroup() {
 
         var limit = as(plan, Limit.class);
         var agg = as(limit.child(), Aggregate.class);
-        assertThat(Expressions.names(agg.groupings()), contains("a", "bar"));
+        assertThat(Expressions.names(agg.groupings()), contains("a", "b"));
         var row = as(agg.child(), Row.class);
     }
 
     /**
      * Expected
      * Limit[1000[INTEGER]]
-     * \_Aggregate[[emp_no{f}#11, bar{r}#4],[MAX(salary{f}#16) AS baz, gender{f}#13 AS bar]]
-     *   \_EsRelation[test][_meta_field{f}#17, emp_no{f}#11, first_name{f}#12, ..]
+     * \_Aggregate[[emp_no{f}#14, gender{f}#16],[MAX(salary{f}#19) AS baz, gender{f}#16 AS bar]]
+     *   \_EsRelation[test][_meta_field{f}#20, emp_no{f}#14, first_name{f}#15, ..]
      */
     public void testMultipleRenameStatsDropGroupMultirow() {
         LogicalPlan plan = optimizedPlan("""
@@ -2765,7 +2790,7 @@ public void testMultipleRenameStatsDropGroupMultirow() {
 
         var limit = as(plan, Limit.class);
         var agg = as(limit.child(), Aggregate.class);
-        assertThat(Expressions.names(agg.groupings()), contains("emp_no", "bar"));
+        assertThat(Expressions.names(agg.groupings()), contains("emp_no", "gender"));
         var row = as(agg.child(), EsRelation.class);
     }
 
@@ -2846,7 +2871,7 @@ public void testIsNotNullConstraintForStatsWithAndOnGrouping() {
     /**
      * Expects
      * Limit[1000[INTEGER]]
-     * \_Aggregate[[x{r}#4],[SUM(salary{f}#13) AS sum(salary), salary{f}#13 AS x]]
+     * \_Aggregate[[salary{f}#13],[SUM(salary{f}#13) AS sum(salary), salary{f}#13 AS x]]
      *   \_EsRelation[test][_meta_field{f}#14, emp_no{f}#8, first_name{f}#9, ge..]
      */
     public void testIsNotNullConstraintForStatsWithAndOnGroupingAlias() {
@@ -2858,7 +2883,7 @@ public void testIsNotNullConstraintForStatsWithAndOnGroupingAlias() {
 
         var limit = as(plan, Limit.class);
         var agg = as(limit.child(), Aggregate.class);
-        assertThat(Expressions.names(agg.groupings()), contains("x"));
+        assertThat(Expressions.names(agg.groupings()), contains("salary"));
         assertThat(Expressions.names(agg.aggregates()), contains("sum(salary)", "x"));
         var from = as(agg.child(), EsRelation.class);
     }
