error on seed in sampling operator

Author: jan-elastic

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/capabilities/PostPhysicalOptimizationVerificationAware.java

@@ -0,0 +1,22 @@
+/*
+ * 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.capabilities;
+
+import org.elasticsearch.xpack.esql.common.Failures;
+
+/**
+ * Interface implemented by expressions that require validation post physical optimization.
+ */
+public interface PostPhysicalOptimizationVerificationAware {
+
+    /**
+     * Validates the implementing expression - discovered failures are reported to the given
+     * {@link Failures} class.
+     */
+    void postPhysicalOptimizationVerification(Failures failures);
+}

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

@@ -8,7 +8,7 @@
 package org.elasticsearch.xpack.esql.optimizer;
 
 import org.elasticsearch.xpack.esql.VerificationException;
-import org.elasticsearch.xpack.esql.common.Failure;
+import org.elasticsearch.xpack.esql.common.Failures;
 import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.EnableSpatialDistancePushdown;
 import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.InsertFieldExtraction;
 import org.elasticsearch.xpack.esql.optimizer.rules.physical.local.PushFiltersToSource;
@@ -24,7 +24,6 @@
 import org.elasticsearch.xpack.esql.rule.Rule;
 
 import java.util.ArrayList;
-import java.util.Collection;
 import java.util.List;
 
 /**
@@ -46,8 +45,8 @@ public PhysicalPlan localOptimize(PhysicalPlan plan) {
     }
 
     PhysicalPlan verify(PhysicalPlan plan) {
-        Collection<Failure> failures = verifier.verify(plan);
-        if (failures.isEmpty() == false) {
+        Failures failures = verifier.verify(plan);
+        if (failures.hasFailures()) {
             throw new VerificationException(failures);
         }
         return plan;

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

@@ -8,14 +8,13 @@
 package org.elasticsearch.xpack.esql.optimizer;
 
 import org.elasticsearch.xpack.esql.VerificationException;
-import org.elasticsearch.xpack.esql.common.Failure;
+import org.elasticsearch.xpack.esql.common.Failures;
 import org.elasticsearch.xpack.esql.optimizer.rules.physical.ProjectAwayColumns;
 import org.elasticsearch.xpack.esql.plan.physical.FragmentExec;
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
 import org.elasticsearch.xpack.esql.rule.ParameterizedRuleExecutor;
 import org.elasticsearch.xpack.esql.rule.RuleExecutor;
 
-import java.util.Collection;
 import java.util.List;
 
 /**
@@ -39,8 +38,8 @@ public PhysicalPlan optimize(PhysicalPlan plan) {
     }
 
     PhysicalPlan verify(PhysicalPlan plan) {
-        Collection<Failure> failures = verifier.verify(plan);
-        if (failures.isEmpty() == false) {
+        Failures failures = verifier.verify(plan);
+        if (failures.hasFailures()) {
             throw new VerificationException(failures);
         }
         return plan;

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

@@ -7,7 +7,7 @@
 
 package org.elasticsearch.xpack.esql.optimizer;
 
-import org.elasticsearch.xpack.esql.common.Failure;
+import org.elasticsearch.xpack.esql.capabilities.PostPhysicalOptimizationVerificationAware;
 import org.elasticsearch.xpack.esql.common.Failures;
 import org.elasticsearch.xpack.esql.core.expression.Attribute;
 import org.elasticsearch.xpack.esql.core.expression.Expressions;
@@ -17,10 +17,6 @@
 import org.elasticsearch.xpack.esql.plan.physical.FieldExtractExec;
 import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
 
-import java.util.Collection;
-import java.util.LinkedHashSet;
-import java.util.Set;
-
 import static org.elasticsearch.xpack.esql.common.Failure.fail;
 
 /** Physical plan verifier. */
@@ -31,8 +27,8 @@ public final class PhysicalVerifier {
     private PhysicalVerifier() {}
 
     /** Verifies the physical plan. */
-    public Collection<Failure> verify(PhysicalPlan plan) {
-        Set<Failure> failures = new LinkedHashSet<>();
+    public Failures verify(PhysicalPlan plan) {
+        Failures failures = new Failures();
         Failures depFailures = new Failures();
 
         // AwaitsFix https://github.com/elastic/elasticsearch/issues/118531
@@ -56,6 +52,17 @@ public Collection<Failure> verify(PhysicalPlan plan) {
                 }
             }
             PlanConsistencyChecker.checkPlan(p, depFailures);
+
+            if (failures.hasFailures() == false) {
+                if (p instanceof PostPhysicalOptimizationVerificationAware va) {
+                    va.postPhysicalOptimizationVerification(failures);
+                }
+                p.forEachExpression(ex -> {
+                    if (ex instanceof PostPhysicalOptimizationVerificationAware va) {
+                        va.postPhysicalOptimizationVerification(failures);
+                    }
+                });
+            }
         });
 
         if (depFailures.hasFailures()) {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/SampleExec.java

@@ -11,6 +11,8 @@
 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;
@@ -19,7 +21,9 @@
 import java.io.IOException;
 import java.util.Objects;
 
-public class SampleExec extends UnaryExec {
+import static org.elasticsearch.xpack.esql.common.Failure.fail;
+
+public class SampleExec extends UnaryExec implements PostPhysicalOptimizationVerificationAware {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
         PhysicalPlan.class,
         "SampleExec",
@@ -96,4 +100,15 @@ public boolean equals(Object 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"));
+        }
+    }
 }

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

@@ -7701,6 +7701,8 @@ public void testPruneRedundantOrderBy() {
      *        \_EsRelation[test][_meta_field{f}#12, emp_no{f}#6, first_name{f}#7, ge..]
      */
     public void testSampleMerged() {
+        assumeTrue("sample must be enabled", EsqlCapabilities.Cap.SAMPLE.isEnabled());
+
         var query = """
             FROM TEST
             | SAMPLE .3 5
@@ -7721,6 +7723,8 @@ public void testSampleMerged() {
     }
 
     public void testSamplePushDown() {
+        assumeTrue("sample must be enabled", EsqlCapabilities.Cap.SAMPLE.isEnabled());
+
         for (var command : List.of(
             "ENRICH languages_idx on first_name",
             "EVAL x = 1",
@@ -7745,6 +7749,8 @@ public void testSamplePushDown() {
     }
 
     public void testSamplePushDown_sort() {
+        assumeTrue("sample must be enabled", EsqlCapabilities.Cap.SAMPLE.isEnabled());
+
         var query = "FROM TEST | WHERE emp_no > 0 | SAMPLE 0.5 | LIMIT 100";
         var optimized = optimizedPlan(query);
 
@@ -7758,6 +7764,8 @@ public void testSamplePushDown_sort() {
     }
 
     public void testSamplePushDown_where() {
+        assumeTrue("sample must be enabled", EsqlCapabilities.Cap.SAMPLE.isEnabled());
+
         var query = "FROM TEST | SORT emp_no | SAMPLE 0.5 | LIMIT 100";
         var optimized = optimizedPlan(query);
 
@@ -7770,6 +7778,8 @@ public void testSamplePushDown_where() {
     }
 
     public void testSampleNoPushDown() {
+        assumeTrue("sample must be enabled", EsqlCapabilities.Cap.SAMPLE.isEnabled());
+
         for (var command : List.of("LIMIT 100", "MV_EXPAND languages", "STATS COUNT()")) {
             var query = "FROM TEST | " + command + " | SAMPLE .5";
             var optimized = optimizedPlan(query);
@@ -7789,7 +7799,9 @@ public void testSampleNoPushDown() {
      *        | \_EsRelation[test][_meta_field{f}#12, emp_no{f}#6, first_name{f}#7, ge..]
      *        \_EsRelation[languages_lookup][LOOKUP][language_code{f}#17, language_name{f}#18]
      */
-    public void testRandomSampleNoPushDownLookupJoin() {
+    public void testSampleNoPushDownLookupJoin() {
+        assumeTrue("sample must be enabled", EsqlCapabilities.Cap.SAMPLE.isEnabled());
+
         var query = """
             FROM TEST
             | EVAL language_code = emp_no
@@ -7814,6 +7826,8 @@ public void testRandomSampleNoPushDownLookupJoin() {
      *            \_EsRelation[test][_meta_field{f}#12, emp_no{f}#6, first_name{f}#7, ge..]
      */
     public void testSampleNoPushDownChangePoint() {
+        assumeTrue("sample must be enabled", EsqlCapabilities.Cap.SAMPLE.isEnabled());
+
         var query = """
             FROM TEST
             | CHANGE_POINT emp_no ON hire_date

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

@@ -7801,7 +7801,9 @@ public void testReductionPlanForAggs() {
      *                  query[{"bool":{"filter":[{"sampling":{"probability":0.1,"seed":234,"hash":0}}],"boost":1.0}}]
      *                  [_doc{f}#24], limit[1000], sort[] estimatedRowSize[332]
      */
-    public void testRandomSamplePushDown() {
+    public void testSamplePushDown() {
+        assumeTrue("sample must be enabled", EsqlCapabilities.Cap.SAMPLE.isEnabled());
+
         var plan = physicalPlan("""
             FROM test
             | SAMPLE +0.1 -234
@@ -7822,6 +7824,20 @@ public void testRandomSamplePushDown() {
         assertThat(randomSampling.hash(), equalTo(0));
     }
 
+    public void testSample_seedNotSupportedInOperator() {
+        assumeTrue("sample must be enabled", EsqlCapabilities.Cap.SAMPLE.isEnabled());
+
+        optimizedPlan(physicalPlan("FROM test | SAMPLE 0.1"));
+        optimizedPlan(physicalPlan("FROM test | SAMPLE 0.1 42"));
+        optimizedPlan(physicalPlan("FROM test | MV_EXPAND first_name | SAMPLE 0.1"));
+
+        VerificationException e = expectThrows(
+            VerificationException.class,
+            () -> optimizedPlan(physicalPlan("FROM test | MV_EXPAND first_name | SAMPLE 0.1 42"))
+        );
+        assertThat(e.getMessage(), equalTo("Found 1 problem\nline 1:47: Seed not supported when sampling can't be pushed down to Lucene"));
+    }
+
     @SuppressWarnings("SameParameterValue")
     private static void assertFilterCondition(
         Filter filter,
@@ -8005,7 +8021,7 @@ private PhysicalPlan physicalPlan(String query, TestDataSource dataSource, boole
         var logical = logicalOptimizer.optimize(dataSource.analyzer.analyze(parser.createStatement(query)));
         // System.out.println("Logical\n" + logical);
         var physical = mapper.map(logical);
-        // System.out.println(physical);
+        // System.out.println("Physical\n" + physical);
         if (assertSerialization) {
             assertSerialization(physical);
         }