add SampleBreaking interface

Author: jan-elastic

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

@@ -12,10 +12,9 @@
 import org.elasticsearch.xpack.esql.expression.function.aggregate.HasSampleCorrection;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Mul;
 import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
-import org.elasticsearch.xpack.esql.plan.logical.Limit;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
-import org.elasticsearch.xpack.esql.plan.logical.MvExpand;
 import org.elasticsearch.xpack.esql.plan.logical.Sample;
+import org.elasticsearch.xpack.esql.plan.logical.SampleBreaking;
 import org.elasticsearch.xpack.esql.rule.Rule;
 
 import java.util.ArrayList;
@@ -30,15 +29,14 @@ public LogicalPlan apply(LogicalPlan logicalPlan) {
             if (plan instanceof Sample sample) {
                 sampleProbabilities.add(sample.probability());
             }
-            if (plan instanceof Limit || plan instanceof MvExpand) {
-                sampleProbabilities.clear();
-            }
             if (plan instanceof Aggregate && sampleProbabilities.isEmpty() == false) {
                 plan = plan.transformExpressionsOnly(
                     e -> e instanceof HasSampleCorrection hsc && hsc.isSampleCorrected() == false
                         ? hsc.sampleCorrection(getSampleProbability(sampleProbabilities, e.source()))
                         : e
                 );
+            }
+            if (plan instanceof SampleBreaking) {
                 sampleProbabilities.clear();
             }
             return plan;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Aggregate.java

@@ -42,7 +42,7 @@
 import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
 import static org.elasticsearch.xpack.esql.plan.logical.Filter.checkFilterConditionDataType;
 
-public class Aggregate extends UnaryPlan implements PostAnalysisVerificationAware, TelemetryAware, SortAgnostic {
+public class Aggregate extends UnaryPlan implements PostAnalysisVerificationAware, TelemetryAware, SortAgnostic, SampleBreaking {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
         LogicalPlan.class,
         "Aggregate",

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Limit.java

@@ -18,7 +18,7 @@
 import java.io.IOException;
 import java.util.Objects;
 
-public class Limit extends UnaryPlan implements TelemetryAware {
+public class Limit extends UnaryPlan implements TelemetryAware, SampleBreaking {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "Limit", Limit::new);
 
     private final Expression limit;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/MvExpand.java

@@ -23,7 +23,7 @@
 import java.util.List;
 import java.util.Objects;
 
-public class MvExpand extends UnaryPlan implements TelemetryAware, SortAgnostic {
+public class MvExpand extends UnaryPlan implements TelemetryAware, SortAgnostic, SampleBreaking {
     public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(LogicalPlan.class, "MvExpand", MvExpand::new);
 
     private final NamedExpression target;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/SampleBreaking.java

@@ -0,0 +1,85 @@
+/*
+ * 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.plan.logical;
+
+/**
+ *  This interface is to check whether a plan breaks the random sampling context.
+ *  <p>
+ *
+ *  Random sampling aims to correct <code>STATS</code> after them, e.g.
+ *  <p>
+ *  <code>
+ *  | SAMPLE 0.1 | STATS SUM(value)
+ *  </code>
+ *  <p>
+ *  gives an estimate of the sum of the values, and not 10% of the sum.
+ *  <p>
+ *
+ *  For many commands inbetween this works fine, because they can be swapped
+ *  with <code>SAMPLE</code>. For example,
+ *  <p>
+ *  <code>
+ *  | SAMPLE 0.1 | SORT value
+ *  </code>
+ *  <p>
+ *  <code>
+ *  | SAMPLE 0.1 | WHERE value > 10
+ *  </code>
+ *  <p>
+ *  are equivalent to
+ *  <p>
+ *  <code>
+ *  | SORT value | SAMPLE 0.1
+ *  </code>
+ *  <p>
+ *  <code>
+ *  | WHERE value > 10 | SAMPLE 0.1
+ *  </code>
+ *  <p>
+ *  (statistically equivalent, not necessary identical), and therefore succeeding
+ *  <code>STATS</code> can be adjusted for the sample size.
+ *  <p>
+ *
+ *  In other cases, commands cannot be swapped with <code>SAMPLE</code>, e.g.
+ *  <p>
+ *  <code>
+ *  | SAMPLE 0.1 | MV_EXPAND value
+ *  </code>
+ *  <p>
+ *  <code>
+ *  | SAMPLE 0.1 | LIMIT 100
+ *  </code>
+ *  <p>
+ *  In those cases, it also makes no sense to correct any succeeding <code>STATS</code>.
+ *  <p>
+ *
+ *  As a rule of thumb, if an operator can be swapped with random sampling if it maps:
+ *  <ul>
+ *      <li>
+ *          one row to one row (e.g. <code>DISSECT</code>, <code>DROP</code>, <code>ENRICH</code>,
+ *          <code>EVAL</code>, <code>GROK</code>, <code>KEEP</code>, <code>RENAME</code>)
+ *      </li>
+ *      <li>
+ *          one row to zero or one row (<code>WHERE</code>)
+ *      </li>
+ *      <li>
+ *          reorders the rows (<code>SORT</code>)
+ *      </li>
+ *  </ul>
+ *  Contrarily, it is sampling breaking (and should implement this interface) if it maps:
+ *  <ul>
+ *      <li>
+ *          one row to many rows (<code>MV_EXPAND</code>)
+ *      </li>
+ *      <li>
+ *          many rows to many rows (<code>LIMIT</code>, <code>STATS</code>)
+ *      </li>
+ *  </ul>
+ *
+ */
+public interface SampleBreaking {}