ES|QL: Fix drop of renamed grouping (#102282)

Fixes #102121

Aggs groupings were not taken into account while merging aggs with
projections, so they were wrongly removed in case of DROP

docs/changelog/102282.yaml

@@ -0,0 +1,6 @@
+pr: 102282
+summary: "ES|QL: Fix drop of renamed grouping"
+area: ES|QL
+type: bug
+issues:
+ - 102121

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

@@ -54,3 +54,36 @@ c:l|mi:i|s:l
 0  |null|null
 ;
 
+
+// see https://github.com/elastic/elasticsearch/issues/102121
+dropGrouping#[skip:-8.11.99, reason:planning bug fixed in v8.12]
+row a = 1 | rename a AS foo | stats bar = count(*) by foo | drop foo;
+
+bar:long
+1
+;
+
+dropGroupingMulti#[skip:-8.11.99]
+row a = 1, b = 2 | rename a AS foo, b as bar | stats baz = count(*) by foo, bar | drop foo;
+
+baz:long  | bar:integer
+1         | 2
+;
+
+dropGroupingMulti2#[skip:-8.11.99]
+row a = 1, b = 2 | rename a AS foo, b as bar | stats baz = count(*) by foo, bar | drop foo, bar;
+
+baz:long
+1
+;
+
+
+dropGroupingMultirow#[skip:-8.11.99]
+from employees | rename gender AS foo | stats bar = count(*) by foo | drop foo | sort bar;
+
+bar:long
+10
+33
+57
+;
+

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

@@ -281,7 +281,10 @@ protected LogicalPlan rule(UnaryPlan plan) {
                     // eliminate lower project but first replace the aliases in the upper one
                     return p.withProjections(combineProjections(project.projections(), p.projections()));
                 } else if (child instanceof Aggregate a) {
-                    return new Aggregate(a.source(), a.child(), a.groupings(), combineProjections(project.projections(), a.aggregates()));
+                    var aggs = a.aggregates();
+                    var newAggs = combineProjections(project.projections(), aggs);
+                    var newGroups = replacePrunedAliasesUsedInGroupBy(a.groupings(), aggs, newAggs);
+                    return new Aggregate(a.source(), a.child(), newGroups, newAggs);
                 }
             }
 
@@ -320,6 +323,39 @@ private List<NamedExpression> combineProjections(List<? extends NamedExpression>
             return replaced;
         }
 
+        /**
+         * Replace grouping alias previously contained in the aggregations that might have been projected away.
+         */
+        private List<Expression> replacePrunedAliasesUsedInGroupBy(
+            List<Expression> groupings,
+            List<? extends NamedExpression> oldAggs,
+            List<? extends NamedExpression> newAggs
+        ) {
+            AttributeMap<Expression> removedAliases = new AttributeMap<>();
+            AttributeSet currentAliases = new AttributeSet(Expressions.asAttributes(newAggs));
+
+            // record only removed aliases
+            for (NamedExpression ne : oldAggs) {
+                if (ne instanceof Alias alias) {
+                    var attr = ne.toAttribute();
+                    if (currentAliases.contains(attr) == false) {
+                        removedAliases.put(attr, alias.child());
+                    }
+                }
+            }
+
+            if (removedAliases.isEmpty()) {
+                return groupings;
+            }
+
+            var newGroupings = new ArrayList<Expression>(groupings.size());
+            for (Expression group : groupings) {
+                newGroupings.add(group.transformUp(Attribute.class, a -> removedAliases.resolve(a, a)));
+            }
+
+            return newGroupings;
+        }
+
         public static Expression trimNonTopLevelAliases(Expression e) {
             if (e instanceof Alias a) {
                 return new Alias(a.source(), a.name(), a.qualifier(), trimAliases(a.child()), a.id());

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

@@ -2454,6 +2454,63 @@ public void testMvExpandFoldable() {
         var row = as(expand.child(), Row.class);
     }
 
+    /**
+     * Expected
+     * Limit[500[INTEGER]]
+     * \_Aggregate[[a{r}#2],[COUNT([2a][KEYWORD]) AS bar]]
+     *   \_Row[[1[INTEGER] AS a]]
+     */
+    public void testRenameStatsDropGroup() {
+        LogicalPlan plan = optimizedPlan("""
+            row a = 1
+            | rename a AS foo
+            | stats bar = count(*) by foo
+            | drop foo""");
+
+        var limit = as(plan, Limit.class);
+        var agg = as(limit.child(), Aggregate.class);
+        assertThat(Expressions.names(agg.groupings()), contains("a"));
+        var row = as(agg.child(), Row.class);
+    }
+
+    /**
+     * Expected
+     * Limit[500[INTEGER]]
+     * \_Aggregate[[a{r}#2, bar{r}#8],[COUNT([2a][KEYWORD]) AS baz, b{r}#4 AS bar]]
+     *   \_Row[[1[INTEGER] AS a, 2[INTEGER] AS b]]
+     */
+    public void testMultipleRenameStatsDropGroup() {
+        LogicalPlan plan = optimizedPlan("""
+            row a = 1, b = 2
+            | rename a AS foo, b as bar
+            | stats baz = count(*) by foo, bar
+            | drop foo""");
+
+        var limit = as(plan, Limit.class);
+        var agg = as(limit.child(), Aggregate.class);
+        assertThat(Expressions.names(agg.groupings()), contains("a", "bar"));
+        var row = as(agg.child(), Row.class);
+    }
+
+    /**
+     * Expected
+     * Limit[500[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, ..]
+     */
+    public void testMultipleRenameStatsDropGroupMultirow() {
+        LogicalPlan plan = optimizedPlan("""
+            from test
+            | rename emp_no AS foo, gender as bar
+            | stats baz = max(salary) by foo, bar
+            | drop foo""");
+
+        var limit = as(plan, Limit.class);
+        var agg = as(limit.child(), Aggregate.class);
+        assertThat(Expressions.names(agg.groupings()), contains("emp_no", "bar"));
+        var row = as(agg.child(), EsRelation.class);
+    }
+
     private <T> T aliased(Expression exp, Class<T> clazz) {
         var alias = as(exp, Alias.class);
         return as(alias.child(), clazz);