Test accuracy of sampling operator

Author: jan-elastic

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

@@ -16,8 +16,11 @@
 import java.util.List;
 import java.util.stream.LongStream;
 
+import static org.hamcrest.Matchers.both;
 import static org.hamcrest.Matchers.closeTo;
 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 {
@@ -37,7 +40,7 @@ protected void assertSimpleOutput(List<Page> input, List<Page> results) {
     }
 
     @Override
-    protected Operator.OperatorFactory simple() {
+    protected SampleOperator.Factory simple() {
         return new SampleOperator.Factory(0.5, randomInt());
     }
 
@@ -50,4 +53,23 @@ protected Matcher<String> expectedDescriptionOfSimple() {
     protected Matcher<String> expectedToStringOfSimple() {
         return equalTo("SampleOperator[sampled = 0/0]");
     }
+
+    public void testAccuracy() {
+        BlockFactory blockFactory = driverContext().blockFactory();
+        int totalPositionCount = 0;
+
+        for (int iter = 0; iter < 10000; iter++) {
+            SampleOperator operator = simple().get(driverContext());
+            operator.addInput(new Page(blockFactory.newConstantNullBlock(20000)));
+            Page output = operator.getOutput();
+            // 10000 expected rows, stddev=sqrt(10000)=100, so this is 10 stddevs.
+            assertThat(output.getPositionCount(), both(greaterThan(9000)).and(lessThan(11000)));
+            totalPositionCount += output.getPositionCount();
+            output.releaseBlocks();
+        }
+
+        int averagePositionCount = totalPositionCount / 10000;
+        // Running 10000 times, so the stddev is divided by sqrt(10000)=100, so this 10 stddevs again.
+        assertThat(averagePositionCount, both(greaterThan(9990)).and(lessThan(10010)));
+    }
 }