Make isSampledCorrected field final

Author: jan-elastic

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

@@ -42,7 +42,7 @@
 public class Count extends AggregateFunction implements ToAggregator, SurrogateExpression, HasSampleCorrection {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Count", Count::new);
 
-    private boolean sampleCorrected = false;
+    private final boolean isSampleCorrected;
 
     @FunctionInfo(
         returnType = "long",
@@ -98,11 +98,17 @@ public Count(
     }
 
     public Count(Source source, Expression field, Expression filter) {
+        this(source, field, filter, false);
+    }
+
+    public Count(Source source, Expression field, Expression filter, boolean isSampleCorrected) {
         super(source, field, filter, emptyList());
+        this.isSampleCorrected = isSampleCorrected;
     }
 
     private Count(StreamInput in) throws IOException {
         super(in);
+        this.isSampleCorrected = false;
     }
 
     @Override
@@ -175,14 +181,12 @@ public Expression surrogate() {
     }
 
     @Override
-    public boolean sampleCorrected() {
-        return sampleCorrected;
+    public boolean isSampleCorrected() {
+        return isSampleCorrected;
     }
 
     @Override
     public Expression sampleCorrection(Expression sampleProbability) {
-        Count count = new Count(source(), field(), filter());
-        count.sampleCorrected = true;
-        return new ToLong(source(), new Div(source(), count, sampleProbability));
+        return new ToLong(source(), new Div(source(), new Count(source(), field(), filter(), true), sampleProbability));
     }
 }

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

@@ -15,7 +15,7 @@
  */
 public interface HasSampleCorrection {
 
-    boolean sampleCorrected();
+    boolean isSampleCorrected();
 
     Expression sampleCorrection(Expression sampleProbability);
 }

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

@@ -48,7 +48,7 @@
 public class Sum extends NumericAggregate implements SurrogateExpression, HasSampleCorrection {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(Expression.class, "Sum", Sum::new);
 
-    private boolean sampleCorrected = false;
+    private final boolean isSampleCorrected;
 
     @FunctionInfo(
         returnType = { "long", "double" },
@@ -69,11 +69,17 @@ 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 isSampleCorrected) {
         super(source, field, filter, emptyList());
+        this.isSampleCorrected = isSampleCorrected;
     }
 
     private Sum(StreamInput in) throws IOException {
         super(in);
+        this.isSampleCorrected = false;
     }
 
     @Override
@@ -153,15 +159,13 @@ public Expression surrogate() {
     }
 
     @Override
-    public boolean sampleCorrected() {
-        return sampleCorrected;
+    public boolean isSampleCorrected() {
+        return isSampleCorrected;
     }
 
     @Override
     public Expression sampleCorrection(Expression sampleProbability) {
-        Sum sum = new Sum(source(), field(), filter());
-        sum.sampleCorrected = true;
-        Expression correctedSum = new Div(source(), sum, sampleProbability);
+        Expression correctedSum = new Div(source(), new Sum(source(), field(), filter(), true), sampleProbability);
         return switch (dataType()) {
             case DOUBLE -> correctedSum;
             case LONG -> new ToLong(source(), correctedSum);

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

@@ -26,7 +26,7 @@ public LogicalPlan apply(LogicalPlan logicalPlan) {
             }
             if (plan instanceof Aggregate && sampleProbability.get() != null) {
                 plan = plan.transformExpressionsOnly(
-                    e -> e instanceof HasSampleCorrection hsc && hsc.sampleCorrected() == false
+                    e -> e instanceof HasSampleCorrection hsc && hsc.isSampleCorrected() == false
                         ? hsc.sampleCorrection(sampleProbability.get())
                         : e
                 );