Include push down limits past sorting

Author: carlosdelest

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

@@ -17,6 +17,7 @@
 import org.elasticsearch.xpack.esql.plan.logical.Limit;
 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.RegexExtract;
 import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
@@ -73,6 +74,17 @@ public LogicalPlan rule(Limit limit, LogicalOptimizerContext ctx) {
                     return limit.replaceChild(filter.with(limitAppliedExpression));
                 }
             }
+            else if (unary instanceof OrderBy orderBy) {
+                return limit.replaceChild(orderBy.transformDown(lp -> {
+                    if (lp instanceof Filter filter) {
+                        Expression limitAppliedExpression = limitFilterExpressions(filter.condition(), limit, ctx);
+                        if (limitAppliedExpression.equals(filter.condition()) == false) {
+                            return filter.with(limitAppliedExpression);
+                        }
+                    }
+                    return lp;
+                }));
+            }
         } else if (limit.child() instanceof Join join && join.config().type() == JoinTypes.LEFT) {
             // Left joins increase the number of rows if any join key has multiple matches from the right hand side.
             // Therefore, we cannot simply push down the limit - but we can add another limit before the join.

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LocalPhysicalPlanOptimizerTests.java

@@ -2073,6 +2073,45 @@ public void testKnnWithoutExplicitLimit() {
         assertThat(knnQuery.k(), is(config.resultTruncationDefaultSize()));
     }
 
+    public void testKnnWithoutExplicitLimitSorted() {
+        var query = """
+            from test metadata _score
+            | where knn(dense_vector, [0, 1, 2])
+            | sort _score desc
+            """;
+        var analyzer = makeAnalyzer("mapping-all-types.json");
+        var plan = plannerOptimizer.plan(query, IS_SV_STATS, analyzer);
+        var topNExec = as(plan, TopNExec.class);
+        assertThat(topNExec.limit().fold(FoldContext.small()), is(config.resultTruncationDefaultSize()));
+        var exchangeExec = as(topNExec.child(), ExchangeExec.class);
+        var projectExec = as(exchangeExec.child(), ProjectExec.class);
+        var fieldExtractExec = as(projectExec.child(), FieldExtractExec.class);
+        var queryExec = as(fieldExtractExec.child(), EsQueryExec.class);
+        assertThat(queryExec.limit().fold(FoldContext.small()), is(config.resultTruncationDefaultSize()));
+        var knnQuery = as(queryExec.query(), KnnVectorQueryBuilder.class);
+        assertThat(knnQuery.k(), is(config.resultTruncationDefaultSize()));
+    }
+
+    public void testKnnWithoExplicitLimitSorted() {
+        var query = """
+            from test metadata _score
+            | where knn(dense_vector, [0, 1, 2])
+            | sort _score desc
+            | limit 10
+            """;
+        var analyzer = makeAnalyzer("mapping-all-types.json");
+        var plan = plannerOptimizer.plan(query, IS_SV_STATS, analyzer);
+        var topNExec = as(plan, TopNExec.class);
+        assertThat(topNExec.limit().fold(FoldContext.small()), is(10));
+        var exchangeExec = as(topNExec.child(), ExchangeExec.class);
+        var projectExec = as(exchangeExec.child(), ProjectExec.class);
+        var fieldExtractExec = as(projectExec.child(), FieldExtractExec.class);
+        var queryExec = as(fieldExtractExec.child(), EsQueryExec.class);
+        assertThat(queryExec.limit().fold(FoldContext.small()), is(10));
+        var knnQuery = as(queryExec.query(), KnnVectorQueryBuilder.class);
+        assertThat(knnQuery.k(), is(10));
+    }
+
     public void testKnnWithExplicitLimit() {
         var query = """
             from test