use samplebreaking in pushdown

Author: jan-elastic

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

@@ -43,7 +43,7 @@ public LogicalPlan apply(LogicalPlan logicalPlan) {
         });
     }
 
-    private Expression getSampleProbability(List<Expression> sampleProbabilities, Source source) {
+    private static Expression getSampleProbability(List<Expression> sampleProbabilities, Source source) {
         Expression result = null;
         for (Expression probability : sampleProbabilities) {
             result = result == null ? probability : new Mul(source, result, probability);

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

@@ -11,15 +11,9 @@
 import org.elasticsearch.xpack.esql.core.expression.Foldables;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
 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.RegexExtract;
 import org.elasticsearch.xpack.esql.plan.logical.Sample;
+import org.elasticsearch.xpack.esql.plan.logical.SampleBreaking;
 import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
 
 public class PushDownAndCombineSample extends OptimizerRules.ParameterizedOptimizerRule<Sample, LogicalOptimizerContext> {
@@ -32,20 +26,13 @@ public PushDownAndCombineSample() {
     protected LogicalPlan rule(Sample sample, LogicalOptimizerContext context) {
         LogicalPlan plan = sample;
         var child = sample.child();
-        if (child instanceof Sample rsChild) {
-            var probability = combinedProbability(context, sample, rsChild);
-            var seed = combinedSeed(context, sample, rsChild);
-            plan = new Sample(sample.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;
-                plan = unaryChild.replaceChild(sample.replaceChild(unaryChild.child()));
-            }
+        if (child instanceof Sample sampleChild) {
+            var probability = combinedProbability(context, sample, sampleChild);
+            var seed = combinedSeed(context, sample, sampleChild);
+            plan = new Sample(sample.source(), probability, seed, sampleChild.child());
+        } else if (child instanceof UnaryPlan unaryChild && child instanceof SampleBreaking == false) {
+            plan = unaryChild.replaceChild(sample.replaceChild(unaryChild.child()));
+        }
         return plan;
     }