fix projection flattening with .toAttribute

Author: bpintea

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java

@@ -253,7 +253,11 @@ public static Range rangeOf(Expression value, Expression lower, boolean includeL
     }
 
     public static EsRelation relation() {
-        return new EsRelation(EMPTY, new EsIndex(randomAlphaOfLength(8), emptyMap()), IndexMode.STANDARD);
+        return relation(IndexMode.STANDARD);
+    }
+
+    public static EsRelation relation(IndexMode mode) {
+        return new EsRelation(EMPTY, new EsIndex(randomAlphaOfLength(8), emptyMap()), mode);
     }
 
     /**

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

@@ -65,7 +65,7 @@ public static Tuple<List<Alias>, List<NamedExpression>> aliasedNulls(
                 }
                 // otherwise point to it since this avoids creating field copies
                 else {
-                    projection = new Alias(attr.source(), attr.name(), nullAlias.toAttribute(), nullAlias.id());
+                    projection = new Alias(attr.source(), attr.name(), nullAlias.toAttribute(), attr.id());
                 }
             } else {
                 projection = attr;

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

@@ -10,7 +10,6 @@
 import org.elasticsearch.xpack.esql.core.expression.AttributeMap;
 import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
-import org.elasticsearch.xpack.esql.core.expression.Expressions;
 import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
 import org.elasticsearch.xpack.esql.optimizer.rules.RuleUtils;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.OptimizerRules;
@@ -59,6 +58,6 @@ private static LogicalPlan replaceJoin(Join join) {
         }
         var aliasedNulls = RuleUtils.aliasedNulls(joinRightOutput, a -> true);
         var eval = new Eval(join.source(), join.left(), aliasedNulls.v1());
-        return new Project(join.source(), eval, join.computeOutput(join.left().output(), Expressions.asAttributes(aliasedNulls.v2())));
+        return new Project(join.source(), eval, join.computeOutputExpressions(join.left().output(), aliasedNulls.v2()));
     }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Lookup.java

@@ -15,6 +15,7 @@
 import org.elasticsearch.xpack.esql.core.capabilities.Resolvables;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.Expressions;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
@@ -139,7 +140,7 @@ public List<Attribute> output() {
             if (localRelation == null) {
                 throw new IllegalStateException("Cannot determine output of LOOKUP with unresolved table");
             }
-            lazyOutput = Join.computeOutput(child().output(), localRelation.output(), joinConfig());
+            lazyOutput = Expressions.asAttributes(Join.computeOutputExpressions(child().output(), localRelation.output(), joinConfig()));
         }
         return lazyOutput;
     }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/join/InlineJoin.java

@@ -200,7 +200,7 @@ public Join replaceChildren(LogicalPlan left, LogicalPlan right) {
     }
 
     @Override
-    public List<Attribute> computeOutput(List<Attribute> left, List<Attribute> right) {
+    List<Attribute> computeOutput(List<Attribute> left, List<Attribute> right) {
         JoinType joinType = config().type();
         List<Attribute> output;
         if (LEFT.equals(joinType)) {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/join/Join.java

@@ -18,6 +18,7 @@
 import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Expressions;
+import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
 import org.elasticsearch.xpack.esql.core.expression.ReferenceAttribute;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
@@ -64,7 +65,7 @@
 import static org.elasticsearch.xpack.esql.core.type.DataType.UNSIGNED_LONG;
 import static org.elasticsearch.xpack.esql.core.type.DataType.UNSUPPORTED;
 import static org.elasticsearch.xpack.esql.core.type.DataType.VERSION;
-import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
+import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputExpressions;
 import static org.elasticsearch.xpack.esql.plan.logical.join.JoinTypes.LEFT;
 import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.commonType;
 
@@ -210,30 +211,38 @@ public List<Attribute> rightOutputFields() {
         return rightOutputFields;
     }
 
-    public List<Attribute> computeOutput(List<Attribute> left, List<Attribute> right) {
-        return computeOutput(left, right, config);
+    List<Attribute> computeOutput(List<Attribute> left, List<Attribute> right) {
+        return Expressions.asAttributes(computeOutputExpressions(left, right, config));
+    }
+
+    public List<NamedExpression> computeOutputExpressions(List<? extends NamedExpression> left, List<? extends NamedExpression> right) {
+        return computeOutputExpressions(left, right, config);
     }
 
     /**
      * Combine the two lists of attributes into one.
      * In case of (name) conflicts, specify which sides wins, that is overrides the other column - the left or the right.
      */
-    public static List<Attribute> computeOutput(List<Attribute> leftOutput, List<Attribute> rightOutput, JoinConfig config) {
+    public static List<NamedExpression> computeOutputExpressions(
+        List<? extends NamedExpression> leftOutput,
+        List<? extends NamedExpression> rightOutput,
+        JoinConfig config
+    ) {
         JoinType joinType = config.type();
-        List<Attribute> output;
+        List<NamedExpression> output;
         // TODO: make the other side nullable
         if (LEFT.equals(joinType)) {
             if (config.joinOnConditions() == null) {
                 // right side becomes nullable and overrides left except for join keys, which we preserve from the left
                 AttributeSet rightKeys = AttributeSet.of(config.rightFields());
-                List<Attribute> rightOutputWithoutMatchFields = rightOutput.stream()
-                    .filter(attr -> rightKeys.contains(attr) == false)
+                List<? extends NamedExpression> rightOutputWithoutMatchFields = rightOutput.stream()
+                    .filter(ne -> rightKeys.contains(ne.toAttribute()) == false)
                     .toList();
-                output = mergeOutputAttributes(rightOutputWithoutMatchFields, leftOutput);
+                output = mergeOutputExpressions(rightOutputWithoutMatchFields, leftOutput);
             } else {
                 // We don't allow any attributes in the joinOnConditions that don't have unique names
                 // so right always overwrites left in case of name clashes
-                output = mergeOutputAttributes(rightOutput, leftOutput);
+                output = mergeOutputExpressions(rightOutput, leftOutput);
             }
 
         } else {

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

@@ -61,6 +61,9 @@
 import org.elasticsearch.xpack.esql.plan.logical.Project;
 import org.elasticsearch.xpack.esql.plan.logical.Row;
 import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
+import org.elasticsearch.xpack.esql.plan.logical.join.Join;
+import org.elasticsearch.xpack.esql.plan.logical.join.JoinConfig;
+import org.elasticsearch.xpack.esql.plan.logical.join.JoinTypes;
 import org.elasticsearch.xpack.esql.plan.logical.local.EmptyLocalSupplier;
 import org.elasticsearch.xpack.esql.plan.logical.local.EsqlProject;
 import org.elasticsearch.xpack.esql.plan.logical.local.LocalRelation;
@@ -96,6 +99,7 @@
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.withDefaultLimitWarning;
 import static org.elasticsearch.xpack.esql.core.tree.Source.EMPTY;
 import static org.elasticsearch.xpack.esql.core.type.DataType.INTEGER;
+import static org.elasticsearch.xpack.esql.core.type.DataType.KEYWORD;
 import static org.elasticsearch.xpack.esql.optimizer.rules.logical.OptimizerRules.TransformDirection.DOWN;
 import static org.elasticsearch.xpack.esql.optimizer.rules.logical.OptimizerRules.TransformDirection.UP;
 import static org.hamcrest.Matchers.contains;
@@ -1019,6 +1023,58 @@ public List<Attribute> output() {
         assertThat(e.getMessage(), containsString("Output has changed from"));
     }
 
+    /**
+     * Input:
+     * Project[[key{f}#2, int{f}#3, field1{f}#7, field2{f}#8]]
+     * \_Join[LEFT,[key{f}#2],[key{f}#6],null]
+     *   |_EsRelation[JLfQlKmn][key{f}#2, int{f}#3, field1{f}#4, field2{f}#5]
+     *   \_EsRelation[HQtEBOWq][LOOKUP][key{f}#6, field1{f}#7, field2{f}#8]
+     *
+     * Output:
+     * Project[[key{r}#2, int{f}#3, field1{r}#7, field1{r}#7 AS field2#8]]
+     * \_Eval[[null[KEYWORD] AS key#2, null[INTEGER] AS field1#7]]
+     *   \_EsRelation[JLfQlKmn][key{f}#2, int{f}#3, field1{f}#4, field2{f}#5]
+     */
+    public void testPruneLeftJoinOnNullMatchingFieldAndShadowingAttributes() {
+        var keyLeft = getFieldAttribute("key", KEYWORD);
+        var intFieldLeft = getFieldAttribute("int");
+        var fieldLeft1 = getFieldAttribute("field1");
+        var fieldLeft2 = getFieldAttribute("field2");
+        var leftRelation = EsqlTestUtils.relation(IndexMode.STANDARD)
+            .withAttributes(List.of(keyLeft, intFieldLeft, fieldLeft1, fieldLeft2));
+
+        var keyRight = getFieldAttribute("key", KEYWORD);
+        var fieldRight1 = getFieldAttribute("field1");
+        var fieldRight2 = getFieldAttribute("field2");
+        var rightRelation = EsqlTestUtils.relation(IndexMode.LOOKUP).withAttributes(List.of(keyRight, fieldRight1, fieldRight2));
+
+        JoinConfig joinConfig = new JoinConfig(JoinTypes.LEFT, List.of(keyLeft), List.of(keyRight), null);
+        var join = new Join(EMPTY, leftRelation, rightRelation, joinConfig);
+        var project = new Project(EMPTY, join, List.of(keyLeft, intFieldLeft, fieldRight1, fieldRight2));
+
+        var testStats = statsForMissingField("key");
+        var localPlan = localPlan(project, testStats);
+
+        var projectOut = as(localPlan, Project.class);
+        var projectionsOut = projectOut.projections();
+        assertThat(Expressions.names(projectionsOut), contains("key", "int", "field1", "field2"));
+        assertThat(projectionsOut.get(0).id(), is(keyLeft.id()));
+        assertThat(projectionsOut.get(1).id(), is(intFieldLeft.id()));
+        assertThat(projectionsOut.get(2).id(), is(fieldRight1.id())); // id must remain from the RHS.
+        var aliasField2 = as(projectionsOut.get(3), Alias.class); // the projection must contain an alias ...
+        assertThat(aliasField2.id(), is(fieldRight2.id())); // ... with the same id as the original field.
+
+        var eval = as(projectOut.child(), Eval.class);
+        assertThat(Expressions.names(eval.fields()), contains("key", "field1"));
+        var keyEval = as(Alias.unwrap(eval.fields().get(0)), Literal.class);
+        assertThat(keyEval.value(), is(nullValue()));
+        assertThat(keyEval.dataType(), is(KEYWORD));
+        var field1Eval = as(Alias.unwrap(eval.fields().get(1)), Literal.class);
+        assertThat(field1Eval.value(), is(nullValue()));
+        assertThat(field1Eval.dataType(), is(INTEGER));
+        var source = as(eval.child(), EsRelation.class);
+    }
+
     private IsNotNull isNotNull(Expression field) {
         return new IsNotNull(EMPTY, field);
     }
@@ -1078,6 +1134,6 @@ protected List<String> filteredWarnings() {
     }
 
     public static EsRelation relation() {
-        return new EsRelation(EMPTY, new EsIndex(randomAlphaOfLength(8), emptyMap()), randomFrom(IndexMode.values()));
+        return EsqlTestUtils.relation(randomFrom(IndexMode.values()));
     }
 }