refactor mapper

Author: smalyshev

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

@@ -16,6 +16,7 @@
 import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+import org.elasticsearch.xpack.esql.plan.logical.OrderBy;
 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;
@@ -97,22 +98,26 @@ public void postAnalysisVerification(Failures failures) {
     private void checkRemoteJoin(Failures failures) {
         boolean[] agg = { false };
         boolean[] enrichCoord = { false };
+        boolean[] sort = { false };
 
         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;
-            }
-            if (u instanceof Enrich enrich && enrich.mode() == Enrich.Mode.REMOTE) {
-                if (agg[0]) {
-                    failures.add(fail(enrich, "LOOKUP JOIN with remote indices can't be executed after STATS"));
-                }
-                if (enrichCoord[0]) {
-                    failures.add(fail(enrich, "LOOKUP JOIN with remote indices can't be executed after ENRICH with coordinator policy"));
-                }
+            } else if (u instanceof OrderBy) {
+                sort[0] = true;
             }
         });
+        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 (sort[0]) {
+            failures.add(fail(this, "LOOKUP JOIN with remote indices can't be executed after SORT"));
+        }
     }
 
     public boolean isRemote() {

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

@@ -29,6 +29,7 @@
 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.physical.AggregateExec;
 import org.elasticsearch.xpack.esql.plan.physical.EnrichExec;
 import org.elasticsearch.xpack.esql.plan.physical.ExchangeExec;
 import org.elasticsearch.xpack.esql.plan.physical.FragmentExec;
@@ -202,8 +203,12 @@ private PhysicalPlan mapUnary(UnaryPlan unary) {
 
     private PhysicalPlan mapToFragmentExec(LogicalPlan logical, PhysicalPlan child) {
         Holder<Boolean> hasFragment = new Holder<>(false);
+        Holder<Boolean> forceLocal = new Holder<>(false);
 
         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);
@@ -216,12 +221,21 @@ private PhysicalPlan mapToFragmentExec(LogicalPlan logical, PhysicalPlan child)
                 }));
             }
             if (f instanceof EnrichExec enrichExec) {
-                // It can only be ANY because COORDINATOR would have errored out earlier, and REMOTE should be under FragmentExec
-                assert enrichExec.mode() == Enrich.Mode.ANY : "enrich must be in ANY mode here";
-                return enrichExec.child();
+                assert enrichExec.mode() != Enrich.Mode.REMOTE : "Unexpected remote ENRICH when looking for fragment";
+                if (enrichExec.mode() == Enrich.Mode.ANY) {
+                    return enrichExec.child();
+                } else {
+                    forceLocal.set(true);
+                    return f;
+                }
             }
             if (f instanceof UnaryExec unaryExec) {
-                if (f instanceof LimitExec || f instanceof ExchangeExec || f instanceof TopNExec) {
+                if (f instanceof AggregateExec || f instanceof TopNExec) {
+                    // We can't make a fragment here...
+                    forceLocal.set(true);
+                    return f;
+                }
+                if (f instanceof LimitExec || f instanceof ExchangeExec) {
                     return f;
                 } else {
                     return unaryExec.child();
@@ -230,10 +244,13 @@ private PhysicalPlan mapToFragmentExec(LogicalPlan logical, PhysicalPlan child)
             if (f instanceof LookupJoinExec lj) {
                 return lj.right();
             }
-            // Currently, it's either UnaryExec or LeafExec. Leaf will either resolve to FragmentExec or we'll ignore it.
             return f;
         });
 
+        if (forceLocal.get()) {
+            return null;
+        }
+
         if (hasFragment.get()) {
             return childTransformed;
         }