diff --git a/docs/changelog/107131.yaml b/docs/changelog/107131.yaml new file mode 100644 index 0000000000000..ebb696931777b --- /dev/null +++ b/docs/changelog/107131.yaml @@ -0,0 +1,6 @@ +pr: 107131 +summary: "ESQL: Fix bug when combining projections" +area: ES|QL +type: bug +issues: + - 107083 diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/stats.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/stats.csv-spec index 547d6be8dec8e..2c7f16e806d88 100644 --- a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/stats.csv-spec +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/stats.csv-spec @@ -1315,3 +1315,37 @@ FROM employees rows:l 6 ; + +evalOverridingKey#[skip:-8.13.1,reason:supported in 8.13.2] +FROM employees +| EVAL k = languages +| STATS c = COUNT() BY languages, k +| DROP k +| SORT languages +; + +c:l| languages:i +15 | 1 +19 | 2 +17 | 3 +18 | 4 +21 | 5 +10 | null +; + +evalMultipleOverridingKeys#[skip:-8.13.1,reason:supported in 8.13.2] +FROM employees +| EVAL k = languages, k1 = k +| STATS c = COUNT() BY languages, k, k1, languages +| DROP k +| SORT languages +; + +c:l | k1:i | languages:i +15 | 1 | 1 +19 | 2 | 2 +17 | 3 | 3 +18 | 4 | 4 +21 | 5 | 5 +10 | null | null +; 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 437453642b686..35f8fce828b79 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 @@ -363,11 +363,6 @@ private static List projectAggregations( List upperProjection, List lowerAggregations ) { - AttributeMap lowerAliases = new AttributeMap<>(); - for (NamedExpression ne : lowerAggregations) { - lowerAliases.put(ne.toAttribute(), Alias.unwrap(ne)); - } - AttributeSet seen = new AttributeSet(); for (NamedExpression upper : upperProjection) { Expression unwrapped = Alias.unwrap(upper); @@ -391,11 +386,18 @@ private static List combineProjections( List lower ) { - // collect aliases in the lower list - AttributeMap aliases = new AttributeMap<>(); + // collect named expressions declaration in the lower list + AttributeMap namedExpressions = new AttributeMap<>(); + // while also collecting the alias map for resolving the source (f1 = 1, f2 = f1, etc..) + AttributeMap aliases = new AttributeMap<>(); for (NamedExpression ne : lower) { - if ((ne instanceof Attribute) == false) { - aliases.put(ne.toAttribute(), ne); + // record the alias + aliases.put(ne.toAttribute(), Alias.unwrap(ne)); + + // record named expression as is + if (ne instanceof Alias as) { + Expression child = as.child(); + namedExpressions.put(ne.toAttribute(), as.replaceChild(aliases.resolve(child, child))); } } List replaced = new ArrayList<>(); @@ -403,7 +405,7 @@ private static List combineProjections( // replace any matching attribute with a lower alias (if there's a match) // but clean-up non-top aliases at the end for (NamedExpression ne : upper) { - NamedExpression replacedExp = (NamedExpression) ne.transformUp(Attribute.class, a -> aliases.resolve(a, a)); + NamedExpression replacedExp = (NamedExpression) ne.transformUp(Attribute.class, a -> namedExpressions.resolve(a, a)); replaced.add((NamedExpression) trimNonTopLevelAliases(replacedExp)); } return replaced; @@ -436,7 +438,10 @@ private List replacePrunedAliasesUsedInGroupBy( var newGroupings = new ArrayList(groupings.size()); for (Expression group : groupings) { - newGroupings.add(group.transformUp(Attribute.class, a -> removedAliases.resolve(a, a))); + var transformed = group.transformUp(Attribute.class, a -> removedAliases.resolve(a, a)); + if (Expressions.anyMatch(newGroupings, g -> Expressions.equalsAsAttribute(g, transformed)) == false) { + newGroupings.add(transformed); + } } return newGroupings; 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 bb38e3843a661..74dc952c2200a 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 @@ -306,6 +306,52 @@ public void testCombineProjections() { var relation = as(limit.child(), EsRelation.class); } + /** + * Expects + * Project[[languages{f}#12 AS f2]] + * \_Limit[1000[INTEGER]] + * \_EsRelation[test][_meta_field{f}#15, emp_no{f}#9, first_name{f}#10, g..] + */ + public void testCombineProjectionsWithEvalAndDrop() { + var plan = plan(""" + from test + | eval f1 = languages, f2 = f1 + | keep f2 + """); + + var keep = as(plan, Project.class); + assertThat(Expressions.names(keep.projections()), contains("f2")); + assertThat(Expressions.name(Alias.unwrap(keep.projections().get(0))), is("languages")); + var limit = as(keep.child(), Limit.class); + var relation = as(limit.child(), EsRelation.class); + + } + + /** + * Expects + * Project[[last_name{f}#26, languages{f}#25 AS f2, f4{r}#13]] + * \_Eval[[languages{f}#25 + 3[INTEGER] AS f4]] + * \_Limit[1000[INTEGER]] + * \_EsRelation[test][_meta_field{f}#28, emp_no{f}#22, first_name{f}#23, ..] + */ + public void testCombineProjectionsWithEval() { + var plan = plan(""" + from test + | eval f1 = languages, f2 = f1, f3 = 1 + 2, f4 = f3 + languages + | keep emp_no, *name, salary, f* + | drop f3 + | keep last_name, f2, f4 + """); + + var keep = as(plan, Project.class); + assertThat(Expressions.names(keep.projections()), contains("last_name", "f2", "f4")); + var eval = as(keep.child(), Eval.class); + assertThat(Expressions.names(eval.fields()), contains("f4")); + var add = as(Alias.unwrap(eval.fields().get(0)), Add.class); + var limit = as(eval.child(), Limit.class); + var relation = as(limit.child(), EsRelation.class); + } + public void testCombineProjectionWithFilterInBetween() { var plan = plan(""" from test @@ -348,6 +394,27 @@ public void testCombineProjectionWithAggregation() { assertThat(Expressions.names(agg.groupings()), contains("last_name", "first_name")); } + /** + * Expects + * Limit[1000[INTEGER]] + * \_Aggregate[[last_name{f}#23, first_name{f}#20, k{r}#4],[SUM(salary{f}#24) AS s, last_name{f}#23, first_name{f}#20, first_n + * ame{f}#20 AS k]] + * \_EsRelation[test][_meta_field{f}#25, emp_no{f}#19, first_name{f}#20, ..] + */ + public void testCombineProjectionWithAggregationAndEval() { + var plan = plan(""" + from test + | eval k = first_name, k1 = k + | stats s = sum(salary) by last_name, first_name, k, k1 + | keep s, last_name, first_name, k + """); + + var limit = as(plan, Limit.class); + var agg = as(limit.child(), Aggregate.class); + assertThat(Expressions.names(agg.aggregates()), contains("s", "last_name", "first_name", "k")); + assertThat(Expressions.names(agg.groupings()), contains("last_name", "first_name", "k")); + } + /** * Project[[s{r}#4 AS d, s{r}#4, last_name{f}#21, first_name{f}#18]] * \_Limit[1000[INTEGER]]