ESQL: improve mv_expand planning

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java

@@ -362,7 +362,8 @@ private LogicalPlan resolveMvExpand(MvExpand p, List<Attribute> childrenOutput)
                             null,
                             false
                         )
-                        : resolved
+                        : resolved,
+                    p.limit()
                 );
             }
             return p;
@@ -701,6 +702,7 @@ public LogicalPlan apply(LogicalPlan logicalPlan, AnalyzerContext context) {
             } else {
                 limit = context.configuration().resultTruncationMaxSize(); // user provided a limit: cap result entries to the max
             }
+
             return new Limit(Source.EMPTY, new Literal(Source.EMPTY, limit, DataTypes.INTEGER), logicalPlan);
         }
     }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanNamedTypes.java

@@ -785,7 +785,7 @@ static void writeLimit(PlanStreamOutput out, Limit limit) throws IOException {
     }
 
     static MvExpand readMvExpand(PlanStreamInput in) throws IOException {
-        return new MvExpand(in.readSource(), in.readLogicalPlanNode(), in.readNamedExpression(), in.readAttribute());
+        return new MvExpand(in.readSource(), in.readLogicalPlanNode(), in.readNamedExpression(), in.readAttribute(), null);
     }
 
     static void writeMvExpand(PlanStreamOutput out, MvExpand mvExpand) throws IOException {

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

@@ -52,7 +52,6 @@
 import org.elasticsearch.xpack.ql.optimizer.OptimizerRules.SetAsOptimized;
 import org.elasticsearch.xpack.ql.optimizer.OptimizerRules.SimplifyComparisonsArithmetics;
 import org.elasticsearch.xpack.ql.plan.logical.Aggregate;
-import org.elasticsearch.xpack.ql.plan.logical.EsRelation;
 import org.elasticsearch.xpack.ql.plan.logical.Filter;
 import org.elasticsearch.xpack.ql.plan.logical.Limit;
 import org.elasticsearch.xpack.ql.plan.logical.LogicalPlan;
@@ -62,7 +61,6 @@
 import org.elasticsearch.xpack.ql.rule.ParameterizedRule;
 import org.elasticsearch.xpack.ql.rule.ParameterizedRuleExecutor;
 import org.elasticsearch.xpack.ql.rule.Rule;
-import org.elasticsearch.xpack.ql.tree.Source;
 import org.elasticsearch.xpack.ql.type.DataType;
 import org.elasticsearch.xpack.ql.type.DataTypes;
 import org.elasticsearch.xpack.ql.util.CollectionUtils;
@@ -137,13 +135,16 @@ protected static List<Batch<LogicalPlan>> rules() {
             new PruneColumns(),
             new PruneLiteralsInOrderBy(),
             new PushDownAndCombineLimits(),
-            new DuplicateLimitAfterMvExpand(),
             new PushDownAndCombineFilters(),
             new PushDownEval(),
             new PushDownRegexExtract(),
             new PushDownEnrich(),
             new PushDownAndCombineOrderBy(),
             new PruneOrderByBeforeStats(),
+            // see https://github.com/elastic/elasticsearch/issues/103543
+            // new PushDownMvExpandPastProject(),
+            new DuplicateLimitAfterMvExpand(),
+            new PushDownMvExpandPastOrderBy(),
             new PruneRedundantSortClauses()
         );
 
@@ -153,12 +154,12 @@ protected static List<Batch<LogicalPlan>> rules() {
             new ReplaceDuplicateAggWithEval(),
             // pushing down limits again, because ReplaceDuplicateAggWithEval could create new Project nodes that can still be optimized
             new PushDownAndCombineLimits(),
-            new ReplaceLimitAndSortAsTopN()
+            new ReplaceLimitAndSortAsTopN(),
+            new CleanUpMvExpand()
         );
-        var defaultTopN = new Batch<>("Add default TopN", new AddDefaultTopN());
         var label = new Batch<>("Set as Optimized", Limiter.ONCE, new SetAsOptimized());
 
-        return asList(substitutions, operators, skip, cleanup, defaultTopN, label);
+        return asList(substitutions, operators, skip, cleanup, label);
     }
 
     // TODO: currently this rule only works for aggregate functions (AVG)
@@ -436,7 +437,7 @@ protected LogicalPlan rule(Limit limit) {
                 var l2 = (int) childLimit.limit().fold();
                 return new Limit(limit.source(), Literal.of(limitSource, Math.min(l1, l2)), childLimit.child());
             } else if (limit.child() instanceof UnaryPlan unary) {
-                if (unary instanceof Eval || unary instanceof Project || unary instanceof RegexExtract || unary instanceof Enrich) {
+                if (isTransparentForLimit(unary)) {
                     return unary.replaceChild(limit.replaceChild(unary.child()));
                 }
                 // check if there's a 'visible' descendant limit lower than the current one
@@ -481,59 +482,33 @@ private static Limit descendantLimit(UnaryPlan unary) {
         }
     }
 
-    static class DuplicateLimitAfterMvExpand extends OptimizerRules.OptimizerRule<Limit> {
+    private static boolean isTransparentForLimit(LogicalPlan plan) {
+        return plan instanceof Eval || plan instanceof Project || plan instanceof RegexExtract || plan instanceof Enrich;
+    }
 
+    protected static class PushDownMvExpandPastOrderBy extends OptimizerRules.OptimizerRule<OrderBy> {
         @Override
-        protected LogicalPlan rule(Limit limit) {
-            var child = limit.child();
-            var shouldSkip = child instanceof Eval
-                || child instanceof Project
-                || child instanceof RegexExtract
-                || child instanceof Enrich
-                || child instanceof Limit;
-
-            if (shouldSkip == false && child instanceof UnaryPlan unary) {
+        protected LogicalPlan rule(OrderBy upperOrderBy) {
+            if (upperOrderBy.child() instanceof UnaryPlan unary) {
                 MvExpand mvExpand = descendantMvExpand(unary);
-                if (mvExpand != null) {
-                    Limit limitBeforeMvExpand = limitBeforeMvExpand(mvExpand);
-                    // if there is no "appropriate" limit before mv_expand, then push down a copy of the one after it so that:
-                    // - a possible TopN is properly built as low as possible in the tree (closed to Lucene)
-                    // - the input of mv_expand is as small as possible before it is expanded (less rows to inflate and occupy memory)
-                    if (limitBeforeMvExpand == null) {
-                        var duplicateLimit = new Limit(limit.source(), limit.limit(), mvExpand.child());
-                        return limit.replaceChild(propagateDuplicateLimitUntilMvExpand(duplicateLimit, mvExpand, unary));
-                    }
+                if (mvExpand != null && mvExpand.child() instanceof OrderBy) {
+                    return upperOrderBy.replaceChild(replaceChildUntilMvExpandAndOrderBy(unary, mvExpand));
                 }
             }
-            return limit;
+
+            return upperOrderBy;
         }
 
         private static MvExpand descendantMvExpand(UnaryPlan unary) {
             UnaryPlan plan = unary;
-            AttributeSet filterReferences = new AttributeSet();
-            while (plan instanceof Aggregate == false) {
-                if (plan instanceof MvExpand mve) {
-                    // don't return the mv_expand that has a filter after it which uses the expanded values
-                    // since this will trigger the use of a potentially incorrect (too restrictive) limit further down in the tree
-                    if (filterReferences.isEmpty() == false) {
-                        if (filterReferences.contains(mve.target()) // the same field or reference attribute is used in mv_expand AND filter
-                            || mve.target() instanceof ReferenceAttribute // or the mv_expand attr hasn't yet been resolved to a field attr
-                            // or not all filter references have been resolved to field attributes
-                            || filterReferences.stream().anyMatch(ref -> ref instanceof ReferenceAttribute)) {
-                            return null;
-                        }
-                    }
+            while (plan instanceof UnaryPlan) {
+                // find the first MvExpand that also has an orderBy in front of it, otherwise keep looking,
+                // otherwise we will skip valid mv_expands further down the tree
+                if (plan instanceof MvExpand mve && mve.child() instanceof OrderBy) {
                     return mve;
-                } else if (plan instanceof Filter filter) {
-                    // gather all the filters' references to be checked later when a mv_expand is found
-                    filterReferences.addAll(filter.references());
                 } else if (plan instanceof OrderBy) {
-                    // ordering after mv_expand COULD break the order of the results, so the limit shouldn't be copied past mv_expand
-                    // something like from test | sort emp_no | mv_expand job_positions | sort first_name | limit 5
-                    // (the sort first_name likely changes the order of the docs after sort emp_no, so "limit 5" shouldn't be copied down
                     return null;
                 }
-
                 if (plan.child() instanceof UnaryPlan unaryPlan) {
                     plan = unaryPlan;
                 } else {
@@ -542,31 +517,97 @@ private static MvExpand descendantMvExpand(UnaryPlan unary) {
             }
             return null;
         }
+    }
 
-        private static Limit limitBeforeMvExpand(MvExpand mvExpand) {
-            UnaryPlan plan = mvExpand;
-            while (plan instanceof Aggregate == false) {
-                if (plan instanceof Limit limit) {
-                    return limit;
+    protected static class PushDownMvExpandPastProject extends OptimizerRules.OptimizerRule<MvExpand> {
+        @Override
+        protected LogicalPlan rule(MvExpand mvExpand) {
+            LogicalPlan child = mvExpand.child();
+            if (child instanceof Project) {
+                return pushDownPastProject(mvExpand);
+            }
+
+            return mvExpand;
+        }
+    }
+
+    static class DuplicateLimitAfterMvExpand extends OptimizerRules.OptimizerRule<Limit> {
+        @Override
+        protected LogicalPlan rule(Limit limit) {
+            var child = limit.child();
+            // let the PushDownAndCombineLimits rule to push Limit further down
+            var shouldSkip = isTransparentForLimit(child) || child instanceof Limit;
+            if (shouldSkip || child instanceof UnaryPlan == false) {
+                return limit;
+            }
+
+            UnaryPlan unary = (UnaryPlan) child;
+            MvExpand mvExpand = null;
+            AttributeSet filterReferences = new AttributeSet();
+            boolean hasOrderBy = false;
+            boolean isFiltered = false;
+            // look for any orderBy between this limit and mv_expand
+            // and, separately, for any filter that is filtering the expanded values of mv_expand
+            while (mvExpand == null && unary instanceof UnaryPlan) {
+                if (unary instanceof MvExpand mve) {
+                    if (filterReferences.isEmpty() == false) {
+                        if (filterReferences.contains(mve.expanded())) {// if the filter is on the mv_expand field
+                            isFiltered = true;
+                        }
+                    }
+                    mvExpand = mve;
+                    break;
+                } else if (unary instanceof Filter filter) {
+                    // gather all the filters' references to be checked later when a mv_expand is found
+                    filterReferences.addAll(filter.references());
+                } else if (unary instanceof OrderBy) {
+                    hasOrderBy = true;
+                    break;
+                } else if (unary instanceof Limit) {
+                    break;
                 }
-                if (plan.child() instanceof UnaryPlan unaryPlan) {
-                    plan = unaryPlan;
+
+                if (unary.child() instanceof UnaryPlan unaryPlan) {
+                    unary = unaryPlan;
                 } else {
                     break;
                 }
             }
-            return null;
+
+            if (mvExpand != null) {
+                // if the filter is on the expanded values and the order of the results doesn't change, push down mv_expand past sort
+                if (isFiltered && hasOrderBy == false) {
+                    if (mvExpand.child() instanceof OrderBy) {
+                        return limit.replaceChild(replaceChildUntilMvExpandAndOrderBy((UnaryPlan) child, mvExpand));
+                    }
+                } else if (mvExpand.limit() == null) { // otherwise push down a copy of the limit before mv_expand
+                    var duplicateLimit = new Limit(limit.source(), limit.limit(), mvExpand.child());
+                    return limit.replaceChild(propagateDuplicateLimitUntilMvExpand(duplicateLimit, mvExpand, (UnaryPlan) child));
+                }
+            }
+
+            return limit;
         }
 
         private LogicalPlan propagateDuplicateLimitUntilMvExpand(Limit duplicateLimit, MvExpand mvExpand, UnaryPlan child) {
             if (child == mvExpand) {
-                return mvExpand.replaceChild(duplicateLimit);
+                return new MvExpand(mvExpand.source(), duplicateLimit, mvExpand.target(), mvExpand.expanded(), duplicateLimit);
             } else {
                 return child.replaceChild(propagateDuplicateLimitUntilMvExpand(duplicateLimit, mvExpand, (UnaryPlan) child.child()));
             }
         }
     }
 
+    private static LogicalPlan replaceChildUntilMvExpandAndOrderBy(UnaryPlan child, MvExpand mvExpand) {
+        if (child == mvExpand) {
+            assert mvExpand.child() instanceof OrderBy;
+            OrderBy orderBy = (OrderBy) mvExpand.child();
+            return orderBy.replaceChild(mvExpand.replaceChild(orderBy.child()));
+        } else {
+            return child.replaceChild(replaceChildUntilMvExpandAndOrderBy((UnaryPlan) child.child(), mvExpand));
+        }
+    }
+
     // 3 in (field, 4, 5) --> 3 in (field) or 3 in (4, 5)
     public static class SplitInWithFoldableValue extends OptimizerRules.OptimizerExpressionRule<In> {
 
@@ -1016,37 +1057,10 @@ protected LogicalPlan rule(Limit plan) {
         }
     }
 
-    /**
-     * This adds an explicit TopN node to a plan that only has an OrderBy right before Lucene.
-     * To date, the only known use case that "needs" this is a query of the form
-     * from test
-     * | sort emp_no
-     * | mv_expand first_name
-     * | rename first_name AS x
-     * | where x LIKE "*a*"
-     * | limit 15
-     *
-     * or
-     *
-     * from test
-     * | sort emp_no
-     * | mv_expand first_name
-     * | sort first_name
-     * | limit 15
-     *
-     * PushDownAndCombineLimits rule will copy the "limit 15" after "sort emp_no" if there is no filter on the expanded values
-     * OR if there is no sort between "limit" and "mv_expand".
-     * But, since this type of query has such a filter, the "sort emp_no" will have no limit when it reaches the current rule.
-     */
-    static class AddDefaultTopN extends ParameterizedOptimizerRule<LogicalPlan, LogicalOptimizerContext> {
-
+    static class CleanUpMvExpand extends OptimizerRules.OptimizerRule<MvExpand> {
         @Override
-        protected LogicalPlan rule(LogicalPlan plan, LogicalOptimizerContext context) {
-            if (plan instanceof UnaryPlan unary && unary.child() instanceof OrderBy order && order.child() instanceof EsRelation relation) {
-                var limit = new Literal(Source.EMPTY, context.configuration().resultTruncationMaxSize(), DataTypes.INTEGER);
-                return unary.replaceChild(new TopN(plan.source(), relation, order.order(), limit));
-            }
-            return plan;
+        protected LogicalPlan rule(MvExpand plan) {
+            return plan.limit() == null ? plan : new MvExpand(plan.source(), plan.child(), plan.target(), plan.expanded(), null);
         }
     }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/LogicalPlanBuilder.java

@@ -151,7 +151,7 @@ public PlanFactory visitDissectCommand(EsqlBaseParser.DissectCommandContext ctx)
     public PlanFactory visitMvExpandCommand(EsqlBaseParser.MvExpandCommandContext ctx) {
         UnresolvedAttribute field = visitQualifiedName(ctx.qualifiedName());
         Source src = source(ctx);
-        return child -> new MvExpand(src, child, field, new UnresolvedAttribute(src, field.name()));
+        return child -> new MvExpand(src, child, field, new UnresolvedAttribute(src, field.name()), null);
 
     }
 

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

@@ -9,6 +9,7 @@
 
 import org.elasticsearch.xpack.ql.expression.Attribute;
 import org.elasticsearch.xpack.ql.expression.NamedExpression;
+import org.elasticsearch.xpack.ql.plan.logical.Limit;
 import org.elasticsearch.xpack.ql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.ql.plan.logical.UnaryPlan;
 import org.elasticsearch.xpack.ql.tree.NodeInfo;
@@ -23,11 +24,13 @@ public class MvExpand extends UnaryPlan {
     private final Attribute expanded;
 
     private List<Attribute> output;
+    private Limit limit; // a temporary "marker" indicating if a limit has been "associated" with this mv_expand
 
-    public MvExpand(Source source, LogicalPlan child, NamedExpression target, Attribute expanded) {
+    public MvExpand(Source source, LogicalPlan child, NamedExpression target, Attribute expanded, Limit limit) {
         super(source, child);
         this.target = target;
         this.expanded = expanded;
+        this.limit = limit;
     }
 
     public static List<Attribute> calculateOutput(List<Attribute> input, NamedExpression target, Attribute expanded) {
@@ -50,14 +53,18 @@ public Attribute expanded() {
         return expanded;
     }
 
+    public Limit limit() {
+        return limit;
+    }
+
     @Override
     public boolean expressionsResolved() {
         return target.resolved();
     }
 
     @Override
     public UnaryPlan replaceChild(LogicalPlan newChild) {
-        return new MvExpand(source(), newChild, target, expanded);
+        return new MvExpand(source(), newChild, target, expanded, limit);
     }
 
     @Override
@@ -70,19 +77,20 @@ public List<Attribute> output() {
 
     @Override
     protected NodeInfo<? extends LogicalPlan> info() {
-        return NodeInfo.create(this, MvExpand::new, child(), target, expanded);
+        return NodeInfo.create(this, MvExpand::new, child(), target, expanded, limit);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(super.hashCode(), target, expanded);
+        return Objects.hash(super.hashCode(), target, expanded, limit);
     }
 
     @Override
     public boolean equals(Object obj) {
-        if (false == super.equals(obj)) {
-            return false;
+        if (super.equals(obj)) {
+            var other = (MvExpand) obj;
+            return Objects.equals(target, other.target) && Objects.equals(expanded, other.expanded) && Objects.equals(limit, other.limit);
         }
-        return Objects.equals(target, ((MvExpand) obj).target) && Objects.equals(expanded, ((MvExpand) obj).expanded);
+        return false;
     }
 }

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/io/stream/PlanNamedTypesTests.java

@@ -488,7 +488,7 @@ public void testEsqlProject() throws IOException {
 
     public void testMvExpand() throws IOException {
         var esRelation = new EsRelation(Source.EMPTY, randomEsIndex(), List.of(randomFieldAttribute()), randomBoolean());
-        var orig = new MvExpand(Source.EMPTY, esRelation, randomFieldAttribute(), randomFieldAttribute());
+        var orig = new MvExpand(Source.EMPTY, esRelation, randomFieldAttribute(), randomFieldAttribute(), null);
         BytesStreamOutput bso = new BytesStreamOutput();
         PlanStreamOutput out = new PlanStreamOutput(bso, planNameRegistry);
         PlanNamedTypes.writeMvExpand(out, orig);