ESQL: Fix injected attributes's IDs in UnionAll branches #141262
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,5 @@ | ||
| pr: 141262 | ||
| summary: Fix injected attributes's IDs in `UnionAll` branches | ||
| area: ES|QL | ||
| type: bug | ||
| issues: [] |
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 |
|---|---|---|
|
|
@@ -19,6 +19,7 @@ | |
| import org.elasticsearch.xpack.esql.core.expression.FieldAttribute; | ||
| import org.elasticsearch.xpack.esql.core.expression.Literal; | ||
| import org.elasticsearch.xpack.esql.core.expression.NameId; | ||
| import org.elasticsearch.xpack.esql.core.expression.NamedExpression; | ||
| import org.elasticsearch.xpack.esql.core.expression.UnresolvedAttribute; | ||
| import org.elasticsearch.xpack.esql.core.expression.UnresolvedPattern; | ||
| import org.elasticsearch.xpack.esql.core.expression.UnresolvedTimestamp; | ||
|
|
@@ -28,7 +29,6 @@ | |
| import org.elasticsearch.xpack.esql.plan.logical.Eval; | ||
| import org.elasticsearch.xpack.esql.plan.logical.Fork; | ||
| import org.elasticsearch.xpack.esql.plan.logical.LeafPlan; | ||
| import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan; | ||
| import org.elasticsearch.xpack.esql.plan.logical.Project; | ||
| import org.elasticsearch.xpack.esql.plan.logical.Row; | ||
|
|
@@ -44,6 +44,7 @@ | |
|
|
||
| import static org.elasticsearch.xpack.esql.analysis.Analyzer.ResolveRefs.insistKeyword; | ||
| import static org.elasticsearch.xpack.esql.core.util.CollectionUtils.combine; | ||
| import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes; | ||
|
|
||
| /** | ||
| * The rule handles fields that don't show up in the index mapping, but are used within the query. These fields can either be missing | ||
|
|
@@ -81,8 +82,9 @@ private static LogicalPlan resolve(LogicalPlan plan, boolean load) { | |
| if (unresolved.isEmpty()) { | ||
| return plan; | ||
| } | ||
| var unresolvedLinkedSet = unresolvedLinkedSet(unresolved); | ||
|
|
||
| var transformed = load ? load(plan, unresolvedLinkedSet) : nullify(plan, unresolvedLinkedSet); | ||
|
Contributor
There was a problem hiding this comment. If the unresolveds being in a linked (thus order-preserving) set is important, should the signature of |
||
|
|
||
| return transformed.equals(plan) ? plan : refreshPlan(transformed, unresolved); | ||
| } | ||
|
|
@@ -91,21 +93,17 @@ private static LogicalPlan resolve(LogicalPlan plan, boolean load) { | |
| * The method introduces {@code EVAL missing_field = NULL}-equivalent into the plan, on top of the source, for every attribute in | ||
| * {@code unresolved}. It also "patches" the introduced attributes through the plan, where needed (like through Fork/UntionAll). | ||
| */ | ||
| private static LogicalPlan nullify(LogicalPlan plan, Set<UnresolvedAttribute> unresolved) { | ||
| // insert an Eval on top of every LeafPlan, if there's a UnaryPlan atop it | ||
| var transformed = plan.transformUp( | ||
| n -> n instanceof UnaryPlan unary && unary.child() instanceof LeafPlan, | ||
| p -> evalUnresolvedAtopUnary((UnaryPlan) p, nullAliases(unresolved)) | ||
| ); | ||
| // insert an Eval on top of those LeafPlan that are children of n-ary plans (could happen with UnionAll) | ||
| return transformed.transformUp( | ||
| n -> n instanceof UnaryPlan == false && n instanceof LeafPlan == false, | ||
| nAry -> evalUnresolvedAtopNary(nAry, nullAliases(unresolved)) | ||
| ); | ||
| } | ||
|
|
||
| /** | ||
|
|
@@ -115,95 +113,73 @@ private static LogicalPlan nullify(LogicalPlan plan, List<UnresolvedAttribute> u | |
| * <p> | ||
| * It also "patches" the introduced attributes through the plan, where needed (like through Fork/UntionAll). | ||
| */ | ||
| private static LogicalPlan load(LogicalPlan plan, Set<UnresolvedAttribute> unresolved) { | ||
| // TODO: this will need to be revisited for non-lookup joining or scenarios where we won't want extraction from specific sources | ||
| return plan.transformUp(EsRelation.class, esr -> { | ||
| if (esr.indexMode() == IndexMode.LOOKUP) { | ||
| return esr; | ||
| } | ||
| List<FieldAttribute> fieldsToLoad = fieldsToLoad(unresolved, Expressions.names(esr.output())); | ||
| // there shouldn't be any duplicates, we can just merge the two lists | ||
| return fieldsToLoad.isEmpty() ? esr : esr.withAttributes(combine(esr.output(), fieldsToLoad)); | ||
| }); | ||
| } | ||
|
|
||
| private static List<FieldAttribute> fieldsToLoad(Set<UnresolvedAttribute> unresolved, List<String> exclude) { | ||
|
Contributor
There was a problem hiding this comment. Isn't this a mistake? I would have expected the
Contributor
Author
There was a problem hiding this comment. It's a Set because the initial collection of UnresolvedAttributes is dedup'd -- this is what What we want here is to exclude those attributes produces by the EsRelation itself into which we would then later inject/ Not sure if it's worth instantiating new collection types to wrap the existing ones. |
||
| List<FieldAttribute> insisted = new ArrayList<>(unresolved.size()); | ||
| for (var ua : unresolved) { | ||
| if (exclude.contains(ua.name()) == false) { | ||
| insisted.add(insistKeyword(ua)); | ||
| } | ||
| } | ||
| return insisted; | ||
| } | ||
|
|
||
| // TODO: would an alternative to this be to have ResolveRefs#resolveFork re-resolve the Fork? | ||
| // We might need some plan delimiters/markers to make it unequivocal which nodes belong to | ||
| // "make Fork work" - like ([Limit -] Project [- Eval])s - and which don't. | ||
| // PruneColumns does the same dance. There's some fragility w.r.t. assuming there to be a top Project and danger of the outputs not | ||
| // being aligned after applying the changes. | ||
| /** | ||
| * Update the Fork's top Projects in the subplans, and correspondingly, its output, to account for newly introduced aliases. | ||
| */ | ||
| private static Fork patchFork(Fork fork) { | ||
| List<LogicalPlan> newChildren = new ArrayList<>(fork.children().size()); | ||
| boolean childrenChanged = false; | ||
| for (var child : fork.children()) { | ||
| Holder<Boolean> patched = new Holder<>(false); | ||
| var transformed = child.transformDown( | ||
| // TODO add a suitable forEachDownMayReturnEarly equivalent | ||
| n -> patched.get() == false && n instanceof Project, // process top Project only (Fork-injected) | ||
| n -> { | ||
| patched.set(true); | ||
| return patchForkProject((Project) n); | ||
| } | ||
| ); | ||
| childrenChanged |= transformed != child; | ||
| newChildren.add(transformed); | ||
| } | ||
| return childrenChanged ? fork.withSubPlans(newChildren) : fork; | ||
| } | ||
|
|
||
| /** | ||
| * Add any missing attributes that are found in the child's output but not in the Project's output. These have been injected before | ||
| * by the evalUnresolvedAtopXXX methods and need to be "let through" the Project. | ||
| */ | ||
| private static Project patchForkProject(Project project) { | ||
| var projectOutput = project.output(); | ||
| var childOutput = project.child().output(); | ||
| if (projectOutput.equals(childOutput) == false) { | ||
|
Comment on lines
+170
to
+172
Contributor
There was a problem hiding this comment.
Contributor
Author
There was a problem hiding this comment.
Since these changes aren't functionally impacting, I'd apply them to a follow-up PR (unless other changes will be required), if ok with you? |
||
| List<Attribute> delta = new ArrayList<>(childOutput); | ||
| delta.removeAll(projectOutput); | ||
| project = project.withProjections(mergeOutputAttributes(delta, projectOutput)); | ||
|
Contributor
There was a problem hiding this comment. Could we please avoid this pattern of renaming the input parameter and then returning it outside the block? Just use an early exit above.
Contributor
Author
There was a problem hiding this comment. This is a pre-existing pattern. Some folks find it easier to read code with fewer returns. (Myself, I don't necessarily, but I don't mind this style either).
...reason being: if the control hasn't visited the block, the input is simply returned with no change.
Contributor
There was a problem hiding this comment. I personally find it unfathomable that it's harder to read code with more early exits than it is to read code with more changes if the input variable, but to each their own I guess 🥲. |
||
| } | ||
| return project; | ||
| } | ||
|
|
||
| private static LogicalPlan refreshPlan(LogicalPlan plan, List<UnresolvedAttribute> unresolved) { | ||
| var refreshed = refreshUnresolved(plan, unresolved); | ||
| return refreshed.transformDown(Fork.class, ResolveUnmapped::patchFork); | ||
| } | ||
|
|
||
| /** | ||
|
|
@@ -222,25 +198,10 @@ private static LogicalPlan refreshUnresolved(LogicalPlan plan, List<UnresolvedAt | |
| }); | ||
| } | ||
|
|
||
| /** | ||
| * Inserts an Eval atop each child of the given {@code nAry}, if the child is a LeafPlan. | ||
| */ | ||
| private static LogicalPlan evalUnresolvedAtopNary(LogicalPlan nAry, List<Alias> nullAliases) { | ||
| List<LogicalPlan> newChildren = new ArrayList<>(nAry.children().size()); | ||
| boolean changed = false; | ||
| for (var child : nAry.children()) { | ||
|
|
@@ -257,7 +218,7 @@ private static LogicalPlan evalUnresolvedNary(LogicalPlan nAry, List<Alias> null | |
| /** | ||
| * Inserts an Eval atop the given {@code unaryAtopSource}, if this isn't an Eval already. Otherwise it merges the nullAliases into it. | ||
| */ | ||
| private static LogicalPlan evalUnresolvedAtopUnary(UnaryPlan unaryAtopSource, List<Alias> nullAliases) { | ||
| assertSourceType(unaryAtopSource.child()); | ||
| if (unaryAtopSource instanceof Eval eval && eval.resolved()) { // if this Eval isn't resolved, insert a new (resolved) one | ||
| List<Alias> pre = new ArrayList<>(nullAliases.size()); | ||
|
|
@@ -291,16 +252,23 @@ private static void assertSourceType(LogicalPlan source) { | |
| } | ||
| } | ||
|
|
||
| private static List<Alias> nullAliases(Set<UnresolvedAttribute> unresolved) { | ||
|
Contributor
There was a problem hiding this comment. Same here: If the output list is supposed to be stable, I think we should explicitly require a |
||
| List<Alias> aliases = new ArrayList<>(unresolved.size()); | ||
|
Contributor
There was a problem hiding this comment. Why not use
Contributor
Author
There was a problem hiding this comment. Not sure I understand how would |
||
| unresolved.forEach(u -> aliases.add(nullAlias(u))); | ||
| return aliases; | ||
| } | ||
|
|
||
| private static Alias nullAlias(NamedExpression attribute) { | ||
| return new Alias(attribute.source(), attribute.name(), NULLIFIED); | ||
| } | ||
|
|
||
| // Some plans may reference the same UA multiple times (Aggregate groupings in aggregates, Eval): dedupe | ||
| private static LinkedHashSet<UnresolvedAttribute> unresolvedLinkedSet(List<UnresolvedAttribute> unresolved) { | ||
| Map<String, UnresolvedAttribute> aliasesMap = new LinkedHashMap<>(unresolved.size()); | ||
| unresolved.forEach(u -> aliasesMap.putIfAbsent(u.name(), u)); | ||
| return new LinkedHashSet<>(aliasesMap.values()); | ||
| } | ||
|
|
||
| /** | ||
| * @return all the {@link UnresolvedAttribute}s in the given node / {@code plan}, but excluding the {@link UnresolvedPattern} and | ||
| * {@link UnresolvedTimestamp} subtypes. | ||
|
|
||
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.
Why not the
containsand.getapproach? I see the code that reaches this part is fairly safe to assume that there won't be anynullsets returned for a key, but we are not sure how this code will evolve in the future.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.
This isn't strictly related, but I thought I might need an update of this code and spotted the pattern.
There's nothing wrong from the functional PoV, but the code, as it was, checks if the key is in the map, then does it again, but fetching the corresponding value. The code as is in the proposed change only does the latter. If the result/value is
null, the key isn't in there. (Well, I guess it could [have] be[en] anullvalue, to be exact, but that would have resulted in a NPE by now(?))