ES|QL: Remove redundant sorts from execution plan

docs/changelog/121156.yaml

@@ -0,0 +1,5 @@
+pr: 121156
+summary: Remove redundant sorts from execution plan
+area: ES|QL
+type: bug
+issues: []

muted-tests.yml

@@ -314,8 +314,6 @@ tests:
 - class: org.elasticsearch.xpack.security.profile.ProfileIntegTests
   method: testSetEnabled
   issue: https://github.com/elastic/elasticsearch/issues/121183
-- class: org.elasticsearch.xpack.esql.optimizer.LogicalPlanOptimizerTests
-  issue: https://github.com/elastic/elasticsearch/issues/121185
 - class: org.elasticsearch.xpack.security.CoreWithSecurityClientYamlTestSuiteIT
   method: test {yaml=cat.aliases/10_basic/Simple alias}
   issue: https://github.com/elastic/elasticsearch/issues/121186

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

@@ -1346,3 +1346,50 @@ language_code:integer | language_name:keyword       | country:text
 1                     | English                     | United States of America
 1                     | English                     | null
 ;
+
+
+sortBeforeAndAfterJoin
+required_capability: join_lookup_v12
+required_capability: remove_redundant_sort
+
+FROM employees
+| sort first_name
+| EVAL language_code = languages
+| LOOKUP JOIN languages_lookup ON language_code
+| WHERE emp_no >= 10091 AND emp_no < 10094
+| SORT emp_no
+| KEEP emp_no, language_code, language_name
+;
+
+emp_no:integer | language_code:integer | language_name:keyword
+10091          | 3                     | Spanish
+10092          | 1                     | English
+10093          | 3                     | Spanish
+;
+
+
+
+sortBeforeAndAfterMultipleJoinAndMvExpand
+required_capability: join_lookup_v12
+required_capability: remove_redundant_sort
+
+FROM employees
+| sort first_name
+| EVAL language_code = languages
+| LOOKUP JOIN languages_lookup ON language_code
+| WHERE emp_no >= 10091 AND emp_no < 10094
+| SORT language_name
+| MV_EXPAND first_name
+| SORT first_name
+| MV_EXPAND last_name
+| SORT last_name
+| LOOKUP JOIN languages_lookup ON language_code
+| SORT emp_no
+| KEEP emp_no, language_code, language_name
+;
+
+emp_no:integer | language_code:integer | language_name:keyword
+10091          | 3                     | Spanish
+10092          | 1                     | English
+10093          | 3                     | Spanish
+;

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

@@ -404,3 +404,17 @@ from employees | where emp_no == 10003  | mv_expand first_name | keep first_name
 first_name:keyword
 Parto
 ;
+
+
+sortBeforeAndAfterMvExpand
+from employees 
+| sort first_name 
+| mv_expand job_positions 
+| sort emp_no, job_positions 
+| keep emp_no, job_positions 
+| limit 2;
+
+emp_no:integer | job_positions:keyword   
+10001          | Accountant
+10001          | Senior Python Developer          
+;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -784,7 +784,14 @@ public enum Cap {
         /**
          * Support for aggregate_metric_double type
          */
-        AGGREGATE_METRIC_DOUBLE;
+        AGGREGATE_METRIC_DOUBLE,
+
+        /**
+         * Fix for https://github.com/elastic/elasticsearch/issues/120817
+         * and https://github.com/elastic/elasticsearch/issues/120803
+         * Support for queries that have multiple SORTs that cannot become TopN
+         */
+        REMOVE_REDUNDANT_SORT;
 
         private final boolean enabled;
 

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

@@ -40,6 +40,7 @@
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownEnrich;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownEval;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PushDownRegexExtract;
+import org.elasticsearch.xpack.esql.optimizer.rules.logical.RemoveRedundantSort;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.RemoveStatsOverride;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.ReplaceAggregateAggExpressionWithEval;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.ReplaceAggregateNestedExpressionWithEval;
@@ -195,6 +196,6 @@ protected static Batch<LogicalPlan> operators() {
     }
 
     protected static Batch<LogicalPlan> cleanup() {
-        return new Batch<>("Clean Up", new ReplaceLimitAndSortAsTopN(), new ReplaceRowAsLocalRelation());
+        return new Batch<>("Clean Up", new ReplaceLimitAndSortAsTopN(), new ReplaceRowAsLocalRelation(), new RemoveRedundantSort());
     }
 }

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

@@ -27,6 +27,9 @@ public Failures verify(LogicalPlan plan) {
             PlanConsistencyChecker.checkPlan(p, dependencyFailures);
 
             if (failures.hasFailures() == false) {
+                if (p instanceof PostOptimizationVerificationAware pova) {
+                    pova.postOptimizationVerification(failures);
+                }
                 p.forEachExpression(ex -> {
                     if (ex instanceof PostOptimizationVerificationAware va) {
                         va.postOptimizationVerification(failures);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/RemoveRedundantSort.java

@@ -0,0 +1,88 @@
+/*
+ * 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.rules.logical;
+
+import org.elasticsearch.xpack.esql.plan.logical.Dissect;
+import org.elasticsearch.xpack.esql.plan.logical.Enrich;
+import org.elasticsearch.xpack.esql.plan.logical.Filter;
+import org.elasticsearch.xpack.esql.plan.logical.Grok;
+import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+import org.elasticsearch.xpack.esql.plan.logical.MvExpand;
+import org.elasticsearch.xpack.esql.plan.logical.OrderBy;
+import org.elasticsearch.xpack.esql.plan.logical.Project;
+import org.elasticsearch.xpack.esql.plan.logical.Rename;
+import org.elasticsearch.xpack.esql.plan.logical.TopN;
+import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
+import org.elasticsearch.xpack.esql.plan.logical.join.Join;
+
+/**
+ * SORT cannot be executed without a LIMIT, as ES|QL doesn't support unbounded sort (yet).
+ *
+ * The planner tries to push down LIMIT and transform all the unbounded sorts into a TopN.
+ * In some cases it's not possible though, eg.
+ *
+ * from test | sort x | lookup join lookup on x | sort y
+ *
+ * from test | sort x | mv_expand x | sort y
+ *
+ * "sort y" will become a TopN, but "sort x" will remain unbounded, so the query could not be executed.
+ *
+ * In most cases though, last SORT make the previous SORTs redundant,
+ * ie. it will re-sort previously sorted results
+ * often with a different order.
+ *
+ * This rule finds and removes redundant SORTs, making the plan executable.
+ */
+public class RemoveRedundantSort extends OptimizerRules.OptimizerRule<TopN> {
+
+    @Override
+    protected LogicalPlan rule(TopN plan) {
+        OrderBy redundant = findRedundantSort(plan);
+        if (redundant == null) {
+            return plan;
+        }
+        return plan.transformDown(p -> {
+            if (p == redundant) {
+                return redundant.child();
+            }
+            return p;
+        });
+    }
+
+    private OrderBy findRedundantSort(TopN plan) {
+        LogicalPlan p = plan.child();
+        while (true) {
+            if (p instanceof OrderBy ob) {
+                return ob;
+            }
+            if (p instanceof UnaryPlan unary) {
+                if (unary instanceof Filter
+                    || unary instanceof Project
+                    || unary instanceof Rename
+                    || unary instanceof MvExpand
+                    || unary instanceof Enrich
+                    || unary instanceof Grok
+                    || unary instanceof Dissect
+                // If we introduce window functions, the previous sort could actually become relevant
+                // so to be sure we don't introduce regressions, we'll have to exclude places where these functions could be used
+                // || unary instanceof Eval
+                // || unary instanceof Aggregate
+                ) {
+                    p = unary.child();
+                    continue;
+                }
+            } else if (p instanceof Join lj) {
+                p = lj.left();
+                // TODO do it also on the right-hand side?
+                continue;
+            }
+            return null;
+        }
+    }
+
+}

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

@@ -10,6 +10,7 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.xpack.esql.capabilities.PostAnalysisVerificationAware;
+import org.elasticsearch.xpack.esql.capabilities.PostOptimizationVerificationAware;
 import org.elasticsearch.xpack.esql.capabilities.TelemetryAware;
 import org.elasticsearch.xpack.esql.common.Failures;
 import org.elasticsearch.xpack.esql.core.capabilities.Resolvables;
@@ -25,7 +26,7 @@
 
 import static org.elasticsearch.xpack.esql.common.Failure.fail;
 
-public class OrderBy extends UnaryPlan implements PostAnalysisVerificationAware, TelemetryAware {
+public class OrderBy extends UnaryPlan implements PostAnalysisVerificationAware, PostOptimizationVerificationAware, TelemetryAware {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "OrderBy", OrderBy::new);
 
     private final List<Order> order;
@@ -109,4 +110,9 @@ public void postAnalysisVerification(Failures failures) {
             }
         });
     }
+
+    @Override
+    public void postOptimizationVerification(Failures failures) {
+        failures.add(fail(this, "The query cannot be executed because it would require unbounded sort"));
+    }
 }