Feedback, pass 1

smalyshev · smalyshev · commit 6df935267f95 · 2025-10-03T12:07:51.000-06:00
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/CombineLimitTopN.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/CombineLimitTopN.java
@@ -7,29 +7,30 @@
 
 package org.elasticsearch.xpack.esql.optimizer.rules.logical;
 
-import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
+import org.elasticsearch.xpack.esql.expression.Foldables;
 import org.elasticsearch.xpack.esql.plan.logical.Limit;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.plan.logical.TopN;
 
 /**
- * Combines a Limit followed by a TopN into a single TopN.
+ * Combines a Limit immediately followed by a TopN into a single TopN.
+ * This is needed because {@link HoistRemoteEnrichTopN} can create new limits that are not covered by the previous rules.
  */
-public final class CombineLimitTopN extends OptimizerRules.ParameterizedOptimizerRule<Limit, LogicalOptimizerContext> {
+public final class CombineLimitTopN extends OptimizerRules.OptimizerRule<Limit> {
 
     public CombineLimitTopN() {
         super(OptimizerRules.TransformDirection.DOWN);
     }
 
     @Override
-    public LogicalPlan rule(Limit limit, LogicalOptimizerContext ctx) {
+    public LogicalPlan rule(Limit limit) {
         if (limit.child() instanceof TopN topn) {
-            int thisLimitValue = (int) limit.limit().fold(ctx.foldCtx());
-            int topNValue = (int) topn.limit().fold(ctx.foldCtx());
+            int thisLimitValue = Foldables.limitValue(limit.limit(), limit.sourceText());
+            int topNValue = Foldables.limitValue(topn.limit(), topn.sourceText());
             if (topNValue <= thisLimitValue) {
                 return topn;
             } else {
-                return new TopN(topn.source(), topn.child(), topn.order(), limit.limit(), topn.isLocal());
+                return new TopN(topn.source(), topn.child(), topn.order(), limit.limit(), topn.local());
             }
         }
         return limit;
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/HoistRemoteEnrichLimit.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/HoistRemoteEnrichLimit.java
@@ -41,7 +41,7 @@ protected LogicalPlan rule(Enrich en, LogicalOptimizerContext ctx) {
             // Enrich.
             Set<Limit> seenLimits = Collections.newSetFromMap(new IdentityHashMap<>());
             en.child().forEachDownMayReturnEarly((p, stop) -> {
-                if (p instanceof Limit l && l.isLocal() == false) {
+                if (p instanceof Limit l && l.local() == false) {
                     seenLimits.add(l);
                     return;
                 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/HoistRemoteEnrichTopN.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/HoistRemoteEnrichTopN.java
@@ -8,7 +8,6 @@
 package org.elasticsearch.xpack.esql.optimizer.rules.logical;
 
 import org.elasticsearch.xpack.esql.core.expression.NamedExpression;
-import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
 import org.elasticsearch.xpack.esql.plan.logical.CardinalityPreserving;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
 import org.elasticsearch.xpack.esql.plan.logical.Eval;
@@ -29,20 +28,18 @@
  * This is the same idea as {@link HoistRemoteEnrichLimit} but for TopN instead of Limit.
  * This must happen after {@link ReplaceLimitAndSortAsTopN}.
  */
-public final class HoistRemoteEnrichTopN extends OptimizerRules.ParameterizedOptimizerRule<Enrich, LogicalOptimizerContext>
-    implements
-        OptimizerRules.CoordinatorOnly {
+public final class HoistRemoteEnrichTopN extends OptimizerRules.OptimizerRule<Enrich> implements OptimizerRules.CoordinatorOnly {
     public HoistRemoteEnrichTopN() {
         super(OptimizerRules.TransformDirection.UP);
     }
 
     @Override
-    protected LogicalPlan rule(Enrich en, LogicalOptimizerContext ctx) {
+    protected LogicalPlan rule(Enrich en) {
         if (en.mode() == Enrich.Mode.REMOTE) {
             LogicalPlan plan = en.child();
             // This loop only takes care of one TopN, repeated application will stack them in correct order.
             while (true) {
-                if (plan instanceof TopN top && top.isLocal() == false) {
+                if (plan instanceof TopN top && top.local() == false) {
                     // Create a fake OrderBy and "push" Enrich through it to generate aliases
                     Enrich topWithEnrich = (Enrich) en.replaceChild(new OrderBy(top.source(), en.child(), top.order()));
                     LogicalPlan pushPlan = PushDownUtils.pushGeneratingPlanPastProjectAndOrderBy(topWithEnrich);
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/OptimizerRules.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/OptimizerRules.java
@@ -110,7 +110,8 @@ public final LogicalPlan apply(LogicalPlan plan, P context) {
         protected abstract LogicalPlan rule(SubPlan plan, P context);
     }
 
-    public interface CoordinatorOnly {
-        // This rule should only be applied on the coordinator plan, not for local plan
-    }
+    /**
+     * Marker interface: This rule should only be applied on the coordinator plan, not for a local plan.
+     */
+    public interface CoordinatorOnly {}
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PushDownAndCombineLimits.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PushDownAndCombineLimits.java
@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.esql.optimizer.rules.logical;
 
+import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
 import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
@@ -34,7 +35,7 @@ public PushDownAndCombineLimits() {
     @Override
     public LogicalPlan rule(Limit limit, LogicalOptimizerContext ctx) {
         if (limit.child() instanceof Limit childLimit) {
-            return limit.combine(childLimit, ctx.foldCtx());
+            return combineLimits(limit, childLimit, ctx.foldCtx());
         } else if (limit.child() instanceof UnaryPlan unary) {
             if (unary instanceof Eval || unary instanceof Project || unary instanceof RegexExtract || unary instanceof InferencePlan<?>) {
                 // Push the limit under unary
@@ -79,6 +80,18 @@ public LogicalPlan rule(Limit limit, LogicalOptimizerContext ctx) {
         return limit;
     }
 
+    private static Limit combineLimits(Limit upper, Limit lower, FoldContext ctx) {
+        // Keep the smallest limit
+        var thisLimitValue = (int) upper.limit().fold(ctx);
+        var otherLimitValue = (int) lower.limit().fold(ctx);
+        // We want to preserve the duplicated() value of the smaller limit.
+        if (otherLimitValue <= thisLimitValue) {
+            return lower.withLocal(upper.local() || lower.local());
+        } else {
+            return new Limit(upper.source(), upper.limit(), lower.child(), upper.duplicated(), upper.local() || lower.local());
+        }
+    }
+
     /**
      * Checks the existence of another 'visible' Limit, that exists behind an operation that doesn't produce output more data than
      * its input (that is not a relation/source nor aggregation).
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/ChangePoint.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/ChangePoint.java
@@ -47,7 +47,6 @@
 public class ChangePoint extends UnaryPlan
     implements
         SurrogateLogicalPlan,
-        CardinalityPreserving,
         PostAnalysisVerificationAware,
         LicenseAware,
         ExecutesOn.Coordinator {
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/InlineStats.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/InlineStats.java
@@ -38,13 +38,7 @@
  *     underlying aggregate.
  * </p>
  */
-public class InlineStats extends UnaryPlan
-    implements
-        NamedWriteable,
-        SurrogateLogicalPlan,
-        TelemetryAware,
-        CardinalityPreserving,
-        SortAgnostic {
+public class InlineStats extends UnaryPlan implements NamedWriteable, SurrogateLogicalPlan, TelemetryAware, SortAgnostic {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
         LogicalPlan.class,
         "InlineStats",
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Limit.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Limit.java
@@ -11,7 +11,6 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.xpack.esql.capabilities.TelemetryAware;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
-import org.elasticsearch.xpack.esql.core.expression.FoldContext;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
@@ -91,18 +90,6 @@ public Limit replaceChild(LogicalPlan newChild) {
         return new Limit(source(), limit, newChild, duplicated, local);
     }
 
-    public Limit combine(Limit other, FoldContext ctx) {
-        // Keep the smallest limit
-        var thisLimitValue = (int) limit.fold(ctx);
-        var otherLimitValue = (int) other.limit.fold(ctx);
-        // We want to preserve the duplicated() value of the smaller limit.
-        if (otherLimitValue <= thisLimitValue) {
-            return other.withLocal(local || other.local);
-        } else {
-            return new Limit(source(), limit, other.child(), duplicated, local || other.local);
-        }
-    }
-
     public Expression limit() {
         return limit;
     }
@@ -115,7 +102,7 @@ public boolean duplicated() {
         return duplicated;
     }
 
-    public boolean isLocal() {
+    public boolean local() {
         return local;
     }
 
@@ -157,6 +144,6 @@ public boolean equals(Object obj) {
     @Override
     public ExecuteLocation executesOn() {
         // Global limit always needs to be on the coordinator
-        return isLocal() ? ExecuteLocation.ANY : ExecuteLocation.COORDINATOR;
+        return local ? ExecuteLocation.ANY : ExecuteLocation.COORDINATOR;
     }
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Lookup.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Lookup.java
@@ -32,7 +32,7 @@
  * Looks up values from the associated {@code tables}.
  * The class is supposed to be substituted by a {@link Join}.
  */
-public class Lookup extends UnaryPlan implements SurrogateLogicalPlan, TelemetryAware, CardinalityPreserving, SortAgnostic {
+public class Lookup extends UnaryPlan implements SurrogateLogicalPlan, TelemetryAware, SortAgnostic {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "Lookup", Lookup::new);
 
     private final Expression tableName;
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/OrderBy.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/OrderBy.java
@@ -33,7 +33,6 @@ public class OrderBy extends UnaryPlan
         PostAnalysisVerificationAware,
         PostOptimizationPlanVerificationAware,
         TelemetryAware,
-        CardinalityPreserving,
         SortAgnostic,
         PipelineBreaker {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "OrderBy", OrderBy::new);
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/TopN.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/TopN.java
@@ -81,7 +81,7 @@ public TopN withLocal(boolean local) {
         return new TopN(source(), child(), order, limit, local);
     }
 
-    public boolean isLocal() {
+    public boolean local() {
         return local;
     }
 
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/mapper/Mapper.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/mapper/Mapper.java
@@ -91,8 +91,8 @@ private PhysicalPlan mapUnary(UnaryPlan unary) {
             }
             // in case of a fragment, push to it any current streaming operator
             if (unary instanceof PipelineBreaker == false
-                || (unary instanceof Limit limit && limit.isLocal())
-                || (unary instanceof TopN topN && topN.isLocal())) {
+                || (unary instanceof Limit limit && limit.local())
+                || (unary instanceof TopN topN && topN.local())) {
                 return new FragmentExec(unary);
             }
         }
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/logical/LimitSerializationTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/logical/LimitSerializationTests.java
@@ -28,7 +28,7 @@ protected Limit mutateInstance(Limit instance) throws IOException {
         Expression limit = instance.limit();
         LogicalPlan child = instance.child();
         boolean duplicated = instance.duplicated();
-        boolean local = instance.isLocal();
+        boolean local = instance.local();
         switch (randomIntBetween(0, 3)) {
             case 0 -> limit = randomValueOtherThan(limit, () -> FieldAttributeTests.createFieldAttribute(0, false));
             case 1 -> child = randomValueOtherThan(child, () -> randomChild(0));
@@ -49,7 +49,7 @@ protected Limit copyInstance(Limit instance, TransportVersion version) throws IO
         // Limit#duplicated() is ALWAYS false when being serialized and we assert that in Limit#writeTo().
         // The same applies to Limit#local.
         // So, we need to manually simulate this situation.
-        Limit deserializedCopy = super.copyInstance(instance.withDuplicated(false).withLocal(false), version);
-        return deserializedCopy.withDuplicated(instance.duplicated()).withLocal(instance.isLocal());
+        Limit deserializedCopy = super.copyInstance(instance, version);
+        return deserializedCopy.withDuplicated(instance.duplicated()).withLocal(instance.local());
     }
 }
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/logical/TopNSerializationTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/logical/TopNSerializationTests.java
@@ -41,7 +41,7 @@ protected TopN mutateInstance(TopN instance) throws IOException {
         LogicalPlan child = instance.child();
         List<Order> order = instance.order();
         Expression limit = instance.limit();
-        boolean local = instance.isLocal();
+        boolean local = instance.local();
         switch (between(0, 3)) {
             case 0 -> child = randomValueOtherThan(child, () -> randomChild(0));
             case 1 -> order = randomValueOtherThan(order, TopNSerializationTests::randomOrders);
@@ -59,8 +59,8 @@ protected boolean alwaysEmptySource() {
     @Override
     protected TopN copyInstance(TopN instance, TransportVersion version) throws IOException {
         // TopN#local is ALWAYS false after serialization.
-        TopN deserializedCopy = super.copyInstance(instance.withLocal(false), version);
-        return deserializedCopy.withLocal(instance.isLocal());
+        TopN deserializedCopy = super.copyInstance(instance, version);
+        return deserializedCopy.withLocal(instance.local());
     }
 
 }

Original file line number	Diff line number	Diff line change
`@@ -41,7 +41,7 @@ protected LogicalPlan rule(Enrich en, LogicalOptimizerContext ctx) {`
`41`	`41`	`// Enrich.`
`42`	`42`	`Set<Limit> seenLimits = Collections.newSetFromMap(new IdentityHashMap<>());`
`43`	`43`	`en.child().forEachDownMayReturnEarly((p, stop) -> {`
`44`		`- if (p instanceof Limit l && l.isLocal() == false) {`
	`44`	`+ if (p instanceof Limit l && l.local() == false) {`
`45`	`45`	`seenLimits.add(l);`
`46`	`46`	`return;`
`47`	`47`	`}`
Original file line number	Diff line number	Diff line change
`@@ -110,7 +110,8 @@ public final LogicalPlan apply(LogicalPlan plan, P context) {`
`110`	`110`	`protected abstract LogicalPlan rule(SubPlan plan, P context);`
`111`	`111`	`}`
`112`	`112`
`113`		`- public interface CoordinatorOnly {`
`114`		`- // This rule should only be applied on the coordinator plan, not for local plan`
`115`		`- }`
	`113`	`+ /**`
	`114`	`+ * Marker interface: This rule should only be applied on the coordinator plan, not for a local plan.`
	`115`	`+ */`
	`116`	`+ public interface CoordinatorOnly {}`
`116`	`117`	`}`