From 7e669831e87220943907da6e731437340683ddc0 Mon Sep 17 00:00:00 2001 From: Luigi Dell'Aquila Date: Mon, 20 Nov 2023 12:43:54 +0100 Subject: [PATCH] ES|QL: Fix drop of renamed grouping (#102282) Fixes https://github.com/elastic/elasticsearch/issues/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 | 6 ++ .../src/main/resources/drop.csv-spec | 33 +++++++++++ .../esql/optimizer/LogicalPlanOptimizer.java | 38 ++++++++++++- .../optimizer/LogicalPlanOptimizerTests.java | 57 +++++++++++++++++++ 4 files changed, 133 insertions(+), 1 deletion(-) create mode 100644 docs/changelog/102282.yaml diff --git a/docs/changelog/102282.yaml b/docs/changelog/102282.yaml new file mode 100644 index 0000000000000..4860d70f99ccc --- /dev/null +++ b/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 diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/drop.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/drop.csv-spec index cd3afa25fc0a6..601b4f329f9d7 100644 --- a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/drop.csv-spec +++ b/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 +; + diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizer.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizer.java index bfa9ea449ad46..29b61949b6778 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizer.java +++ b/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 combineProjections(List return replaced; } + /** + * Replace grouping alias previously contained in the aggregations that might have been projected away. + */ + private List replacePrunedAliasesUsedInGroupBy( + List groupings, + List oldAggs, + List newAggs + ) { + AttributeMap 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(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()); 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 e825f1f96a8b3..b82bb46ec103e 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 @@ -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 aliased(Expression exp, Class clazz) { var alias = as(exp, Alias.class); return as(alias.child(), clazz);