Skip to content

Commit b042096

Browse files
committed
fix(esql): refactor MvExpand pushdown logic to directly manipulate projections
- Replace AttributeMap import with Attribute for more direct attribute handling - Inline MvExpand pushdown logic instead of delegating to utility methods - Directly update projections to point to expanded attributes during pushdown - Handle both Alias and direct NamedExpression cases when replacing target references - Simplify the overall flow by constructing Project directly with updated projections - Improve code clarity and maintainability by making the transformation explicit
1 parent dd86f32 commit b042096

File tree

1 file changed

+19
-5
lines changed

1 file changed

+19
-5
lines changed

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

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
package org.elasticsearch.xpack.esql.optimizer.rules.logical;
99

1010
import org.elasticsearch.xpack.esql.core.expression.Alias;
11-
import org.elasticsearch.xpack.esql.core.expression.AttributeMap;
11+
import org.elasticsearch.xpack.esql.core.expression.Attribute;
1212
import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
1313
import org.elasticsearch.xpack.esql.plan.logical.Eval;
1414
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
@@ -70,7 +70,7 @@ protected LogicalPlan rule(MvExpand mvExpand) {
7070
alias.child()
7171
);
7272
projections.set(i, mvExpand.expanded());
73-
pj = new Project(pj.source(), new Eval(aliasAlias.source(), pj.child(), List.of(aliasAlias)), projections);
73+
pj = pj.replaceChild(new Eval(aliasAlias.source(), pj.child(), List.of(aliasAlias)));
7474
mvExpand = new MvExpand(mvExpand.source(), pj, aliasAlias.toAttribute(), mvExpand.expanded());
7575
break;
7676
} else if (alias.child().semanticEquals(mvExpand.target().toAttribute())) {
@@ -81,15 +81,29 @@ protected LogicalPlan rule(MvExpand mvExpand) {
8181
alias.child()
8282
);
8383
projections.set(i, alias.replaceChild(aliasAlias.toAttribute()));
84-
pj = new Project(pj.source(), new Eval(aliasAlias.source(), pj.child(), List.of(aliasAlias)), projections);
84+
pj = pj.replaceChild(new Eval(aliasAlias.source(), pj.child(), List.of(aliasAlias)));
8585
mvExpand = mvExpand.replaceChild(pj);
8686
break;
8787
}
8888
}
8989
}
9090

91-
Project project = PushDownUtils.pushDownPastProject(mvExpand);
92-
return PushDownUtils.resolveRenamesFromMap(project, AttributeMap.of(mvExpand.target().toAttribute(), mvExpand.expanded()));
91+
// Push down the MvExpand past the Project
92+
MvExpand pushedDownMvExpand = mvExpand.replaceChild(pj.child());
93+
94+
// Update projections to point to the expanded attribute
95+
Attribute target = mvExpand.target().toAttribute();
96+
Attribute expanded = mvExpand.expanded();
97+
for (int i = 0; i < projections.size(); i++) {
98+
NamedExpression ne = projections.get(i);
99+
if (ne instanceof Alias alias && alias.child().semanticEquals(target)) {
100+
projections.set(i, alias.replaceChild(expanded));
101+
} else if (ne.semanticEquals(target)) {
102+
projections.set(i, expanded);
103+
}
104+
}
105+
106+
return new Project(pj.source(), pushedDownMvExpand, projections);
93107
}
94108
return mvExpand;
95109
}

0 commit comments

Comments
 (0)