rulesListener = SubscribableListener.newSucceeded(plan);
+ for (LogicalPlanPreOptimizerRule preOptimizerRule : preOptimizerRules) {
+ rulesListener = rulesListener.andThen((l, p) -> preOptimizerRule.apply(p, l));
+ }
+ rulesListener.addListener(listener);
}
}
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPreOptimizerContext.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPreOptimizerContext.java
index d082bd56fc46d..d08c639d0470b 100644
--- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPreOptimizerContext.java
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPreOptimizerContext.java
@@ -8,36 +8,11 @@
package org.elasticsearch.xpack.esql.optimizer;
import org.elasticsearch.xpack.esql.core.expression.FoldContext;
+import org.elasticsearch.xpack.esql.inference.InferenceService;
-import java.util.Objects;
-
-public class LogicalPreOptimizerContext {
-
- private final FoldContext foldCtx;
-
- public LogicalPreOptimizerContext(FoldContext foldCtx) {
- this.foldCtx = foldCtx;
- }
-
- 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);
- }
+/**
+ * Context passed to logical pre-optimizer rules.
+ */
+public record LogicalPreOptimizerContext(FoldContext foldCtx, InferenceService inferenceService) {
- @Override
- public String toString() {
- return "LogicalPreOptimizerContext[foldCtx=" + foldCtx + ']';
- }
}
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/preoptimizer/FoldInferenceFunctions.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/preoptimizer/FoldInferenceFunctions.java
new file mode 100644
index 0000000000000..4a3fc87c43795
--- /dev/null
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/preoptimizer/FoldInferenceFunctions.java
@@ -0,0 +1,131 @@
+/*
+ * 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.optimizer.rules.logical.preoptimizer;
+
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.CountDownActionListener;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.expression.function.inference.InferenceFunction;
+import org.elasticsearch.xpack.esql.inference.InferenceFunctionEvaluator;
+import org.elasticsearch.xpack.esql.optimizer.LogicalPreOptimizerContext;
+import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Pre-optimizer rule that performs constant folding for inference functions in logical plans.
+ *
+ * This rule identifies inference functions with constant parameters and evaluates them at optimization time,
+ * replacing them with their computed results.
+ *
+ * The folding process is recursive and handles nested inference functions by processing them in multiple
+ * passes until no more foldable functions remain.
+ *
+ * Example transformation:
+ * {@code TEXT_EMBEDDING("hello world", "model1")} → {@code [0.1, 0.2, 0.3, ...]}
+ */
+public class FoldInferenceFunctions implements LogicalPlanPreOptimizerRule {
+
+ private final InferenceFunctionEvaluator inferenceFunctionEvaluator;
+
+ public FoldInferenceFunctions(LogicalPreOptimizerContext preOptimizerContext) {
+ this(InferenceFunctionEvaluator.factory().create(preOptimizerContext.foldCtx(), preOptimizerContext.inferenceService()));
+ }
+
+ protected FoldInferenceFunctions(InferenceFunctionEvaluator inferenceFunctionEvaluator) {
+ this.inferenceFunctionEvaluator = inferenceFunctionEvaluator;
+ }
+
+ @Override
+ public void apply(LogicalPlan plan, ActionListener listener) {
+ foldInferenceFunctions(plan, listener);
+ }
+
+ /**
+ * Recursively folds inference functions in the logical plan.
+ *
+ * This method collects all foldable inference functions, evaluates them in parallel,
+ * and then replaces them with their computed results. If new foldable inference functions are remaining
+ * after the first round of folding (due to nested function resolution), it recursively processes
+ * them until no more foldable functions remain.
+ *
+ *
+ * @param plan the logical plan to fold inference functions in
+ * @param listener the listener to notify when the folding is complete
+ */
+ private void foldInferenceFunctions(LogicalPlan plan, ActionListener listener) {
+ // Collect all foldable inference functions from the current plan
+ List> inferenceFunctions = collectFoldableInferenceFunctions(plan);
+
+ if (inferenceFunctions.isEmpty()) {
+ // No foldable inference functions were found - return the original plan unchanged
+ listener.onResponse(plan);
+ return;
+ }
+
+ // Map to store the computed results for each inference function
+ Map, Expression> inferenceFunctionsToResults = new HashMap<>();
+
+ // Create a countdown listener that will be triggered when all inference functions complete
+ // Once all are done, replace the functions in the plan with their results and recursively
+ // process any remaining foldable inference functions
+ CountDownActionListener completionListener = new CountDownActionListener(
+ inferenceFunctions.size(),
+ listener.delegateFailureIgnoreResponseAndWrap(l -> {
+ // Transform the plan by replacing inference functions with their computed results
+ LogicalPlan transformedPlan = plan.transformExpressionsUp(
+ InferenceFunction.class,
+ f -> inferenceFunctionsToResults.getOrDefault(f, f)
+ );
+
+ // Recursively process the transformed plan to handle any remaining inference functions
+ foldInferenceFunctions(transformedPlan, l);
+ })
+ );
+
+ // Evaluate each inference function asynchronously
+ for (InferenceFunction inferenceFunction : inferenceFunctions) {
+ inferenceFunctionEvaluator.fold(inferenceFunction, completionListener.delegateFailureAndWrap((l, result) -> {
+ inferenceFunctionsToResults.put(inferenceFunction, result);
+ l.onResponse(null);
+ }));
+ }
+ }
+
+ /**
+ * Collects all foldable inference functions from the logical plan.
+ *
+ * A function is considered foldable if it meets all of the following criteria:
+ *
+ * - It's an instance of {@link InferenceFunction}
+ * - It's marked as foldable (all parameters are constants)
+ * - It doesn't contain nested inference functions (to avoid dependency issues)
+ *
+ *
+ * Functions with nested inference functions are excluded to ensure proper evaluation order.
+ * They will be considered for folding in subsequent recursive passes after their nested
+ * functions have been resolved.
+ *
+ * @param plan the logical plan to collect inference functions from
+ * @return a list of foldable inference functions, may be empty if none are found
+ */
+ private List> collectFoldableInferenceFunctions(LogicalPlan plan) {
+ List> inferenceFunctions = new ArrayList<>();
+
+ plan.forEachExpressionUp(InferenceFunction.class, f -> {
+ if (f.foldable() && f.hasNestedInferenceFunction() == false) {
+ inferenceFunctions.add(f);
+ }
+ });
+
+ return inferenceFunctions;
+ }
+}
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/preoptimizer/LogicalPlanPreOptimizerRule.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/preoptimizer/LogicalPlanPreOptimizerRule.java
new file mode 100644
index 0000000000000..29d9be564f1bf
--- /dev/null
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/preoptimizer/LogicalPlanPreOptimizerRule.java
@@ -0,0 +1,25 @@
+/*
+ * 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.optimizer.rules.logical.preoptimizer;
+
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+
+/**
+ * A rule that can be applied to an analyzed logical plan before it is optimized.
+ */
+public interface LogicalPlanPreOptimizerRule {
+
+ /**
+ * Apply the rule to the logical plan.
+ *
+ * @param plan the analyzed logical plan to pre-optimize
+ * @param listener the listener returning the pre-optimized plan when pre-optimization rule is applied
+ */
+ void apply(LogicalPlan plan, ActionListener listener);
+}
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/CsvTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/CsvTests.java
index 56b97759db83c..c3eece1f85b4d 100644
--- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/CsvTests.java
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/CsvTests.java
@@ -589,7 +589,7 @@ private ActualResults executePlan(BigArrays bigArrays) throws Exception {
null,
null,
null,
- new LogicalPlanPreOptimizer(new LogicalPreOptimizerContext(foldCtx)),
+ new LogicalPlanPreOptimizer(new LogicalPreOptimizerContext(foldCtx, mock(InferenceService.class))),
functionRegistry,
new LogicalPlanOptimizer(new LogicalOptimizerContext(configuration, foldCtx)),
mapper,
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/TextEmbeddingErrorTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/TextEmbeddingErrorTests.java
new file mode 100644
index 0000000000000..3864c8ccc8685
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/TextEmbeddingErrorTests.java
@@ -0,0 +1,76 @@
+/*
+ * 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.expression.function.inference;
+
+import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.TypeResolutions;
+import org.elasticsearch.xpack.esql.core.tree.Source;
+import org.elasticsearch.xpack.esql.core.type.DataType;
+import org.elasticsearch.xpack.esql.expression.function.AbstractFunctionTestCase;
+import org.elasticsearch.xpack.esql.expression.function.ErrorsForCasesWithoutExamplesTestCase;
+import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier;
+import org.hamcrest.Matcher;
+import org.junit.Before;
+
+import java.util.List;
+import java.util.Locale;
+import java.util.Set;
+
+import static org.hamcrest.Matchers.equalTo;
+
+/**
+ * Tests error conditions and type validation for TEXT_EMBEDDING function.
+ */
+public class TextEmbeddingErrorTests extends ErrorsForCasesWithoutExamplesTestCase {
+
+ @Before
+ public void checkCapability() {
+ assumeTrue("TEXT_EMBEDDING is not enabled", EsqlCapabilities.Cap.TEXT_EMBEDDING_FUNCTION.isEnabled());
+ }
+
+ @Override
+ protected List cases() {
+ return paramsToSuppliers(TextEmbeddingTests.parameters());
+ }
+
+ @Override
+ protected Expression build(Source source, List args) {
+ return new TextEmbedding(source, args.get(0), args.get(1));
+ }
+
+ @Override
+ protected Matcher expectedTypeErrorMatcher(List> validPerPosition, List signature) {
+ return equalTo(typeErrorMessage(true, validPerPosition, signature, (v, p) -> "string"));
+ }
+
+ protected static String typeErrorMessage(
+ boolean includeOrdinal,
+ List> validPerPosition,
+ List signature,
+ AbstractFunctionTestCase.PositionalErrorMessageSupplier positionalErrorMessageSupplier
+ ) {
+ for (int i = 0; i < signature.size(); i++) {
+ if (signature.get(i) == DataType.NULL) {
+ String ordinal = includeOrdinal ? TypeResolutions.ParamOrdinal.fromIndex(i).name().toLowerCase(Locale.ROOT) + " " : "";
+ return ordinal + "argument of [" + sourceForSignature(signature) + "] cannot be null, received []";
+ }
+
+ if (validPerPosition.get(i).contains(signature.get(i)) == false) {
+ break;
+ }
+ }
+
+ return ErrorsForCasesWithoutExamplesTestCase.typeErrorMessage(
+ includeOrdinal,
+ validPerPosition,
+ signature,
+ positionalErrorMessageSupplier
+ );
+ }
+}
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/TextEmbeddingTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/TextEmbeddingTests.java
new file mode 100644
index 0000000000000..f7a1ed4cf2025
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/TextEmbeddingTests.java
@@ -0,0 +1,69 @@
+/*
+ * 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.expression.function.inference;
+
+import com.carrotsearch.randomizedtesting.annotations.Name;
+import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
+
+import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.tree.Source;
+import org.elasticsearch.xpack.esql.expression.function.AbstractFunctionTestCase;
+import org.elasticsearch.xpack.esql.expression.function.FunctionName;
+import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier;
+import org.hamcrest.Matchers;
+import org.junit.Before;
+
+import java.util.List;
+import java.util.function.Supplier;
+
+import static org.elasticsearch.xpack.esql.core.type.DataType.DENSE_VECTOR;
+import static org.elasticsearch.xpack.esql.core.type.DataType.KEYWORD;
+import static org.hamcrest.Matchers.equalTo;
+
+@FunctionName("text_embedding")
+public class TextEmbeddingTests extends AbstractFunctionTestCase {
+ @Before
+ public void checkCapability() {
+ assumeTrue("TEXT_EMBEDDING is not enabled", EsqlCapabilities.Cap.TEXT_EMBEDDING_FUNCTION.isEnabled());
+ }
+
+ public TextEmbeddingTests(@Name("TestCase") Supplier testCaseSupplier) {
+ this.testCase = testCaseSupplier.get();
+ }
+
+ @ParametersFactory
+ public static Iterable parameters() {
+ return parameterSuppliersFromTypedData(
+ List.of(
+ new TestCaseSupplier(
+ List.of(KEYWORD, KEYWORD),
+ () -> new TestCaseSupplier.TestCase(
+ List.of(
+ new TestCaseSupplier.TypedData(randomBytesReference(10).toBytesRef(), KEYWORD, "text"),
+ new TestCaseSupplier.TypedData(randomBytesReference(10).toBytesRef(), KEYWORD, "inference_id")
+ ),
+ Matchers.blankOrNullString(),
+ DENSE_VECTOR,
+ equalTo(true)
+ )
+ )
+ )
+ );
+ }
+
+ @Override
+ protected Expression build(Source source, List args) {
+ return new TextEmbedding(source, args.get(0), args.get(1));
+ }
+
+ @Override
+ protected boolean canSerialize() {
+ return false;
+ }
+}
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceFunctionEvaluatorTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceFunctionEvaluatorTests.java
new file mode 100644
index 0000000000000..b4c63e6553b95
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceFunctionEvaluatorTests.java
@@ -0,0 +1,186 @@
+/*
+ * 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.compute.data.FloatBlock;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.operator.Operator;
+import org.elasticsearch.compute.test.ComputeTestCase;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.xpack.esql.core.expression.Attribute;
+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.core.tree.Source;
+import org.elasticsearch.xpack.esql.core.type.DataType;
+import org.elasticsearch.xpack.esql.expression.function.inference.InferenceFunction;
+import org.elasticsearch.xpack.esql.expression.function.inference.TextEmbedding;
+import org.junit.After;
+import org.junit.Before;
+
+import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.elasticsearch.xpack.esql.EsqlTestUtils.as;
+import static org.hamcrest.Matchers.containsString;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.instanceOf;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+public class InferenceFunctionEvaluatorTests extends ComputeTestCase {
+
+ private ThreadPool threadPool;
+
+ @Before
+ public void setupThreadPool() {
+ this.threadPool = createThreadPool();
+ }
+
+ @After
+ public void tearDownThreadPool() {
+ terminate(threadPool);
+ }
+
+ public void testFoldTextEmbeddingFunction() throws Exception {
+ // Create a mock TextEmbedding function
+ TextEmbedding textEmbeddingFunction = new TextEmbedding(
+ Source.EMPTY,
+ Literal.keyword(Source.EMPTY, "test-model"),
+ Literal.keyword(Source.EMPTY, "test input")
+ );
+
+ // Create a mock operator that returns a result
+ Operator operator = mock(Operator.class);
+
+ Float[] embedding = randomArray(1, 100, Float[]::new, ESTestCase::randomFloat);
+
+ when(operator.getOutput()).thenAnswer(i -> {
+ FloatBlock.Builder outputBlockBuilder = blockFactory().newFloatBlockBuilder(1).beginPositionEntry();
+
+ for (int j = 0; j < embedding.length; j++) {
+ outputBlockBuilder.appendFloat(embedding[j]);
+ }
+
+ outputBlockBuilder.endPositionEntry();
+
+ return new Page(outputBlockBuilder.build());
+ });
+
+ InferenceFunctionEvaluator.InferenceOperatorProvider inferenceOperatorProvider = (f, driverContext) -> operator;
+
+ // Execute the fold operation
+ InferenceFunctionEvaluator evaluator = new InferenceFunctionEvaluator(FoldContext.small(), inferenceOperatorProvider);
+
+ AtomicReference resultExpression = new AtomicReference<>();
+ evaluator.fold(textEmbeddingFunction, ActionListener.wrap(resultExpression::set, ESTestCase::fail));
+
+ assertBusy(() -> {
+ assertNotNull(resultExpression.get());
+ Literal result = as(resultExpression.get(), Literal.class);
+ assertThat(result.dataType(), equalTo(DataType.DENSE_VECTOR));
+ assertThat(as(result.value(), List.class).toArray(), equalTo(embedding));
+ });
+
+ // Check all breakers are empty after the operation is executed
+ allBreakersEmpty();
+ }
+
+ public void testFoldWithNonFoldableFunction() {
+ // A function with a non-literal argument is not foldable.
+ TextEmbedding textEmbeddingFunction = new TextEmbedding(
+ Source.EMPTY,
+ mock(Attribute.class),
+ Literal.keyword(Source.EMPTY, "test input")
+ );
+
+ InferenceFunctionEvaluator evaluator = new InferenceFunctionEvaluator(
+ FoldContext.small(),
+ (f, driverContext) -> mock(Operator.class)
+ );
+
+ AtomicReference error = new AtomicReference<>();
+ evaluator.fold(textEmbeddingFunction, ActionListener.wrap(r -> fail("should have failed"), error::set));
+
+ assertNotNull(error.get());
+ assertThat(error.get(), instanceOf(IllegalArgumentException.class));
+ assertThat(error.get().getMessage(), equalTo("Inference function must be foldable"));
+ }
+
+ public void testFoldWithAsyncFailure() throws Exception {
+ TextEmbedding textEmbeddingFunction = new TextEmbedding(
+ Source.EMPTY,
+ Literal.keyword(Source.EMPTY, "test-model"),
+ Literal.keyword(Source.EMPTY, "test input")
+ );
+
+ // Mock an operator that will trigger an async failure
+ Operator operator = mock(Operator.class);
+ doAnswer(invocation -> {
+ // Simulate the operator finishing and then immediately calling the failure listener.
+ // In that case getOutput() will replay the failure when called allowing us to catch the error.
+ throw new RuntimeException("async failure");
+ }).when(operator).getOutput();
+
+ InferenceFunctionEvaluator.InferenceOperatorProvider inferenceOperatorProvider = (f, driverContext) -> operator;
+ InferenceFunctionEvaluator evaluator = new InferenceFunctionEvaluator(FoldContext.small(), inferenceOperatorProvider);
+
+ AtomicReference error = new AtomicReference<>();
+ evaluator.fold(textEmbeddingFunction, ActionListener.wrap(r -> fail("should have failed"), error::set));
+
+ assertBusy(() -> assertNotNull(error.get()));
+ assertThat(error.get(), instanceOf(RuntimeException.class));
+ assertThat(error.get().getMessage(), equalTo("async failure"));
+
+ allBreakersEmpty();
+ }
+
+ public void testFoldWithNullOutputPage() throws Exception {
+ TextEmbedding textEmbeddingFunction = new TextEmbedding(
+ Source.EMPTY,
+ Literal.keyword(Source.EMPTY, "test-model"),
+ Literal.keyword(Source.EMPTY, "test input")
+ );
+
+ Operator operator = mock(Operator.class);
+ when(operator.getOutput()).thenReturn(null);
+
+ InferenceFunctionEvaluator.InferenceOperatorProvider inferenceOperatorProvider = (f, driverContext) -> operator;
+ InferenceFunctionEvaluator evaluator = new InferenceFunctionEvaluator(FoldContext.small(), inferenceOperatorProvider);
+
+ AtomicReference error = new AtomicReference<>();
+ evaluator.fold(textEmbeddingFunction, ActionListener.wrap(r -> fail("should have failed"), error::set));
+
+ assertBusy(() -> assertNotNull(error.get()));
+ assertThat(error.get(), instanceOf(IllegalStateException.class));
+ assertThat(error.get().getMessage(), equalTo("Expected output page from inference operator"));
+
+ allBreakersEmpty();
+ }
+
+ public void testFoldWithUnsupportedFunction() throws Exception {
+ InferenceFunction unsupported = mock(InferenceFunction.class);
+ when(unsupported.foldable()).thenReturn(true);
+
+ InferenceFunctionEvaluator evaluator = new InferenceFunctionEvaluator(FoldContext.small(), (f, driverContext) -> {
+ throw new IllegalArgumentException("Unknown inference function: " + f.getClass().getName());
+ });
+
+ AtomicReference error = new AtomicReference<>();
+ evaluator.fold(unsupported, ActionListener.wrap(r -> fail("should have failed"), error::set));
+
+ assertNotNull(error.get());
+ assertThat(error.get(), instanceOf(IllegalArgumentException.class));
+ assertThat(error.get().getMessage(), containsString("Unknown inference function"));
+
+ allBreakersEmpty();
+ }
+}
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanPreOptimizerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanPreOptimizerTests.java
index 8e573dd1cf3c9..ad7c1dbea1a50 100644
--- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanPreOptimizerTests.java
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanPreOptimizerTests.java
@@ -30,6 +30,7 @@
import static org.elasticsearch.xpack.esql.EsqlTestUtils.of;
import static org.hamcrest.Matchers.equalTo;
import static org.hamcrest.Matchers.notNullValue;
+import static org.mockito.Mockito.mock;
public class LogicalPlanPreOptimizerTests extends ESTestCase {
@@ -72,7 +73,8 @@ public LogicalPlan preOptimizedPlan(LogicalPlan plan) throws Exception {
}
private LogicalPlanPreOptimizer preOptimizer() {
- LogicalPreOptimizerContext preOptimizerContext = new LogicalPreOptimizerContext(FoldContext.small());
+ var inferenceService = mock(org.elasticsearch.xpack.esql.inference.InferenceService.class);
+ LogicalPreOptimizerContext preOptimizerContext = new LogicalPreOptimizerContext(FoldContext.small(), inferenceService);
return new LogicalPlanPreOptimizer(preOptimizerContext);
}