Init inference function evaluator.

Author: afoucret

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

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

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/InferenceFunctionEvaluator.java

@@ -0,0 +1,57 @@
+/*
+ * 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.inference;
+
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.common.lucene.BytesRefs;
+import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
+import org.elasticsearch.compute.operator.Operator;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.FoldContext;
+import org.elasticsearch.xpack.esql.core.expression.Literal;
+import org.elasticsearch.xpack.esql.evaluator.EvalMapper;
+import org.elasticsearch.xpack.esql.expression.function.inference.InferenceFunction;
+import org.elasticsearch.xpack.esql.expression.function.inference.TextEmbedding;
+import org.elasticsearch.xpack.esql.inference.textembedding.TextEmbeddingOperator;
+
+public class InferenceFunctionEvaluator {
+
+    private final FoldContext foldContext;
+    private final InferenceService inferenceService;
+
+    public InferenceFunctionEvaluator(FoldContext foldContext, InferenceService inferenceService) {
+        this.foldContext = foldContext;
+        this.inferenceService = inferenceService;
+    }
+
+    public void fold(InferenceFunction<?> f, ActionListener<Object> listener) {
+        assert f.foldable() : "Inference function must be foldable";
+
+
+    }
+
+    private Operator.OperatorFactory createInferenceOperatorFactory(InferenceFunction<?> f) {
+        return switch (f) {
+            case TextEmbedding textEmbedding -> new TextEmbeddingOperator.Factory(
+                inferenceService,
+                inferenceId(f),
+                expressionEvaluatorFactory(textEmbedding.inputText())
+            );
+            default -> throw new IllegalArgumentException("Unknown inference function: " + f.getClass().getName());
+        };
+    }
+
+    private String inferenceId(InferenceFunction<?> f) {
+        return BytesRefs.toString(f.inferenceId().fold(foldContext));
+    }
+
+    private ExpressionEvaluator.Factory expressionEvaluatorFactory(Expression e) {
+        assert e.foldable() : "Input expression must be foldable";
+        return EvalMapper.toEvaluator(foldContext, Literal.of(foldContext, e), null);
+    }
+}

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

@@ -8,8 +8,12 @@
 package org.elasticsearch.xpack.esql.optimizer;
 
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.SubscribableListener;
+import org.elasticsearch.xpack.esql.inference.InferenceFunctionEvaluator;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 
+import java.util.List;
+
 /**
  * The class is responsible for invoking any steps that need to be applied to the logical plan,
  * before this is being optimized.
@@ -25,6 +29,8 @@ public LogicalPlanPreOptimizer(LogicalPreOptimizerContext preOptimizerContext) {
         this.preOptimizerContext = preOptimizerContext;
     }
 
+    private static final List<Rule> RULES = List.of();
+
     /**
      * Pre-optimize a logical plan.
      *
@@ -44,7 +50,27 @@ 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);
+        SubscribableListener<LogicalPlan> ruleChainListener = SubscribableListener.newSucceeded(plan);
+        for (Rule rule : RULES) {
+            ruleChainListener = ruleChainListener.andThen((l, p) -> rule.apply(p, l));
+        }
+        ruleChainListener.addListener(listener);
+    }
+
+    public interface Rule {
+        void apply(LogicalPlan plan, ActionListener<LogicalPlan> listener);
+    }
+
+    private static class FoldInferenceFunction implements Rule {
+        private final InferenceFunctionEvaluator inferenceEvaluator;
+
+        private FoldInferenceFunction(LogicalPreOptimizerContext preOptimizerContext) {
+            this.inferenceEvaluator = new InferenceFunctionEvaluator(preOptimizerContext.foldCtx(), preOptimizerContext.inferenceService());
+        }
+
+        @Override
+        public void apply(LogicalPlan plan, ActionListener<LogicalPlan> listener) {
+
+        }
     }
 }

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

@@ -8,36 +8,29 @@
 package org.elasticsearch.xpack.esql.optimizer;
 
 import org.elasticsearch.xpack.esql.core.expression.FoldContext;
-
-import java.util.Objects;
+import org.elasticsearch.xpack.esql.inference.InferenceService;
 
 public class LogicalPreOptimizerContext {
 
     private final FoldContext foldCtx;
 
-    public LogicalPreOptimizerContext(FoldContext foldCtx) {
+    private final InferenceService inferenceService;
+
+    public LogicalPreOptimizerContext(FoldContext foldCtx, InferenceService inferenceService) {
         this.foldCtx = foldCtx;
+        this.inferenceService = inferenceService;
     }
 
     public FoldContext foldCtx() {
         return foldCtx;
     }
 
-    @Override
-    public boolean equals(Object obj) {
-        if (obj == this) return true;
-        if (obj == null || obj.getClass() != this.getClass()) return false;
-        var that = (LogicalPreOptimizerContext) obj;
-        return this.foldCtx.equals(that.foldCtx);
-    }
-
-    @Override
-    public int hashCode() {
-        return Objects.hash(foldCtx);
-    }
-
     @Override
     public String toString() {
         return "LogicalPreOptimizerContext[foldCtx=" + foldCtx + ']';
     }
+
+    public InferenceService inferenceService() {
+        return inferenceService;
+    }
 }

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

@@ -588,7 +588,7 @@ private ActualResults executePlan(BigArrays bigArrays) throws Exception {
             null,
             null,
             null,
-            new LogicalPlanPreOptimizer(new LogicalPreOptimizerContext(foldCtx)),
+            new LogicalPlanPreOptimizer(new LogicalPreOptimizerContext(foldCtx, null)),
             functionRegistry,
             new LogicalPlanOptimizer(new LogicalOptimizerContext(configuration, foldCtx)),
             mapper,

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

@@ -72,7 +72,7 @@ public LogicalPlan preOptimizedPlan(LogicalPlan plan) throws Exception {
     }
 
     private LogicalPlanPreOptimizer preOptimizer() {
-        LogicalPreOptimizerContext preOptimizerContext = new LogicalPreOptimizerContext(FoldContext.small());
+        LogicalPreOptimizerContext preOptimizerContext = new LogicalPreOptimizerContext(FoldContext.small(), null);
         return new LogicalPlanPreOptimizer(preOptimizerContext);
     }