diff --git a/docs/changelog/126598.yaml b/docs/changelog/126598.yaml new file mode 100644 index 0000000000000..39187f0d60977 --- /dev/null +++ b/docs/changelog/126598.yaml @@ -0,0 +1,6 @@ +pr: 126598 +summary: "ESQL: Retain aggregate when grouping" +area: ES|QL +type: bug +issues: + - 126026 diff --git a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeRestTest.java b/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeRestTest.java index 844c06e8ac17e..6850753a865de 100644 --- a/x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/generative/GenerativeRestTest.java +++ b/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\\[\\[.* 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 diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/eval.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/eval.csv-spec index ec2e63aad659f..787495e204d90 100644 --- a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/eval.csv-spec +++ b/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 +; diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java index cde9e98c9edef..7a85e91ecd982 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java +++ b/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 ES|QL: columns not projected away despite KEEP #126026 + */ + RETAIN_AGGREGATE_WHEN_GROUPING, + /** * Fix for union-types when some indexes are missing the required field. Done in #111932. */ diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PruneColumns.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PruneColumns.java index 1e00411863771..58e20030adb40 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PruneColumns.java +++ b/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 { 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 index 23ab2576a4d71..cbd3998f5a850 100644 --- 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 @@ -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