ESQL: Retain aggregate when grouping (#126598)

Author: kanoshiou

docs/changelog/126598.yaml

@@ -0,0 +1,6 @@
+pr: 126598
+summary: "ESQL: Retain aggregate when grouping"
+area: ES|QL
+type: bug
+issues:
+  - 126026

x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeRestTest.java

@@ -52,7 +52,6 @@ public abstract class GenerativeRestTest extends ESRestTestCase {
         "Plan \\[ProjectExec\\[\\[<no-fields>.* optimized incorrectly due to missing references", // https://github.com/elastic/elasticsearch/issues/125866
         "token recognition error at: '``", // https://github.com/elastic/elasticsearch/issues/125870
                                            // https://github.com/elastic/elasticsearch/issues/127167
-        "Unknown column \\[.*\\]", // https://github.com/elastic/elasticsearch/issues/126026
         "optimized incorrectly due to missing references", // https://github.com/elastic/elasticsearch/issues/116781
         "No matches found for pattern", // https://github.com/elastic/elasticsearch/issues/126418
         "The incoming YAML document exceeds the limit:" // still to investigate, but it seems to be specific to the test framework

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

@@ -673,3 +673,82 @@ ROW foo = [10, 11, 12]
 max:integer
 13
 ;
+
+evalAfterAvgGroupingUsingSameName
+required_capability: retain_aggregate_when_grouping
+from employees
+| stats max = max(salary), avg = avg(salary) by g = gender 
+| keep max, g
+| eval max = 12, g = 3
+;
+
+max:integer | g:integer
+12          | 3
+12          | 3
+12          | 3
+;
+
+evalAfterGroupingUsingSameName
+required_capability: retain_aggregate_when_grouping
+row foo = [10,11,9], bar = [1,2,3]
+| mv_expand foo
+| mv_expand bar 
+| stats this = max(foo) by bar 
+| keep this 
+| eval this = 12
+;
+
+this:integer
+12
+12
+12
+;
+
+evalAfterGroupingUsingSameName2
+required_capability: retain_aggregate_when_grouping
+from employees 
+| stats count = count(emp_no) by gender, is_rehired 
+| keep count 
+| rename count as x 
+| keep x 
+| eval x = 12
+;
+
+x:integer
+12
+12
+12
+12
+12
+12
+12
+12
+12
+;
+
+evalAfterGroupingUsingSameName3
+required_capability: retain_aggregate_when_grouping
+row foo = 10 
+| stats field_1 = max(foo), this = count(*), field_2 = max(foo) by foo  
+| rename this AS field_2 
+| keep field_1, field_2  
+| eval field_2 = 1, field_1 = 1
+;
+
+field_2:integer | field_1:integer
+1               | 1
+;
+
+evalAfterGroupingUsingSameName4
+required_capability: retain_aggregate_when_grouping
+from employees
+| stats sum = sum(salary), max = max(salary) by g = gender
+| keep max, g
+| eval g = 12, max = 2
+;
+
+g:integer   | max:integer
+12          | 2
+12          | 2
+12          | 2
+;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -418,6 +418,12 @@ public enum Cap {
          */
         REMOVE_EMPTY_ATTRIBUTE_IN_MERGING_OUTPUT,
 
+        /**
+         * Support for retain aggregate when grouping.
+         * See <a href="https://github.com/elastic/elasticsearch/issues/126026"> ES|QL: columns not projected away despite KEEP #126026 </a>
+         */
+        RETAIN_AGGREGATE_WHEN_GROUPING,
+
         /**
          * Fix for union-types when some indexes are missing the required field. Done in #111932.
          */

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

@@ -10,6 +10,8 @@
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockUtils;
 import org.elasticsearch.index.IndexMode;
+import org.elasticsearch.xpack.esql.core.expression.Alias;
+import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
 import org.elasticsearch.xpack.esql.core.expression.Expressions;
 import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
@@ -83,7 +85,11 @@ public LogicalPlan apply(LogicalPlan plan) {
                                 );
                             } else {
                                 // Aggs cannot produce pages with 0 columns, so retain one grouping.
-                                remaining = List.of(Expressions.attribute(aggregate.groupings().get(0)));
+                                Attribute attribute = Expressions.attribute(aggregate.groupings().getFirst());
+                                NamedExpression firstAggregate = aggregate.aggregates().getFirst();
+                                remaining = List.of(
+                                    new Alias(firstAggregate.source(), firstAggregate.name(), attribute, firstAggregate.id())
+                                );
                                 p = aggregate.with(aggregate.groupings(), remaining);
                             }
                         } else {

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

@@ -423,7 +423,7 @@ public void testEmptyProjectInStatWithGroupAndEval() {
         var relation = as(filter.child(), EsRelation.class);
 
         assertThat(Expressions.names(agg.groupings()), contains("emp_no"));
-        assertThat(Expressions.names(agg.aggregates()), contains("emp_no"));
+        assertThat(Expressions.names(agg.aggregates()), contains("c"));
 
         var exprs = eval.fields();
         assertThat(exprs.size(), equalTo(1));
@@ -2802,6 +2802,27 @@ public void testEvalAfterStats() {
         assertThat(Expressions.names(eval.output()), contains("x"));
     }
 
+    /**
+     * Expects
+     * Eval[[2[INTEGER] AS x]]
+     * \_Limit[1000[INTEGER],false]
+     *   \_Aggregate[[foo{r}#3],[foo{r}#3 AS x]]
+     *     \_LocalRelation[[foo{r}#3],[IntVectorBlock[vector=ConstantIntVector[positions=1, value=1]]]]
+     */
+    public void testEvalAfterGroupBy() {
+        var plan = optimizedPlan("""
+            ROW foo = 1
+            | STATS x = max(foo) by foo
+            | KEEP x
+            | EVAL x = 2
+            """);
+        var eval = as(plan, Eval.class);
+        var limit = as(eval.child(), Limit.class);
+        var aggregate = as(limit.child(), Aggregate.class);
+        var localRelation = as(aggregate.child(), LocalRelation.class);
+        assertThat(Expressions.names(eval.output()), contains("x"));
+    }
+
     public void testCombineLimitWithOrderByThroughFilterAndEval() {
         LogicalPlan plan = optimizedPlan("""
             from test