ESQL: Replace any Attribute type when pushing down past Project (elastic#135295) (elastic#136275)

This will allow any attribute type be replaced in node that's being pushed down past a `Project`. So far only `ReferenceAttribute`s were replaced. Some nodes reference other types, like `FieldAttribute` (in `OrderBy`), which would otherwise be pushed down unchanged, and thus possibly wrongly (such as when the `Project` renames _to_ the field attribute name, thus showdowing it).

Closes elastic#134407

Author: bpintea

docs/changelog/135295.yaml

@@ -0,0 +1,6 @@
+pr: 135295
+summary: Replace any Attribute type when pushing down past Project
+area: ES|QL
+type: bug
+issues:
+ - 134407

x-pack/plugin/esql/qa/testFixtures/src/main/resources/lookup-join.csv-spec

@@ -5610,3 +5610,22 @@ id_int:integer | name_str:keyword | extra1:keyword | other1:keyword | other2:int
 13             | Mia              | thud           | xi             | 14000
 14             | Nina             | foo2           | omicron        | 15000
 ;
+
+// https://github.com/elastic/elasticsearch/issues/134407
+SortOnFieldNullifiedDueToJoinKeyMissingInSomeOfTheIndices
+required_capability: join_lookup_v12
+required_capability: fixed_pushdown_past_project_with_attributes_resolution
+
+from sample_data,languages_mixed_numerics,apps
+| rename language_code_byte as language_code
+| lookup join languages_lookup on language_code
+| rename id as language_code // `id` can be missing locally, so "joined" `language_code` is eval'd to null hereafter
+| lookup join languages_lookup on language_code
+| sort language_code_short DESC NULLS LAST, name ASC NULLS LAST, language_name
+| limit 2
+;
+
+@timestamp:datetime|client_ip:ip|event_duration:l| language_code:i | language_code_double:d | language_code_float:d  |language_code_half_float:d|language_code_integer:i| language_code_long:l |language_code_scaled_float:d|language_code_short:i|       message:keyword       |     name:keyword      |    version:version    | language_name:keyword
+null           |null           |null           |null           |32767.0               |32767.0               |32768.0                 |32767                |32767               |32767.0                   |32767              |null                 |null           |null           |null
+null           |null           |null           |null           |128.0                 |128.0                 |128.0                   |128                  |128                 |128.0                     |128                |null                 |null           |null           |null
+;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -544,6 +544,11 @@ public enum Cap {
          */
         FIXED_PUSHDOWN_PAST_PROJECT,
 
+        /**
+         * When resolving renames, consider all {@code Attribute}s in the plan, not just the {@code ReferenceAttribute}s.
+         */
+        FIXED_PUSHDOWN_PAST_PROJECT_WITH_ATTRIBUTES_RESOLUTION,
+
         /**
          * Adds the {@code MV_PSERIES_WEIGHTED_SUM} function for converting sorted lists of numbers into
          * a bounded score. This is a generalization of the

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

@@ -30,9 +30,11 @@ public final class RuleUtils {
     private RuleUtils() {}
 
     /**
-     * Returns a tuple of two lists:
-     * 1. A list of aliases to null literals for those data types in the {@param outputAttributes} that {@param shouldBeReplaced}.
-     * 2. A list of named expressions where attributes that match the predicate are replaced with their corresponding null alias.
+     * @return a tuple of two lists:
+     * <ol>
+     * <li>A list of aliases to null literals for those data types in the {@code outputAttributes} that {@code shouldBeReplaced}.</li>
+     * <li>A list of named expressions where attributes that match the predicate are replaced with their corresponding null alias.</li>
+     * </ol>
      *
      * @param outputAttributes The original output attributes.
      * @param shouldBeReplaced A predicate to determine which attributes should be replaced with null aliases.

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

@@ -14,7 +14,6 @@
 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.expression.Order;
 import org.elasticsearch.xpack.esql.plan.GeneratingPlan;
 import org.elasticsearch.xpack.esql.plan.logical.Eval;
@@ -206,7 +205,7 @@ private static <P extends LogicalPlan> P resolveRenamesFromProject(P plan, Proje
 
     @SuppressWarnings("unchecked")
     public static <P extends LogicalPlan> P resolveRenamesFromMap(P plan, AttributeMap<Expression> map) {
-        return (P) plan.transformExpressionsOnly(ReferenceAttribute.class, r -> map.resolve(r, r));
+        return (P) plan.transformExpressionsOnly(Attribute.class, r -> map.resolve(r, r));
     }
 
     private record AttributeReplacement(List<Expression> rewrittenExpressions, AttributeMap<Alias> replacedAttributes) {}

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

@@ -97,6 +97,7 @@
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.OptimizerRules;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneRedundantOrderBy;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownAndCombineLimits;
+import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownAndCombineOrderBy;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownEnrich;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownEval;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownInferencePlan;
@@ -165,6 +166,7 @@
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.localSource;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.randomLiteral;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.referenceAttribute;
+import static org.elasticsearch.xpack.esql.EsqlTestUtils.relation;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.singleValue;
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.withDefaultLimitWarning;
 import static org.elasticsearch.xpack.esql.analysis.Analyzer.NO_FIELDS;
@@ -5737,6 +5739,51 @@ public void testPushShadowingGeneratingPlanPastRenamingProjectWithResolution() {
         }
     }
 
+    /*
+     * Test for: https://github.com/elastic/elasticsearch/issues/134407
+     *
+     * Input:
+     * OrderBy[[Order[a{f}#2,ASC,ANY]]]
+     * \_Project[[aTemp{r}#3 AS a#2]]
+     *   \_Eval[[null[INTEGER] AS aTemp#3]]
+     *     \_EsRelation[uYiPqAFD][LOOKUP][a{f}#2]
+     *
+     * Output:
+     * Project[[aTemp{r}#3 AS a#2]]
+     * \_OrderBy[[Order[aTemp{r}#3,ASC,ANY]]]
+     *   \_Eval[[null[INTEGER] AS aTemp#3]]
+     *     \_EsRelation[uYiPqAFD][LOOKUP][a{f}#2]
+     */
+    public void testPushDownOrderByPastRename() {
+        FieldAttribute a = getFieldAttribute("a");
+        EsRelation relation = relation().withAttributes(List.of(a));
+
+        Alias aTemp = new Alias(EMPTY, "aTemp", new Literal(EMPTY, null, a.dataType()));
+        Eval eval = new Eval(EMPTY, relation, List.of(aTemp));
+
+        // Rename the null literal to "a" so that the OrderBy can refer to it. Requires re-using the id of original "a" attribute.
+        Alias aliasA = new Alias(EMPTY, "a", aTemp.toAttribute(), a.id());
+        Project project = new Project(EMPTY, eval, List.of(aliasA));
+
+        // OrderBy sorts on original `a` attribute; after pushing down it should sort on aTemp.
+        OrderBy orderBy = new OrderBy(EMPTY, project, List.of(new Order(EMPTY, a, Order.OrderDirection.ASC, Order.NullsPosition.ANY)));
+
+        LogicalPlan optimized = new PushDownAndCombineOrderBy().apply(orderBy);
+
+        var projectOut = as(optimized, Project.class);
+        assertThat(projectOut.projections(), equalTo(project.projections()));
+        var orderByOutput = as(projectOut.child(), OrderBy.class);
+        var orderAttr = as(orderByOutput.order().getFirst().child(), ReferenceAttribute.class);
+
+        // the actual fix test
+        assertThat(orderAttr.name(), equalTo("aTemp"));
+        assertThat(orderAttr.id(), equalTo(aTemp.id()));
+
+        var evalOutput = as(orderByOutput.child(), Eval.class);
+        assertThat(evalOutput.fields(), equalTo(eval.fields()));
+        assertThat(evalOutput.child(), equalTo(relation));
+    }
+
     /**
      * Expects
      * <pre>{@code