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

Author: smalyshev

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

@@ -310,20 +310,28 @@ private static void checkForPlansForbiddenBeforeRemoteEnrich(Enrich enrich, Fail
             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.Coordinator
+                || u instanceof Enrich upstreamEnrich && upstreamEnrich.mode() == Enrich.Mode.COORDINATOR
+                || u instanceof LookupJoin) {
+                badCommands.add(u.source());
             }
+
+            // 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");
+            // }
         });
 
-        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()))
+        );
     }
 }

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

@@ -0,0 +1,40 @@
+/*
+ * 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
+    }
+
+    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/LookupJoin.java

@@ -14,6 +14,7 @@
 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;
@@ -107,7 +108,7 @@ private void checkRemoteJoin(Failures failures) {
         List<Source> fails = new LinkedList<>();
 
         this.forEachUp(UnaryPlan.class, u -> {
-            if (u instanceof PipelineBreaker) {
+            if (u instanceof PipelineBreaker || u instanceof ExecutesOn.Coordinator) {
                 fails.add(u.source());
             }
             if (u instanceof Enrich enrich && enrich.mode() == Enrich.Mode.COORDINATOR) {

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) {