From b808bba29635eb1719b37dd1eb48d7064d86dc79 Mon Sep 17 00:00:00 2001 From: kanoshiou Date: Thu, 10 Apr 2025 19:43:41 +0800 Subject: [PATCH 01/10] Retain aggregates --- .../xpack/esql/optimizer/rules/logical/PruneColumns.java | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 c5115aeda944c..652bec22a9484 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.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 { From 46faa0725017e62ea0b032d558428ada06ed1eb9 Mon Sep 17 00:00:00 2001 From: kanoshiou Date: Thu, 10 Apr 2025 20:12:48 +0800 Subject: [PATCH 02/10] Update csv tests --- .../src/main/resources/eval.csv-spec | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) 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 72660c11d8b73..760271b637457 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 @@ -636,4 +636,52 @@ row foo = "Bar" | where foo rlike "(?i)Ba.*" foo:keyword ; +evalAfterAvgGroupingUsingSameName +from employees +| stats avg = avg(salary) by gender +| keep avg +| eval avg = 12 +; + +avg:integer +12 +12 +12 +; + +evalAfterGroupingUsingSameName +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 +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 +; From cab7c3ed3d1ad30bc21f4d274a190c0949e00cc1 Mon Sep 17 00:00:00 2001 From: kanoshiou Date: Thu, 10 Apr 2025 20:19:26 +0800 Subject: [PATCH 03/10] Update docs/changelog/126598.yaml --- docs/changelog/126598.yaml | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 docs/changelog/126598.yaml 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 From 7745195dfabdd3df6ec12506c2b5d4b7f1db659d Mon Sep 17 00:00:00 2001 From: kanoshiou Date: Thu, 10 Apr 2025 20:47:28 +0800 Subject: [PATCH 04/10] Fix bwc --- .../esql/qa/testFixtures/src/main/resources/eval.csv-spec | 3 +++ .../elasticsearch/xpack/esql/action/EsqlCapabilities.java | 6 ++++++ 2 files changed, 9 insertions(+) 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 760271b637457..ee32caea71081 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 @@ -637,6 +637,7 @@ foo:keyword ; evalAfterAvgGroupingUsingSameName +required_capability: retain_aggregate_when_grouping from employees | stats avg = avg(salary) by gender | keep avg @@ -650,6 +651,7 @@ avg:integer ; evalAfterGroupingUsingSameName +required_capability: retain_aggregate_when_grouping row foo = [10,11,9], bar = [1,2,3] | mv_expand foo | mv_expand bar @@ -665,6 +667,7 @@ this:integer ; evalAfterGroupingUsingSameName2 +required_capability: retain_aggregate_when_grouping from employees | stats count = count(emp_no) by gender, is_rehired | keep count 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 ba9fd985a7877..c454a183d65d9 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 @@ -413,6 +413,12 @@ public enum Cap { */ RENAME_SEQUENTIAL_PROCESSING, + /** + * 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. */ From 65791cccb82752620942ded88b4e54669ebbddd7 Mon Sep 17 00:00:00 2001 From: kanoshiou Date: Thu, 10 Apr 2025 23:29:43 +0800 Subject: [PATCH 05/10] Update LogicalPlanOptimizerTests --- .../optimizer/LogicalPlanOptimizerTests.java | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) 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 47ca4ec31ee5e..914bae4046181 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 @@ -2779,6 +2779,28 @@ private static List 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 From a18e16703c944f8039efa5efd104c673c53a8c89 Mon Sep 17 00:00:00 2001 From: elasticsearchmachine Date: Wed, 16 Apr 2025 13:34:30 +0000 Subject: [PATCH 06/10] [CI] Auto commit changes from spotless --- .../xpack/esql/optimizer/LogicalPlanOptimizerTests.java | 1 - 1 file changed, 1 deletion(-) 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 914bae4046181..fd5fec0e2695f 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 @@ -2779,7 +2779,6 @@ private static List orderNames(TopN topN) { return topN.order().stream().map(o -> as(o.child(), NamedExpression.class).name()).toList(); } - /** * Expects * Eval[[2[INTEGER] AS x]] From ca6f19b760a1a7f6c3a3b3b2d63f5177c4f56ff7 Mon Sep 17 00:00:00 2001 From: kanoshiou Date: Fri, 18 Apr 2025 00:37:57 +0800 Subject: [PATCH 07/10] Remove 126026 from `ALLOWED_ERRORS` --- .../xpack/esql/qa/rest/generative/GenerativeRestTest.java | 1 - 1 file changed, 1 deletion(-) 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 f609fc2bec403..6e8a3e7aeecb5 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 "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 "JOIN left field .* is incompatible with right field", // https://github.com/elastic/elasticsearch/issues/126419 From df357aa2dde0088f147cdbca292f212dab2f769e Mon Sep 17 00:00:00 2001 From: kanoshiou Date: Fri, 18 Apr 2025 19:39:38 +0800 Subject: [PATCH 08/10] Update --- .../qa/testFixtures/src/main/resources/eval.csv-spec | 12 ++++++++++++ .../esql/optimizer/rules/logical/PruneColumns.java | 5 ++++- .../esql/optimizer/LogicalPlanOptimizerTests.java | 2 +- 3 files changed, 17 insertions(+), 2 deletions(-) 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 ee32caea71081..a9192c6d7f94a 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 @@ -688,3 +688,15 @@ x:integer 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 +; 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 652bec22a9484..c0d4751becace 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 @@ -87,7 +87,10 @@ public LogicalPlan apply(LogicalPlan plan) { } else { // Aggs cannot produce pages with 0 columns, so retain one grouping. Attribute attribute = Expressions.attribute(aggregate.groupings().getFirst()); - remaining = aggregate.aggregates().stream().map(v -> new Alias(v.source(), v.name(), attribute)).toList(); + 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 aebb6d15bf424..019e3db8b2999 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)); From 8b93e2a5cb8bdfe68a698d65e1cb1717ea6a07bd Mon Sep 17 00:00:00 2001 From: kanoshiou Date: Wed, 23 Apr 2025 09:41:24 +0800 Subject: [PATCH 09/10] Update tests --- .../src/main/resources/eval.csv-spec | 28 ++++++++++++++----- 1 file changed, 21 insertions(+), 7 deletions(-) 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 f0aa4e301ff01..7a633fb580cb4 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 @@ -677,15 +677,15 @@ max:integer evalAfterAvgGroupingUsingSameName required_capability: retain_aggregate_when_grouping from employees -| stats avg = avg(salary) by gender -| keep avg -| eval avg = 12 +| stats max = max(salary), avg = avg(salary) by g = gender +| keep max, g +| eval max = 12, g = 3 ; -avg:integer -12 -12 -12 +avg:integer | g:integer +12 | 3 +12 | 3 +12 | 3 ; evalAfterGroupingUsingSameName @@ -738,3 +738,17 @@ row foo = 10 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 +; From f1b55246cb4a8595a160343b5a47dda00915de07 Mon Sep 17 00:00:00 2001 From: kanoshiou Date: Wed, 23 Apr 2025 17:31:46 +0800 Subject: [PATCH 10/10] Fix --- .../esql/qa/testFixtures/src/main/resources/eval.csv-spec | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 7a633fb580cb4..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 @@ -682,7 +682,7 @@ from employees | eval max = 12, g = 3 ; -avg:integer | g:integer +max:integer | g:integer 12 | 3 12 | 3 12 | 3