Skip to content

Commit b5ede88

Browse files
kanoshioukanoshiou
committed
Create a temporary attribute for the UnionAll case
1 parent d8c221a commit b5ede88

File tree

1 file changed

+75
-23
lines changed

1 file changed

+75
-23
lines changed

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PushDownMvExpandPastProject.java‎

Lines changed: 75 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@
1010
import org.elasticsearch.xpack.esql.core.expression.Alias;
1111
import org.elasticsearch.xpack.esql.core.expression.Attribute;
1212
import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
13+
import org.elasticsearch.xpack.esql.core.expression.Nullability;
14+
import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
1315
import org.elasticsearch.xpack.esql.plan.logical.Eval;
1416
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
1517
import org.elasticsearch.xpack.esql.plan.logical.MvExpand;
@@ -24,34 +26,82 @@ public final class PushDownMvExpandPastProject extends OptimizerRules.OptimizerR
2426
@Override
2527
protected LogicalPlan rule(MvExpand mvExpand) {
2628
if (mvExpand.child() instanceof Project pj) {
27-
List<NamedExpression> projections = new ArrayList<>(pj.projections());
2829
LogicalPlan finalChild = pj.child();
2930
NamedExpression finalTarget = mvExpand.target();
30-
Attribute expanded = mvExpand.expanded();
31+
Attribute finalExpanded = mvExpand.expanded();
3132

32-
// Skip if the expanded field has the same name as a field in the projection's input set, and
33-
// the projection shadows that specific field from the projection input set.
34-
// Pushing down the MvExpand in such cases would cause duplicate output attributes.
35-
// This can happen with aliases generated by ResolveUnionTypesInUnionAll.
36-
// Example:
37-
// MvExpand[salary{r}#168,salary{r}#175]
38-
// \_Project[[$$salary$converted_to$keyword{r$}#178 AS salary#168]]
39-
// \_UnionAll[[salary{r}#174, $$salary$converted_to$keyword{r$}#178]]
40-
String expandedFieldName = expanded.name();
33+
String expandedFieldName = finalExpanded.name();
34+
List<NamedExpression> projections = new ArrayList<>(pj.projections());
4135
Set<String> inputNames = pj.inputSet().stream().map(NamedExpression::name).collect(Collectors.toSet());
42-
if (projections.stream()
43-
.anyMatch(
44-
e -> e instanceof Alias alias
45-
&& inputNames.contains(expandedFieldName)
46-
&& inputNames.contains(alias.toAttribute().name())
47-
)) {
48-
return mvExpand;
49-
}
5036

5137
// Find if the target is aliased in the project and create an alias with temporary names for it.
5238
for (int i = 0; i < projections.size(); i++) {
5339
if (projections.get(i) instanceof Alias alias) {
40+
boolean replaced = false;
41+
/*
42+
* If the expanded field has the same name as a field in the projection's input set,
43+
* and the projection shadows that specific field from the projection input set.
44+
* Pushing down the MvExpand in such cases would cause duplicate output attributes.
45+
* To avoid this case, we create a temporary attribute for the expanded field and
46+
* update the projection to alias this temporary attribute back to the original name.
47+
* This can happen with aliases generated by ResolveUnionTypesInUnionAll.
48+
*
49+
* Example query:
50+
* from employees, (from employees | keep salary)
51+
* | eval salary = salary::keyword
52+
* | keep salary
53+
* | mv_expand salary
54+
*
55+
* From plan:
56+
* MvExpand[language_code{r}#4,language_code{r}#17]
57+
* \_Project[[$$language_code$converted_to$keyword{r$}#20 AS language_code#4]]
58+
* \_UnionAll[[language_code{r}#15, $$language_code$converted_to$keyword{r$}#20, language_name{r}#16]]
59+
*
60+
* To plan:
61+
* Project[[$$language_code$temp_name$21{r$}#22 AS language_code#17]]
62+
* \_MvExpand[$$language_code$converted_to$keyword{r$}#20,$$language_code$temp_name$21{r$}#22]
63+
* \_UnionAll[[language_code{r}#15, $$language_code$converted_to$keyword{r$}#20, language_name{r}#16]]
64+
*
65+
*
66+
* If the original mv_expand target field is referenced elsewhere in the projections,
67+
* a defensive eval will also be injected.
68+
*
69+
* Example query:
70+
* from languages, (from languages | keep language_code)
71+
* | eval language_code = language_code::keyword
72+
* | eval tmp = language_code
73+
* | keep language_code, tmp
74+
* | mv_expand language_code
75+
*
76+
* From plan:
77+
* MvExpand[language_code{r}#4,language_code{r}#22]
78+
* \_Project[[$$language_code$converted_to$keyword{r$}#25 AS language_code#4,$$language_code$converted_to$keyword{r$}#25
79+
* AS tmp#7]]
80+
* \_UnionAll[[language_code{r}#20, $$language_code$converted_to$keyword{r$}#25, language_name{r}#21]]
81+
*
82+
* To plan:
83+
* Project[[$$language_code$temp_name$26{r$}#27 AS language_code#22, $$language_code$converted_to$keyword{r$}#25
84+
* AS tmp#7]]
85+
* \_MvExpand[$$language_code$converted_to$keyword$language_code$0{r}#28,$$language_code$temp_name$26{r$}#27]
86+
* \_Eval[[$$language_code$converted_to$keyword{r$}#25 AS $$language_code$converted_to$keyword$language_code$0#28]]
87+
* \_UnionAll[[language_code{r}#20, $$language_code$converted_to$keyword{r$}#25, language_name{r}#21]]
88+
*/
5489
if (alias.toAttribute().semanticEquals(finalTarget.toAttribute())) {
90+
if (inputNames.contains(expandedFieldName) && inputNames.contains(alias.toAttribute().name())) {
91+
ReferenceAttribute tempAttribute = new ReferenceAttribute(
92+
alias.source(),
93+
null,
94+
TemporaryNameUtils.locallyUniqueTemporaryName(alias.name()),
95+
alias.dataType(),
96+
Nullability.FALSE,
97+
null,
98+
true
99+
);
100+
projections.set(i, new Alias(alias.source(), expandedFieldName, tempAttribute, finalExpanded.id()));
101+
finalExpanded = tempAttribute;
102+
replaced = true;
103+
}
104+
55105
// Check if the alias's original field (child) is referenced elsewhere in the projections.
56106
// If the original field is not referenced by any other projection or alias,
57107
// we don't need to inject an Eval to preserve it, and can safely resolve renames and push down.
@@ -71,7 +121,9 @@ protected LogicalPlan rule(MvExpand mvExpand) {
71121
TemporaryNameUtils.temporaryName(alias.child(), alias.toAttribute(), 0),
72122
alias.child()
73123
);
74-
projections.set(i, expanded);
124+
if (replaced == false) {
125+
projections.set(i, finalExpanded);
126+
}
75127
finalChild = new Eval(aliasAlias.source(), finalChild, List.of(aliasAlias));
76128
finalTarget = aliasAlias.toAttribute();
77129
break;
@@ -90,16 +142,16 @@ protected LogicalPlan rule(MvExpand mvExpand) {
90142
}
91143

92144
// Push down the MvExpand past the Project
93-
MvExpand pushedDownMvExpand = new MvExpand(mvExpand.source(), finalChild, finalTarget, expanded);
145+
MvExpand pushedDownMvExpand = new MvExpand(mvExpand.source(), finalChild, finalTarget, finalExpanded);
94146

95147
// Update projections to point to the expanded attribute
96148
Attribute target = finalTarget.toAttribute();
97149
for (int i = 0; i < projections.size(); i++) {
98150
NamedExpression ne = projections.get(i);
99151
if (ne instanceof Alias alias && alias.child().semanticEquals(target)) {
100-
projections.set(i, alias.replaceChild(expanded));
152+
projections.set(i, alias.replaceChild(finalExpanded));
101153
} else if (ne.semanticEquals(target)) {
102-
projections.set(i, expanded);
154+
projections.set(i, finalExpanded);
103155
}
104156
}
105157

0 commit comments

Comments
 (0)