Refactor sample correction

Author: jan-elastic

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

@@ -40,7 +40,6 @@ public abstract class AggregateFunction extends Function implements PostAnalysis
     private final Expression field;
     private final List<? extends Expression> parameters;
     private final Expression filter;
-    private final boolean isCorrectedForSampling;
 
     protected AggregateFunction(Source source, Expression field) {
         this(source, field, Literal.TRUE, emptyList());
@@ -51,21 +50,10 @@ protected AggregateFunction(Source source, Expression field, List<? extends Expr
     }
 
     protected AggregateFunction(Source source, Expression field, Expression filter, List<? extends Expression> parameters) {
-        this(source, field, filter, parameters, false);
-    }
-
-    protected AggregateFunction(
-        Source source,
-        Expression field,
-        Expression filter,
-        List<? extends Expression> parameters,
-        boolean isCorrectedForSampling
-    ) {
         super(source, CollectionUtils.combine(asList(field, filter), parameters));
         this.field = field;
         this.filter = filter;
         this.parameters = parameters;
-        this.isCorrectedForSampling = isCorrectedForSampling;
     }
 
     protected AggregateFunction(StreamInput in) throws IOException {
@@ -130,10 +118,6 @@ public AggregateFunction withParameters(List<? extends Expression> parameters) {
         return (AggregateFunction) replaceChildren(CollectionUtils.combine(asList(field, filter), parameters));
     }
 
-    public boolean isCorrectedForSampling() {
-        return isCorrectedForSampling;
-    }
-
     /**
      * Corrects the aggregation in the context of random sampling. By default,
      * nothing is done, but subclasses can override this method if some correction
@@ -156,8 +140,7 @@ public boolean equals(Object obj) {
             AggregateFunction other = (AggregateFunction) obj;
             return Objects.equals(other.field(), field())
                 && Objects.equals(other.filter(), filter())
-                && Objects.equals(other.parameters(), parameters())
-                && other.isCorrectedForSampling() == isCorrectedForSampling();
+                && Objects.equals(other.parameters(), parameters());
         }
         return false;
     }

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

@@ -27,7 +27,6 @@
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.FromAggregateMetricDouble;
 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;
 
@@ -38,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 {
+public class Count extends AggregateFunction implements ToAggregator, SurrogateExpression, HasSampleCorrection {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Count", Count::new);
 
     @FunctionInfo(
@@ -98,10 +97,6 @@ public Count(Source source, Expression field, Expression filter) {
         super(source, field, filter, emptyList());
     }
 
-    public Count(Source source, Expression field, Expression filter, boolean isCorrectedForSampling) {
-        super(source, field, filter, emptyList(), isCorrectedForSampling);
-    }
-
     private Count(StreamInput in) throws IOException {
         super(in);
     }
@@ -176,7 +171,7 @@ public Expression surrogate() {
     }
 
     @Override
-    public Expression correctForSampling(Expression samplingProbability) {
-        return isCorrectedForSampling() ? this : new Div(source(), new Count(source(), field(), filter(), true), samplingProbability);
+    public Expression sampleCorrection(Expression sampleProbability) {
+        return new CountSampleCorrection(source(), field(), filter(), sampleProbability);
     }
 }

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

@@ -0,0 +1,28 @@
+/*
+ * 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.expression.function.aggregate;
+
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.tree.Source;
+import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToLong;
+import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Div;
+
+public class CountSampleCorrection extends Count {
+
+    private final Expression sampleProbability;
+
+    public CountSampleCorrection(Source source, Expression field, Expression filter, Expression sampleProbability) {
+        super(source, field, filter);
+        this.sampleProbability = sampleProbability;
+    }
+
+    @Override
+    public Expression surrogate() {
+         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/HasSampleCorrection.java

@@ -0,0 +1,15 @@
+/*
+ * 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.expression.function.aggregate;
+
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+
+public interface HasSampleCorrection {
+
+    Expression sampleCorrection(Expression sampleProbability);
+}

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

@@ -53,10 +53,6 @@ public abstract class NumericAggregate extends AggregateFunction implements ToAg
         super(source, field, filter, parameters);
     }
 
-    NumericAggregate(Source source, Expression field, Expression filter, List<Expression> parameters, boolean isCorrectedForSampling) {
-        super(source, field, filter, parameters, isCorrectedForSampling);
-    }
-
     NumericAggregate(Source source, Expression field) {
         super(source, field);
     }

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

@@ -27,7 +27,6 @@
 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.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;
@@ -44,7 +43,7 @@
 /**
  * Sum all values of a field in matching documents.
  */
-public class Sum extends NumericAggregate implements SurrogateExpression {
+public class Sum extends NumericAggregate implements SurrogateExpression, HasSampleCorrection {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Sum", Sum::new);
 
     @FunctionInfo(
@@ -66,11 +65,7 @@ public Sum(Source source, @Param(name = "number", type = { "aggregate_metric_dou
     }
 
     public Sum(Source source, Expression field, Expression filter) {
-        this(source, field, filter, false);
-    }
-
-    public Sum(Source source, Expression field, Expression filter, boolean isCorrectedForSampling) {
-        super(source, field, filter, emptyList(), isCorrectedForSampling);
+        super(source, field, filter, emptyList());
     }
 
     private Sum(StreamInput in) throws IOException {
@@ -154,7 +149,7 @@ public Expression surrogate() {
     }
 
     @Override
-    public Expression correctForSampling(Expression samplingProbability) {
-        return isCorrectedForSampling() ? this : new Div(source(), new Sum(source(), field(), filter(), true), samplingProbability);
+    public Expression sampleCorrection(Expression sampleProbability) {
+        return new SumSampleCorrection(source(), field(), filter(), sampleProbability);
     }
 }

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

@@ -0,0 +1,32 @@
+/*
+ * 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.expression.function.aggregate;
+
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.tree.Source;
+import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToLong;
+import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Div;
+
+public class SumSampleCorrection extends Sum {
+
+    private final Expression sampleProbability;
+
+    public SumSampleCorrection(Source source, Expression field, Expression filter, Expression sampleProbability) {
+        super(source, field, filter);
+        this.sampleProbability = sampleProbability;
+    }
+
+    @Override
+    public Expression surrogate() {
+        return switch (dataType()) {
+            case DOUBLE -> new Div(source(), new Sum(source(), field(), filter()), sampleProbability);
+            case LONG -> new ToLong(source(), new Div(source(), new Sum(source(), field(), filter()), sampleProbability));
+            default -> throw new IllegalStateException("unexpected data type [" + dataType() + "]");
+        };
+    }
+}

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

@@ -127,8 +127,8 @@ protected static Batch<LogicalPlan> substitutions() {
         return new Batch<>(
             "Substitutions",
             Limiter.ONCE,
-            new SubstituteSurrogatePlans(),
             new PropagateSampleFrequencyToAggs(),
+            new SubstituteSurrogatePlans(),
             // Translate filtered expressions into aggregate with filters - can't use surrogate expressions because it was
             // retrofitted for constant folding - this needs to be fixed.
             // Needs to occur before ReplaceAggregateAggExpressionWithEval, which will update the functions, losing the filter.

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/PropagateSampleFrequencyToAggs.java

@@ -9,13 +9,14 @@
 
 import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.util.Holder;
-import org.elasticsearch.xpack.esql.expression.function.aggregate.AggregateFunction;
+import org.elasticsearch.xpack.esql.expression.function.aggregate.HasSampleCorrection;
 import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.plan.logical.RandomSample;
 import org.elasticsearch.xpack.esql.rule.Rule;
 
 public class PropagateSampleFrequencyToAggs extends Rule<LogicalPlan, LogicalPlan> {
+
     @Override
     public LogicalPlan apply(LogicalPlan logicalPlan) {
         Holder<Expression> sampleProbability = new Holder<>(null);
@@ -24,7 +25,8 @@ public LogicalPlan apply(LogicalPlan logicalPlan) {
                 sampleProbability.set(randomSample.probability());
             }
             if (plan instanceof Aggregate && sampleProbability.get() != null) {
-                plan = plan.transformExpressionsOnly(AggregateFunction.class, af -> af.correctForSampling(sampleProbability.get()));
+                plan = plan.transformExpressionsOnly(
+                    e -> e instanceof HasSampleCorrection hsc ? hsc.sampleCorrection(sampleProbability.get()) : e);
                 sampleProbability.set(null);
             }
             return plan;