Implement algo

Author: alex-spies

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

@@ -7,12 +7,23 @@
 
 package org.elasticsearch.xpack.esql.optimizer.rules.logical;
 
+import org.elasticsearch.xpack.esql.core.expression.Alias;
+import org.elasticsearch.xpack.esql.core.expression.Attribute;
+import org.elasticsearch.xpack.esql.core.expression.AttributeMap;
+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.plan.logical.Eval;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.plan.logical.Project;
-import org.elasticsearch.xpack.esql.plan.logical.Eval;
 import org.elasticsearch.xpack.esql.plan.logical.join.Join;
 import org.elasticsearch.xpack.esql.plan.logical.join.JoinTypes;
 
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
 /**
  * If a {@link Project} is found in the left child of a left {@link Join}, perform it after. Due to requiring the projected attributes
  * later, field extractions can also happen later, making joins cheapter to execute on data nodes.
@@ -26,20 +37,76 @@
 public final class PushDownJoinPastProject extends OptimizerRules.OptimizerRule<Join> {
     @Override
     protected LogicalPlan rule(Join join) {
-        if (join.left() instanceof Project project && join.config().type() == JoinTypes.LEFT)
-        {
-            // 1. Propagate any renames into the Join, as we will remove the upstream Project.
-            //    E.g. `RENAME field AS key | LOOKUP JOIN idx ON key` -> `LOOKUP JOIN idx ON field | ...`
-            // 2. Construct the downstream Project using the Join's output.
-            //    Use trivial aliases for now, so we can easily update the expressions.
-            //    This is like adding a `RENAME field1 AS field1, field2 AS field2, ...` after the Join, where we can easily change the
-            //    referenced field names to deal with name conflicts.
-            // 3. Propagate any renames from the upstream Project into the new downstream Project.
-            // 4. Look for name conflicts: any name shadowed by the `LOOKUP JOIN` that was used in the original Project needs to be
-            //    aliased temporarily. Add an Eval upstream from the `LOOKUP JOIN` and propagate its renames into the new downstream
-            //    Project.
-            // 5. Remove any trivial aliases from the new downstream Project - `RENAME field AS field` just becomes (the equivalent of)
-            //    `KEEP field`.
+        if (join.left() instanceof Project project && join.config().type() == JoinTypes.LEFT) {
+            AttributeMap.Builder<Expression> aliasBuilder = AttributeMap.builder();
+            project.forEachExpression(Alias.class, a -> aliasBuilder.put(a.toAttribute(), a.child()));
+            var aliasesFromProject = aliasBuilder.build();
+
+            // Propagate any renames into the Join, as we will remove the upstream Project.
+            // E.g. `RENAME field AS key | LOOKUP JOIN idx ON key` -> `LOOKUP JOIN idx ON field | ...`
+            Join updatedJoin = PushDownUtils.resolveRenamesFromMap(join, aliasesFromProject);
+
+            // Construct the expressions for the new downstream Project using the Join's output.
+            // We need to carry over RENAMEs/aliases from the original upstream Project.
+            List<Attribute> originalOutput = join.output();
+            List<NamedExpression> newProjections = new ArrayList<>(originalOutput.size());
+            for (Attribute attr : originalOutput) {
+                Attribute resolved = (Attribute) aliasesFromProject.resolve(attr, attr);
+                if (attr.semanticEquals(resolved)) {
+                    newProjections.add(attr);
+                } else {
+                    Alias renamed = new Alias(attr.source(), attr.name(), resolved, attr.id(), attr.synthetic());
+                    newProjections.add(renamed);
+                }
+            }
+
+            // This doesn't deal with name conflicts yet. Any name shadowed by a lookup field from the `LOOKUP JOIN` could still have been
+            // used in the original Project; any such conflict needs to be resolved by copying the attribute under a temporary name via an
+            // Eval - and using the attribute from said Eval in the new downstream Project.
+            Set<String> lookupFieldNames = new HashSet<>(Expressions.names(join.rightOutputFields()));
+            PushDownUtils.AttributeReplacement replacement = PushDownUtils.renameAttributesInExpressions(lookupFieldNames, newProjections);
+
+            List<Expression> conflictFreeProjections = replacement.rewrittenExpressions();
+            List<Alias> evalAliases = new ArrayList<>(replacement.replacedAttributes().values());
+
+            if (evalAliases.isEmpty()) {
+                // No name conflicts, so no eval needed.
+                return new Project(project.source(), updatedJoin.replaceLeft(project.child()), newProjections);
+            }
+
+            // The conflict free projections replaced any shadowed attributes by temporary attributes that we'll create in an Eval.
+            // That's good if the replaced attribute was in an Alias; if the projection was a mere attribute to begin with, we need to
+            // alias it back to the name/id that's expected in the original output.
+            List<NamedExpression> finalProjections = new ArrayList<>(conflictFreeProjections.size());
+            for (int i = 0; i < newProjections.size(); i++) {
+                Expression conflictFreeProj = conflictFreeProjections.get(i);
+                if (conflictFreeProj instanceof Alias as) {
+                    // Already aliased - keep it, it's fine if the child was rewritten.
+                    finalProjections.add(as);
+                } else if (conflictFreeProj instanceof Attribute conflictFreeAttr) {
+                    Attribute expectedOutputAttr = (Attribute) newProjections.get(i);
+                    if (expectedOutputAttr.semanticEquals(conflictFreeAttr)) {
+                        // no conflict, wasn't rewritten
+                        finalProjections.add(expectedOutputAttr);
+                    } else {
+                        Alias renameBack = new Alias(
+                            expectedOutputAttr.source(),
+                            expectedOutputAttr.name(),
+                            conflictFreeAttr,
+                            expectedOutputAttr.id(),
+                            expectedOutputAttr.synthetic()
+                        );
+                        finalProjections.add(renameBack);
+                    }
+                } else {
+                    throw new IllegalStateException("Resolving a list of projections must yield a list of projections again");
+                }
+            }
+
+            Eval eval = new Eval(project.source(), project.child(), evalAliases);
+            Join finalJoin = new Join(join.source(), eval, updatedJoin.right(), updatedJoin.config());
+
+            return new Project(project.source(), finalJoin, finalProjections);
         }
 
         return join;

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

@@ -26,7 +26,6 @@
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -55,7 +54,7 @@ class PushDownUtils {
     public static <Plan extends UnaryPlan & GeneratingPlan<Plan>> LogicalPlan pushGeneratingPlanPastProjectAndOrderBy(Plan generatingPlan) {
         LogicalPlan child = generatingPlan.child();
         if (child instanceof OrderBy orderBy) {
-            Set<String> generatedFieldNames = new LinkedHashSet<>(Expressions.names(generatingPlan.generatedAttributes()));
+            Set<String> generatedFieldNames = new HashSet<>(Expressions.names(generatingPlan.generatedAttributes()));
 
             // Look for attributes in the OrderBy's expressions and create aliases with temporary names for them.
             AttributeReplacement nonShadowedOrders = renameAttributesInExpressions(generatedFieldNames, orderBy.order());
@@ -91,8 +90,7 @@ public static <Plan extends UnaryPlan & GeneratingPlan<Plan>> LogicalPlan pushGe
 
             List<Attribute> generatedAttributes = generatingPlan.generatedAttributes();
 
-            @SuppressWarnings("unchecked")
-            Plan generatingPlanWithResolvedExpressions = (Plan) resolveRenamesFromProject(generatingPlan, project);
+            Plan generatingPlanWithResolvedExpressions = resolveRenamesFromProject(generatingPlan, project);
 
             Set<String> namesReferencedInRenames = new HashSet<>();
             for (NamedExpression ne : project.projections()) {
@@ -145,7 +143,7 @@ public static <Plan extends UnaryPlan & GeneratingPlan<Plan>> LogicalPlan pushGe
      * Returns the rewritten expressions and a map with an alias for each replaced attribute; the rewritten expressions reference
      * these aliases.
      */
-    private static AttributeReplacement renameAttributesInExpressions(
+    public static AttributeReplacement renameAttributesInExpressions(
         Set<String> attributeNamesToRename,
         List<? extends Expression> expressions
     ) {
@@ -198,13 +196,18 @@ public static Project pushDownPastProject(UnaryPlan parent) {
         }
     }
 
-    private static UnaryPlan resolveRenamesFromProject(UnaryPlan plan, Project project) {
+    public static <P extends LogicalPlan> P resolveRenamesFromProject(P plan, Project project) {
         AttributeMap.Builder<Expression> aliasBuilder = AttributeMap.builder();
         project.forEachExpression(Alias.class, a -> aliasBuilder.put(a.toAttribute(), a.child()));
         var aliases = aliasBuilder.build();
 
-        return (UnaryPlan) plan.transformExpressionsOnly(ReferenceAttribute.class, r -> aliases.resolve(r, r));
+        return resolveRenamesFromMap(plan, aliases);
+    }
+
+    @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));
     }
 
-    private record AttributeReplacement(List<Expression> rewrittenExpressions, AttributeMap<Alias> replacedAttributes) {}
+    public record AttributeReplacement(List<Expression> rewrittenExpressions, AttributeMap<Alias> replacedAttributes) {}
 }