push down through where and sort

Author: jan-elastic

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

@@ -13,8 +13,10 @@
 import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
 import org.elasticsearch.xpack.esql.plan.logical.Eval;
+import org.elasticsearch.xpack.esql.plan.logical.Filter;
 import org.elasticsearch.xpack.esql.plan.logical.Insist;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+import org.elasticsearch.xpack.esql.plan.logical.OrderBy;
 import org.elasticsearch.xpack.esql.plan.logical.Project;
 import org.elasticsearch.xpack.esql.plan.logical.RandomSample;
 import org.elasticsearch.xpack.esql.plan.logical.RegexExtract;
@@ -36,7 +38,9 @@ protected LogicalPlan rule(RandomSample randomSample, LogicalOptimizerContext co
             plan = new RandomSample(randomSample.source(), probability, seed, rsChild.child());
         } else if (child instanceof Enrich
             || child instanceof Eval
+            || child instanceof Filter
             || child instanceof Insist
+            || child instanceof OrderBy
             || child instanceof Project
             || child instanceof RegexExtract) {
                 var unaryChild = (UnaryPlan) child;

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

@@ -7672,14 +7672,33 @@ public void testRandomSamplePushDown() {
         }
     }
 
+    public void testRandomSamplePushDown_sort() {
+        var query = "FROM TEST | WHERE emp_no > 0 | RANDOM_SAMPLE 0.5 | LIMIT 100";
+        var optimized = optimizedPlan(query);
+
+        var limit = as(optimized, Limit.class);
+        var filter = as(limit.child(), Filter.class);
+        var randomSample = as(filter.child(), RandomSample.class);
+        var source = as(randomSample.child(), EsRelation.class);
+
+        assertThat(randomSample.probability().fold(FoldContext.small()), equalTo(0.5));
+        assertNull(randomSample.seed());
+    }
+
+    public void testRandomSamplePushDown_where() {
+        var query = "FROM TEST | SORT emp_no | RANDOM_SAMPLE 0.5 | LIMIT 100";
+        var optimized = optimizedPlan(query);
+
+        var topN = as(optimized, TopN.class);
+        var randomSample = as(topN.child(), RandomSample.class);
+        var source = as(randomSample.child(), EsRelation.class);
+
+        assertThat(randomSample.probability().fold(FoldContext.small()), equalTo(0.5));
+        assertNull(randomSample.seed());
+    }
+
     public void testRandomSampleNoPushDown() {
-        for (var command : List.of(
-            "LIMIT 100",
-            "MV_EXPAND languages",
-            "SORT emp_no | LIMIT 100", // the limit avoids an unbounded sort
-            "STATS COUNT()",
-            "WHERE emp_no > 1"
-        )) {
+        for (var command : List.of("LIMIT 100", "MV_EXPAND languages", "STATS COUNT()")) {
             var query = "FROM TEST | " + command + " | RANDOM_SAMPLE .5";
             var optimized = optimizedPlan(query);