One of the ways to fix the limit issue

Author: smalyshev

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);
+        LogicalPlan newFirstGrandChild = limit.replaceChild(firstGrandChild).withImplied(true);
 
         List<LogicalPlan> newGrandChildren = new ArrayList<>();
         newGrandChildren.add(newFirstGrandChild);

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

@@ -29,19 +29,21 @@ 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);
+        this(source, limit, child, false, false);
     }
 
-    public Limit(Source source, Expression limit, LogicalPlan child, boolean duplicated) {
+    public Limit(Source source, Expression limit, LogicalPlan child, boolean duplicated, boolean implied) {
         super(source, child);
         this.limit = limit;
         this.duplicated = duplicated;
+        this.implied = implied;
     }
 
     /**
@@ -52,6 +54,7 @@ private Limit(StreamInput in) throws IOException {
             Source.readFrom((PlanStreamInput) in),
             in.readNamedWriteable(Expression.class),
             in.readNamedWriteable(LogicalPlan.class),
+            false,
             false
         );
     }
@@ -77,28 +80,36 @@ public String getWriteableName() {
 
     @Override
     protected NodeInfo<Limit> info() {
-        return NodeInfo.create(this, Limit::new, limit, child(), duplicated);
+        return NodeInfo.create(this, Limit::new, limit, child(), duplicated, implied);
     }
 
     @Override
     public Limit replaceChild(LogicalPlan newChild) {
-        return new Limit(source(), limit, newChild, duplicated);
+        return new Limit(source(), limit, newChild, duplicated, implied);
     }
 
     public Expression limit() {
         return limit;
     }
 
     public Limit withLimit(Expression limit) {
-        return new Limit(source(), limit, child(), duplicated);
+        return new Limit(source(), limit, child(), duplicated, implied);
     }
 
     public boolean duplicated() {
         return duplicated;
     }
 
+    public boolean implied() {
+        return implied;
+    }
+
     public Limit withDuplicated(boolean duplicated) {
-        return new Limit(source(), limit, child(), duplicated);
+        return new Limit(source(), limit, child(), duplicated, implied);
+    }
+
+    public Limit withImplied(boolean implied) {
+        return new Limit(source(), limit, child(), duplicated, implied);
     }
 
     @Override
@@ -108,7 +119,7 @@ public boolean expressionsResolved() {
 
     @Override
     public int hashCode() {
-        return Objects.hash(limit, child(), duplicated);
+        return Objects.hash(limit, child(), duplicated, implied);
     }
 
     @Override

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

@@ -205,20 +205,19 @@ 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;
             }
             // Once we reached FragmentExec, we stuff our Enrich under it
             if (f instanceof FragmentExec) {
                 hasFragment.set(true);
-                // FIXME: hack to remove duplicate limits. This is probably not the right way to do it.
-                return new FragmentExec(logical.transformUp(Limit.class, l -> {
-                    if (l.duplicated()) {
-                        return l.child();
-                    }
-                    return l;
-                }));
+                return new FragmentExec(logical);
             }
             if (f instanceof EnrichExec enrichExec) {
                 assert enrichExec.mode() != Enrich.Mode.REMOTE : "Unexpected remote ENRICH when looking for fragment";
@@ -236,14 +235,19 @@ private PhysicalPlan mapToFragmentExec(Join logical, PhysicalPlan child) {
                     forceLocal.set(true);
                     return f;
                 }
-                if (f instanceof LimitExec || f instanceof ExchangeExec || f instanceof TopNExec) {
+                // No need to include LimitExec because:
+                // 1. If it's an implied limit, we have a copy of it already on the top
+                // 2. If it's a forced limit, we can't execute this as a remote join anyway.
+                if (f instanceof ExchangeExec || f instanceof TopNExec) {
                     return f;
                 } else {
                     return unaryExec.child();
                 }
             }
-            if (f instanceof LookupJoinExec lj) {
-                return lj.right();
+            if (f instanceof LookupJoinExec) {
+                // We shouldn't be meeting LookupJoinExec unless we can't execute this as a remote join
+                forceLocal.set(true);
+                return f;
             }
             if (f instanceof MergeExec) {
                 forceLocal.set(true);
@@ -252,14 +256,7 @@ private PhysicalPlan mapToFragmentExec(Join logical, PhysicalPlan child) {
             return f;
         });
 
-        if (forceLocal.get()) {
-            if (logical.isRemote()) {
-                throw new EsqlIllegalArgumentException("Remote joins are not supported in this context");
-            }
-            return null;
-        }
-
-        if (hasFragment.get()) {
+        if (forceLocal.get() == false && hasFragment.get()) {
             return childTransformed;
         }
         return null;
@@ -288,7 +285,11 @@ private PhysicalPlan mapBinary(BinaryPlan bp) {
 
             if (FRAGMENT_EXEC_HACK_ENABLED) {
                 var leftPlan = mapToFragmentExec(join, left);
-                if (leftPlan != null) {
+                if (leftPlan == null) {
+                    if (join.isRemote()) {
+                        throw new EsqlIllegalArgumentException("Remote joins are not supported in this context: [" + join + "]");
+                    }
+                } else {
                     return leftPlan;
                 }
             }