Sketch out another approach to planning remote join

Author: alex-spies

x-pack/plugin/esql/qa/server/multi-clusters/src/javaRestTest/java/org/elasticsearch/xpack/esql/ccq/MultiClusterSpecIT.java

@@ -119,11 +119,11 @@ public MultiClusterSpecIT(
         "StatsAndLookupMessageFromIndex",
         "MvJoinKeyOnTheLookupIndexAfterStats",
         "MvJoinKeyOnFromAfterStats",
-        // Lookup join after SORT+LIMIT is not supported in CCS yet
-        "NullifiedJoinKeyToPurgeTheJoin"
-    // "SortBeforeAndAfterJoin",
-    // "SortEvalBeforeLookup",
-    // "SortBeforeAndAfterMultipleJoinAndMvExpand"
+        // Lookup join after SORT is not supported in CCS yet
+        "NullifiedJoinKeyToPurgeTheJoin",
+        "SortBeforeAndAfterJoin",
+        "SortEvalBeforeLookup",
+        "SortBeforeAndAfterMultipleJoinAndMvExpand"
     );
 
     @Override

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

@@ -331,7 +331,7 @@ emp_no:integer
 10001
 ;
 
-dropAllFieldsUsedInLookupOnMaybeTheCoordinator
+dropAllFieldsUsedInLookupOnTheCoordinator
 required_capability: join_lookup_v12
 
 FROM employees
@@ -367,7 +367,7 @@ language_code:integer | language_name:keyword       | country:text
 1                     | English                     | null
 ;
 
-nonUniqueRightKeyOnMaybeTheCoordinatorLateLimit
+nonUniqueRightKeyOnTheCoordinatorLateLimit
 required_capability: join_lookup_v12
 required_capability: join_lookup_fix_limit_pushdown
 
@@ -714,7 +714,7 @@ emp_no:integer | language_code:integer | language_name:keyword
 10003          | 4                     | German
 ;
 
-filterOnTheDataNodeThenFilterOnMaybeTheCoordinator
+filterOnTheDataNodeThenFilterOnTheCoordinator
 required_capability: join_lookup_v12
 
 FROM employees

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

@@ -108,7 +108,7 @@ private static Limit duplicateLimitAsFirstGrandchild(Limit limit) {
 
         List<LogicalPlan> grandChildren = limit.child().children();
         LogicalPlan firstGrandChild = grandChildren.getFirst();
-        LogicalPlan newFirstGrandChild = limit.replaceChild(firstGrandChild).withImplied(true);
+        LogicalPlan newFirstGrandChild = limit.replaceChild(firstGrandChild);
 
         List<LogicalPlan> newGrandChildren = new ArrayList<>();
         newGrandChildren.add(newFirstGrandChild);

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

@@ -45,7 +45,7 @@
 import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
 import static org.elasticsearch.xpack.esql.plan.logical.Filter.checkFilterConditionDataType;
 
-public class Aggregate extends UnaryPlan implements PostAnalysisVerificationAware, TelemetryAware, SortAgnostic {
+public class Aggregate extends UnaryPlan implements PostAnalysisVerificationAware, TelemetryAware, SortAgnostic, PipelineBreaker {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
         LogicalPlan.class,
         "Aggregate",

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

@@ -18,7 +18,7 @@
 import java.io.IOException;
 import java.util.Objects;
 
-public class Limit extends UnaryPlan implements TelemetryAware {
+public class Limit extends UnaryPlan implements TelemetryAware, PipelineBreaker {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "Limit", Limit::new);
 
     private final Expression limit;
@@ -29,21 +29,19 @@ public class Limit extends UnaryPlan implements TelemetryAware {
      * infinite loops from adding a duplicate of the limit past the child over and over again.
      */
     private final transient boolean duplicated;
-    private final transient boolean implied;
 
     /**
      * Default way to create a new instance. Do not use this to copy an existing instance, as this sets {@link Limit#duplicated} to
      * {@code false}.
      */
     public Limit(Source source, Expression limit, LogicalPlan child) {
-        this(source, limit, child, false, false);
+        this(source, limit, child, false);
     }
 
-    public Limit(Source source, Expression limit, LogicalPlan child, boolean duplicated, boolean implied) {
+    public Limit(Source source, Expression limit, LogicalPlan child, boolean duplicated) {
         super(source, child);
         this.limit = limit;
         this.duplicated = duplicated;
-        this.implied = implied;
     }
 
     /**
@@ -54,24 +52,20 @@ private Limit(StreamInput in) throws IOException {
             Source.readFrom((PlanStreamInput) in),
             in.readNamedWriteable(Expression.class),
             in.readNamedWriteable(LogicalPlan.class),
-            false,
             false
         );
     }
 
     /**
-     * Omits serializing {@link Limit#duplicated} because when sent to a data node, this should always be {@code false}.
-     * That's because if it's true, this means a copy of this limit was pushed down below an MvExpand or Join, and thus there's
-     * another pipeline breaker further upstream - we're already on the coordinator node.
+     * Omits serializing {@link Limit#duplicated} because this is only required to avoid pushing duplicating a limit past
+     * {@link org.elasticsearch.xpack.esql.plan.logical.join.Join} or {@link MvExpand} in an infinite loop, see
+     * {@link org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownAndCombineLimits}.
      */
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         Source.EMPTY.writeTo(out);
         out.writeNamedWriteable(limit());
         out.writeNamedWriteable(child());
-        // Let's make sure we notice during tests if we ever serialize a duplicated Limit.
-        // FIXME: verify that this condition is correct and fix the comments
-        assert implied || duplicated == false;
     }
 
     @Override
@@ -81,36 +75,28 @@ public String getWriteableName() {
 
     @Override
     protected NodeInfo<Limit> info() {
-        return NodeInfo.create(this, Limit::new, limit, child(), duplicated, implied);
+        return NodeInfo.create(this, Limit::new, limit, child(), duplicated);
     }
 
     @Override
     public Limit replaceChild(LogicalPlan newChild) {
-        return new Limit(source(), limit, newChild, duplicated, implied);
+        return new Limit(source(), limit, newChild, duplicated);
     }
 
     public Expression limit() {
         return limit;
     }
 
     public Limit withLimit(Expression limit) {
-        return new Limit(source(), limit, child(), duplicated, implied);
+        return new Limit(source(), limit, child(), duplicated);
     }
 
     public boolean duplicated() {
         return duplicated;
     }
 
-    public boolean implied() {
-        return implied;
-    }
-
     public Limit withDuplicated(boolean duplicated) {
-        return new Limit(source(), limit, child(), duplicated, implied);
-    }
-
-    public Limit withImplied(boolean implied) {
-        return new Limit(source(), limit, child(), duplicated, implied);
+        return new Limit(source(), limit, child(), duplicated);
     }
 
     @Override
@@ -120,7 +106,7 @@ public boolean expressionsResolved() {
 
     @Override
     public int hashCode() {
-        return Objects.hash(limit, child(), duplicated, implied);
+        return Objects.hash(limit, child(), duplicated);
     }
 
     @Override

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

@@ -31,7 +31,8 @@ public class OrderBy extends UnaryPlan
         PostAnalysisVerificationAware,
         PostOptimizationVerificationAware,
         TelemetryAware,
-        SortAgnostic {
+        SortAgnostic,
+        PipelineBreaker {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "OrderBy", OrderBy::new);
 
     private final List<Order> order;

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

@@ -0,0 +1,10 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.plan.logical;
+
+public interface PipelineBreaker {}

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

@@ -21,7 +21,7 @@
 import java.util.List;
 import java.util.Objects;
 
-public class TopN extends UnaryPlan {
+public class TopN extends UnaryPlan implements PipelineBreaker {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "TopN", TopN::new);
 
     private final List<Order> order;

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

@@ -13,16 +13,16 @@
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
-import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
-import org.elasticsearch.xpack.esql.plan.logical.Limit;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
-import org.elasticsearch.xpack.esql.plan.logical.OrderBy;
+import org.elasticsearch.xpack.esql.plan.logical.PipelineBreaker;
 import org.elasticsearch.xpack.esql.plan.logical.SurrogateLogicalPlan;
 import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
 import org.elasticsearch.xpack.esql.plan.logical.join.JoinTypes.UsingJoinType;
 
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 
 import static java.util.Collections.emptyList;
 import static org.elasticsearch.xpack.esql.common.Failure.fail;
@@ -97,28 +97,18 @@ public void postAnalysisVerification(Failures failures) {
     }
 
     private void checkRemoteJoin(Failures failures) {
-        boolean[] agg = { false };
-        boolean[] enrichCoord = { false };
-        boolean[] limit = { false };
+        Set<String> fails = new HashSet<>();
 
         this.forEachUp(UnaryPlan.class, u -> {
-            if (u instanceof Aggregate) {
-                agg[0] = true;
-            } else if (u instanceof Enrich enrich && enrich.mode() == Enrich.Mode.COORDINATOR) {
-                enrichCoord[0] = true;
-            } else if (u instanceof Limit) {
-                limit[0] = true;
+            if (u instanceof PipelineBreaker) {
+                fails.add(u.nodeName());
+            }
+            if (u instanceof Enrich enrich && enrich.mode() == Enrich.Mode.COORDINATOR) {
+                fails.add("ENRICH with coordinator policy");
             }
         });
-        if (agg[0]) {
-            failures.add(fail(this, "LOOKUP JOIN with remote indices can't be executed after STATS"));
-        }
-        if (enrichCoord[0]) {
-            failures.add(fail(this, "LOOKUP JOIN with remote indices can't be executed after ENRICH with coordinator policy"));
-        }
-        if (limit[0]) {
-            failures.add(fail(this, "LOOKUP JOIN with remote indices can't be executed after LIMIT"));
-        }
+
+        fails.forEach(f -> failures.add(fail(this, "LOOKUP JOIN with remote indices can't be executed after " + f)));
     }
 
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/mapper/Mapper.java

@@ -20,7 +20,7 @@
 import org.elasticsearch.xpack.esql.plan.logical.LeafPlan;
 import org.elasticsearch.xpack.esql.plan.logical.Limit;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
-import org.elasticsearch.xpack.esql.plan.logical.OrderBy;
+import org.elasticsearch.xpack.esql.plan.logical.PipelineBreaker;
 import org.elasticsearch.xpack.esql.plan.logical.Sample;
 import org.elasticsearch.xpack.esql.plan.logical.TopN;
 import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
@@ -205,11 +205,6 @@ private PhysicalPlan mapToFragmentExec(Join logical, PhysicalPlan child) {
         Holder<Boolean> hasFragment = new Holder<>(false);
         Holder<Boolean> forceLocal = new Holder<>(false);
 
-        // If we have any forced limits inside, we can't execute this as a remote join
-        if (logical.left().anyMatch(pl -> pl instanceof Limit limit && limit.implied() == false)) {
-            return null;
-        }
-
         var childTransformed = child.transformUp(f -> {
             if (forceLocal.get()) {
                 return f;
@@ -276,23 +271,31 @@ private PhysicalPlan mapBinary(BinaryPlan bp) {
                 return new FragmentExec(bp);
             }
 
+            if (join.isRemote()) {
+                // This is generally wrong in case of pipeline breakers upstream from the join, but we validate against these.
+                // The only potential pipeline breakers upstream should be limits duplicated past the join from PushdownAndCombineLimits,
+                // but they are okay to perform on the data nodes because they only serve to reduce the number of rows processed and
+                // don't affect correctness due to another limit being downstream.
+                return new FragmentExec(bp);
+            }
+
             PhysicalPlan left = map(bp.left());
 
             // only broadcast joins supported for now - hence push down as a streaming operator
             if (left instanceof FragmentExec fragment) {
                 return new FragmentExec(bp);
             }
 
-            if (FRAGMENT_EXEC_HACK_ENABLED) {
-                var leftPlan = mapToFragmentExec(join, left);
-                if (leftPlan == null) {
-                    if (join.isRemote()) {
-                        throw new EsqlIllegalArgumentException("Remote joins are not supported in this context: [" + join + "]");
-                    }
-                } else {
-                    return leftPlan;
-                }
-            }
+            // if (FRAGMENT_EXEC_HACK_ENABLED) {
+            // var leftPlan = mapToFragmentExec(join, left);
+            // if (leftPlan == null) {
+            // if (join.isRemote()) {
+            // throw new EsqlIllegalArgumentException("Remote joins are not supported in this context: [" + join + "]");
+            // }
+            // } else {
+            // return leftPlan;
+            // }
+            // }
 
             PhysicalPlan right = map(bp.right());
             // if the right is data we can use a hash join directly
@@ -322,7 +325,7 @@ private PhysicalPlan mapFork(Fork fork) {
     }
 
     public static boolean isPipelineBreaker(LogicalPlan p) {
-        return p instanceof Aggregate || p instanceof TopN || p instanceof Limit || p instanceof OrderBy;
+        return p instanceof PipelineBreaker;
     }
 
     private PhysicalPlan addExchangeForFragment(LogicalPlan logical, PhysicalPlan child) {