Introduce a new PRE_OPTIMIZED to the LogicalPlan

Author: afoucret

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

@@ -36,7 +36,14 @@ public PreOptimizer(TransportActionServices services, FoldContext foldContext) {
     }
 
     public void preOptimize(LogicalPlan plan, ActionListener<LogicalPlan> listener) {
-        inferencePreOptimizer.foldInferenceFunctions(plan, listener);
+        if (plan.analyzed() == false) {
+            throw new IllegalStateException("Expected analyzed plan");
+        }
+
+        inferencePreOptimizer.foldInferenceFunctions(plan, listener.safeMap(p -> {
+            p.setPreOptimized();
+            return p;
+        }));
     }
 
     private static class InferencePreOptimizer {

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

@@ -65,6 +65,14 @@ 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/test/java/org/elasticsearch/xpack/esql/optimizer/PreOptimizerTests.java

@@ -27,6 +27,7 @@
 import org.elasticsearch.xpack.esql.plan.logical.Eval;
 import org.elasticsearch.xpack.esql.plan.logical.Filter;
 
+import java.util.ArrayList;
 import java.util.List;
 
 import static org.elasticsearch.xpack.esql.EsqlTestUtils.as;
@@ -76,18 +77,20 @@ private void testEvalFunctionEmbedding(TextEmbeddingModelMock textEmbeddingModel
             relation,
             List.of(new Alias(Source.EMPTY, fieldName, new TextEmbedding(Source.EMPTY, of(query), of(inferenceId))))
         );
+        eval.setAnalyzed();
 
         SetOnce<Object> preOptimizedPlanHolder = new SetOnce<>();
         preOptimizer.preOptimize(eval, ActionListener.wrap(preOptimizedPlanHolder::set, ESTestCase::fail));
 
         assertBusy(() -> {
             assertThat(preOptimizedPlanHolder.get(), notNullValue());
             Eval preOptimizedEval = as(preOptimizedPlanHolder.get(), Eval.class);
+            assertThat(preOptimizedEval.preOptimized(), equalTo(true));
             assertThat(preOptimizedEval.fields(), hasSize(1));
             assertThat(preOptimizedEval.fields().get(0).name(), equalTo(fieldName));
             Literal preOptimizedQuery = as(preOptimizedEval.fields().get(0).child(), Literal.class);
             assertThat(preOptimizedQuery.dataType(), equalTo(DENSE_VECTOR));
-            assertThat(preOptimizedQuery.value(), equalTo(textEmbeddingModel.embedding(query)));
+            assertThat(preOptimizedQuery.value(), equalTo(textEmbeddingModel.embeddingList(query)));
         });
     }
 
@@ -102,6 +105,7 @@ private void testKnnFunctionEmbedding(TextEmbeddingModelMock textEmbeddingModel)
             relation,
             new Knn(Source.EMPTY, getFieldAttribute("a"), new TextEmbedding(Source.EMPTY, of(query), of(inferenceId)), of(10), null)
         );
+        filter.setAnalyzed();
         Knn knn = as(filter.condition(), Knn.class);
 
         SetOnce<Object> preOptimizedHolder = new SetOnce<>();
@@ -110,14 +114,15 @@ private void testKnnFunctionEmbedding(TextEmbeddingModelMock textEmbeddingModel)
         assertBusy(() -> {
             assertThat(preOptimizedHolder.get(), notNullValue());
             Filter preOptimizedFilter = as(preOptimizedHolder.get(), Filter.class);
+            assertThat(preOptimizedFilter.preOptimized(), equalTo(true));
             Knn preOptimizedKnn = as(preOptimizedFilter.condition(), Knn.class);
             assertThat(preOptimizedKnn.field(), equalTo(knn.field()));
             assertThat(preOptimizedKnn.k(), equalTo(knn.k()));
             assertThat(preOptimizedKnn.options(), equalTo(knn.options()));
 
             Literal preOptimizedQuery = as(preOptimizedKnn.query(), Literal.class);
             assertThat(preOptimizedQuery.dataType(), equalTo(DENSE_VECTOR));
-            assertThat(preOptimizedQuery.value(), equalTo(textEmbeddingModel.embedding(query)));
+            assertThat(preOptimizedQuery.value(), equalTo(textEmbeddingModel.embeddingList(query)));
         });
     }
 
@@ -141,6 +146,15 @@ private interface TextEmbeddingModelMock {
         TextEmbeddingResults<?> embeddingResults(String input);
 
         float[] embedding(String input);
+
+        default List<Float> embeddingList(String input) {
+            float[] embedding = embedding(input);
+            List<Float> embeddingList = new ArrayList<>(embedding.length);
+            for (float value : embedding) {
+                embeddingList.add(value);
+            }
+            return embeddingList;
+        }
     }
 
     private static final TextEmbeddingModelMock FLOAT_EMBEDDING_MODEL = new TextEmbeddingModelMock() {