Revert "Reuse prod code and reduce EsqlSession public surface"

muted-tests.yml

@@ -510,93 +510,18 @@ tests:
 - class: org.elasticsearch.gradle.internal.transport.TransportVersionManagementPluginFuncTest
   method: cannot change committed ids to a branch
   issue: https://github.com/elastic/elasticsearch/issues/132790
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:spatial.ConvertFromStringParseError}
-  issue: https://github.com/elastic/elasticsearch/issues/132805
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:floats.ScalbWithHugeScaleFactor}
-  issue: https://github.com/elastic/elasticsearch/issues/132806
 - class: org.elasticsearch.reservedstate.service.FileSettingsServiceIT
   method: testSettingsAppliedOnStart
   issue: https://github.com/elastic/elasticsearch/issues/131210
 - class: org.elasticsearch.index.mapper.vectors.SparseVectorFieldMapperTests
   method: testPruningScenarios
   issue: https://github.com/elastic/elasticsearch/issues/132810
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:spatial_shapes.ConvertFromStringParseError}
-  issue: https://github.com/elastic/elasticsearch/issues/132812
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:string.MvStringNotEqualsLong}
-  issue: https://github.com/elastic/elasticsearch/issues/132813
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:spatial.ConvertCartesianFromStringParseError}
-  issue: https://github.com/elastic/elasticsearch/issues/132814
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:spatial_shapes.ConvertCartesianShapeFromStringParseError}
-  issue: https://github.com/elastic/elasticsearch/issues/132815
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:string.RepeatNegative}
-  issue: https://github.com/elastic/elasticsearch/issues/132818
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:stats.WeightedAvgWeightMvWarning}
-  issue: https://github.com/elastic/elasticsearch/issues/132821
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:ip.PushDownIPWithIn}
-  issue: https://github.com/elastic/elasticsearch/issues/132822
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:floats.ScalbWithHugeconstantFirstArgument}
-  issue: https://github.com/elastic/elasticsearch/issues/132823
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:boolean.MvSliceWarnings}
-  issue: https://github.com/elastic/elasticsearch/issues/132824
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:change_point.Values null column}
-  issue: https://github.com/elastic/elasticsearch/issues/132825
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:string.MvStringEqualsLongString}
-  issue: https://github.com/elastic/elasticsearch/issues/132826
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:stats.WeightedAvgFieldMvWarning}
-  issue: https://github.com/elastic/elasticsearch/issues/132827
 - class: org.elasticsearch.test.rest.yaml.CcsCommonYamlTestSuiteIT
   method: test {p0=search/160_exists_query/Test exists query on mapped date field with no doc values}
   issue: https://github.com/elastic/elasticsearch/issues/132828
 - class: org.elasticsearch.test.rest.yaml.RcsCcsCommonYamlTestSuiteIT
   method: test {p0=search/160_exists_query/Test exists query on keyword field in empty index}
   issue: https://github.com/elastic/elasticsearch/issues/132829
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:ip.CidrMatchFieldArg}
-  issue: https://github.com/elastic/elasticsearch/issues/132834
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:change_point.Null values}
-  issue: https://github.com/elastic/elasticsearch/issues/132835
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:ip.Conditional}
-  issue: https://github.com/elastic/elasticsearch/issues/132836
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:ip.CdirMatchOrsIPs}
-  issue: https://github.com/elastic/elasticsearch/issues/132837
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:string.SpaceNegative}
-  issue: https://github.com/elastic/elasticsearch/issues/132838
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:string.MvStringNotEqualsFound}
-  issue: https://github.com/elastic/elasticsearch/issues/132839
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:string.ConcatOfText}
-  issue: https://github.com/elastic/elasticsearch/issues/132840
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:floats.EqualToOrEqualToMultivalue}
-  issue: https://github.com/elastic/elasticsearch/issues/132841
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:mv_percentile.FromIndexPercentile}
-  issue: https://github.com/elastic/elasticsearch/issues/132846
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:ip.CidrMatchFunctionArg}
-  issue: https://github.com/elastic/elasticsearch/issues/132850
-- class: org.elasticsearch.xpack.esql.CsvTests
-  method: test {csv-spec:conditional.CaseOnMv}
-  issue: https://github.com/elastic/elasticsearch/issues/132851
 - class: org.elasticsearch.xpack.esql.CsvTests
   method: test {csv-spec:floats.EqualToMultivalue}
   issue: https://github.com/elastic/elasticsearch/issues/132852

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

@@ -180,31 +180,19 @@ public void execute(EsqlQueryRequest request, EsqlExecutionInfo executionInfo, P
         analyzedPlan(parsed, executionInfo, request.filter(), new EsqlCCSUtils.CssPartialErrorsActionListener(executionInfo, listener) {
             @Override
             public void onResponse(LogicalPlan analyzedPlan) {
-                optimizeAndExecute(request, executionInfo, planRunner, analyzedPlan, listener);
+                SubscribableListener.<LogicalPlan>newForked(l -> preOptimizedPlan(analyzedPlan, l))
+                    .<LogicalPlan>andThen((l, p) -> preMapper.preMapper(optimizedPlan(p), l))
+                    .<Result>andThen((l, p) -> executeOptimizedPlan(request, executionInfo, planRunner, p, l))
+                    .addListener(listener);
             }
         });
     }
 
-    // visible for testing in CsvTests
-    public void optimizeAndExecute(
-        EsqlQueryRequest request,
-        EsqlExecutionInfo executionInfo,
-        PlanRunner planRunner,
-        LogicalPlan analyzedPlan,
-        ActionListener<Result> listener
-    ) {
-        SubscribableListener.<LogicalPlan>newForked(l -> logicalPlanPreOptimizer.preOptimize(analyzedPlan, l))
-            .andThenApply(this::optimizedPlan)
-            .<LogicalPlan>andThen((l, p) -> preMapper.preMapper(p, l))
-            .<Result>andThen((l, p) -> executeOptimizedPlan(request, executionInfo, planRunner, p, l))
-            .addListener(listener);
-    }
-
     /**
      * Execute an analyzed plan. Most code should prefer calling {@link #execute} but
      * this is public for testing.
      */
-    private void executeOptimizedPlan(
+    public void executeOptimizedPlan(
         EsqlQueryRequest request,
         EsqlExecutionInfo executionInfo,
         PlanRunner planRunner,
@@ -804,7 +792,7 @@ private PhysicalPlan logicalPlanToPhysicalPlan(LogicalPlan optimizedPlan, EsqlQu
         return EstimatesRowSize.estimateRowSize(0, physicalPlan);
     }
 
-    private LogicalPlan optimizedPlan(LogicalPlan logicalPlan) {
+    public LogicalPlan optimizedPlan(LogicalPlan logicalPlan) {
         if (logicalPlan.preOptimized() == false) {
             throw new IllegalStateException("Expected pre-optimized plan");
         }
@@ -813,7 +801,11 @@ private LogicalPlan optimizedPlan(LogicalPlan logicalPlan) {
         return plan;
     }
 
-    private PhysicalPlan physicalPlan(LogicalPlan optimizedPlan) {
+    public void preOptimizedPlan(LogicalPlan logicalPlan, ActionListener<LogicalPlan> listener) {
+        logicalPlanPreOptimizer.preOptimize(logicalPlan, listener);
+    }
+
+    public PhysicalPlan physicalPlan(LogicalPlan optimizedPlan) {
         if (optimizedPlan.optimized() == false) {
             throw new IllegalStateException("Expected optimized plan");
         }
@@ -823,7 +815,7 @@ private PhysicalPlan physicalPlan(LogicalPlan optimizedPlan) {
         return plan;
     }
 
-    private PhysicalPlan optimizedPhysicalPlan(LogicalPlan optimizedPlan) {
+    public PhysicalPlan optimizedPhysicalPlan(LogicalPlan optimizedPlan) {
         var plan = physicalPlanOptimizer.optimize(physicalPlan(optimizedPlan));
         LOGGER.debug("Optimized physical plan:\n{}", plan);
         return plan;

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

@@ -596,21 +596,28 @@ private ActualResults executePlan(BigArrays bigArrays) throws Exception {
         TestPhysicalOperationProviders physicalOperationProviders = testOperationProviders(foldCtx, testDatasets);
 
         PlainActionFuture<ActualResults> listener = new PlainActionFuture<>();
-        session.optimizeAndExecute(
-            new EsqlQueryRequest(),
-            new EsqlExecutionInfo(randomBoolean()),
-            planRunner(bigArrays, foldCtx, physicalOperationProviders),
-            analyzed,
-            listener.map(
-                result -> new ActualResults(
-                    result.schema().stream().map(Attribute::name).toList(),
-                    result.schema().stream().map(a -> Type.asType(a.dataType().nameUpper())).toList(),
-                    result.schema().stream().map(Attribute::dataType).toList(),
-                    result.pages(),
-                    threadPool.getThreadContext().getResponseHeaders()
+
+        session.preOptimizedPlan(analyzed, listener.delegateFailureAndWrap((l, preOptimized) -> {
+            session.executeOptimizedPlan(
+                new EsqlQueryRequest(),
+                new EsqlExecutionInfo(randomBoolean()),
+                planRunner(bigArrays, foldCtx, physicalOperationProviders),
+                session.optimizedPlan(preOptimized),
+                listener.delegateFailureAndWrap(
+                    // Wrap so we can capture the warnings in the calling thread
+                    (next, result) -> next.onResponse(
+                        new ActualResults(
+                            result.schema().stream().map(Attribute::name).toList(),
+                            result.schema().stream().map(a -> Type.asType(a.dataType().nameUpper())).toList(),
+                            result.schema().stream().map(Attribute::dataType).toList(),
+                            result.pages(),
+                            threadPool.getThreadContext().getResponseHeaders()
+                        )
+                    )
                 )
-            )
-        );
+            );
+        }));
+
         return listener.get();
     }