correct aggregations for random sampling

Author: jan-elastic

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

@@ -40,6 +40,7 @@ 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());
@@ -50,10 +51,21 @@ 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 {
@@ -118,6 +130,19 @@ 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
+     * is needed. See {@link org.elasticsearch.xpack.esql.expression.function.aggregate.Sum} for an example.
+     */
+    public Expression correctForSampling(Expression samplingProbability) {
+        return this;
+    }
+
     @Override
     public int hashCode() {
         // NB: the hashcode is currently used for key generation so
@@ -131,7 +156,8 @@ 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());
+                && Objects.equals(other.parameters(), parameters())
+                && other.isCorrectedForSampling() == isCorrectedForSampling();
         }
         return false;
     }

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

@@ -27,6 +27,7 @@
 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;
 
@@ -97,6 +98,10 @@ 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);
     }
@@ -169,4 +174,9 @@ public Expression surrogate() {
 
         return null;
     }
+
+    @Override
+    public Expression correctForSampling(Expression samplingProbability) {
+        return isCorrectedForSampling() ? this : new Div(source(), new Count(source(), field(), filter(), true), samplingProbability);
+    }
 }

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

@@ -53,6 +53,10 @@ 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,6 +27,7 @@
 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;
@@ -65,7 +66,11 @@ public Sum(Source source, @Param(name = "number", type = { "aggregate_metric_dou
     }
 
     public Sum(Source source, Expression field, Expression filter) {
-        super(source, field, filter, emptyList());
+        this(source, field, filter, false);
+    }
+
+    public Sum(Source source, Expression field, Expression filter, boolean isCorrectedForSampling) {
+        super(source, field, filter, emptyList(), isCorrectedForSampling);
     }
 
     private Sum(StreamInput in) throws IOException {
@@ -147,4 +152,9 @@ public Expression surrogate() {
             ? new Mul(s, new MvSum(s, field), new Count(s, new Literal(s, StringUtils.WILDCARD, DataType.KEYWORD)))
             : null;
     }
+
+    @Override
+    public Expression correctForSampling(Expression samplingProbability) {
+        return isCorrectedForSampling() ? this : new Div(source(), new Sum(source(), field(), filter(), true), samplingProbability);
+    }
 }

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

@@ -27,6 +27,7 @@
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PropagateEvalFoldables;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PropagateInlineEvals;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PropagateNullable;
+import org.elasticsearch.xpack.esql.optimizer.rules.logical.PropagateSampleFrequencyToAggs;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PropgateUnmappedFields;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneColumns;
 import org.elasticsearch.xpack.esql.optimizer.rules.logical.PruneEmptyPlans;
@@ -127,6 +128,7 @@ protected static Batch<LogicalPlan> substitutions() {
             "Substitutions",
             Limiter.ONCE,
             new SubstituteSurrogatePlans(),
+            new PropagateSampleFrequencyToAggs(),
             // 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

@@ -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.optimizer.rules.logical;
+
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.expression.function.aggregate.AggregateFunction;
+import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+import org.elasticsearch.xpack.esql.plan.logical.RandomSample;
+import org.elasticsearch.xpack.esql.rule.Rule;
+
+import java.util.concurrent.atomic.AtomicReference;
+
+public class PropagateSampleFrequencyToAggs extends Rule<LogicalPlan, LogicalPlan> {
+    @Override
+    public LogicalPlan apply(LogicalPlan logicalPlan) {
+        AtomicReference<Expression> sampleProbability = new AtomicReference<>(null);
+        return logicalPlan.transformUp(plan -> {
+            if (plan instanceof RandomSample randomSample) {
+                sampleProbability.set(randomSample.probability());
+            }
+            if (sampleProbability.get() != null) {
+                plan = plan.transformExpressionsOnly(AggregateFunction.class, af -> af.correctForSampling(sampleProbability.get()));
+            }
+            return plan;
+        });
+    }
+}