ESQL: Retain aggregate when grouping

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
         "only supports KEYWORD or TEXT values, found expression", // https://github.com/elastic/elasticsearch/issues/126017
         "token recognition error at: '``", // https://github.com/elastic/elasticsearch/issues/125870
-        "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

@@ -636,4 +636,55 @@ row foo = "Bar" | where foo rlike "(?i)Ba.*"
 foo:keyword
 ;
 
+evalAfterAvgGroupingUsingSameName
+required_capability: retain_aggregate_when_grouping
+from employees
+| stats avg = avg(salary) by gender 
+| keep avg 
+| eval avg = 12
+;
+
+avg:integer
+12
+12
+12
+;
+
+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
+;
 

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

@@ -413,6 +413,12 @@ public enum Cap {
          */
         RENAME_SEQUENTIAL_PROCESSING,
 
+        /**
+         * 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.EmptyAttribute;
 import org.elasticsearch.xpack.esql.core.expression.Expressions;
@@ -84,7 +86,8 @@ 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());
+                                remaining = aggregate.aggregates().stream().map(v -> new Alias(v.source(), v.name(), attribute)).toList();
                                 p = aggregate.with(aggregate.groupings(), remaining);
                             }
                         } else {

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

@@ -2784,6 +2784,27 @@ private static List<String> orderNames(TopN topN) {
         return topN.order().stream().map(o -> as(o.child(), NamedExpression.class).name()).toList();
     }
 
+    /**
+     * 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