Introduce execution location marker for better handling of remote/local compatibility

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

@@ -45,7 +45,13 @@
 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, PipelineBreaker {
+public class Aggregate extends UnaryPlan
+    implements
+        PostAnalysisVerificationAware,
+        TelemetryAware,
+        SortAgnostic,
+        PipelineBreaker,
+        ExecutesOn.Coordinator {
     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/ChangePoint.java

@@ -44,7 +44,12 @@
  * enforced by the Limit in the surrogate plan.
  */
 @SupportsObservabilityTier(tier = COMPLETE)
-public class ChangePoint extends UnaryPlan implements SurrogateLogicalPlan, PostAnalysisVerificationAware, LicenseAware {
+public class ChangePoint extends UnaryPlan
+    implements
+        SurrogateLogicalPlan,
+        PostAnalysisVerificationAware,
+        LicenseAware,
+        ExecutesOn.Coordinator {
 
     private final Attribute value;
     private final Attribute key;

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

@@ -17,7 +17,7 @@
 import org.elasticsearch.index.IndexMode;
 import org.elasticsearch.transport.RemoteClusterAware;
 import org.elasticsearch.xpack.core.enrich.EnrichPolicy;
-import org.elasticsearch.xpack.esql.capabilities.PostAnalysisPlanVerificationAware;
+import org.elasticsearch.xpack.esql.capabilities.PostAnalysisVerificationAware;
 import org.elasticsearch.xpack.esql.capabilities.TelemetryAware;
 import org.elasticsearch.xpack.esql.common.Failures;
 import org.elasticsearch.xpack.esql.core.capabilities.Resolvables;
@@ -35,7 +35,6 @@
 import org.elasticsearch.xpack.esql.index.EsIndex;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
 import org.elasticsearch.xpack.esql.plan.GeneratingPlan;
-import org.elasticsearch.xpack.esql.plan.logical.join.LookupJoin;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -45,13 +44,18 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
-import java.util.function.BiConsumer;
 
 import static org.elasticsearch.xpack.esql.common.Failure.fail;
 import static org.elasticsearch.xpack.esql.core.expression.Expressions.asAttributes;
 import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
 
-public class Enrich extends UnaryPlan implements GeneratingPlan<Enrich>, PostAnalysisPlanVerificationAware, TelemetryAware, SortAgnostic {
+public class Enrich extends UnaryPlan
+    implements
+        GeneratingPlan<Enrich>,
+        PostAnalysisVerificationAware,
+        TelemetryAware,
+        SortAgnostic,
+        ExecutesOn {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
         LogicalPlan.class,
         "Enrich",
@@ -68,6 +72,16 @@ public class Enrich extends UnaryPlan implements GeneratingPlan<Enrich>, PostAna
 
     private final Mode mode;
 
+    @Override
+    public ExecuteLocation executesOn() {
+        if (mode == Mode.REMOTE) {
+            return ExecuteLocation.REMOTE;
+        } else if (mode == Mode.COORDINATOR) {
+            return ExecuteLocation.COORDINATOR;
+        }
+        return ExecuteLocation.ANY;
+    }
+
     public enum Mode {
         ANY,
         COORDINATOR,
@@ -279,11 +293,6 @@ public int hashCode() {
         return Objects.hash(super.hashCode(), mode, policyName, matchField, policy, concreteIndices, enrichFields);
     }
 
-    @Override
-    public BiConsumer<LogicalPlan, Failures> postAnalysisPlanVerification() {
-        return Enrich::checkRemoteEnrich;
-    }
-
     /**
      * Ensure that no remote enrich is allowed after a reduction or an enrich with coordinator mode.
      * <p>
@@ -296,34 +305,26 @@ public BiConsumer<LogicalPlan, Failures> postAnalysisPlanVerification() {
      * We might consider implementing the actual remote enrich on the coordinating cluster, however, this requires
      * retaining the originating cluster and restructing pages for routing, which might be complicated.
      */
-    private static void checkRemoteEnrich(LogicalPlan plan, Failures failures) {
-        // First look for remote ENRICH, and then look at its children. Going over the whole plan once is trickier as remote ENRICHs can be
-        // in separate FORK branches which are valid by themselves.
-        plan.forEachUp(Enrich.class, enrich -> checkForPlansForbiddenBeforeRemoteEnrich(enrich, failures));
-    }
-
-    /**
-     * For a given remote {@link Enrich}, check if there are any forbidden plans upstream.
-     */
     private static void checkForPlansForbiddenBeforeRemoteEnrich(Enrich enrich, Failures failures) {
         if (enrich.mode != Mode.REMOTE) {
             return;
         }
 
-        Set<String> badCommands = new HashSet<>();
+        Set<Source> badCommands = new HashSet<>();
 
         enrich.forEachUp(LogicalPlan.class, u -> {
-            if (u instanceof Aggregate) {
-                badCommands.add("STATS");
-            } else if (u instanceof Enrich upstreamEnrich && upstreamEnrich.mode() == Enrich.Mode.COORDINATOR) {
-                badCommands.add("another ENRICH with coordinator policy");
-            } else if (u instanceof LookupJoin) {
-                badCommands.add("LOOKUP JOIN");
-            } else if (u instanceof Fork) {
-                badCommands.add("FORK");
+            if (u instanceof ExecutesOn ex && ex.executesOn() == ExecuteLocation.COORDINATOR) {
+                badCommands.add(u.source());
             }
         });
 
-        badCommands.forEach(c -> failures.add(fail(enrich, "ENRICH with remote policy can't be executed after " + c)));
+        badCommands.forEach(
+            f -> failures.add(fail(enrich, "ENRICH with remote policy can't be executed after [" + f.text() + "]" + f.source()))
+        );
+    }
+
+    @Override
+    public void postAnalysisVerification(Failures failures) {
+        checkForPlansForbiddenBeforeRemoteEnrich(this, failures);
     }
 }

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

@@ -0,0 +1,41 @@
+/*
+ * 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;
+
+/**
+ * Mark nodes that execute only in a specific way, either on the coordinator or on a remote node.
+ */
+public interface ExecutesOn {
+    enum ExecuteLocation {
+        COORDINATOR,
+        REMOTE,
+        ANY; // Can be executed on either coordinator or remote nodes
+    }
+
+    ExecuteLocation executesOn();
+
+    /**
+     * Executes on the remote nodes only (note that may include coordinator, but not on the aggregation stage).
+     */
+    interface Remote extends ExecutesOn {
+        @Override
+        default ExecuteLocation executesOn() {
+            return ExecuteLocation.REMOTE;
+        }
+    }
+
+    /**
+     * Executes on the coordinator only. Can not be run on remote nodes.
+     */
+    interface Coordinator extends ExecutesOn {
+        @Override
+        default ExecuteLocation executesOn() {
+            return ExecuteLocation.COORDINATOR;
+        }
+    }
+}

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

@@ -33,7 +33,12 @@
  * A Fork is a n-ary {@code Plan} where each child is a sub plan, e.g.
  * {@code FORK [WHERE content:"fox" ] [WHERE content:"dog"] }
  */
-public class Fork extends LogicalPlan implements PostAnalysisPlanVerificationAware, TelemetryAware, PipelineBreaker {
+public class Fork extends LogicalPlan
+    implements
+        PostAnalysisPlanVerificationAware,
+        TelemetryAware,
+        PipelineBreaker,
+        ExecutesOn.Coordinator {
 
     public static final String FORK_FIELD = "_fork";
     public static final int MAX_BRANCHES = 8;

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

@@ -23,6 +23,7 @@
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
+import org.elasticsearch.xpack.esql.plan.logical.ExecutesOn;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 
 import java.io.IOException;
@@ -33,7 +34,7 @@
 import static org.elasticsearch.xpack.esql.core.type.DataType.TEXT;
 import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
 
-public class Completion extends InferencePlan<Completion> implements TelemetryAware, PostAnalysisVerificationAware {
+public class Completion extends InferencePlan<Completion> implements TelemetryAware, PostAnalysisVerificationAware, ExecutesOn.Coordinator {
 
     public static final String DEFAULT_OUTPUT_FIELD_NAME = "completion";
 

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

@@ -23,6 +23,7 @@
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
+import org.elasticsearch.xpack.esql.plan.logical.ExecutesOn;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
 
@@ -33,7 +34,7 @@
 import static org.elasticsearch.xpack.esql.core.expression.Expressions.asAttributes;
 import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
 
-public class Rerank extends InferencePlan<Rerank> implements TelemetryAware {
+public class Rerank extends InferencePlan<Rerank> implements TelemetryAware, ExecutesOn.Coordinator {
 
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "Rerank", Rerank::new);
     public static final String DEFAULT_INFERENCE_ID = ".rerank-v1-elasticsearch";

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

@@ -21,6 +21,7 @@
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
 import org.elasticsearch.xpack.esql.plan.logical.BinaryPlan;
+import org.elasticsearch.xpack.esql.plan.logical.ExecutesOn;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.plan.logical.SortAgnostic;
 
@@ -56,7 +57,7 @@
 import static org.elasticsearch.xpack.esql.plan.logical.join.JoinTypes.LEFT;
 import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.commonType;
 
-public class Join extends BinaryPlan implements PostAnalysisVerificationAware, SortAgnostic {
+public class Join extends BinaryPlan implements PostAnalysisVerificationAware, SortAgnostic, ExecutesOn {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "Join", Join::new);
     public static final DataType[] UNSUPPORTED_TYPES = {
         TEXT,
@@ -309,4 +310,9 @@ private static boolean comparableTypes(Attribute left, Attribute right) {
     public boolean isRemote() {
         return isRemote;
     }
+
+    @Override
+    public ExecuteLocation executesOn() {
+        return isRemote ? ExecuteLocation.REMOTE : ExecuteLocation.COORDINATOR;
+    }
 }

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

@@ -13,15 +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.Enrich;
+import org.elasticsearch.xpack.esql.plan.logical.ExecutesOn;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 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.LinkedList;
+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;
@@ -104,13 +105,10 @@ public void postAnalysisVerification(Failures failures) {
     }
 
     private void checkRemoteJoin(Failures failures) {
-        List<Source> fails = new LinkedList<>();
+        Set<Source> fails = new HashSet<>();
 
         this.forEachUp(UnaryPlan.class, u -> {
-            if (u instanceof PipelineBreaker) {
-                fails.add(u.source());
-            }
-            if (u instanceof Enrich enrich && enrich.mode() == Enrich.Mode.COORDINATOR) {
+            if (u instanceof PipelineBreaker || (u instanceof ExecutesOn ex && ex.executesOn() == ExecuteLocation.COORDINATOR)) {
                 fails.add(u.source());
             }
         });

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/VerifierTests.java

@@ -2358,7 +2358,7 @@ public void testRemoteEnrichAfterLookupJoin() {
             | %s
             | ENRICH _remote:languages ON language_code
             """, lookupCommand), analyzer);
-        assertThat(err, containsString("4:3: ENRICH with remote policy can't be executed after LOOKUP JOIN"));
+        assertThat(err, containsString("4:3: ENRICH with remote policy can't be executed after [" + lookupCommand + "]@3:3"));
 
         err = error(Strings.format("""
             FROM test
@@ -2367,7 +2367,7 @@ public void testRemoteEnrichAfterLookupJoin() {
             | ENRICH _remote:languages ON language_code
             | %s
             """, lookupCommand, lookupCommand), analyzer);
-        assertThat(err, containsString("4:3: ENRICH with remote policy can't be executed after LOOKUP JOIN"));
+        assertThat(err, containsString("4:3: ENRICH with remote policy can't be executed after [" + lookupCommand + "]@3:3"));
 
         err = error(Strings.format("""
             FROM test
@@ -2377,7 +2377,7 @@ public void testRemoteEnrichAfterLookupJoin() {
             | MV_EXPAND language_code
             | ENRICH _remote:languages ON language_code
             """, lookupCommand), analyzer);
-        assertThat(err, containsString("6:3: ENRICH with remote policy can't be executed after LOOKUP JOIN"));
+        assertThat(err, containsString("6:3: ENRICH with remote policy can't be executed after [" + lookupCommand + "]@3:3"));
     }
 
     public void testRemoteEnrichAfterCoordinatorOnlyPlans() {
@@ -2420,7 +2420,7 @@ public void testRemoteEnrichAfterCoordinatorOnlyPlans() {
             | STATS count(*) BY language_code
             | ENRICH _remote:languages ON language_code
             """, analyzer);
-        assertThat(err, containsString("4:3: ENRICH with remote policy can't be executed after STATS"));
+        assertThat(err, containsString("4:3: ENRICH with remote policy can't be executed after [STATS count(*) BY language_code]@3:3"));
 
         err = error("""
             FROM test
@@ -2430,7 +2430,7 @@ public void testRemoteEnrichAfterCoordinatorOnlyPlans() {
             | MV_EXPAND language_code
             | ENRICH _remote:languages ON language_code
             """, analyzer);
-        assertThat(err, containsString("6:3: ENRICH with remote policy can't be executed after STATS"));
+        assertThat(err, containsString("6:3: ENRICH with remote policy can't be executed after [STATS count(*) BY language_code]@3:3"));
 
         query("""
             FROM test
@@ -2445,7 +2445,10 @@ public void testRemoteEnrichAfterCoordinatorOnlyPlans() {
             | ENRICH _coordinator:languages ON language_code
             | ENRICH _remote:languages ON language_code
             """, analyzer);
-        assertThat(err, containsString("4:3: ENRICH with remote policy can't be executed after another ENRICH with coordinator policy"));
+        assertThat(
+            err,
+            containsString("4:3: ENRICH with remote policy can't be executed after [ENRICH _coordinator:languages ON language_code]@3:3")
+        );
 
         err = error("""
             FROM test
@@ -2456,15 +2459,23 @@ public void testRemoteEnrichAfterCoordinatorOnlyPlans() {
             | DISSECT language_name "%{foo}"
             | ENRICH _remote:languages ON language_code
             """, analyzer);
-        assertThat(err, containsString("7:3: ENRICH with remote policy can't be executed after another ENRICH with coordinator policy"));
+        assertThat(
+            err,
+            containsString("7:3: ENRICH with remote policy can't be executed after [ENRICH _coordinator:languages ON language_code]@3:3")
+        );
 
         err = error("""
             FROM test
             | FORK (WHERE languages == 1) (WHERE languages == 2)
             | EVAL language_code = languages
             | ENRICH _remote:languages ON language_code
             """, analyzer);
-        assertThat(err, containsString("4:3: ENRICH with remote policy can't be executed after FORK"));
+        assertThat(
+            err,
+            containsString(
+                "4:3: ENRICH with remote policy can't be executed after [FORK (WHERE languages == 1) (WHERE languages == 2)]@2:3"
+            )
+        );
     }
 
     private void checkFullTextFunctionsInStats(String functionInvocation) {

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

@@ -7004,7 +7004,7 @@ public void testAggThenEnrichRemote() {
             | eval employee_id = to_str(emp_no)
             | ENRICH _remote:departments
             """));
-        assertThat(error.getMessage(), containsString("line 4:3: ENRICH with remote policy can't be executed after STATS"));
+        assertThat(error.getMessage(), containsString("line 4:3: ENRICH with remote policy can't be executed after [STATS size=count(*) BY emp_no]@2:3"));
     }
 
     public void testEnrichBeforeLimit() {
@@ -7354,7 +7354,7 @@ public void testRejectRemoteEnrichAfterCoordinatorEnrich() {
             """));
         assertThat(
             error.getMessage(),
-            containsString("ENRICH with remote policy can't be executed after another ENRICH with coordinator policy")
+            containsString("ENRICH with remote policy can't be executed after [ENRICH _coordinator:departments]@3:3")
         );
     }