don't sample with prob > 10%

Author: jan-elastic

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/approximate/Approximate.java

@@ -243,6 +243,12 @@ public record QueryProperties(boolean canDecreaseRowCount, boolean canIncreaseRo
     // TODO: set via a query setting
     private static final double CONFIDENCE_LEVEL = 0.90;
 
+    /**
+     * Don't sample with a probability higher than this threshold. The cost of
+     * tracking confidence intervals doesn't outweigh the benefits of sampling.
+     */
+    private static final double SAMPLE_PROBABILITY_THRESHOLD = 0.1;
+
     /**
      * The number of times (trials) the sampled rows are divided into buckets.
      */
@@ -442,8 +448,8 @@ private ActionListener<Result> sourceCountListener(ActionListener<Result> listen
             sourceRowCount = rowCount(countResult);
             logger.debug("sourceCountPlan result: {} rows", sourceRowCount);
             double sampleProbability = sourceRowCount <= SAMPLE_ROW_COUNT ? 1.0 : (double) SAMPLE_ROW_COUNT / sourceRowCount;
-            if (queryProperties.canIncreaseRowCount == false && sampleProbability == 1.0) {
-                // If the query cannot increase the number of rows, and the sample probability is 1.0,
+            if (queryProperties.canIncreaseRowCount == false && sampleProbability > SAMPLE_PROBABILITY_THRESHOLD) {
+                // If the query cannot increase the number of rows, and the sample probability is large,
                 // we can directly run the original query without sampling.
                 runner.run(toPhysicalPlan.apply(logicalPlan), configuration, foldContext, listener);
             } else if (queryProperties.canIncreaseRowCount == false && queryProperties.canDecreaseRowCount == false) {
@@ -510,7 +516,7 @@ private ActionListener<Result> countListener(double sampleProbability, ActionLis
             long rowCount = rowCount(countResult);
             logger.debug("countPlan result (p={}): {} rows", sampleProbability, rowCount);
             double newSampleProbability = Math.min(1.0, sampleProbability * SAMPLE_ROW_COUNT / Math.max(1, rowCount));
-            if (rowCount <= SAMPLE_ROW_COUNT / 2 && newSampleProbability < 1.0) {
+            if (rowCount <= SAMPLE_ROW_COUNT / 2 && newSampleProbability < SAMPLE_PROBABILITY_THRESHOLD) {
                 runner.run(
                     toPhysicalPlan.apply(countPlan(newSampleProbability)),
                     configuration,
@@ -583,7 +589,7 @@ private long rowCount(Result countResult) {
      * </pre>
      */
     private LogicalPlan approximatePlan(double sampleProbability) {
-        if (sampleProbability == 1.0) {
+        if (sampleProbability > SAMPLE_PROBABILITY_THRESHOLD) {
             logger.debug("using original plan (too few rows)");
             return logicalPlan;
         }

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/approximate/ApproximateTests.java

@@ -311,7 +311,7 @@ public void testCountPlan_smallDataAfterMvExpanding() throws Exception {
         // - one pass to get the number of expanded rows (which determines the sample probability)
         // - one pass to execute the original query
         assertThat(runner.invocations, hasSize(4));
-        assertThat(runner.invocations.get(0), allOf(hasMvExpand("emp_no")));
+        assertThat(runner.invocations.get(0), allOf(hasMvExpand("emp_no"), not(hasSample())));
         assertThat(runner.invocations.get(1), allOf(not(hasMvExpand("emp_no")), not(hasSample())));
         assertThat(runner.invocations.get(2), allOf(hasMvExpand("emp_no"), not(hasSample())));
         assertThat(runner.invocations.get(3), allOf(hasMvExpand("emp_no"), not(hasSample())));
@@ -327,7 +327,7 @@ public void testCountPlan_largeDataAfterMvExpanding() throws Exception {
         // - one pass to get the number of expanded rows (which determines the sample probability)
         // - one pass to approximate the query
         assertThat(runner.invocations, hasSize(4));
-        assertThat(runner.invocations.get(0), allOf(hasMvExpand("emp_no")));
+        assertThat(runner.invocations.get(0), allOf(hasMvExpand("emp_no"), not(hasSample())));
         assertThat(runner.invocations.get(1), allOf(not(hasMvExpand("emp_no")), not(hasSample())));
         assertThat(runner.invocations.get(2), allOf(hasMvExpand("emp_no"), not(hasSample())));
         assertThat(runner.invocations.get(3), allOf(hasMvExpand("emp_no"), hasSample(1e-4)));
@@ -349,6 +349,37 @@ public void testCountPlan_largeDataBeforeMvExpanding() throws Exception {
         assertThat(runner.invocations.get(3), allOf(hasMvExpand("emp_no"), hasSample(1e-7)));
     }
 
+    public void testCountPlan_sampleProbabilityThreshold_noFilter() throws Exception {
+        TestRunner runner = new TestRunner(500_000, 500_000);
+        Approximate approximate = createApproximate("FROM test | STATS SUM(emp_no)", runner);
+        approximate.approximate(TestRunner.resultCloser);
+        // This plan needs three passes:
+        // - one pass to check whether it's an ES stats query (always no for the test runner)
+        // - one pass to get the total number of rows
+        // - one pass to execute the original query (because the sample probability is 20%)
+        assertThat(runner.invocations, hasSize(3));
+        assertThat(runner.invocations.get(0), allOf(not(hasSample())));
+        assertThat(runner.invocations.get(1), allOf(not(hasSample())));
+        assertThat(runner.invocations.get(1), allOf(not(hasSample())));
+    }
+
+    public void testCountPlan_sampleProbabilityThreshold_withFilter() throws Exception {
+        TestRunner runner = new TestRunner(1_000_000_000_000L, 200_000);
+        Approximate approximate = createApproximate("FROM test | WHERE emp_no > 1 | STATS SUM(emp_no)", runner);
+        approximate.approximate(TestRunner.resultCloser);
+        // This plan needs three passes:
+        // - one pass to check whether it's an ES stats query (always no for the test runner)
+        // - one pass to get the total number of rows
+        // - two passes to get the number of filtered rows (which determines the sample probability)
+        // - one pass to execute the original query (because the sample probability is 50%)
+        assertThat(runner.invocations, hasSize(5));
+        assertThat(runner.invocations.get(0), allOf(not(hasSample()), hasFilter("emp_no")));
+        assertThat(runner.invocations.get(1), allOf(not(hasSample()), not(hasFilter("emp_no"))));
+        assertThat(runner.invocations.get(2), allOf(hasSample(1e-7), hasFilter("emp_no")));
+        assertThat(runner.invocations.get(3), allOf(hasSample(1e-2), hasFilter("emp_no")));
+        assertThat(runner.invocations.get(4), allOf(not(hasSample()), hasFilter("emp_no")));
+    }
+
     public void testApproximatePlan_createsConfidenceInterval_withoutGrouping() throws Exception {
         TestRunner runner = new TestRunner(1_000_000_000, 1_000_000_000);
         Approximate approximate = createApproximate("FROM test | STATS COUNT(), SUM(emp_no)", runner);