Replace remote Enrich hack with proper handling of remote Enrich

docs/changelog/135051.yaml

@@ -0,0 +1,5 @@
+pr: 135051
+summary: Ban Limit + `MvExpand` before remote Enrich
+area: ES|QL
+type: enhancement
+issues: []

server/src/main/resources/transport/upper_bounds/8.18.csv

@@ -1 +1 @@
-initial_elasticsearch_8_18_6,8840008
+initial_elasticsearch_8_18_8,8840010

server/src/main/resources/transport/upper_bounds/8.19.csv

@@ -1 +1 @@
-initial_elasticsearch_8_19_3,8841067
+initial_elasticsearch_8_19_5,8841069

server/src/main/resources/transport/upper_bounds/9.0.csv

@@ -1 +1 @@
-initial_elasticsearch_9_0_6,9000015
+initial_elasticsearch_9_0_8,9000017

server/src/main/resources/transport/upper_bounds/9.1.csv

@@ -1 +1 @@
-initial_elasticsearch_9_1_4,9112007
+initial_elasticsearch_9_1_5,9112008

server/src/main/resources/transport/upper_bounds/9.2.csv

@@ -1 +1 @@
-ml_inference_endpoint_cache,9157000
+inference_results_map_with_cluster_alias,9166000

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

@@ -9,6 +9,7 @@
 
 import org.elasticsearch.core.Tuple;
 import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.test.junit.annotations.TestLogging;
 import org.elasticsearch.xpack.esql.VerificationException;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
 
@@ -27,6 +28,7 @@
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 
+@TestLogging(value = "org.elasticsearch.xpack.esql:TRACE", reason = "debug")
 public class CrossClusterEnrichIT extends AbstractEnrichBasedCrossClusterTestCase {
 
     @Override
@@ -402,6 +404,33 @@ 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 testLimitThenEnrichRemote() {
@@ -466,6 +495,19 @@ public void testEnrichCoordinatorThenEnrichRemote() {
         assertThat(error.getMessage(), containsString("ENRICH with remote policy can't be executed after [ENRICH  _COORDINATOR"));
     }
 
+    public void testEnrichAfterMvExpandLimit() {
+        String query = String.format(Locale.ROOT, """
+            FROM *:events,events
+            | SORT timestamp
+            | LIMIT 2
+            | eval ip= TO_STR(host)
+            | MV_EXPAND host
+            | %s
+            """, enrichHosts(Enrich.Mode.REMOTE));
+        var error = expectThrows(VerificationException.class, () -> runQuery(query, randomBoolean()).close());
+        assertThat(error.getMessage(), containsString("MV_EXPAND after LIMIT is incompatible with remote ENRICH"));
+    }
+
     private static void assertCCSExecutionInfoDetails(EsqlExecutionInfo executionInfo) {
         assertThat(executionInfo.overallTook().millis(), greaterThanOrEqualTo(0L));
         assertTrue(executionInfo.isCrossClusterSearch());

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

@@ -35,11 +35,11 @@
  * {@link org.elasticsearch.xpack.esql.stats.SearchStats} which provides access to metadata about the index.
  *
  * <p>NB: This class also reapplies all the rules from {@link LogicalPlanOptimizer#operators(boolean)}
- * and {@link LogicalPlanOptimizer#cleanup()}
+ * and {@link LogicalPlanOptimizer#cleanup(boolean)}
  */
 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 +53,7 @@ public class LocalLogicalPlanOptimizer extends ParameterizedRuleExecutor<Logical
             new ReplaceDateTruncBucketWithRoundTo()
         ),
         localOperators(),
-        cleanup()
+        cleanup(true)
     );
 
     public LocalLogicalPlanOptimizer(LocalLogicalOptimizerContext localLogicalOptimizerContext) {
@@ -90,7 +90,7 @@ private static Batch<LogicalPlan> localOperators() {
 
     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

@@ -19,6 +19,8 @@
 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;
@@ -89,19 +91,19 @@
  *     time), combining expressions, dropping redundant clauses, and some normalization such as putting literals on the right whenever
  *     possible.  These rules are run in a loop until none of the rules make any changes to the plan (there is also a safety shut off
  *     after many iterations, although hitting that is considered a bug)</li>
- *     <li>{@link LogicalPlanOptimizer#cleanup()}  Which can replace sorts+limit with a TopN</li>
+ *     <li>{@link LogicalPlanOptimizer#cleanup(boolean)}  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(boolean)} steps are reapplied
+ * at the {@link LocalLogicalPlanOptimizer} layer.</p>
  */
 public class LogicalPlanOptimizer extends ParameterizedRuleExecutor<LogicalPlan, LogicalOptimizerContext> {
 
     private static final List<RuleExecutor.Batch<LogicalPlan>> RULES = List.of(
         substitutions(),
         operators(false),
         new Batch<>("Skip Compute", new SkipQueryOnLimitZero()),
-        cleanup(),
+        cleanup(false),
         new Batch<>("Set as Optimized", Limiter.ONCE, new SetAsOptimized())
     );
 
@@ -114,7 +116,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);
         }
@@ -165,6 +167,7 @@ protected static Batch<LogicalPlan> substitutions() {
     protected static Batch<LogicalPlan> operators(boolean local) {
         return new Batch<>(
             "Operator Optimization",
+            new HoistRemoteEnrichLimit(local),
             new CombineProjections(local),
             new CombineEvals(),
             new PruneEmptyPlans(),
@@ -209,7 +212,13 @@ protected static Batch<LogicalPlan> operators(boolean local) {
         );
     }
 
-    protected static Batch<LogicalPlan> cleanup() {
-        return new Batch<>("Clean Up", new ReplaceLimitAndSortAsTopN(), new ReplaceRowAsLocalRelation(), new PropgateUnmappedFields());
+    protected static Batch<LogicalPlan> cleanup(boolean local) {
+        return new Batch<>(
+            "Clean Up",
+            new ReplaceLimitAndSortAsTopN(),
+            new HoistRemoteEnrichTopN(local),
+            new ReplaceRowAsLocalRelation(),
+            new PropgateUnmappedFields()
+        );
     }
 }

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

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 general 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.