ESQL: mv_expand pushes down a limit copy and keeps the limit after it…

… untouched (#100782) (#101268)

- allow mv_expand to push down limit and project past it
- accept a limit after mv_expand when there is also a second limit before the mv_expand
- adds a default TopN for cases when there is only a sort at Lucene level
- adds OrderBy node type to the exceptions for duplicating the limit after mv_expand

(cherry picked from commit 4679b09)

docs/changelog/100782.yaml

@@ -0,0 +1,8 @@
+pr: 100782
+summary: "ESQL: `mv_expand` pushes down limit and project and keep the limit after\
+  \ it untouched"
+area: ES|QL
+type: bug
+issues:
+ - 99971
+ - 100774

x-pack/plugin/esql/qa/testFixtures/src/main/resources/mv_expand.csv-spec

@@ -98,3 +98,130 @@ sum_a:long    | b:integer
 12555000      | 29
 12555000      | 30
 ;
+
+expandAfterSort1
+from employees | keep job_positions, emp_no | sort emp_no | mv_expand job_positions | limit 10 | sort job_positions;
+
+   job_positions:keyword  |emp_no:integer     
+Accountant                |10001          
+Head Human Resources      |10004          
+Principal Support Engineer|10006          
+Reporting Analyst         |10004          
+Senior Python Developer   |10001          
+Senior Team Lead          |10002          
+Support Engineer          |10004          
+Tech Lead                 |10004          
+null                      |10005          
+null                      |10003          
+;
+
+expandAfterSort2
+from employees | sort emp_no | mv_expand job_positions | keep job_positions, emp_no | limit 5;
+
+   job_positions:keyword  |emp_no:integer     
+Accountant             |10001          
+Senior Python Developer|10001          
+Senior Team Lead       |10002          
+null                   |10003          
+Head Human Resources   |10004
+;
+
+expandWithMultiSort
+from employees | keep emp_no, job_positions | sort emp_no | mv_expand job_positions | limit 10 | where emp_no <= 10006 | sort job_positions nulls first;
+
+emp_no:integer |  job_positions:keyword       
+10003          |null                      
+10005          |null                      
+10001          |Accountant                
+10004          |Head Human Resources      
+10006          |Principal Support Engineer
+10004          |Reporting Analyst         
+10001          |Senior Python Developer   
+10002          |Senior Team Lead          
+10004          |Support Engineer          
+10004          |Tech Lead
+;
+
+filterMvExpanded
+from employees | keep emp_no, job_positions | mv_expand job_positions | where job_positions like "A*" | sort job_positions, emp_no;
+
+emp_no:integer |  job_positions:keyword     
+10001          |Accountant     
+10012          |Accountant     
+10016          |Accountant     
+10023          |Accountant     
+10025          |Accountant     
+10028          |Accountant     
+10034          |Accountant     
+10037          |Accountant     
+10044          |Accountant     
+10045          |Accountant     
+10050          |Accountant     
+10051          |Accountant     
+10066          |Accountant     
+10081          |Accountant     
+10085          |Accountant     
+10089          |Accountant     
+10092          |Accountant     
+10094          |Accountant     
+10010          |Architect      
+10011          |Architect      
+10031          |Architect      
+10032          |Architect      
+10042          |Architect      
+10047          |Architect      
+10059          |Architect      
+10068          |Architect      
+10072          |Architect      
+10076          |Architect      
+10078          |Architect      
+10096          |Architect      
+10098          |Architect      
+;
+
+doubleSort_OnDifferentThan_MvExpandedFields
+from employees | sort emp_no | mv_expand job_positions | keep emp_no, job_positions, salary | sort salary, job_positions | limit 5;
+
+emp_no:integer |  job_positions:keyword   |salary:integer
+10015          |Head Human Resources      |25324          
+10015          |Junior Developer          |25324          
+10015          |Principal Support Engineer|25324          
+10015          |Support Engineer          |25324          
+10035          |Data Scientist            |25945
+;
+
+doubleLimit_expandLimitLowerThanAvailable
+from employees | where emp_no == 10004 | limit 1 | keep emp_no, job_positions | mv_expand job_positions | limit 2;
+
+emp_no:integer |  job_positions:keyword
+10004          |Head Human Resources
+10004          |Reporting Analyst
+;
+
+doubleLimit_expandLimitGreaterThanAvailable
+from employees | where emp_no == 10004 | limit 1 | keep emp_no, job_positions | mv_expand job_positions | limit 5;
+
+emp_no:integer |  job_positions:keyword
+10004          |Head Human Resources
+10004          |Reporting Analyst   
+10004          |Support Engineer    
+10004          |Tech Lead
+;
+
+doubleLimitWithSort
+from employees | where emp_no == 10004 | limit 1 | keep emp_no, job_positions | mv_expand job_positions | limit 5 | sort job_positions desc;
+
+emp_no:integer |  job_positions:keyword
+10004          |Tech Lead           
+10004          |Support Engineer    
+10004          |Reporting Analyst   
+10004          |Head Human Resources
+;
+
+tripleLimit_WithWhere_InBetween_MvExpand_And_Limit
+from employees | where emp_no == 10004 | limit 1 | keep emp_no, job_positions | mv_expand job_positions | where job_positions LIKE "*a*" | limit 2 | where job_positions LIKE "*a*" | limit 3;
+
+emp_no:integer |  job_positions:keyword
+10004          |Head Human Resources
+10004          |Reporting Analyst
+;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/execution/PlanExecutor.java

@@ -13,6 +13,7 @@
 import org.elasticsearch.xpack.esql.analysis.Verifier;
 import org.elasticsearch.xpack.esql.enrich.EnrichPolicyResolver;
 import org.elasticsearch.xpack.esql.expression.function.EsqlFunctionRegistry;
+import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
 import org.elasticsearch.xpack.esql.optimizer.LogicalPlanOptimizer;
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
 import org.elasticsearch.xpack.esql.planner.Mapper;
@@ -30,7 +31,6 @@ public class PlanExecutor {
     private final IndexResolver indexResolver;
     private final PreAnalyzer preAnalyzer;
     private final FunctionRegistry functionRegistry;
-    private final LogicalPlanOptimizer logicalPlanOptimizer;
     private final Mapper mapper;
     private final Metrics metrics;
     private final Verifier verifier;
@@ -39,7 +39,6 @@ public PlanExecutor(IndexResolver indexResolver) {
         this.indexResolver = indexResolver;
         this.preAnalyzer = new PreAnalyzer();
         this.functionRegistry = new EsqlFunctionRegistry();
-        this.logicalPlanOptimizer = new LogicalPlanOptimizer();
         this.mapper = new Mapper(functionRegistry);
         this.metrics = new Metrics();
         this.verifier = new Verifier(metrics);
@@ -59,7 +58,7 @@ public void esql(
             enrichPolicyResolver,
             preAnalyzer,
             functionRegistry,
-            logicalPlanOptimizer,
+            new LogicalPlanOptimizer(new LogicalOptimizerContext(cfg)),
             mapper,
             verifier
         );

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

@@ -10,4 +10,36 @@
 import org.elasticsearch.xpack.esql.session.EsqlConfiguration;
 import org.elasticsearch.xpack.esql.stats.SearchStats;
 
-public record LocalLogicalOptimizerContext(EsqlConfiguration configuration, SearchStats searchStats) {}
+import java.util.Objects;
+
+public final class LocalLogicalOptimizerContext extends LogicalOptimizerContext {
+    private final SearchStats searchStats;
+
+    public LocalLogicalOptimizerContext(EsqlConfiguration configuration, SearchStats searchStats) {
+        super(configuration);
+        this.searchStats = searchStats;
+    }
+
+    public SearchStats searchStats() {
+        return searchStats;
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (super.equals(obj)) {
+            var that = (LocalLogicalOptimizerContext) obj;
+            return Objects.equals(this.searchStats, that.searchStats);
+        }
+        return false;
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(super.hashCode(), searchStats);
+    }
+
+    @Override
+    public String toString() {
+        return "LocalLogicalOptimizerContext[" + "configuration=" + configuration() + ", " + "searchStats=" + searchStats + ']';
+    }
+}

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

@@ -116,10 +116,7 @@ else if (plan instanceof Project project) {
         }
     }
 
-    public abstract static class ParameterizedOptimizerRule<SubPlan extends LogicalPlan, P> extends ParameterizedRule<
-        SubPlan,
-        LogicalPlan,
-        P> {
+    abstract static class ParameterizedOptimizerRule<SubPlan extends LogicalPlan, P> extends ParameterizedRule<SubPlan, LogicalPlan, P> {
 
         public final LogicalPlan apply(LogicalPlan plan, P context) {
             return plan.transformUp(typeToken(), t -> rule(t, context));

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

@@ -0,0 +1,43 @@
+/*
+ * 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.optimizer;
+
+import org.elasticsearch.xpack.esql.session.EsqlConfiguration;
+
+import java.util.Objects;
+
+public class LogicalOptimizerContext {
+    private final EsqlConfiguration configuration;
+
+    public LogicalOptimizerContext(EsqlConfiguration configuration) {
+        this.configuration = configuration;
+    }
+
+    public EsqlConfiguration configuration() {
+        return configuration;
+    }
+
+    @Override
+    public boolean equals(Object obj) {
+        if (obj == this) return true;
+        if (obj == null || obj.getClass() != this.getClass()) return false;
+        var that = (LogicalOptimizerContext) obj;
+        return Objects.equals(this.configuration, that.configuration);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(configuration);
+    }
+
+    @Override
+    public String toString() {
+        return "LogicalOptimizerContext[" + "configuration=" + configuration + ']';
+    }
+
+}