refactor(esql): Improve MvExpand push-down logic in logical optimizer

- Enhance PushDownMvExpandPastProject rule to handle more complex query scenarios
- Add checks to prevent duplicate output attributes when pushing down MvExpand
- Support pushing down MvExpand for queries with aliased expressions
- Handle cases where target is an alias or referenced in an alias
- Add test case to verify complex MvExpand push-down behavior
Resolves edge cases in multi-stage query optimization where MvExpand interactions with Project and Eval nodes were previously limited.

Author: kanoshiou

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

@@ -17,26 +17,54 @@
 
 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]]
+            //   \_Limit[1000[INTEGER],false,false]
+            //     \_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 && alias.toAttribute().equals(mvExpand.target().toAttribute())) {
-                    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;
+                if (projection instanceof Alias alias) {
+                    if (alias.toAttribute().semanticEquals(mvExpand.target().toAttribute())) {
+                        // 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;
+                    }
                 }
             }
 

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizerTests.java

@@ -3149,13 +3149,14 @@ public void testPushDownMvExpandPastProject() {
      * <pre>{@code
      * Project[[a{r}#13, b{r}#14]]
      * \_Limit[1000[INTEGER],true,false]
-     *   \_MvExpand[$$a$b$0{r}#15,b{r}#14]
-     *     \_Eval[[a{r}#13 AS $$a$b$0#15]]
+     *   \_MvExpand[$$a$b$0$b$0{r}#16,b{r}#14]
+     *     \_Eval[[$$a$b$0{r}#15 AS $$a$b$0$b$0#16]]
      *       \_Limit[1000[INTEGER],true,false]
      *         \_MvExpand[a{r}#6,a{r}#13]
-     *           \_Limit[1000[INTEGER],false,false]
-     *             \_Aggregate[[],[COUNT(*[KEYWORD],true[BOOLEAN]) AS a#6]]
-     *               \_LocalRelation[[a{r}#4],Page{blocks=[IntVectorBlock[vector=ConstantIntVector[positions=1, value=1]]]}]
+     *           \_Eval[[a{r}#6 AS $$a$b$0#15]]
+     *             \_Limit[1000[INTEGER],false,false]
+     *               \_Aggregate[[],[COUNT(*[KEYWORD],true[BOOLEAN],PT0S[TIME_DURATION]) AS a#6]]
+     *                 \_LocalRelation[[a{r}#4],Page{blocks=[IntVectorBlock[vector=ConstantIntVector[positions=1, value=1]]]}]]
      * }</pre>
      */
     public void testPushDownMvExpandPastProject2() {
@@ -3171,7 +3172,8 @@ public void testPushDownMvExpandPastProject2() {
         var eval = as(mvExpand.child(), Eval.class);
         limit = asLimit(eval.child(), 1000, true);
         mvExpand = as(limit.child(), MvExpand.class);
-        limit = asLimit(mvExpand.child(), 1000, false);
+        eval = as(mvExpand.child(), Eval.class);
+        limit = asLimit(eval.child(), 1000, false);
         var agg = as(limit.child(), Aggregate.class);
         as(agg.child(), LocalRelation.class);
     }