fixes

Author: jan-elastic

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

@@ -229,7 +229,7 @@ public interface LogicalPlanRunner {
 
     public Approximate(LogicalPlan logicalPlan) {
         this.logicalPlan = logicalPlan;
-        this.preservesRows = verifyPlan();
+        this.preservesRows = verifyPlan(logicalPlan);
     }
 
     /**
@@ -245,7 +245,7 @@ public void approximate(LogicalPlanRunner runner, ActionListener<Result> listene
      * @return whether the part of the query until the STATS command preserves all rows
      * @throws VerificationException if the plan is not suitable for approximation
      */
-    private boolean verifyPlan() throws VerificationException {
+    public static boolean verifyPlan(LogicalPlan logicalPlan) throws VerificationException {
         if (logicalPlan.preOptimized() == false) {
             throw new IllegalStateException("Expected pre-optimized plan");
         }

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

@@ -68,13 +68,13 @@ public class ConfidenceInterval extends EsqlScalarFunction {
     @FunctionInfo(returnType = { "double", }, description = "...")
     public ConfidenceInterval(
         Source source,
-        @Param(name = "bestEstimate", type = { "double", "int", "long" }) Expression bestEstimate,
-        @Param(name = "estimates", type = { "double", "int", "long" }) Expression estimates,
-        @Param(name = "trialCount", type = { "int" }) Expression trialCount,
-        @Param(name = "bucketCount", type = { "int" }) Expression bucketCount,
+        @Param(name = "bestEstimate", type = { "double", "integer", "long" }) Expression bestEstimate,
+        @Param(name = "estimates", type = { "double", "integer", "long" }) Expression estimates,
+        @Param(name = "trialCount", type = { "integer" }) Expression trialCount,
+        @Param(name = "bucketCount", type = { "integer" }) Expression bucketCount,
         @Param(name = "confidenceLevel", type = { "double" }) Expression confidenceLevel
     ) {
-        super(source, List.of(bestEstimate, estimates));
+        super(source, List.of(bestEstimate, estimates, trialCount, bucketCount, confidenceLevel));
         this.bestEstimate = bestEstimate;
         this.estimates = estimates;
         this.trialCount = trialCount;

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

@@ -58,11 +58,11 @@ public class Reliable extends EsqlScalarFunction {
     @FunctionInfo(returnType = { "boolean", }, description = "...")
     public Reliable(
         Source source,
-        @Param(name = "estimates", type = { "double", "int", "long" }) Expression estimates,
-        @Param(name = "trialCount", type = { "int" }) Expression trialCount,
-        @Param(name = "bucketCount", type = { "int" }) Expression bucketCount
+        @Param(name = "estimates", type = { "double", "integer", "long" }) Expression estimates,
+        @Param(name = "trialCount", type = { "integer" }) Expression trialCount,
+        @Param(name = "bucketCount", type = { "integer" }) Expression bucketCount
     ) {
-        super(source, List.of(estimates));
+        super(source, List.of(estimates, trialCount, bucketCount));
         this.estimates = estimates;
         this.trialCount = trialCount;
         this.bucketCount = bucketCount;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/EsqlSession.java

@@ -212,7 +212,6 @@ public void execute(EsqlQueryRequest request, EsqlExecutionInfo executionInfo, P
             plan = explain.query();
             parsedPlanString = plan.toString();
         }
-<<<<<<< HEAD
         analyzedPlan(
             plan,
             configuration,
@@ -229,7 +228,7 @@ public void onResponse(LogicalPlan analyzedPlan) {
                     SubscribableListener.<LogicalPlan>newForked(l -> preOptimizedPlan(analyzedPlan, logicalPlanPreOptimizer, l))
                         .<LogicalPlan>andThen((l, p) -> {
                             if (request.approximate()) {
-                                new Approximate(p);  // to verify whether the pre-optimized plan is suitable for approximation
+                                Approximate.verifyPlan(p);
                             }
                             l.onResponse(p);
                         })
@@ -292,7 +291,17 @@ public void executeOptimizedPlan(
             // TODO: this could be snuck into the underlying listener
             EsqlCCSUtils.updateExecutionInfoAtEndOfPlanning(executionInfo);
             // execute any potential subplans
-            executeSubPlans(optimizedPlan, configuration, foldContext, planRunner, executionInfo, request, logicalPlanOptimizer, physicalPlanOptimizer, listener);
+            executeSubPlans(
+                optimizedPlan,
+                configuration,
+                foldContext,
+                planRunner,
+                executionInfo,
+                request,
+                logicalPlanOptimizer,
+                physicalPlanOptimizer,
+                listener
+            );
         }
     }
 
@@ -329,7 +338,12 @@ private void executeSubPlans(
             );
         } else if (request.approximate()) {
             new Approximate(optimizedPlan).approximate(
-                (p, l) -> runner.run(logicalPlanToPhysicalPlan(optimizedPlan(p, logicalPlanOptimizer), request, physicalPlanOptimizer), configuration, foldContext, l),
+                (p, l) -> runner.run(
+                    logicalPlanToPhysicalPlan(optimizedPlan(p, logicalPlanOptimizer), request, physicalPlanOptimizer),
+                    configuration,
+                    foldContext,
+                    l
+                ),
                 listener
             );
         } else {

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/approximate/ApproximateTests.java

@@ -30,6 +30,7 @@
 import org.elasticsearch.xpack.esql.optimizer.LogicalPlanPreOptimizer;
 import org.elasticsearch.xpack.esql.optimizer.LogicalPreOptimizerContext;
 import org.elasticsearch.xpack.esql.parser.EsqlParser;
+import org.elasticsearch.xpack.esql.parser.QueryParams;
 import org.elasticsearch.xpack.esql.plan.logical.Eval;
 import org.elasticsearch.xpack.esql.plan.logical.Filter;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
@@ -161,8 +162,8 @@ public void testVerify_incompatibleProcessingCommand() {
 
     public void testVerify_incompatibleAggregation() {
         assertError(
-            "FROM test | SORT emp_no STATS MIN(emp_no) | LIMIT 100",
-            equalTo("line 1:19: aggregation function [MIN] cannot be approximated")
+            "FROM test | SORT emp_no | STATS MIN(emp_no) | LIMIT 100",
+            equalTo("line 1:33: aggregation function [MIN] cannot be approximated")
         );
         assertError(
             "FROM test | STATS SUM(emp_no), VALUES(emp_no), TOP(emp_no, 2, \"ASC\"), COUNT()",
@@ -309,7 +310,7 @@ private void assertError(String esql, Matcher<String> matcher) {
     private Approximate createApproximate(String query) throws Exception {
         SetOnce<LogicalPlan> resultHolder = new SetOnce<>();
         SetOnce<Exception> exceptionHolder = new SetOnce<>();
-        LogicalPlan plan = parser.createStatement(query, TEST_CFG);
+        LogicalPlan plan = parser.createStatement(query, new QueryParams());
         plan = analyzer.analyze(plan);
         plan.setAnalyzed();
         preOptimizer.preOptimize(plan, ActionListener.wrap(resultHolder::set, exceptionHolder::set));