Move sample correction to approximate class

Author: jan-elastic

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

@@ -30,12 +30,14 @@
 import org.elasticsearch.xpack.esql.expression.function.aggregate.StdDev;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Sum;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.WeightedAvg;
+import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToLong;
 import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.ConfidenceInterval;
 import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvAppend;
 import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvContains;
 import org.elasticsearch.xpack.esql.expression.function.scalar.random.Random;
 import org.elasticsearch.xpack.esql.expression.predicate.logical.And;
 import org.elasticsearch.xpack.esql.expression.predicate.nulls.IsNotNull;
+import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Div;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.NotEquals;
 import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
 import org.elasticsearch.xpack.esql.plan.logical.ChangePoint;
@@ -131,6 +133,11 @@ public interface LogicalPlanRunner {
         WeightedAvg.class
     );
 
+    private static final Set<Class<? extends AggregateFunction>> SAMPLE_CORRECTED_AGGS = Set.of(
+        Count.class,
+        Sum.class
+    );
+
     private static final Set<Class<? extends AggregateFunction>> SUPPORTED_MULTI_VALUED_AGGS = Set.of(
         org.elasticsearch.xpack.esql.expression.function.aggregate.Sample.class
     );
@@ -362,11 +369,7 @@ private LogicalPlan approximatePlan(double sampleProbability) {
                     for (int bucketId = -1; bucketId < bucketCount; bucketId++) {
                         AggregateFunction bucketedAgg = agg.withFilter(
                             new MvContains(Source.EMPTY, bucketIdField.toAttribute(), Literal.integer(Source.EMPTY, bucketId)));
-                        Expression correctedAgg = bucketedAgg instanceof NeedsSampleCorrection nsc
-                            ? nsc.sampleCorrection(
-                                Literal.fromDouble(Source.EMPTY, bucketId == -1 ? sampleProbability : sampleProbability / BUCKET_COUNT)
-                              )
-                            : bucketedAgg;
+                        Expression correctedAgg = correctForSampling(bucketedAgg, bucketId == -1 ? sampleProbability : sampleProbability / BUCKET_COUNT);
                         Alias correctedAggAlias = bucketId == -1
                             ? aggAlias.replaceChild(correctedAgg)
                             : new Alias(
@@ -477,4 +480,16 @@ private LogicalPlan approximatePlan(double sampleProbability) {
         approximatePlan.setPreOptimized();
         return approximatePlan;
     }
+
+    private static Expression correctForSampling(AggregateFunction agg, double sampleProbability) {
+        if (SAMPLE_CORRECTED_AGGS.contains(agg.getClass()) == false) {
+            return agg;
+        }
+        Expression correctedAgg = new Div(agg.source(), agg, Literal.fromDouble(Source.EMPTY, sampleProbability));
+        return switch (agg.dataType()) {
+            case DOUBLE -> correctedAgg;
+            case LONG -> new ToLong(agg.source(), correctedAgg);
+            default -> throw new IllegalStateException("unexpected data type [" + agg.dataType() + "]");
+        };
+    }
 }

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

@@ -12,7 +12,6 @@
 import org.elasticsearch.compute.aggregation.AggregatorFunctionSupplier;
 import org.elasticsearch.compute.aggregation.CountAggregatorFunction;
 import org.elasticsearch.compute.data.AggregateMetricDoubleBlockBuilder;
-import org.elasticsearch.xpack.esql.approximate.NeedsSampleCorrection;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
 import org.elasticsearch.xpack.esql.core.expression.Nullability;
@@ -26,10 +25,8 @@
 import org.elasticsearch.xpack.esql.expression.function.FunctionType;
 import org.elasticsearch.xpack.esql.expression.function.Param;
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.FromAggregateMetricDouble;
-import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToLong;
 import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvCount;
 import org.elasticsearch.xpack.esql.expression.function.scalar.nulls.Coalesce;
-import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Div;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Mul;
 import org.elasticsearch.xpack.esql.planner.ToAggregator;
 
@@ -40,7 +37,7 @@
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.DEFAULT;
 import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isType;
 
-public class Count extends AggregateFunction implements ToAggregator, SurrogateExpression, NeedsSampleCorrection {
+public class Count extends AggregateFunction implements ToAggregator, SurrogateExpression {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Count", Count::new);
 
     @FunctionInfo(
@@ -188,9 +185,4 @@ public Expression surrogate() {
 
         return null;
     }
-
-    @Override
-    public Expression sampleCorrection(Expression sampleProbability) {
-        return new ToLong(source(), new Div(source(), new Count(source(), field(), filter()), sampleProbability));
-    }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/Count.java

@@ -14,7 +14,6 @@
 import org.elasticsearch.compute.aggregation.SumIntAggregatorFunctionSupplier;
 import org.elasticsearch.compute.aggregation.SumLongAggregatorFunctionSupplier;
 import org.elasticsearch.compute.data.AggregateMetricDoubleBlockBuilder;
-import org.elasticsearch.xpack.esql.approximate.NeedsSampleCorrection;
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
 import org.elasticsearch.xpack.esql.core.expression.TypeResolutions;
@@ -28,9 +27,7 @@
 import org.elasticsearch.xpack.esql.expression.function.FunctionType;
 import org.elasticsearch.xpack.esql.expression.function.Param;
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.FromAggregateMetricDouble;
-import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToLong;
 import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvSum;
-import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Div;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Mul;
 
 import java.io.IOException;
@@ -46,7 +43,7 @@
 /**
  * Sum all values of a field in matching documents.
  */
-public class Sum extends NumericAggregate implements SurrogateExpression, NeedsSampleCorrection {
+public class Sum extends NumericAggregate implements SurrogateExpression {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Sum", Sum::readFrom);
 
     private final Expression summationMode;
@@ -176,14 +173,4 @@ public Expression surrogate() {
         // SUM(const) is equivalent to MV_SUM(const)*COUNT(*).
         return field.foldable() ? new Mul(s, new MvSum(s, field), new Count(s, Literal.keyword(s, StringUtils.WILDCARD), filter())) : null;
     }
-
-    @Override
-    public Expression sampleCorrection(Expression sampleProbability) {
-        Expression correctedSum = new Div(source(), new Sum(source(), field(), filter(), summationMode()), sampleProbability);
-        return switch (dataType()) {
-            case DOUBLE -> correctedSum;
-            case LONG -> new ToLong(source(), correctedSum);
-            default -> throw new IllegalStateException("unexpected data type [" + dataType() + "]");
-        };
-    }
 }