Replace remote Enrich hack with proper handling of remote Enrich

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/CrossClusterEnrichIT.java

@@ -402,6 +402,61 @@ public void testTopNThenEnrichRemote() {
             assertThat(executionInfo.clusterAliases(), equalTo(Set.of("", "c1", "c2")));
             assertCCSExecutionInfoDetails(executionInfo);
         }
+
+        // No renames, no KEEP - this is required to verify that ENRICH does not break sort with fields it overrides
+        query = """
+            FROM *:events,events
+            | eval ip= TO_STR(host)
+            | SORT timestamp, user, ip
+            | LIMIT 5
+            | ENRICH _remote:hosts
+            """;
+        try (EsqlQueryResponse resp = runQuery(query, requestIncludeMeta)) {
+            assertThat(
+                getValuesList(resp),
+                equalTo(
+                    List.of(
+                        List.of("192.168.1.2", 1L, "andres", "192.168.1.2", "Windows"),
+                        List.of("192.168.1.3", 1L, "matthew", "192.168.1.3", "MacOS"),
+                        Arrays.asList("192.168.1.25", 1L, "park", (String) null, (String) null),
+                        List.of("192.168.1.5", 2L, "akio", "192.168.1.5", "Android"),
+                        List.of("192.168.1.6", 2L, "sergio", "192.168.1.6", "iOS")
+                    )
+                )
+            );
+            EsqlExecutionInfo executionInfo = resp.getExecutionInfo();
+            assertThat(executionInfo.includeCCSMetadata(), equalTo(responseExpectMeta));
+            assertThat(executionInfo.clusterAliases(), equalTo(Set.of("", "c1", "c2")));
+            assertCCSExecutionInfoDetails(executionInfo);
+        }
+    }
+
+    public void testLimitWithCardinalityChange() {
+        String query = String.format(Locale.ROOT, """
+            FROM *:events,events
+            | eval ip= TO_STR(host)
+            | LIMIT 10
+            | WHERE user != "andres"
+            | %s
+            """, enrichHosts(Enrich.Mode.REMOTE));
+        // This is currently not supported, because WHERE is not cardinality preserving
+        var error = expectThrows(VerificationException.class, () -> runQuery(query, randomBoolean()).close());
+        assertThat(error.getMessage(), containsString("ENRICH with remote policy can't be executed after [LIMIT 10]@3:3"));
+    }
+
+    public void testTopNTwiceThenEnrichRemote() {
+        String query = String.format(Locale.ROOT, """
+            FROM *:events,events
+            | eval ip= TO_STR(host)
+            | SORT timestamp
+            | LIMIT 9
+            | SORT ip, user
+            | LIMIT 5
+            | ENRICH _remote:hosts
+            """, enrichHosts(Enrich.Mode.REMOTE));
+        // This is currently not supported, because we can not handle double topN with remote enrich
+        var error = expectThrows(VerificationException.class, () -> runQuery(query, randomBoolean()).close());
+        assertThat(error.getMessage(), containsString("ENRICH with remote policy can't be executed after [SORT timestamp]"));
     }
 
     public void testLimitThenEnrichRemote() {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/capabilities/PostOptimizationVerificationAware.java

@@ -15,6 +15,11 @@
  */
 public interface PostOptimizationVerificationAware {
 
+    /**
+     * Marker interface for verifiers that should only be run on Coordinator
+     */
+    interface CoordinatorOnly extends PostOptimizationVerificationAware {}
+
     /**
      * Validates the implementing expression - discovered failures are reported to the given
      * {@link Failures} class.

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LocalLogicalPlanOptimizer.java

@@ -9,8 +9,8 @@
 
 import org.elasticsearch.xpack.esql.VerificationException;
 import org.elasticsearch.xpack.esql.common.Failures;
+import org.elasticsearch.xpack.esql.optimizer.rules.logical.OptimizerRules;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PropagateEmptyRelation;
-import org.elasticsearch.xpack.esql.optimizer.rules.logical.ReplaceStatsFilteredAggWithEval;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.ReplaceStringCasingWithInsensitiveRegexMatch;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.local.IgnoreNullMetrics;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.local.InferIsNotNull;
@@ -24,6 +24,7 @@
 import org.elasticsearch.xpack.esql.rule.Rule;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 
 import static org.elasticsearch.common.util.CollectionUtils.arrayAsArrayList;
@@ -39,7 +40,7 @@
  */
 public class LocalLogicalPlanOptimizer extends ParameterizedRuleExecutor<LogicalPlan, LocalLogicalOptimizerContext> {
 
-    private final LogicalVerifier verifier = LogicalVerifier.INSTANCE;
+    private final LogicalVerifier verifier = LogicalVerifier.LOCAL_INSTANCE;
 
     private static final List<Batch<LogicalPlan>> RULES = arrayAsArrayList(
         new Batch<>(
@@ -53,7 +54,7 @@ public class LocalLogicalPlanOptimizer extends ParameterizedRuleExecutor<Logical
             new ReplaceDateTruncBucketWithRoundTo()
         ),
         localOperators(),
-        cleanup()
+        localCleanup()
     );
 
     public LocalLogicalPlanOptimizer(LocalLogicalOptimizerContext localLogicalOptimizerContext) {
@@ -75,8 +76,7 @@ private static Batch<LogicalPlan> localOperators() {
         for (var r : rules) {
             switch (r) {
                 case PropagateEmptyRelation ignoredPropagate -> newRules.add(new LocalPropagateEmptyRelation());
-                // skip it: once a fragment contains an Agg, this can no longer be pruned, which the rule can do
-                case ReplaceStatsFilteredAggWithEval ignoredReplace -> {
+                case OptimizerRules.CoordinatorOnly ignored -> {
                 }
                 default -> newRules.add(r);
             }
@@ -88,9 +88,18 @@ private static Batch<LogicalPlan> localOperators() {
         return operators.with(newRules.toArray(Rule[]::new));
     }
 
+    @SuppressWarnings("unchecked")
+    private static Batch<LogicalPlan> localCleanup() {
+        var cleanup = cleanup();
+        return cleanup.with(
+            // Remove CoordinatorOnly rules
+            Arrays.stream(cleanup.rules()).filter(r -> r instanceof OptimizerRules.CoordinatorOnly == false).toArray(Rule[]::new)
+        );
+    }
+
     public LogicalPlan localOptimize(LogicalPlan plan) {
         LogicalPlan optimized = execute(plan);
-        Failures failures = verifier.verify(optimized, true, plan.output());
+        Failures failures = verifier.verify(optimized, plan.output());
         if (failures.hasFailures()) {
             throw new VerificationException(failures);
         }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizer.java

@@ -36,7 +36,7 @@ public class LocalPhysicalPlanOptimizer extends ParameterizedRuleExecutor<Physic
 
     private static final List<Batch<PhysicalPlan>> RULES = rules(true);
 
-    private final PhysicalVerifier verifier = PhysicalVerifier.INSTANCE;
+    private final PhysicalVerifier verifier = PhysicalVerifier.LOCAL_INSTANCE;
 
     public LocalPhysicalPlanOptimizer(LocalPhysicalOptimizerContext context) {
         super(context);
@@ -47,7 +47,7 @@ public PhysicalPlan localOptimize(PhysicalPlan plan) {
     }
 
     PhysicalPlan verify(PhysicalPlan optimizedPlan, List<Attribute> expectedOutputAttributes) {
-        Failures failures = verifier.verify(optimizedPlan, true, expectedOutputAttributes);
+        Failures failures = verifier.verify(optimizedPlan, expectedOutputAttributes);
         if (failures.hasFailures()) {
             throw new VerificationException(failures);
         }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizer.java

@@ -16,10 +16,13 @@
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.CombineBinaryComparisons;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.CombineDisjunctions;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.CombineEvals;
+import org.elasticsearch.xpack.esql.optimizer.rules.logical.CombineLimitTopN;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.CombineProjections;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.ConstantFolding;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.ExtractAggregateCommonFilter;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.FoldNull;
+import org.elasticsearch.xpack.esql.optimizer.rules.logical.HoistRemoteEnrichLimit;
+import org.elasticsearch.xpack.esql.optimizer.rules.logical.HoistRemoteEnrichTopN;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.LiteralsOnTheRight;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PartiallyFoldCase;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PropagateEmptyRelation;
@@ -93,8 +96,8 @@
  *     <li>{@link LogicalPlanOptimizer#cleanup()}  Which can replace sorts+limit with a TopN</li>
  * </ul>
  *
- * <p>Note that the {@link LogicalPlanOptimizer#operators(boolean)} and {@link LogicalPlanOptimizer#cleanup()} steps are reapplied at the
- * {@link LocalLogicalPlanOptimizer} layer.</p>
+ * <p>Note that the {@link LogicalPlanOptimizer#operators(boolean)} and {@link LogicalPlanOptimizer#cleanup()} steps are reapplied
+ * at the {@link LocalLogicalPlanOptimizer} layer.</p>
  */
 public class LogicalPlanOptimizer extends ParameterizedRuleExecutor<LogicalPlan, LogicalOptimizerContext> {
 
@@ -115,7 +118,7 @@ public LogicalPlanOptimizer(LogicalOptimizerContext optimizerContext) {
     public LogicalPlan optimize(LogicalPlan verified) {
         var optimized = execute(verified);
 
-        Failures failures = verifier.verify(optimized, false, verified.output());
+        Failures failures = verifier.verify(optimized, verified.output());
         if (failures.hasFailures()) {
             throw new VerificationException(failures);
         }
@@ -166,6 +169,7 @@ protected static Batch<LogicalPlan> substitutions() {
     protected static Batch<LogicalPlan> operators(boolean local) {
         return new Batch<>(
             "Operator Optimization",
+            new HoistRemoteEnrichLimit(),
             new CombineProjections(local),
             new CombineEvals(),
             new PropagateEmptyRelation(),
@@ -212,6 +216,13 @@ protected static Batch<LogicalPlan> operators(boolean local) {
     }
 
     protected static Batch<LogicalPlan> cleanup() {
-        return new Batch<>("Clean Up", new ReplaceLimitAndSortAsTopN(), new ReplaceRowAsLocalRelation(), new PropgateUnmappedFields());
+        return new Batch<>(
+            "Clean Up",
+            new ReplaceLimitAndSortAsTopN(),
+            new HoistRemoteEnrichTopN(),
+            new ReplaceRowAsLocalRelation(),
+            new PropgateUnmappedFields(),
+            new CombineLimitTopN()
+        );
     }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalVerifier.java

@@ -11,29 +11,18 @@
 import org.elasticsearch.xpack.esql.capabilities.PostOptimizationVerificationAware;
 import org.elasticsearch.xpack.esql.common.Failures;
 import org.elasticsearch.xpack.esql.optimizer.rules.PlanConsistencyChecker;
-import org.elasticsearch.xpack.esql.plan.logical.Enrich;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 
 import java.util.ArrayList;
 import java.util.List;
 import java.util.function.BiConsumer;
 
 public final class LogicalVerifier extends PostOptimizationPhasePlanVerifier<LogicalPlan> {
+    public static final LogicalVerifier LOCAL_INSTANCE = new LogicalVerifier(true);
+    public static final LogicalVerifier INSTANCE = new LogicalVerifier(false);
 
-    public static final LogicalVerifier INSTANCE = new LogicalVerifier();
-
-    private LogicalVerifier() {}
-
-    @Override
-    boolean skipVerification(LogicalPlan optimizedPlan, boolean skipRemoteEnrichVerification) {
-        if (skipRemoteEnrichVerification) {
-            // AwaitsFix https://github.com/elastic/elasticsearch/issues/118531
-            var enriches = optimizedPlan.collectFirstChildren(Enrich.class::isInstance);
-            if (enriches.isEmpty() == false && ((Enrich) enriches.get(0)).mode() == Enrich.Mode.REMOTE) {
-                return true;
-            }
-        }
-        return false;
+    private LogicalVerifier(boolean isLocal) {
+        super(isLocal);
     }
 
     @Override
@@ -44,14 +33,16 @@ void checkPlanConsistency(LogicalPlan optimizedPlan, Failures failures, Failures
             PlanConsistencyChecker.checkPlan(p, depFailures);
 
             if (failures.hasFailures() == false) {
-                if (p instanceof PostOptimizationVerificationAware pova) {
+                if (p instanceof PostOptimizationVerificationAware pova
+                    && (pova instanceof PostOptimizationVerificationAware.CoordinatorOnly && isLocal) == false) {
                     pova.postOptimizationVerification(failures);
                 }
                 if (p instanceof PostOptimizationPlanVerificationAware popva) {
                     checkers.add(popva.postOptimizationPlanVerification());
                 }
                 p.forEachExpression(ex -> {
-                    if (ex instanceof PostOptimizationVerificationAware va) {
+                    if (ex instanceof PostOptimizationVerificationAware va
+                        && (va instanceof PostOptimizationVerificationAware.CoordinatorOnly && isLocal) == false) {
                         va.postOptimizationVerification(failures);
                     }
                 });

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/PhysicalPlanOptimizer.java

@@ -39,7 +39,7 @@ public PhysicalPlan optimize(PhysicalPlan plan) {
     }
 
     PhysicalPlan verify(PhysicalPlan optimizedPlan, List<Attribute> expectedOutputAttributes) {
-        Failures failures = verifier.verify(optimizedPlan, false, expectedOutputAttributes);
+        Failures failures = verifier.verify(optimizedPlan, expectedOutputAttributes);
         if (failures.hasFailures()) {
             throw new VerificationException(failures);
         }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/PhysicalVerifier.java

@@ -12,30 +12,19 @@
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.Expressions;
 import org.elasticsearch.xpack.esql.optimizer.rules.PlanConsistencyChecker;
-import org.elasticsearch.xpack.esql.plan.logical.Enrich;
-import org.elasticsearch.xpack.esql.plan.physical.EnrichExec;
+import org.elasticsearch.xpack.esql.plan.logical.ExecutesOn;
 import org.elasticsearch.xpack.esql.plan.physical.FieldExtractExec;
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
 
 import static org.elasticsearch.xpack.esql.common.Failure.fail;
 
 /** Physical plan verifier. */
 public final class PhysicalVerifier extends PostOptimizationPhasePlanVerifier<PhysicalPlan> {
+    public static final PhysicalVerifier LOCAL_INSTANCE = new PhysicalVerifier(true);
+    public static final PhysicalVerifier INSTANCE = new PhysicalVerifier(false);
 
-    public static final PhysicalVerifier INSTANCE = new PhysicalVerifier();
-
-    private PhysicalVerifier() {}
-
-    @Override
-    boolean skipVerification(PhysicalPlan optimizedPlan, boolean skipRemoteEnrichVerification) {
-        if (skipRemoteEnrichVerification) {
-            // AwaitsFix https://github.com/elastic/elasticsearch/issues/118531
-            var enriches = optimizedPlan.collectFirstChildren(EnrichExec.class::isInstance);
-            if (enriches.isEmpty() == false && ((EnrichExec) enriches.get(0)).mode() == Enrich.Mode.REMOTE) {
-                return true;
-            }
-        }
-        return false;
+    private PhysicalVerifier(boolean isLocal) {
+        super(isLocal);
     }
 
     @Override
@@ -54,6 +43,19 @@ void checkPlanConsistency(PhysicalPlan optimizedPlan, Failures failures, Failure
                     );
                 }
             }
+
+            // This check applies only for coordinator physical plans (isLocal == false)
+            if (isLocal == false && p instanceof ExecutesOn ex && ex.executesOn() == ExecutesOn.ExecuteLocation.REMOTE) {
+                failures.add(
+                    fail(
+                        p,
+                        "Physical plan contains remote executing operation [{}] in local part. "
+                            + "This usually means this command is incompatible with some of the preceding commands.",
+                        p.nodeName()
+                    )
+                );
+            }
+
             PlanConsistencyChecker.checkPlan(p, depFailures);
 
             if (failures.hasFailures() == false) {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/PostOptimizationPhasePlanVerifier.java

@@ -28,13 +28,17 @@
  */
 public abstract class PostOptimizationPhasePlanVerifier<P extends QueryPlan<P>> {
 
+    // Are we verifying the global plan (coordinator) or a local plan (data node)?
+    protected final boolean isLocal;
+
+    protected PostOptimizationPhasePlanVerifier(boolean isLocal) {
+        this.isLocal = isLocal;
+    }
+
     /** Verifies the optimized plan */
-    public Failures verify(P optimizedPlan, boolean skipRemoteEnrichVerification, List<Attribute> expectedOutputAttributes) {
+    public Failures verify(P optimizedPlan, List<Attribute> expectedOutputAttributes) {
         Failures failures = new Failures();
         Failures depFailures = new Failures();
-        if (skipVerification(optimizedPlan, skipRemoteEnrichVerification)) {
-            return failures;
-        }
 
         checkPlanConsistency(optimizedPlan, failures, depFailures);
 
@@ -47,8 +51,6 @@ public Failures verify(P optimizedPlan, boolean skipRemoteEnrichVerification, Li
         return failures;
     }
 
-    abstract boolean skipVerification(P optimizedPlan, boolean skipRemoteEnrichVerification);
-
     abstract void checkPlanConsistency(P optimizedPlan, Failures failures, Failures depFailures);
 
     private static void verifyOutputNotChanged(QueryPlan<?> optimizedPlan, List<Attribute> expectedOutputAttributes, Failures failures) {
@@ -67,7 +69,7 @@ private static void verifyOutputNotChanged(QueryPlan<?> optimizedPlan, List<Attr
                 .stream()
                 .anyMatch(x -> x.name().equals(ProjectAwayColumns.ALL_FIELDS_PROJECTED));
             // LookupJoinExec represents the lookup index with EsSourceExec and this is turned into EsQueryExec by
-            // ReplaceSourceAttributes. Because InsertFieldExtractions doesn't apply to lookup indices, the
+            // ReplaceSourceAttributes. Because InsertFieldExtraction doesn't apply to lookup indices, the
             // right hand side will only have the EsQueryExec providing the _doc attribute and nothing else.
             // We perform an optimizer run on every fragment. LookupJoinExec also contains such a fragment,
             // and currently it only contains an EsQueryExec after optimization.