merge confidence interval and reliable

Author: jan-elastic

x-pack/plugin/esql/src/main/generated/org/elasticsearch/xpack/esql/expression/function/scalar/approximate/ReliableEvaluator.java

@@ -36,7 +36,6 @@
 import org.elasticsearch.xpack.esql.expression.function.scalar.EsqlScalarFunction;
 import org.elasticsearch.xpack.esql.expression.function.scalar.approximate.ConfidenceInterval;
 import org.elasticsearch.xpack.esql.expression.function.scalar.approximate.Random;
-import org.elasticsearch.xpack.esql.expression.function.scalar.approximate.Reliable;
 import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Case;
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToDouble;
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToInteger;
@@ -744,13 +743,19 @@ private LogicalPlan approximatePlan(double sampleProbability) {
                     default -> throw new IllegalStateException("unexpected data type [" + output.dataType() + "]");
                 };
                 confidenceIntervalsAndReliable.add(
-                    new Alias(Source.EMPTY, "CONFIDENCE_INTERVAL(" + output.name() + ")", confidenceInterval)
+                    new Alias(Source.EMPTY, "CONFIDENCE_INTERVAL(" + output.name() + ")",
+                        new MvSlice(Source.EMPTY, confidenceInterval, Literal.integer(Source.EMPTY, 0), Literal.integer(Source.EMPTY, 1))
+                    )
                 );
                 confidenceIntervalsAndReliable.add(
                     new Alias(
                         Source.EMPTY,
                         "RELIABLE(" + output.name() + ")",
-                        new Reliable(Source.EMPTY, bucketsMv, trialCount, bucketCount)
+                        new GreaterThanOrEqual(
+                            Source.EMPTY,
+                            new MvSlice(Source.EMPTY, confidenceInterval, Literal.integer(Source.EMPTY, 2), Literal.integer(Source.EMPTY, 2)),
+                            Literal.fromDouble(Source.EMPTY, 0.5)
+                        )
                     )
                 );
             }

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

@@ -10,7 +10,6 @@
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.xpack.esql.expression.function.grouping.GroupingWritables;
 import org.elasticsearch.xpack.esql.expression.function.scalar.approximate.ConfidenceInterval;
-import org.elasticsearch.xpack.esql.expression.function.scalar.approximate.Reliable;
 import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Case;
 import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.ClampMax;
 import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.ClampMin;
@@ -108,7 +107,6 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
         entries.add(Or.ENTRY);
         entries.add(Pi.ENTRY);
         entries.add(Pow.ENTRY);
-        entries.add(Reliable.ENTRY);
         entries.add(Right.ENTRY);
         entries.add(Repeat.ENTRY);
         entries.add(Replace.ENTRY);

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/ScalarFunctionWritables.java

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.esql.expression.function.scalar.approximate;
 
 import org.apache.commons.math3.distribution.NormalDistribution;
+import org.apache.commons.math3.stat.descriptive.moment.Kurtosis;
 import org.apache.commons.math3.stat.descriptive.moment.Mean;
 import org.apache.commons.math3.stat.descriptive.moment.Skewness;
 import org.apache.commons.math3.stat.descriptive.moment.StandardDeviation;
@@ -32,6 +33,7 @@
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
 
@@ -193,16 +195,16 @@ static void process(
             return;
         }
         double bestEstimate = bestEstimateBlock.getDouble(bestEstimateBlock.getFirstValueIndex(position));
-        double[] estimates = new double[estimatesBlock.getValueCount(position)];
-        for (int i = 0; i < estimatesBlock.getValueCount(position); i++) {
-            estimates[i] = estimatesBlock.getDouble(estimatesBlock.getFirstValueIndex(position) + i);
-        }
         int trialCount = trialCountBlock.getInt(trialCountBlock.getFirstValueIndex(position));
         int bucketCount = bucketCountBlock.getInt(bucketCountBlock.getFirstValueIndex(position));
-        if (estimates.length != trialCount * bucketCount) {
+        if (estimatesBlock.getValueCount(position) != trialCount * bucketCount) {
             builder.appendNull();
             return;
         }
+        double[] estimates = new double[estimatesBlock.getValueCount(position)];
+        for (int i = 0; i < estimatesBlock.getValueCount(position); i++) {
+            estimates[i] = estimatesBlock.getDouble(estimatesBlock.getFirstValueIndex(position) + i);
+        }
         double confidenceLevel = confidenceLevelBlock.getDouble(confidenceLevelBlock.getFirstValueIndex(position));
         double[] confidenceInterval = computeConfidenceInterval(bestEstimate, estimates, trialCount, bucketCount, confidenceLevel);
         if (confidenceInterval == null) {
@@ -237,13 +239,17 @@ static double[] computeConfidenceInterval(
                     meanZeroNaN.increment(0.0);
                 }
             }
-            if (meanIgnoreNaN.getN() >= 3) {
-                meansIgnoreNaN.increment(meanIgnoreNaN.getResult());
+            double value;
+            if (Double.isNaN(value = meanIgnoreNaN.getResult()) == false) {
+                meansIgnoreNaN.increment(value);
             }
-            if (meanZeroNaN.getN() >= 3) {
-                meansZeroNaN.increment(meanZeroNaN.getResult());
+            if (Double.isNaN(value = meanZeroNaN.getResult()) == false) {
+                meansZeroNaN.increment(value);
             }
         }
+        if (Double.isNaN(meansIgnoreNaN.getResult()) || Double.isNaN(meansZeroNaN.getResult())) {
+            return null;
+        }
 
         double meanIgnoreNan = meansIgnoreNaN.getResult();
         double meanZeroNan = meansZeroNaN.getResult();
@@ -253,42 +259,73 @@ static double[] computeConfidenceInterval(
 
         Mean stddevs = new Mean();
         Mean skews = new Mean();
+        Mean kurtoses = new Mean();
+        int reliableCount = 0;
         for (int trial = 0; trial < trialCount; trial++) {
-            StandardDeviation stdDev = new StandardDeviation(false);
+            StandardDeviation stddev = new StandardDeviation(false);
             Skewness skew = new Skewness();
+            Kurtosis kurtosis = new Kurtosis();
+            boolean hasNans = false;
             for (int bucket = 0; bucket < bucketCount; bucket++) {
                 double estimate = estimates[trial * bucketCount + bucket];
                 if (Double.isNaN(estimate)) {
+                    hasNans = true;
                     if (ignoreNaNs) {
                         continue;
                     } else {
                         estimate = 0.0;
                     }
                 }
-                stdDev.increment(estimate);
+                stddev.increment(estimate);
                 skew.increment(estimate);
+                kurtosis.increment(estimate);
             }
-            if (skew.getN() >= 3) {
-                stddevs.increment(stdDev.getResult());
-                skews.increment(skew.getResult());
+            double stddevResult = stddev.getResult();
+            if (Double.isNaN(stddevResult) == false) {
+                stddevs.increment(stddevResult);
+            }
+            double skewResult = skew.getResult();
+            if (Double.isNaN(skewResult) == false) {
+                skews.increment(skewResult);
+            }
+            double kurtosisResult = kurtosis.getResult();
+            if (Double.isNaN(kurtosisResult) == false) {
+                kurtoses.increment(kurtosisResult);
+            }
+            if (hasNans == false && computeReliable(skewResult, kurtosisResult, bucketCount)) {
+                reliableCount++;
             }
         }
 
         double sm = stddevs.getResult();
+        double skew = skews.getResult();
+        if (Double.isNaN(sm) || Double.isNaN(skew)) {
+            return null;
+        }
         if (sm == 0.0) {
-            return new double[] { bestEstimate, bestEstimate };
+            return new double[] { bestEstimate, bestEstimate, 1.0 };
         }
 
         // Scale the acceleration to account for the dependence of skewness on sample size.
         double scale = 1 / Math.sqrt(bucketCount);
-        double a = scale * skews.getResult() / 6.0;
+        double a = scale * skew / 6.0;
         double z0 = (bestEstimate - mm) / sm;
         double dz = normal.inverseCumulativeProbability((1.0 + confidenceLevel) / 2.0);
         double zl = z0 + (z0 - dz) / (1.0 - Math.min(a * (z0 - dz), 0.9));
         double zu = z0 + (z0 + dz) / (1.0 - Math.min(a * (z0 + dz), 0.9));
         double lower = mm + scale * sm * zl;
         double upper = mm + scale * sm * zu;
-        return lower <= bestEstimate && bestEstimate <= upper ? new double[] { lower, upper } : null;
+
+        return lower <= bestEstimate && bestEstimate <= upper ? new double[] { lower, upper, (double) reliableCount / trialCount } : null;
+    }
+
+    static boolean computeReliable(double skew, double kurtosis, int B) {
+        if (Double.isNaN(skew) || Double.isNaN(kurtosis) || B < 4) {
+            return false;
+        }
+        double maxSkew = Math.sqrt(6.0 * B * (B - 1) / ((B - 2) * (B + 1) * (B + 3))) * 1.96;
+        double maxKurtosis = Math.sqrt(24.0 * B * (B - 1) * (B - 1) / ((B - 3) * (B - 2) * (B + 3) * (B + 5))) * 1.96;
+        return Math.abs(skew) < maxSkew && Math.abs(kurtosis) < maxKurtosis;
     }
 
     @Override