refactor a bit

Author: jan-elastic

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

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.esql.approximate;
 
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.compute.data.LongBlock;
 import org.elasticsearch.xpack.esql.core.InvalidArgumentException;
 import org.elasticsearch.xpack.esql.core.expression.Alias;
@@ -36,6 +37,10 @@
 
 public class Approximate {
 
+    public interface LogicalPlanRunner {
+        void run(LogicalPlan plan, ActionListener<Result> listener);
+    }
+
     private static final Set<Class<? extends LogicalPlan>> SWAPPABLE_WITH_SAMPLE = Set.of(
         Dissect.class,
         Drop.class,
@@ -58,12 +63,29 @@ public Approximate(LogicalPlan logicalPlan) {
         verifyPlan();
     }
 
+    /**
+     * Computes approximate results for the given logical plan.
+     *
+     * This works by first executing a plan that counts the number of rows
+     * getting to the aggregation. That count is used to compute a sample
+     * probability, which is then used to sample approximately 1000 rows
+     * to aggregate over and approximate the aggregation.
+     */
+    public void approximate(LogicalPlanRunner runner, ActionListener<Result> listener) {
+        runner.run(
+            countPlan(),
+            listener.delegateFailureAndWrap(
+                (countListener, countResult) -> runner.run(approximatePlan(sampleProbability(countResult)), listener)
+            )
+        );
+    }
+
     /**
      * Verifies that a plan is suitable for approximation.
      *
      * To be so, the plan must contain at least one STATS function, and all
      * functions between the source and the leftmost STATS function must be
-     * swappable with STATS.
+     * swappable with SAMPLE.
      *
      * In that case, the STATS can be replaced by SAMPLE, STATS with sample
      * correction terms, and the SAMPLE can be moved to the source and
@@ -101,7 +123,7 @@ private void verifyPlan() {
      * off at the leftmost STATS function, followed by "| STATS COUNT(*)".
      * This value can be used to pick a good sample probability.
      */
-    public LogicalPlan countPlan() {
+    private LogicalPlan countPlan() {
         Holder<Boolean> encounteredStats = new Holder<>(false);
         LogicalPlan countPlan = logicalPlan.transformUp(plan -> {
             if (plan instanceof LeafPlan) {
@@ -126,21 +148,20 @@ public LogicalPlan countPlan() {
         return countPlan;
     }
 
+    /**
+     * Returns a sample probability based on the total number of rows.
+     */
+    private double sampleProbability(Result countResult) {
+        long rowCount = ((LongBlock) (countResult.pages().getFirst().getBlock(0))).getLong(0);
+        return rowCount <= SAMPLE_ROW_COUNT ? 1.0 : (double) SAMPLE_ROW_COUNT / rowCount;
+    }
+
     /**
      * Returns a plan that approximates the original plan. It consists of the
      * original plan, with the leftmost STATS function replaced by:
      * "SAMPLE probability | STATS sample_corrected_aggs".
-     *
-     * The sample probability is based on the total row count that would reach
-     * the STATS function, which is obtained by executing the countPlan.
      */
-    public LogicalPlan approximatePlan(Result countResult) {
-        long rowCount = ((LongBlock) (countResult.pages().getFirst().getBlock(0))).getLong(0);
-        if (rowCount <= SAMPLE_ROW_COUNT) {
-            return logicalPlan;
-        }
-        double sampleProbability = (double) SAMPLE_ROW_COUNT / rowCount;
-
+    private LogicalPlan approximatePlan(double sampleProbability) {
         Holder<Boolean> encounteredStats = new Holder<>(false);
         LogicalPlan approximatePlan = logicalPlan.transformUp(plan -> {
             if (plan instanceof LeafPlan) {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/EsqlSession.java

@@ -242,23 +242,15 @@ private void executeSubPlans(
         if (subPlan != null) {
             // code-path to execute subplans
             executeSubPlan(new DriverCompletionInfo.Accumulator(), optimizedPlan, subPlan, executionInfo, runner, request, listener);
+        } else if (request.approximate()) {
+            new Approximate(optimizedPlan).approximate(
+                (p, l) -> runner.run(logicalPlanToPhysicalPlan(optimizedPlan(p), request), l),
+                listener
+            );
         } else {
-            if (request.approximate()) {
-                Approximate approximate = new Approximate(optimizedPlan);
-                runner.run(
-                    logicalPlanToPhysicalPlan(optimizedPlan(approximate.countPlan()), request),
-                    listener.delegateFailureAndWrap(
-                        (countListener, countResult) -> runner.run(
-                            logicalPlanToPhysicalPlan(optimizedPlan(approximate.approximatePlan(countResult)), request),
-                            listener
-                        )
-                    )
-                );
-            } else {
-                PhysicalPlan physicalPlan = logicalPlanToPhysicalPlan(optimizedPlan, request);
-                // execute main plan
-                runner.run(physicalPlan, listener);
-            }
+            PhysicalPlan physicalPlan = logicalPlanToPhysicalPlan(optimizedPlan, request);
+            // execute main plan
+            runner.run(physicalPlan, listener);
         }
     }