Remove optional seed from ES|QL SAMPLE (elastic#128887)

* Remove optional seed from ES|QL SAMPLE

* make it clear that seed is for testing

Author: jan-elastic

x-pack/plugin/esql/compute/src/main/java/org/elasticsearch/compute/operator/SampleOperator.java

@@ -9,6 +9,7 @@
 
 import org.elasticsearch.TransportVersion;
 import org.elasticsearch.TransportVersions;
+import org.elasticsearch.common.Randomness;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -27,16 +28,29 @@
 
 public class SampleOperator implements Operator {
 
-    public record Factory(double probability, int seed) implements OperatorFactory {
+    public static class Factory implements OperatorFactory {
+
+        private final double probability;
+        private final Integer seed;
+
+        public Factory(double probability) {
+            this(probability, null);
+        }
+
+        // visible for testing
+        Factory(double probability, Integer seed) {
+            this.probability = probability;
+            this.seed = seed;
+        }
 
         @Override
         public SampleOperator get(DriverContext driverContext) {
-            return new SampleOperator(probability, seed);
+            return new SampleOperator(probability, seed == null ? Randomness.get().nextInt() : seed);
         }
 
         @Override
         public String describe() {
-            return "SampleOperator[probability = " + probability + ", seed = " + seed + "]";
+            return "SampleOperator[probability = " + probability + "]";
         }
     }
 

x-pack/plugin/esql/compute/src/test/java/org/elasticsearch/compute/operator/SampleOperatorTests.java

@@ -21,7 +21,6 @@
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.lessThan;
-import static org.hamcrest.Matchers.matchesPattern;
 
 public class SampleOperatorTests extends OperatorTestCase {
 
@@ -46,7 +45,7 @@ protected SampleOperator.Factory simple() {
 
     @Override
     protected Matcher<String> expectedDescriptionOfSimple() {
-        return matchesPattern("SampleOperator\\[probability = 0.5, seed = -?\\d+]");
+        return equalTo("SampleOperator[probability = 0.5]");
     }
 
     @Override

x-pack/plugin/esql/src/main/antlr/EsqlBaseParser.g4

@@ -389,5 +389,5 @@ completionCommand
     ;
 
 sampleCommand
-    : DEV_SAMPLE probability=decimalValue seed=integerValue?
+    : DEV_SAMPLE probability=decimalValue
     ;

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

@@ -61,8 +61,7 @@ protected LogicalPlan rule(Sample sample, LogicalOptimizerContext context) {
         var child = sample.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());
+            plan = new Sample(sample.source(), probability, sampleChild.child());
         } else if (child instanceof Enrich
             || child instanceof Eval
             || child instanceof Filter
@@ -80,22 +79,4 @@ private static Expression combinedProbability(LogicalOptimizerContext context, S
         var childProbability = (double) Foldables.valueOf(context.foldCtx(), child.probability());
         return Literal.of(parent.probability(), parentProbability * childProbability);
     }
-
-    private static Expression combinedSeed(LogicalOptimizerContext context, Sample parent, Sample child) {
-        var parentSeed = parent.seed();
-        var childSeed = child.seed();
-        Expression seed;
-        if (parentSeed != null) {
-            if (childSeed != null) {
-                var seedValue = (int) Foldables.valueOf(context.foldCtx(), parentSeed);
-                seedValue ^= (int) Foldables.valueOf(context.foldCtx(), childSeed);
-                seed = Literal.of(parentSeed, seedValue);
-            } else {
-                seed = parentSeed;
-            }
-        } else {
-            seed = childSeed;
-        }
-        return seed;
-    }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/physical/local/PushSampleToSource.java

@@ -33,9 +33,6 @@ protected PhysicalPlan rule(SampleExec sample, LocalPhysicalOptimizerContext ctx
             }
 
             var sampleQuery = new RandomSamplingQueryBuilder((double) Foldables.valueOf(ctx.foldCtx(), sample.probability()));
-            if (sample.seed() != null) {
-                sampleQuery.seed((int) Foldables.valueOf(ctx.foldCtx(), sample.seed()));
-            }
 
             fullQuery.filter(sampleQuery);
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/EsqlBaseParser.interp

@@ -708,20 +708,6 @@ public Literal inferenceId(EsqlBaseParser.IdentifierOrParameterContext ctx) {
 
     public PlanFactory visitSampleCommand(EsqlBaseParser.SampleCommandContext ctx) {
         var probability = visitDecimalValue(ctx.probability);
-        Literal seed;
-        if (ctx.seed != null) {
-            seed = visitIntegerValue(ctx.seed);
-            if (seed.dataType() != DataType.INTEGER) {
-                throw new ParsingException(
-                    seed.source(),
-                    "seed must be an integer, provided [{}] of type [{}]",
-                    ctx.seed.getText(),
-                    seed.dataType()
-                );
-            }
-        } else {
-            seed = null;
-        }
-        return plan -> new Sample(source(ctx), probability, seed, plan);
+        return plan -> new Sample(source(ctx), probability, plan);
     }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/EsqlBaseParser.java

@@ -9,7 +9,6 @@
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
-import org.elasticsearch.core.Nullable;
 import org.elasticsearch.search.aggregations.bucket.sampler.random.RandomSamplingQuery;
 import org.elasticsearch.xpack.esql.capabilities.PostAnalysisVerificationAware;
 import org.elasticsearch.xpack.esql.capabilities.TelemetryAware;
@@ -30,19 +29,16 @@ public class Sample extends UnaryPlan implements TelemetryAware, PostAnalysisVer
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "Sample", Sample::new);
 
     private final Expression probability;
-    private final Expression seed;
 
-    public Sample(Source source, Expression probability, @Nullable Expression seed, LogicalPlan child) {
+    public Sample(Source source, Expression probability, LogicalPlan child) {
         super(source, child);
         this.probability = probability;
-        this.seed = seed;
     }
 
     private Sample(StreamInput in) throws IOException {
         this(
             Source.readFrom((PlanStreamInput) in),
             in.readNamedWriteable(Expression.class), // probability
-            in.readOptionalNamedWriteable(Expression.class), // seed
             in.readNamedWriteable(LogicalPlan.class) // child
         );
     }
@@ -51,7 +47,6 @@ private Sample(StreamInput in) throws IOException {
     public void writeTo(StreamOutput out) throws IOException {
         source().writeTo(out);
         out.writeNamedWriteable(probability);
-        out.writeOptionalNamedWriteable(seed);
         out.writeNamedWriteable(child());
     }
 
@@ -62,30 +57,26 @@ public String getWriteableName() {
 
     @Override
     protected NodeInfo<Sample> info() {
-        return NodeInfo.create(this, Sample::new, probability, seed, child());
+        return NodeInfo.create(this, Sample::new, probability, child());
     }
 
     @Override
     public Sample replaceChild(LogicalPlan newChild) {
-        return new Sample(source(), probability, seed, newChild);
+        return new Sample(source(), probability, newChild);
     }
 
     public Expression probability() {
         return probability;
     }
 
-    public Expression seed() {
-        return seed;
-    }
-
     @Override
     public boolean expressionsResolved() {
-        return probability.resolved() && (seed == null || seed.resolved());
+        return probability.resolved();
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(probability, seed, child());
+        return Objects.hash(probability, child());
     }
 
     @Override
@@ -99,7 +90,7 @@ public boolean equals(Object obj) {
 
         var other = (Sample) obj;
 
-        return Objects.equals(probability, other.probability) && Objects.equals(seed, other.seed) && Objects.equals(child(), other.child());
+        return Objects.equals(probability, other.probability) && Objects.equals(child(), other.child());
     }
 
     @Override

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/LogicalPlanBuilder.java

@@ -10,9 +10,6 @@
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
-import org.elasticsearch.core.Nullable;
-import org.elasticsearch.xpack.esql.capabilities.PostPhysicalOptimizationVerificationAware;
-import org.elasticsearch.xpack.esql.common.Failures;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
@@ -21,30 +18,25 @@
 import java.io.IOException;
 import java.util.Objects;
 
-import static org.elasticsearch.xpack.esql.common.Failure.fail;
-
-public class SampleExec extends UnaryExec implements PostPhysicalOptimizationVerificationAware {
+public class SampleExec extends UnaryExec {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
         PhysicalPlan.class,
         "SampleExec",
         SampleExec::new
     );
 
     private final Expression probability;
-    private final Expression seed;
 
-    public SampleExec(Source source, PhysicalPlan child, Expression probability, @Nullable Expression seed) {
+    public SampleExec(Source source, PhysicalPlan child, Expression probability) {
         super(source, child);
         this.probability = probability;
-        this.seed = seed;
     }
 
     public SampleExec(StreamInput in) throws IOException {
         this(
             Source.readFrom((PlanStreamInput) in),
             in.readNamedWriteable(PhysicalPlan.class), // child
-            in.readNamedWriteable(Expression.class), // probability
-            in.readOptionalNamedWriteable(Expression.class) // seed
+            in.readNamedWriteable(Expression.class) // probability
         );
     }
 
@@ -53,17 +45,16 @@ public void writeTo(StreamOutput out) throws IOException {
         source().writeTo(out);
         out.writeNamedWriteable(child());
         out.writeNamedWriteable(probability);
-        out.writeOptionalNamedWriteable(seed);
     }
 
     @Override
     public UnaryExec replaceChild(PhysicalPlan newChild) {
-        return new SampleExec(source(), newChild, probability, seed);
+        return new SampleExec(source(), newChild, probability);
     }
 
     @Override
     protected NodeInfo<? extends PhysicalPlan> info() {
-        return NodeInfo.create(this, SampleExec::new, child(), probability, seed);
+        return NodeInfo.create(this, SampleExec::new, child(), probability);
     }
 
     /**
@@ -78,13 +69,9 @@ public Expression probability() {
         return probability;
     }
 
-    public Expression seed() {
-        return seed;
-    }
-
     @Override
     public int hashCode() {
-        return Objects.hash(child(), probability, seed);
+        return Objects.hash(child(), probability);
     }
 
     @Override
@@ -98,17 +85,6 @@ public boolean equals(Object obj) {
 
         var other = (SampleExec) obj;
 
-        return Objects.equals(child(), other.child()) && Objects.equals(probability, other.probability) && Objects.equals(seed, other.seed);
-    }
-
-    @Override
-    public void postPhysicalOptimizationVerification(Failures failures) {
-        // It's currently impossible in ES|QL to handle all data in deterministic order, therefore
-        // a fixed random seed in the sample operator doesn't work as intended and is disallowed.
-        // TODO: fix this.
-        if (seed != null) {
-            // TODO: what should the error message here be? This doesn't seem right.
-            failures.add(fail(seed, "Seed not supported when sampling can't be pushed down to Lucene"));
-        }
+        return Objects.equals(child(), other.child()) && Objects.equals(probability, other.probability);
     }
 }