ESQL: Push down MvExpand past Project
#136398
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,8 @@ | ||
| pr: 136398 | ||
| summary: "ESQL: Push down `MvExpand` past `Project`" | ||
| area: ES|QL | ||
| type: enhancement | ||
| issues: | ||
| - 136292 | ||
| - 136596 | ||
| - 119074 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,112 @@ | ||
| /* | ||
| * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one | ||
| * or more contributor license agreements. Licensed under the Elastic License | ||
| * 2.0; you may not use this file except in compliance with the Elastic License | ||
| * 2.0. | ||
| */ | ||
|
|
||
| package org.elasticsearch.xpack.esql.optimizer.rules.logical; | ||
|
|
||
| import org.elasticsearch.xpack.esql.core.expression.Alias; | ||
| import org.elasticsearch.xpack.esql.core.expression.Attribute; | ||
| import org.elasticsearch.xpack.esql.core.expression.AttributeMap; | ||
| import org.elasticsearch.xpack.esql.core.expression.NamedExpression; | ||
| import org.elasticsearch.xpack.esql.plan.logical.Eval; | ||
| import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan; | ||
| import org.elasticsearch.xpack.esql.plan.logical.MvExpand; | ||
| import org.elasticsearch.xpack.esql.plan.logical.Project; | ||
|
|
||
| import java.util.ArrayList; | ||
| import java.util.List; | ||
| import java.util.Set; | ||
| import java.util.stream.Collectors; | ||
|
|
||
| public final class PushDownMvExpandPastProject extends OptimizerRules.OptimizerRule<MvExpand> { | ||
| @Override | ||
| protected LogicalPlan rule(MvExpand mvExpand) { | ||
| if (mvExpand.child() instanceof Project pj) { | ||
| List<NamedExpression> projections = new ArrayList<>(pj.projections()); | ||
|
|
||
| // Skip if any projection alias shadows an input field name, as injecting an Eval would cause | ||
| // duplicate output attributes. This can happen with aliases generated by ResolveUnionTypesInUnionAll. | ||
| // Example: | ||
| // MvExpand[salary{r}#168,salary{r}#175] | ||
| // \_Project[[$$salary$converted_to$keyword{r$}#178 AS salary#168]] | ||
| // \_UnionAll[[salary{r}#174, $$salary$converted_to$keyword{r$}#178]] | ||
|
||
| Set<String> inputNames = pj.inputSet().stream().map(NamedExpression::name).collect(Collectors.toSet()); | ||
| if (projections.stream().anyMatch(e -> e instanceof Alias alias && inputNames.contains(alias.toAttribute().name()))) { | ||
| return mvExpand; | ||
| } | ||
|
||
|
|
||
| // Find if the target is aliased in the project and create an alias with temporary names for it. | ||
| for (int i = 0; i < projections.size(); i++) { | ||
| NamedExpression projection = projections.get(i); | ||
| if (projection instanceof Alias alias) { | ||
| if (alias.toAttribute().semanticEquals(mvExpand.target().toAttribute())) { | ||
| // Check if the alias's original field (child) is referenced elsewhere in the projections. | ||
| // If the original field is not referenced by any other projection or alias, | ||
| // we don't need to inject an Eval to preserve it, and can safely resolve renames and push down. | ||
|
Comment on lines
+105
to
+107
Contributor
There was a problem hiding this comment. Very helpful comments, btw! |
||
| if (projections.stream() | ||
| .anyMatch( | ||
| ne -> ne.semanticEquals(alias.child()) | ||
| || ne instanceof Alias as && as.child().semanticEquals(alias.child()) && as != alias | ||
| ) == false) { | ||
| // The alias's original field is not referenced elsewhere, no need to preserve it, | ||
| mvExpand = PushDownUtils.resolveRenamesFromProject(mvExpand, pj); | ||
| break; | ||
| } | ||
|
|
||
| // for query like: row a = 2 | eval b = a ｜ keep * | mv_expand b | ||
| Alias aliasAlias = new Alias( | ||
| alias.source(), | ||
| TemporaryNameUtils.temporaryName(alias.child(), alias.toAttribute(), 0), | ||
| alias.child() | ||
| ); | ||
| projections.set(i, mvExpand.expanded()); | ||
| pj = new Project(pj.source(), new Eval(aliasAlias.source(), pj.child(), List.of(aliasAlias)), projections); | ||
|
||
| mvExpand = new MvExpand(mvExpand.source(), pj, aliasAlias.toAttribute(), mvExpand.expanded()); | ||
| break; | ||
| } else if (alias.child().semanticEquals(mvExpand.target().toAttribute())) { | ||
| // for query like: row a = 2 | eval b = a ｜ keep * | mv_expand a | ||
| Alias aliasAlias = new Alias( | ||
| alias.source(), | ||
| TemporaryNameUtils.temporaryName(alias.child(), alias.toAttribute(), 0), | ||
| alias.child() | ||
| ); | ||
| projections.set(i, alias.replaceChild(aliasAlias.toAttribute())); | ||
| pj = new Project(pj.source(), new Eval(aliasAlias.source(), pj.child(), List.of(aliasAlias)), projections); | ||
| mvExpand = mvExpand.replaceChild(pj); | ||
| break; | ||
| } | ||
| } | ||
| } | ||
|
|
||
| // Build a map of aliases from the original projection | ||
|
||
| AttributeMap.Builder<Alias> aliasBuilder = AttributeMap.builder(); | ||
| pj.forEachExpression(Alias.class, alias -> aliasBuilder.put(alias.toAttribute(), alias)); | ||
| AttributeMap<Alias> aliases = aliasBuilder.build(); | ||
|
|
||
| // Push down the MvExpand past the Project | ||
| MvExpand pushedDownMvExpand = mvExpand.replaceChild(pj.child()); | ||
|
|
||
| // Create a new projection at the top based on mvExpand.output(), plugging back in the aliases | ||
| List<NamedExpression> newProjections = new ArrayList<>(); | ||
| for (Attribute outputExpr : mvExpand.output()) { | ||
|
||
| Alias alias = aliases.resolve(outputExpr, null); | ||
|
||
| if (alias == null) { | ||
| // No alias, use the output expression directly | ||
| newProjections.add(outputExpr); | ||
| } else if (alias.child().semanticEquals(mvExpand.target().toAttribute())) { | ||
| // Alias child is the target attribute, replace it with the expanded attribute | ||
| newProjections.add(alias.replaceChild(mvExpand.expanded())); | ||
| } else { | ||
| // Keep the alias as is | ||
| newProjections.add(alias); | ||
| } | ||
| } | ||
|
|
||
| return new Project(pj.source(), pushedDownMvExpand, newProjections); | ||
| } | ||
| return mvExpand; | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nit: maybe we could add some tests where the expanded value is actually an MV, like
eval = [1,2].