Renamed PreOptimizer into LogicalPlanPreOptimizer

Author: afoucret

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/execution/PlanExecutor.java

@@ -87,7 +87,7 @@ public void esql(
             indexResolver,
             enrichPolicyResolver,
             preAnalyzer,
-            new LogicalPlanPreOptimizer(new LogicalPreOptimizerContext(foldContext)),
+            new LogicalPlanPreOptimizer(services, new LogicalPreOptimizerContext(foldContext)),
             functionRegistry,
             new LogicalPlanOptimizer(new LogicalOptimizerContext(cfg, foldContext)),
             mapper,

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanPreOptimizer.java

@@ -8,7 +8,19 @@
 package org.elasticsearch.xpack.esql.optimizer;
 
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.CountDownActionListener;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.FoldContext;
+import org.elasticsearch.xpack.esql.expression.function.inference.InferenceFunction;
+import org.elasticsearch.xpack.esql.inference.InferenceFunctionEvaluator;
+import org.elasticsearch.xpack.esql.inference.InferenceRunner;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+import org.elasticsearch.xpack.esql.plugin.TransportActionServices;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 
 /**
  * The class is responsible for invoking any steps that need to be applied to the logical plan,
@@ -19,10 +31,10 @@
  */
 public class LogicalPlanPreOptimizer {
 
-    private final LogicalPreOptimizerContext preOptimizerContext;
+    private final InferenceFunctionFolding inferenceFunctionFolding;
 
-    public LogicalPlanPreOptimizer(LogicalPreOptimizerContext preOptimizerContext) {
-        this.preOptimizerContext = preOptimizerContext;
+    public LogicalPlanPreOptimizer(TransportActionServices services, LogicalPreOptimizerContext preOptimizerContext) {
+        this.inferenceFunctionFolding = new InferenceFunctionFolding(services.inferenceRunner(), preOptimizerContext.foldCtx());
     }
 
     /**
@@ -44,7 +56,53 @@ public void preOptimize(LogicalPlan plan, ActionListener<LogicalPlan> listener)
     }
 
     private void doPreOptimize(LogicalPlan plan, ActionListener<LogicalPlan> listener) {
-        // this is where we will be executing async tasks
-        listener.onResponse(plan);
+        inferenceFunctionFolding.foldInferenceFunctions(plan, listener);
+    }
+
+    private static class InferenceFunctionFolding {
+        private final InferenceRunner inferenceRunner;
+        private final FoldContext foldContext;
+
+        private InferenceFunctionFolding(InferenceRunner inferenceRunner, FoldContext foldContext) {
+            this.inferenceRunner = inferenceRunner;
+            this.foldContext = foldContext;
+        }
+
+        private void foldInferenceFunctions(LogicalPlan plan, ActionListener<LogicalPlan> listener) {
+            // First let's collect all the inference functions
+            List<InferenceFunction<?>> inferenceFunctions = new ArrayList<>();
+            plan.forEachExpressionUp(InferenceFunction.class, inferenceFunctions::add);
+
+            if (inferenceFunctions.isEmpty()) {
+                // No inference functions found. Return the original plan.
+                listener.onResponse(plan);
+                return;
+            }
+
+            // This is a map of inference functions to their results.
+            // We will use this map to replace the inference functions in the plan.
+            Map<InferenceFunction<?>, Expression> inferenceFunctionsToResults = new HashMap<>();
+
+            // Prepare a listener that will be called when all inference functions are done.
+            // This listener will replace the inference functions in the plan with their results.
+            CountDownActionListener completionListener = new CountDownActionListener(
+                inferenceFunctions.size(),
+                listener.map(
+                    ignored -> plan.transformExpressionsUp(InferenceFunction.class, f -> inferenceFunctionsToResults.getOrDefault(f, f))
+                )
+            );
+
+            // Try to compute the result for each inference function.
+            for (InferenceFunction<?> inferenceFunction : inferenceFunctions) {
+                foldInferenceFunction(inferenceFunction, completionListener.map(e -> {
+                    inferenceFunctionsToResults.put(inferenceFunction, e);
+                    return null;
+                }));
+            }
+        }
+
+        private void foldInferenceFunction(InferenceFunction<?> inferenceFunction, ActionListener<Expression> listener) {
+            InferenceFunctionEvaluator.get(inferenceFunction, inferenceRunner).eval(foldContext, listener);
+        }
     }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/PreOptimizer.java

@@ -65,14 +65,6 @@ public boolean optimized() {
         return stage.ordinal() >= Stage.OPTIMIZED.ordinal();
     }
 
-    public void setPreOptimized() {
-        stage = Stage.PRE_OPTIMIZED;
-    }
-
-    public boolean preOptimized() {
-        return stage.ordinal() >= Stage.PRE_OPTIMIZED.ordinal();
-    }
-
     public void setOptimized() {
         stage = Stage.OPTIMIZED;
     }

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

@@ -28,6 +28,6 @@ public record TransportActionServices(
     InferenceServices inferenceServices
 ) {
     public InferenceRunner inferenceRunner() {
-        return inferenceServices.inferenceRunner();
+        return inferenceServices == null ? null : inferenceServices.inferenceRunner();
     }
 }