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..233bf9c2945be 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,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;
+ }
}
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/main/java/org/elasticsearch/xpack/esql/session/EsqlSession.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/EsqlSession.java
index 1d3b710baac35..4d47df6789e8c 100644
--- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/EsqlSession.java
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/EsqlSession.java
@@ -754,7 +754,7 @@ private void analyzeWithRetry(
}
private void resolveInferences(LogicalPlan plan, PreAnalysisResult preAnalysisResult, ActionListener l) {
- inferenceService.inferenceResolver().resolveInferenceIds(plan, l.map(preAnalysisResult::withInferenceResolution));
+ inferenceService.inferenceResolver(functionRegistry).resolveInferenceIds(plan, l.map(preAnalysisResult::withInferenceResolution));
}
private PhysicalPlan logicalPlanToPhysicalPlan(LogicalPlan optimizedPlan, EsqlQueryRequest request) {
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 7a4de299caf79..4a63aea448f9e 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
@@ -303,6 +303,10 @@ public final void test() throws Throwable {
"can't use KQL function in csv tests",
testCase.requiredCapabilities.contains(EsqlCapabilities.Cap.KQL_FUNCTION.capabilityName())
);
+ assumeFalse(
+ "can't use TEXT_EMBEDDING function in csv tests",
+ testCase.requiredCapabilities.contains(EsqlCapabilities.Cap.TEXT_EMBEDDING_FUNCTION.capabilityName())
+ );
assumeFalse(
"can't use KNN function in csv tests",
testCase.requiredCapabilities.contains(EsqlCapabilities.Cap.KNN_FUNCTION_V5.capabilityName())
@@ -584,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,
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTestUtils.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTestUtils.java
index fbfa18dccc477..78a8ca5483246 100644
--- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTestUtils.java
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTestUtils.java
@@ -9,6 +9,7 @@
import org.elasticsearch.index.IndexMode;
import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.test.ESTestCase;
import org.elasticsearch.xpack.core.enrich.EnrichPolicy;
import org.elasticsearch.xpack.esql.EsqlTestUtils;
import org.elasticsearch.xpack.esql.core.type.EsField;
@@ -26,6 +27,7 @@
import org.elasticsearch.xpack.esql.session.Configuration;
import java.util.ArrayList;
+import java.util.Arrays;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
@@ -195,14 +197,39 @@ public static EnrichResolution defaultEnrichResolution() {
return enrichResolution;
}
+ public static final String RERANKING_INFERENCE_ID = "reranking-inference-id";
+ public static final String COMPLETION_INFERENCE_ID = "completion-inference-id";
+ public static final String TEXT_EMBEDDING_INFERENCE_ID = "text-embedding-inference-id";
+ public static final String CHAT_COMPLETION_INFERENCE_ID = "chat-completion-inference-id";
+ public static final String SPARSE_EMBEDDING_INFERENCE_ID = "sparse-embedding-inference-id";
+ public static final List VALID_INFERENCE_IDS = List.of(
+ RERANKING_INFERENCE_ID,
+ COMPLETION_INFERENCE_ID,
+ TEXT_EMBEDDING_INFERENCE_ID,
+ CHAT_COMPLETION_INFERENCE_ID,
+ SPARSE_EMBEDDING_INFERENCE_ID
+ );
+ public static final String ERROR_INFERENCE_ID = "error-inference-id";
+
public static InferenceResolution defaultInferenceResolution() {
return InferenceResolution.builder()
- .withResolvedInference(new ResolvedInference("reranking-inference-id", TaskType.RERANK))
- .withResolvedInference(new ResolvedInference("completion-inference-id", TaskType.COMPLETION))
- .withError("error-inference-id", "error with inference resolution")
+ .withResolvedInference(new ResolvedInference(RERANKING_INFERENCE_ID, TaskType.RERANK))
+ .withResolvedInference(new ResolvedInference(COMPLETION_INFERENCE_ID, TaskType.COMPLETION))
+ .withResolvedInference(new ResolvedInference(TEXT_EMBEDDING_INFERENCE_ID, TaskType.TEXT_EMBEDDING))
+ .withResolvedInference(new ResolvedInference(CHAT_COMPLETION_INFERENCE_ID, TaskType.CHAT_COMPLETION))
+ .withResolvedInference(new ResolvedInference(SPARSE_EMBEDDING_INFERENCE_ID, TaskType.SPARSE_EMBEDDING))
+ .withError(ERROR_INFERENCE_ID, "error with inference resolution")
.build();
}
+ public static String randomInferenceId() {
+ return ESTestCase.randomFrom(VALID_INFERENCE_IDS);
+ }
+
+ public static String randomInferenceId(String... excludes) {
+ return ESTestCase.randomValueOtherThanMany(Arrays.asList(excludes)::contains, AnalyzerTestUtils::randomInferenceId);
+ }
+
public static void loadEnrichPolicyResolution(
EnrichResolution enrich,
String policyType,
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTests.java
index 662e1bdabcab2..79aa821ce4040 100644
--- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTests.java
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/AnalyzerTests.java
@@ -56,6 +56,7 @@
import org.elasticsearch.xpack.esql.expression.function.fulltext.QueryString;
import org.elasticsearch.xpack.esql.expression.function.grouping.Bucket;
import org.elasticsearch.xpack.esql.expression.function.grouping.TBucket;
+import org.elasticsearch.xpack.esql.expression.function.inference.TextEmbedding;
import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToDateNanos;
import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToDatetime;
import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToDenseVector;
@@ -123,6 +124,7 @@
import static org.elasticsearch.xpack.esql.EsqlTestUtils.referenceAttribute;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.withDefaultLimitWarning;
import static org.elasticsearch.xpack.esql.analysis.Analyzer.NO_FIELDS;
+import static org.elasticsearch.xpack.esql.analysis.AnalyzerTestUtils.TEXT_EMBEDDING_INFERENCE_ID;
import static org.elasticsearch.xpack.esql.analysis.AnalyzerTestUtils.analyze;
import static org.elasticsearch.xpack.esql.analysis.AnalyzerTestUtils.analyzer;
import static org.elasticsearch.xpack.esql.analysis.AnalyzerTestUtils.analyzerDefaultMapping;
@@ -130,6 +132,7 @@
import static org.elasticsearch.xpack.esql.analysis.AnalyzerTestUtils.defaultInferenceResolution;
import static org.elasticsearch.xpack.esql.analysis.AnalyzerTestUtils.indexWithDateDateNanosUnionType;
import static org.elasticsearch.xpack.esql.analysis.AnalyzerTestUtils.loadMapping;
+import static org.elasticsearch.xpack.esql.analysis.AnalyzerTestUtils.randomInferenceId;
import static org.elasticsearch.xpack.esql.analysis.AnalyzerTestUtils.tsdbIndexResolution;
import static org.elasticsearch.xpack.esql.core.plugin.EsqlCorePlugin.DENSE_VECTOR_FEATURE_FLAG;
import static org.elasticsearch.xpack.esql.core.tree.Source.EMPTY;
@@ -3765,6 +3768,112 @@ private void assertEmptyEsRelation(LogicalPlan plan) {
assertThat(esRelation.output(), equalTo(NO_FIELDS));
}
+ public void testTextEmbeddingResolveInferenceId() {
+ assumeTrue("TEXT_EMBEDDING function required", EsqlCapabilities.Cap.TEXT_EMBEDDING_FUNCTION.isEnabled());
+
+ LogicalPlan plan = analyze(
+ String.format(Locale.ROOT, """
+ FROM books METADATA _score | EVAL embedding = TEXT_EMBEDDING("italian food recipe", "%s")""", TEXT_EMBEDDING_INFERENCE_ID),
+ "mapping-books.json"
+ );
+
+ Eval eval = as(as(plan, Limit.class).child(), Eval.class);
+ assertThat(eval.fields(), hasSize(1));
+ Alias alias = as(eval.fields().get(0), Alias.class);
+ assertThat(alias.name(), equalTo("embedding"));
+ TextEmbedding function = as(alias.child(), TextEmbedding.class);
+
+ assertThat(function.inputText(), equalTo(string("italian food recipe")));
+ assertThat(function.inferenceId(), equalTo(string(TEXT_EMBEDDING_INFERENCE_ID)));
+ }
+
+ public void testTextEmbeddingFunctionResolveType() {
+ assumeTrue("TEXT_EMBEDDING function required", EsqlCapabilities.Cap.TEXT_EMBEDDING_FUNCTION.isEnabled());
+
+ LogicalPlan plan = analyze(
+ String.format(Locale.ROOT, """
+ FROM books METADATA _score| EVAL embedding = TEXT_EMBEDDING("italian food recipe", "%s")""", TEXT_EMBEDDING_INFERENCE_ID),
+ "mapping-books.json"
+ );
+
+ Eval eval = as(as(plan, Limit.class).child(), Eval.class);
+ assertThat(eval.fields(), hasSize(1));
+ Alias alias = as(eval.fields().get(0), Alias.class);
+ assertThat(alias.name(), equalTo("embedding"));
+
+ TextEmbedding function = as(alias.child(), TextEmbedding.class);
+
+ assertThat(function.foldable(), equalTo(true));
+ assertThat(function.dataType(), equalTo(DENSE_VECTOR));
+ }
+
+ public void testTextEmbeddingFunctionMissingInferenceIdError() {
+ assumeTrue("TEXT_EMBEDDING function required", EsqlCapabilities.Cap.TEXT_EMBEDDING_FUNCTION.isEnabled());
+
+ VerificationException ve = expectThrows(
+ VerificationException.class,
+ () -> analyze(
+ String.format(Locale.ROOT, """
+ FROM books METADATA _score| EVAL embedding = TEXT_EMBEDDING("italian food recipe", "%s")""", "unknow-inference-id"),
+ "mapping-books.json"
+ )
+ );
+
+ assertThat(ve.getMessage(), containsString("unresolved inference [unknow-inference-id]"));
+ }
+
+ public void testTextEmbeddingFunctionInvalidInferenceIdError() {
+ assumeTrue("TEXT_EMBEDDING function required", EsqlCapabilities.Cap.TEXT_EMBEDDING_FUNCTION.isEnabled());
+
+ String inferenceId = randomInferenceId(TEXT_EMBEDDING_INFERENCE_ID);
+ VerificationException ve = expectThrows(
+ VerificationException.class,
+ () -> analyze(
+ String.format(Locale.ROOT, """
+ FROM books METADATA _score| EVAL embedding = TEXT_EMBEDDING("italian food recipe", "%s")""", inferenceId),
+ "mapping-books.json"
+ )
+ );
+
+ assertThat(
+ ve.getMessage(),
+ containsString(String.format(Locale.ROOT, "cannot use inference endpoint [%s] with task type", inferenceId))
+ );
+ }
+
+ public void testTextEmbeddingFunctionWithoutModel() {
+ assumeTrue("TEXT_EMBEDDING function required", EsqlCapabilities.Cap.TEXT_EMBEDDING_FUNCTION.isEnabled());
+
+ ParsingException ve = expectThrows(ParsingException.class, () -> analyze("""
+ FROM books METADATA _score| EVAL embedding = TEXT_EMBEDDING("italian food recipe")""", "mapping-books.json"));
+
+ assertThat(
+ ve.getMessage(),
+ containsString(" error building [text_embedding]: function [text_embedding] expects exactly two arguments")
+ );
+ }
+
+ public void testKnnFunctionWithTextEmbedding() {
+ assumeTrue("KNN function capability required", EsqlCapabilities.Cap.KNN_FUNCTION_V5.isEnabled());
+ assumeTrue("TEXT_EMBEDDING function required", EsqlCapabilities.Cap.TEXT_EMBEDDING_FUNCTION.isEnabled());
+
+ LogicalPlan plan = analyze(
+ String.format(Locale.ROOT, """
+ from test | where KNN(float_vector, TEXT_EMBEDDING("italian food recipe", "%s"))""", TEXT_EMBEDDING_INFERENCE_ID),
+ "mapping-dense_vector.json"
+ );
+
+ Limit limit = as(plan, Limit.class);
+ Filter filter = as(limit.child(), Filter.class);
+ Knn knn = as(filter.condition(), Knn.class);
+ assertThat(knn.field(), instanceOf(FieldAttribute.class));
+ assertThat(((FieldAttribute) knn.field()).name(), equalTo("float_vector"));
+
+ TextEmbedding textEmbedding = as(knn.query(), TextEmbedding.class);
+ assertThat(textEmbedding.inputText(), equalTo(string("italian food recipe")));
+ assertThat(textEmbedding.inferenceId(), equalTo(string(TEXT_EMBEDDING_INFERENCE_ID)));
+ }
+
public void testResolveRerankInferenceId() {
{
LogicalPlan plan = analyze("""
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/VerifierTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/VerifierTests.java
index ba2ef6d152e45..66c8b2d6e463c 100644
--- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/VerifierTests.java
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/VerifierTests.java
@@ -41,6 +41,7 @@
import static org.elasticsearch.xpack.esql.EsqlTestUtils.TEST_CFG;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.paramAsConstant;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.withDefaultLimitWarning;
+import static org.elasticsearch.xpack.esql.analysis.AnalyzerTestUtils.TEXT_EMBEDDING_INFERENCE_ID;
import static org.elasticsearch.xpack.esql.analysis.AnalyzerTestUtils.loadMapping;
import static org.elasticsearch.xpack.esql.core.type.DataType.BOOLEAN;
import static org.elasticsearch.xpack.esql.core.type.DataType.CARTESIAN_POINT;
@@ -2675,6 +2676,53 @@ public void testSortInTimeSeries() {
and the first aggregation [STATS avg(network.connections)] is not allowed"""));
}
+ public void testTextEmbeddingFunctionInvalidQuery() {
+ assertThat(
+ error("from test | EVAL embedding = TEXT_EMBEDDING(null, ?)", defaultAnalyzer, TEXT_EMBEDDING_INFERENCE_ID),
+ equalTo("1:30: first argument of [TEXT_EMBEDDING(null, ?)] cannot be null, received [null]")
+ );
+
+ assertThat(
+ error("from test | EVAL embedding = TEXT_EMBEDDING(42, ?)", defaultAnalyzer, TEXT_EMBEDDING_INFERENCE_ID),
+ equalTo("1:30: first argument of [TEXT_EMBEDDING(42, ?)] must be [string], found value [42] type [integer]")
+ );
+
+ assertThat(
+ error("from test | EVAL embedding = TEXT_EMBEDDING(last_name, ?)", defaultAnalyzer, TEXT_EMBEDDING_INFERENCE_ID),
+ equalTo("1:30: first argument of [TEXT_EMBEDDING(last_name, ?)] must be a constant, received [last_name]")
+ );
+ }
+
+ public void testTextEmbeddingFunctionInvalidInferenceId() {
+ assertThat(
+ error("from test | EVAL embedding = TEXT_EMBEDDING(?, null)", defaultAnalyzer, "query text"),
+ equalTo("1:30: second argument of [TEXT_EMBEDDING(?, null)] cannot be null, received [null]")
+ );
+
+ assertThat(
+ error("from test | EVAL embedding = TEXT_EMBEDDING(?, 42)", defaultAnalyzer, "query text"),
+ equalTo("1:30: second argument of [TEXT_EMBEDDING(?, 42)] must be [string], found value [42] type [integer]")
+ );
+
+ assertThat(
+ error("from test | EVAL embedding = TEXT_EMBEDDING(?, last_name)", defaultAnalyzer, "query text"),
+ equalTo("1:30: second argument of [TEXT_EMBEDDING(?, last_name)] must be a constant, received [last_name]")
+ );
+ }
+
+ // public void testTextEmbeddingFunctionInvalidInferenceId() {
+ // assumeTrue("TEXT_EMBEDDING function required", EsqlCapabilities.Cap.TEXT_EMBEDDING_FUNCTION.isEnabled());
+ //
+ // ParsingException ve = expectThrows(ParsingException.class, () -> analyze("""
+ // FROM books METADATA _score| EVAL embedding = TEXT_EMBEDDING("italian food recipe", CONCAT("machin", title))""",
+ // "mapping-books.json"));
+ //
+ // assertThat(
+ // ve.getMessage(),
+ // containsString(" error building [text_embedding]: function [text_embedding] expects exactly two arguments")
+ // );
+ // }
+
private void checkVectorFunctionsNullArgs(String functionInvocation) throws Exception {
query("from test | eval similarity = " + functionInvocation, fullTextAnalyzer);
}
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..9af017bd5207f
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/TextEmbeddingErrorTests.java
@@ -0,0 +1,74 @@
+/*
+ * 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/TextEmbeddingSerializationTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/TextEmbeddingSerializationTests.java
new file mode 100644
index 0000000000000..5d7e1dfa4301a
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/TextEmbeddingSerializationTests.java
@@ -0,0 +1,46 @@
+/*
+ * 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.tree.Source;
+import org.elasticsearch.xpack.esql.expression.AbstractExpressionSerializationTests;
+import org.junit.Before;
+
+import java.io.IOException;
+
+/** Tests serialization/deserialization of TEXT_EMBEDDING function instances. */
+public class TextEmbeddingSerializationTests extends AbstractExpressionSerializationTests {
+
+ @Before
+ public void checkCapability() {
+ assumeTrue("TEXT_EMBEDDING is not enabled", EsqlCapabilities.Cap.TEXT_EMBEDDING_FUNCTION.isEnabled());
+ }
+
+ @Override
+ protected TextEmbedding createTestInstance() {
+ Source source = randomSource();
+ Expression inputText = randomChild();
+ Expression inferenceId = randomChild();
+ return new TextEmbedding(source, inputText, inferenceId);
+ }
+
+ @Override
+ protected TextEmbedding mutateInstance(TextEmbedding instance) throws IOException {
+ Source source = instance.source();
+ Expression inputText = instance.inputText();
+ Expression inferenceId = instance.inferenceId();
+ if (randomBoolean()) {
+ inputText = randomValueOtherThan(inputText, AbstractExpressionSerializationTests::randomChild);
+ } else {
+ inferenceId = randomValueOtherThan(inferenceId, AbstractExpressionSerializationTests::randomChild);
+ }
+ return new TextEmbedding(source, inputText, inferenceId);
+ }
+}
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..b6fdc7addf984
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/TextEmbeddingTests.java
@@ -0,0 +1,72 @@
+/*
+ * 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.core.type.DataType;
+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.ArrayList;
+import java.util.List;
+import java.util.function.Supplier;
+
+import static org.elasticsearch.xpack.esql.core.type.DataType.DENSE_VECTOR;
+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() {
+ List suppliers = new ArrayList<>();
+
+ // Test all string type combinations for text input and inference endpoint ID
+ for (DataType inputTextDataType : DataType.stringTypes()) {
+ for (DataType inferenceIdDataType : DataType.stringTypes()) {
+ suppliers.add(
+ new TestCaseSupplier(
+ List.of(inputTextDataType, inferenceIdDataType),
+ () -> new TestCaseSupplier.TestCase(
+ List.of(
+ new TestCaseSupplier.TypedData(randomBytesReference(10).toBytesRef(), inputTextDataType, "text"),
+ new TestCaseSupplier.TypedData(randomBytesReference(10).toBytesRef(), inferenceIdDataType, "inference_id")
+ ),
+ Matchers.blankOrNullString(),
+ DENSE_VECTOR,
+ equalTo(true)
+ )
+ )
+ );
+ }
+ }
+
+ return parameterSuppliersFromTypedData(suppliers);
+ }
+
+ @Override
+ protected Expression build(Source source, List args) {
+ return new TextEmbedding(source, args.get(0), args.get(1));
+ }
+}
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/inference/InferenceResolverTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceResolverTests.java
index 39917f849e6f2..13196977f611f 100644
--- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceResolverTests.java
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceResolverTests.java
@@ -23,6 +23,7 @@
import org.elasticsearch.threadpool.TestThreadPool;
import org.elasticsearch.threadpool.ThreadPool;
import org.elasticsearch.xpack.core.inference.action.GetInferenceModelAction;
+import org.elasticsearch.xpack.esql.expression.function.EsqlFunctionRegistry;
import org.elasticsearch.xpack.esql.parser.EsqlParser;
import org.junit.After;
import org.junit.Before;
@@ -44,6 +45,7 @@
public class InferenceResolverTests extends ESTestCase {
private TestThreadPool threadPool;
+ private EsqlFunctionRegistry functionRegistry;
@Before
public void setThreadPool() {
@@ -60,6 +62,11 @@ public void setThreadPool() {
);
}
+ @Before
+ public void setUpFunctionRegistry() {
+ functionRegistry = new EsqlFunctionRegistry();
+ }
+
@After
public void shutdownThreadPool() {
terminate(threadPool);
@@ -78,6 +85,18 @@ public void testCollectInferenceIds() {
List.of("completion-inference-id")
);
+ // Test inference ID collection from an inference function
+ assertCollectInferenceIds(
+ "FROM books METADATA _score | EVAL embedding = TEXT_EMBEDDING(\"description\", \"text-embedding-inference-id\")",
+ List.of("text-embedding-inference-id")
+ );
+
+ // Test inference ID collection with nested functions
+ assertCollectInferenceIds(
+ "FROM books METADATA _score | EVAL embedding = TEXT_EMBEDDING(TEXT_EMBEDDING(\"nested\", \"nested-id\"), \"outer-id\")",
+ List.of("nested-id", "outer-id")
+ );
+
// Multiple inference plans
assertCollectInferenceIds("""
FROM books METADATA _score
@@ -145,7 +164,7 @@ public void testResolveMultipleInferenceIds() throws Exception {
public void testResolveMissingInferenceIds() throws Exception {
InferenceResolver inferenceResolver = inferenceResolver();
- List inferenceIds = List.of("missing-plan");
+ List inferenceIds = List.of("missing-inference-id");
SetOnce inferenceResolutionSetOnce = new SetOnce<>();
@@ -162,7 +181,7 @@ public void testResolveMissingInferenceIds() throws Exception {
assertThat(inferenceResolution.resolvedInferences(), empty());
assertThat(inferenceResolution.hasError(), equalTo(true));
- assertThat(inferenceResolution.getError("missing-plan"), equalTo("inference endpoint not found"));
+ assertThat(inferenceResolution.getError("missing-inference-id"), equalTo("inference endpoint not found"));
});
}
@@ -210,7 +229,7 @@ private static ActionResponse getInferenceModelResponse(GetInferenceModelAction.
}
private InferenceResolver inferenceResolver() {
- return new InferenceResolver(mockClient(), threadPool);
+ return new InferenceResolver(mockClient(), functionRegistry, threadPool);
}
private static ModelConfigurations mockModelConfig(String inferenceId, TaskType taskType) {
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InputTextReaderTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InputTextReaderTests.java
new file mode 100644
index 0000000000000..64dd6a02928e8
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InputTextReaderTests.java
@@ -0,0 +1,231 @@
+/*
+ * 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.apache.lucene.util.BytesRef;
+import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.compute.test.ComputeTestCase;
+
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.nullValue;
+
+public class InputTextReaderTests extends ComputeTestCase {
+
+ public void testReadSingleValuePositions() throws Exception {
+ String[] texts = { "hello", "world", "test" };
+ BytesRefBlock block = createSingleValueBlock(texts);
+
+ try (InputTextReader reader = new InputTextReader(block)) {
+ assertThat(reader.estimatedSize(), equalTo(texts.length));
+
+ for (int i = 0; i < texts.length; i++) {
+ assertThat(reader.readText(i), equalTo(texts[i]));
+ }
+ }
+
+ allBreakersEmpty();
+ }
+
+ public void testReadMultiValuePositions() throws Exception {
+ BytesRefBlock block = createMultiValueBlock();
+
+ try (InputTextReader reader = new InputTextReader(block)) {
+ assertThat(reader.estimatedSize(), equalTo(2));
+
+ // First position has multiple values that should be concatenated with newlines
+ assertThat(reader.readText(0), equalTo("first\nsecond\nthird"));
+
+ // Second position has a single value
+ assertThat(reader.readText(1), equalTo("single"));
+ }
+
+ allBreakersEmpty();
+ }
+
+ public void testReadMultiValuePositionsWithLimit() throws Exception {
+ BytesRefBlock block = createMultiValueBlock();
+
+ try (InputTextReader reader = new InputTextReader(block)) {
+ // Test limiting to first 2 values out of 3
+ assertThat(reader.readText(0, 2), equalTo("first\nsecond"));
+
+ // Test limiting to first 1 value out of 3
+ assertThat(reader.readText(0, 1), equalTo("first"));
+
+ // Test limit larger than available values
+ assertThat(reader.readText(0, 10), equalTo("first\nsecond\nthird"));
+
+ // Test limit of 0
+ assertThat(reader.readText(0, 0), equalTo(""));
+
+ // Test single value position with limit
+ assertThat(reader.readText(1, 1), equalTo("single"));
+ }
+
+ allBreakersEmpty();
+ }
+
+ public void testReadNullValues() throws Exception {
+ BytesRefBlock block = createBlockWithNulls();
+
+ try (InputTextReader reader = new InputTextReader(block)) {
+ assertThat(reader.estimatedSize(), equalTo(3));
+
+ assertThat(reader.readText(0), equalTo("before"));
+ assertThat(reader.readText(1), nullValue());
+ assertThat(reader.readText(2), equalTo("after"));
+ }
+
+ allBreakersEmpty();
+ }
+
+ public void testReadNullValuesWithLimit() throws Exception {
+ BytesRefBlock block = createBlockWithNulls();
+
+ try (InputTextReader reader = new InputTextReader(block)) {
+ // Null values should return null regardless of limit
+ assertThat(reader.readText(0, 1), equalTo("before"));
+ assertThat(reader.readText(1, 1), nullValue());
+ assertThat(reader.readText(1, 10), nullValue());
+ assertThat(reader.readText(2, 1), equalTo("after"));
+ }
+
+ allBreakersEmpty();
+ }
+
+ public void testReadEmptyStrings() throws Exception {
+ String[] texts = { "", "non-empty", "" };
+ BytesRefBlock block = createSingleValueBlock(texts);
+
+ try (InputTextReader reader = new InputTextReader(block)) {
+ for (int i = 0; i < texts.length; i++) {
+ assertThat(reader.readText(i), equalTo(texts[i]));
+ assertThat(reader.readText(i, 1), equalTo(texts[i]));
+ }
+ }
+
+ allBreakersEmpty();
+ }
+
+ public void testReadLargeInput() throws Exception {
+ int size = between(1000, 5000);
+ String[] texts = new String[size];
+ for (int i = 0; i < size; i++) {
+ texts[i] = "text_" + i + "_" + randomAlphaOfLength(10);
+ }
+
+ BytesRefBlock block = createSingleValueBlock(texts);
+
+ try (InputTextReader reader = new InputTextReader(block)) {
+ assertThat(reader.estimatedSize(), equalTo(size));
+
+ for (int i = 0; i < size; i++) {
+ assertThat(reader.readText(i), equalTo(texts[i]));
+ assertThat(reader.readText(i, 1), equalTo(texts[i]));
+ }
+ }
+
+ allBreakersEmpty();
+ }
+
+ public void testReadUnicodeText() throws Exception {
+ String[] texts = { "café", "naïve", "résumé", "🚀 rocket", "多语言支持" };
+ BytesRefBlock block = createSingleValueBlock(texts);
+
+ try (InputTextReader reader = new InputTextReader(block)) {
+ for (int i = 0; i < texts.length; i++) {
+ assertThat(reader.readText(i), equalTo(texts[i]));
+ assertThat(reader.readText(i, 1), equalTo(texts[i]));
+ }
+ }
+
+ allBreakersEmpty();
+ }
+
+ public void testReadMultipleTimesFromSamePosition() throws Exception {
+ String[] texts = { "consistent" };
+ BytesRefBlock block = createSingleValueBlock(texts);
+
+ try (InputTextReader reader = new InputTextReader(block)) {
+ // Reading the same position multiple times should return the same result
+ assertThat(reader.readText(0), equalTo("consistent"));
+ assertThat(reader.readText(0), equalTo("consistent"));
+ assertThat(reader.readText(0, 1), equalTo("consistent"));
+ assertThat(reader.readText(0, 10), equalTo("consistent"));
+ }
+
+ allBreakersEmpty();
+ }
+
+ public void testLimitBoundaryConditions() throws Exception {
+ BytesRefBlock block = createLargeMultiValueBlock();
+
+ try (InputTextReader reader = new InputTextReader(block)) {
+ // Test various limit values on a position with 5 values
+ assertThat(reader.readText(0, 0), equalTo(""));
+ assertThat(reader.readText(0, 1), equalTo("value0"));
+ assertThat(reader.readText(0, 2), equalTo("value0\nvalue1"));
+ assertThat(reader.readText(0, 3), equalTo("value0\nvalue1\nvalue2"));
+ assertThat(reader.readText(0, 4), equalTo("value0\nvalue1\nvalue2\nvalue3"));
+ assertThat(reader.readText(0, 5), equalTo("value0\nvalue1\nvalue2\nvalue3\nvalue4"));
+
+ // Test limit beyond available values
+ assertThat(reader.readText(0, 10), equalTo("value0\nvalue1\nvalue2\nvalue3\nvalue4"));
+ assertThat(reader.readText(0, Integer.MAX_VALUE), equalTo("value0\nvalue1\nvalue2\nvalue3\nvalue4"));
+ }
+
+ allBreakersEmpty();
+ }
+
+ private BytesRefBlock createSingleValueBlock(String[] texts) {
+ try (BytesRefBlock.Builder builder = blockFactory().newBytesRefBlockBuilder(texts.length)) {
+ for (String text : texts) {
+ builder.appendBytesRef(new BytesRef(text));
+ }
+ return builder.build();
+ }
+ }
+
+ private BytesRefBlock createMultiValueBlock() {
+ try (BytesRefBlock.Builder builder = blockFactory().newBytesRefBlockBuilder(2)) {
+ // First position: multiple values
+ builder.beginPositionEntry();
+ builder.appendBytesRef(new BytesRef("first"));
+ builder.appendBytesRef(new BytesRef("second"));
+ builder.appendBytesRef(new BytesRef("third"));
+ builder.endPositionEntry();
+
+ // Second position: single value
+ builder.appendBytesRef(new BytesRef("single"));
+
+ return builder.build();
+ }
+ }
+
+ private BytesRefBlock createBlockWithNulls() {
+ try (BytesRefBlock.Builder builder = blockFactory().newBytesRefBlockBuilder(3)) {
+ builder.appendBytesRef(new BytesRef("before"));
+ builder.appendNull();
+ builder.appendBytesRef(new BytesRef("after"));
+ return builder.build();
+ }
+ }
+
+ private BytesRefBlock createLargeMultiValueBlock() {
+ try (BytesRefBlock.Builder builder = blockFactory().newBytesRefBlockBuilder(1)) {
+ // Single position with 5 values for testing limits
+ builder.beginPositionEntry();
+ for (int i = 0; i < 5; i++) {
+ builder.appendBytesRef(new BytesRef("value" + i));
+ }
+ builder.endPositionEntry();
+
+ return builder.build();
+ }
+ }
+}
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/textembedding/TextEmbeddingOperatorOutputBuilderTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/textembedding/TextEmbeddingOperatorOutputBuilderTests.java
new file mode 100644
index 0000000000000..ea77c6bed3c38
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/textembedding/TextEmbeddingOperatorOutputBuilderTests.java
@@ -0,0 +1,247 @@
+/*
+ * 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.textembedding;
+
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.ElementType;
+import org.elasticsearch.compute.data.FloatBlock;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.test.ComputeTestCase;
+import org.elasticsearch.compute.test.RandomBlock;
+import org.elasticsearch.core.Releasables;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingByteResults;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingFloatResults;
+
+import java.util.List;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class TextEmbeddingOperatorOutputBuilderTests extends ComputeTestCase {
+
+ public void testBuildSmallOutputWithFloatEmbeddings() throws Exception {
+ assertBuildOutputWithFloatEmbeddings(between(1, 100));
+ }
+
+ public void testBuildLargeOutputWithFloatEmbeddings() throws Exception {
+ assertBuildOutputWithFloatEmbeddings(between(1_000, 10_000));
+ }
+
+ public void testBuildSmallOutputWithByteEmbeddings() throws Exception {
+ assertBuildOutputWithByteEmbeddings(between(1, 100));
+ }
+
+ public void testBuildLargeOutputWithByteEmbeddings() throws Exception {
+ assertBuildOutputWithByteEmbeddings(between(1_000, 10_000));
+ }
+
+ public void testHandleNullResponses() throws Exception {
+ final int size = between(10, 100);
+ final Page inputPage = randomInputPage(size, between(1, 5));
+
+ try (
+ TextEmbeddingOperatorOutputBuilder outputBuilder = new TextEmbeddingOperatorOutputBuilder(
+ blockFactory().newFloatBlockBuilder(size),
+ inputPage
+ )
+ ) {
+ // Add some null responses
+ for (int currentPos = 0; currentPos < inputPage.getPositionCount(); currentPos++) {
+ if (randomBoolean()) {
+ outputBuilder.addInferenceResponse(null);
+ } else {
+ float[] embedding = randomFloatEmbedding(randomIntBetween(50, 200));
+ outputBuilder.addInferenceResponse(createFloatEmbeddingResponse(embedding));
+ }
+ }
+
+ final Page outputPage = outputBuilder.buildOutput();
+ assertThat(outputPage.getPositionCount(), equalTo(inputPage.getPositionCount()));
+ assertThat(outputPage.getBlockCount(), equalTo(inputPage.getBlockCount() + 1));
+
+ FloatBlock outputBlock = (FloatBlock) outputPage.getBlock(outputPage.getBlockCount() - 1);
+ assertThat(outputBlock.getPositionCount(), equalTo(size));
+
+ outputPage.releaseBlocks();
+ }
+
+ allBreakersEmpty();
+ }
+
+ public void testHandleEmptyEmbeddings() throws Exception {
+ final int size = between(5, 50);
+ final Page inputPage = randomInputPage(size, between(1, 3));
+
+ try (
+ TextEmbeddingOperatorOutputBuilder outputBuilder = new TextEmbeddingOperatorOutputBuilder(
+ blockFactory().newFloatBlockBuilder(size),
+ inputPage
+ )
+ ) {
+ // Add responses with empty embeddings
+ for (int currentPos = 0; currentPos < inputPage.getPositionCount(); currentPos++) {
+ outputBuilder.addInferenceResponse(createEmptyFloatEmbeddingResponse());
+ }
+
+ final Page outputPage = outputBuilder.buildOutput();
+ FloatBlock outputBlock = (FloatBlock) outputPage.getBlock(outputPage.getBlockCount() - 1);
+
+ // All positions should be null due to empty embeddings
+ for (int pos = 0; pos < outputBlock.getPositionCount(); pos++) {
+ assertThat(outputBlock.isNull(pos), equalTo(true));
+ }
+
+ outputPage.releaseBlocks();
+ }
+
+ allBreakersEmpty();
+ }
+
+ private void assertBuildOutputWithFloatEmbeddings(int size) throws Exception {
+ final Page inputPage = randomInputPage(size, between(1, 10));
+ final int embeddingDim = randomIntBetween(50, 1536); // Common embedding dimensions
+ final float[][] expectedEmbeddings = new float[size][];
+
+ try (
+ TextEmbeddingOperatorOutputBuilder outputBuilder = new TextEmbeddingOperatorOutputBuilder(
+ blockFactory().newFloatBlockBuilder(size),
+ inputPage
+ )
+ ) {
+ for (int currentPos = 0; currentPos < inputPage.getPositionCount(); currentPos++) {
+ float[] embedding = randomFloatEmbedding(embeddingDim);
+ expectedEmbeddings[currentPos] = embedding;
+ outputBuilder.addInferenceResponse(createFloatEmbeddingResponse(embedding));
+ }
+
+ final Page outputPage = outputBuilder.buildOutput();
+ assertThat(outputPage.getPositionCount(), equalTo(inputPage.getPositionCount()));
+ assertThat(outputPage.getBlockCount(), equalTo(inputPage.getBlockCount() + 1));
+
+ assertFloatEmbeddingContent((FloatBlock) outputPage.getBlock(outputPage.getBlockCount() - 1), expectedEmbeddings);
+
+ outputPage.releaseBlocks();
+ }
+
+ allBreakersEmpty();
+ }
+
+ private void assertBuildOutputWithByteEmbeddings(int size) throws Exception {
+ final Page inputPage = randomInputPage(size, between(1, 10));
+ final int embeddingDim = randomIntBetween(50, 1536);
+ final byte[][] expectedByteEmbeddings = new byte[size][];
+
+ try (
+ TextEmbeddingOperatorOutputBuilder outputBuilder = new TextEmbeddingOperatorOutputBuilder(
+ blockFactory().newFloatBlockBuilder(size),
+ inputPage
+ )
+ ) {
+ for (int currentPos = 0; currentPos < inputPage.getPositionCount(); currentPos++) {
+ byte[] embedding = randomByteEmbedding(embeddingDim);
+ expectedByteEmbeddings[currentPos] = embedding;
+ outputBuilder.addInferenceResponse(createByteEmbeddingResponse(embedding));
+ }
+
+ final Page outputPage = outputBuilder.buildOutput();
+ assertThat(outputPage.getPositionCount(), equalTo(inputPage.getPositionCount()));
+ assertThat(outputPage.getBlockCount(), equalTo(inputPage.getBlockCount() + 1));
+
+ assertByteEmbeddingContent((FloatBlock) outputPage.getBlock(outputPage.getBlockCount() - 1), expectedByteEmbeddings);
+
+ outputPage.releaseBlocks();
+ }
+
+ allBreakersEmpty();
+ }
+
+ private void assertFloatEmbeddingContent(FloatBlock block, float[][] expectedEmbeddings) {
+ for (int currentPos = 0; currentPos < block.getPositionCount(); currentPos++) {
+ assertThat(block.isNull(currentPos), equalTo(false));
+ assertThat(block.getValueCount(currentPos), equalTo(expectedEmbeddings[currentPos].length));
+
+ int firstValueIndex = block.getFirstValueIndex(currentPos);
+ for (int i = 0; i < expectedEmbeddings[currentPos].length; i++) {
+ float actualValue = block.getFloat(firstValueIndex + i);
+ float expectedValue = expectedEmbeddings[currentPos][i];
+ assertThat(actualValue, equalTo(expectedValue));
+ }
+ }
+ }
+
+ private void assertByteEmbeddingContent(FloatBlock block, byte[][] expectedByteEmbeddings) {
+ for (int currentPos = 0; currentPos < block.getPositionCount(); currentPos++) {
+ assertThat(block.isNull(currentPos), equalTo(false));
+ assertThat(block.getValueCount(currentPos), equalTo(expectedByteEmbeddings[currentPos].length));
+
+ int firstValueIndex = block.getFirstValueIndex(currentPos);
+ for (int i = 0; i < expectedByteEmbeddings[currentPos].length; i++) {
+ float actualValue = block.getFloat(firstValueIndex + i);
+ // Convert byte to float the same way as TextEmbeddingByteResults.Embedding.toFloatArray()
+ float expectedValue = expectedByteEmbeddings[currentPos][i];
+ assertThat(actualValue, equalTo(expectedValue));
+ }
+ }
+ }
+
+ private float[] randomFloatEmbedding(int dimension) {
+ float[] embedding = new float[dimension];
+ for (int i = 0; i < dimension; i++) {
+ embedding[i] = randomFloat();
+ }
+ return embedding;
+ }
+
+ private byte[] randomByteEmbedding(int dimension) {
+ byte[] embedding = new byte[dimension];
+ for (int i = 0; i < dimension; i++) {
+ embedding[i] = randomByte();
+ }
+ return embedding;
+ }
+
+ private static InferenceAction.Response createFloatEmbeddingResponse(float[] embedding) {
+ var embeddingResult = new TextEmbeddingFloatResults.Embedding(embedding);
+ var textEmbeddingResults = new TextEmbeddingFloatResults(List.of(embeddingResult));
+ return new InferenceAction.Response(textEmbeddingResults);
+ }
+
+ private static InferenceAction.Response createByteEmbeddingResponse(byte[] embedding) {
+ var embeddingResult = new TextEmbeddingByteResults.Embedding(embedding);
+ var textEmbeddingResults = new TextEmbeddingByteResults(List.of(embeddingResult));
+ return new InferenceAction.Response(textEmbeddingResults);
+ }
+
+ private static InferenceAction.Response createEmptyFloatEmbeddingResponse() {
+ var textEmbeddingResults = new TextEmbeddingFloatResults(List.of());
+ return new InferenceAction.Response(textEmbeddingResults);
+ }
+
+ private Page randomInputPage(int positionCount, int columnCount) {
+ final Block[] blocks = new Block[columnCount];
+ try {
+ for (int i = 0; i < columnCount; i++) {
+ blocks[i] = RandomBlock.randomBlock(
+ blockFactory(),
+ RandomBlock.randomElementExcluding(List.of(ElementType.AGGREGATE_METRIC_DOUBLE)),
+ positionCount,
+ randomBoolean(),
+ 0,
+ 0,
+ randomInt(10),
+ randomInt(10)
+ ).block();
+ }
+
+ return new Page(blocks);
+ } catch (Exception e) {
+ Releasables.close(blocks);
+ throw (e);
+ }
+ }
+}
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/textembedding/TextEmbeddingOperatorRequestIteratorTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/textembedding/TextEmbeddingOperatorRequestIteratorTests.java
new file mode 100644
index 0000000000000..a27ee95b941d7
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/textembedding/TextEmbeddingOperatorRequestIteratorTests.java
@@ -0,0 +1,167 @@
+/*
+ * 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.textembedding;
+
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.compute.test.ComputeTestCase;
+import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.nullValue;
+
+public class TextEmbeddingOperatorRequestIteratorTests extends ComputeTestCase {
+
+ public void testIterateSmallInput() throws Exception {
+ assertIterate(between(1, 100));
+ }
+
+ public void testIterateLargeInput() throws Exception {
+ assertIterate(between(10_000, 100_000));
+ }
+
+ public void testIterateWithNullValues() throws Exception {
+ final String inferenceId = randomIdentifier();
+ final BytesRefBlock inputBlock = createBlockWithNulls();
+
+ try (TextEmbeddingOperatorRequestIterator requestIterator = new TextEmbeddingOperatorRequestIterator(inputBlock, inferenceId)) {
+ BytesRef scratch = new BytesRef();
+
+ // First position: "before"
+ InferenceAction.Request request1 = requestIterator.next();
+ assertThat(request1.getInferenceEntityId(), equalTo(inferenceId));
+ assertThat(request1.getTaskType(), equalTo(TaskType.TEXT_EMBEDDING));
+ scratch = inputBlock.getBytesRef(inputBlock.getFirstValueIndex(0), scratch);
+ assertThat(request1.getInput().get(0), equalTo(scratch.utf8ToString()));
+
+ // Second position: null
+ InferenceAction.Request request2 = requestIterator.next();
+ assertThat(request2, nullValue());
+
+ // Third position: "after"
+ InferenceAction.Request request3 = requestIterator.next();
+ assertThat(request3.getInferenceEntityId(), equalTo(inferenceId));
+ assertThat(request3.getTaskType(), equalTo(TaskType.TEXT_EMBEDDING));
+ scratch = inputBlock.getBytesRef(inputBlock.getFirstValueIndex(2), scratch);
+ assertThat(request3.getInput().get(0), equalTo(scratch.utf8ToString()));
+ }
+
+ allBreakersEmpty();
+ }
+
+ public void testIterateWithMultiValuePositions() throws Exception {
+ final String inferenceId = randomIdentifier();
+ final BytesRefBlock inputBlock = createMultiValueBlock();
+
+ try (TextEmbeddingOperatorRequestIterator requestIterator = new TextEmbeddingOperatorRequestIterator(inputBlock, inferenceId)) {
+ // First position: multi-value concatenated with newlines
+ InferenceAction.Request request1 = requestIterator.next();
+ assertThat(request1.getInferenceEntityId(), equalTo(inferenceId));
+ assertThat(request1.getTaskType(), equalTo(TaskType.TEXT_EMBEDDING));
+ assertThat(request1.getInput().get(0), equalTo("first"));
+
+ // Second position: single value
+ InferenceAction.Request request2 = requestIterator.next();
+ assertThat(request2.getInferenceEntityId(), equalTo(inferenceId));
+ assertThat(request2.getTaskType(), equalTo(TaskType.TEXT_EMBEDDING));
+ assertThat(request2.getInput().get(0), equalTo("single"));
+ }
+
+ allBreakersEmpty();
+ }
+
+ public void testEstimatedSize() throws Exception {
+ final String inferenceId = randomIdentifier();
+ final int size = randomIntBetween(10, 1000);
+ final BytesRefBlock inputBlock = randomInputBlock(size);
+
+ try (TextEmbeddingOperatorRequestIterator requestIterator = new TextEmbeddingOperatorRequestIterator(inputBlock, inferenceId)) {
+ assertThat(requestIterator.estimatedSize(), equalTo(size));
+ }
+
+ allBreakersEmpty();
+ }
+
+ public void testHasNextAndIteration() throws Exception {
+ final String inferenceId = randomIdentifier();
+ final int size = randomIntBetween(5, 50);
+ final BytesRefBlock inputBlock = randomInputBlock(size);
+
+ try (TextEmbeddingOperatorRequestIterator requestIterator = new TextEmbeddingOperatorRequestIterator(inputBlock, inferenceId)) {
+ int count = 0;
+ while (requestIterator.hasNext()) {
+ requestIterator.next();
+ count++;
+ }
+ assertThat(count, equalTo(size));
+
+ // Verify hasNext returns false after iteration is complete
+ assertThat(requestIterator.hasNext(), equalTo(false));
+ }
+
+ allBreakersEmpty();
+ }
+
+ private void assertIterate(int size) throws Exception {
+ final String inferenceId = randomIdentifier();
+ final BytesRefBlock inputBlock = randomInputBlock(size);
+
+ try (TextEmbeddingOperatorRequestIterator requestIterator = new TextEmbeddingOperatorRequestIterator(inputBlock, inferenceId)) {
+ BytesRef scratch = new BytesRef();
+
+ for (int currentPos = 0; requestIterator.hasNext(); currentPos++) {
+ InferenceAction.Request request = requestIterator.next();
+
+ assertThat(request.getInferenceEntityId(), equalTo(inferenceId));
+ assertThat(request.getTaskType(), equalTo(TaskType.TEXT_EMBEDDING));
+
+ // Verify the input text matches what's in the block
+ scratch = inputBlock.getBytesRef(inputBlock.getFirstValueIndex(currentPos), scratch);
+ assertThat(request.getInput().get(0), equalTo(scratch.utf8ToString()));
+ }
+ }
+
+ allBreakersEmpty();
+ }
+
+ private BytesRefBlock randomInputBlock(int size) {
+ try (BytesRefBlock.Builder blockBuilder = blockFactory().newBytesRefBlockBuilder(size)) {
+ for (int i = 0; i < size; i++) {
+ blockBuilder.appendBytesRef(new BytesRef(randomAlphaOfLength(10)));
+ }
+
+ return blockBuilder.build();
+ }
+ }
+
+ private BytesRefBlock createBlockWithNulls() {
+ try (BytesRefBlock.Builder builder = blockFactory().newBytesRefBlockBuilder(3)) {
+ builder.appendBytesRef(new BytesRef("before"));
+ builder.appendNull();
+ builder.appendBytesRef(new BytesRef("after"));
+ return builder.build();
+ }
+ }
+
+ private BytesRefBlock createMultiValueBlock() {
+ try (BytesRefBlock.Builder builder = blockFactory().newBytesRefBlockBuilder(2)) {
+ // First position: multiple values
+ builder.beginPositionEntry();
+ builder.appendBytesRef(new BytesRef("first"));
+ builder.appendBytesRef(new BytesRef("second"));
+ builder.appendBytesRef(new BytesRef("third"));
+ builder.endPositionEntry();
+
+ // Second position: single value
+ builder.appendBytesRef(new BytesRef("single"));
+
+ return builder.build();
+ }
+ }
+}
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/textembedding/TextEmbeddingOperatorTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/textembedding/TextEmbeddingOperatorTests.java
new file mode 100644
index 0000000000000..6ff9a90b70b16
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/textembedding/TextEmbeddingOperatorTests.java
@@ -0,0 +1,137 @@
+/*
+ * 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.textembedding;
+
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.compute.data.FloatBlock;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.operator.Operator;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingFloatResults;
+import org.elasticsearch.xpack.esql.inference.InferenceOperatorTestCase;
+import org.hamcrest.Matcher;
+import org.junit.Before;
+
+import java.util.List;
+
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.hasSize;
+
+public class TextEmbeddingOperatorTests extends InferenceOperatorTestCase {
+ private static final String SIMPLE_INFERENCE_ID = "test_text_embedding";
+ private static final int EMBEDDING_DIMENSION = 384; // Common embedding dimension
+
+ private int inputChannel;
+
+ @Before
+ public void initTextEmbeddingChannels() {
+ inputChannel = between(0, inputsCount - 1);
+ }
+
+ @Override
+ protected Operator.OperatorFactory simple(SimpleOptions options) {
+ return new TextEmbeddingOperator.Factory(mockedInferenceService(), SIMPLE_INFERENCE_ID, evaluatorFactory(inputChannel));
+ }
+
+ @Override
+ protected void assertSimpleOutput(List input, List results) {
+ assertThat(results, hasSize(input.size()));
+
+ for (int curPage = 0; curPage < input.size(); curPage++) {
+ Page inputPage = input.get(curPage);
+ Page resultPage = results.get(curPage);
+
+ assertEquals(inputPage.getPositionCount(), resultPage.getPositionCount());
+ assertEquals(inputPage.getBlockCount() + 1, resultPage.getBlockCount());
+
+ for (int channel = 0; channel < inputPage.getBlockCount(); channel++) {
+ Block inputBlock = inputPage.getBlock(channel);
+ Block resultBlock = resultPage.getBlock(channel);
+ assertBlockContentEquals(inputBlock, resultBlock);
+ }
+
+ assertTextEmbeddingResults(inputPage, resultPage);
+ }
+ }
+
+ private void assertTextEmbeddingResults(Page inputPage, Page resultPage) {
+ BytesRefBlock inputBlock = resultPage.getBlock(inputChannel);
+ FloatBlock resultBlock = (FloatBlock) resultPage.getBlock(inputPage.getBlockCount());
+
+ BlockStringReader blockReader = new InferenceOperatorTestCase.BlockStringReader();
+
+ for (int curPos = 0; curPos < inputPage.getPositionCount(); curPos++) {
+ if (inputBlock.isNull(curPos)) {
+ assertThat(resultBlock.isNull(curPos), equalTo(true));
+ } else {
+ // Verify that we have an embedding vector at this position
+ assertThat(resultBlock.isNull(curPos), equalTo(false));
+ assertThat(resultBlock.getValueCount(curPos), equalTo(EMBEDDING_DIMENSION));
+
+ // Get the input text to verify our mock embedding generation
+ String inputText = blockReader.readString(inputBlock, curPos);
+
+ // Verify the embedding values match our mock generation pattern
+ int firstValueIndex = resultBlock.getFirstValueIndex(curPos);
+ for (int i = 0; i < EMBEDDING_DIMENSION; i++) {
+ float expectedValue = generateMockEmbeddingValue(inputText, i);
+ float actualValue = resultBlock.getFloat(firstValueIndex + i);
+ assertThat(actualValue, equalTo(expectedValue));
+ }
+ }
+ }
+ }
+
+ @Override
+ protected TextEmbeddingFloatResults mockInferenceResult(InferenceAction.Request request) {
+ // For text embedding, we expect one input text per request
+ String inputText = request.getInput().get(0);
+
+ // Generate a deterministic mock embedding based on the input text
+ float[] mockEmbedding = generateMockEmbedding(inputText, EMBEDDING_DIMENSION);
+
+ var embeddingResult = new TextEmbeddingFloatResults.Embedding(mockEmbedding);
+ return new TextEmbeddingFloatResults(List.of(embeddingResult));
+ }
+
+ @Override
+ protected Matcher expectedDescriptionOfSimple() {
+ return expectedToStringOfSimple();
+ }
+
+ @Override
+ protected Matcher expectedToStringOfSimple() {
+ return equalTo("TextEmbeddingOperator[inference_id=[" + SIMPLE_INFERENCE_ID + "]]");
+ }
+
+ /**
+ * Generates a deterministic mock embedding vector based on the input text.
+ * This ensures our tests are repeatable and verifiable.
+ */
+ private float[] generateMockEmbedding(String inputText, int dimension) {
+ float[] embedding = new float[dimension];
+ int textHash = inputText.hashCode();
+
+ for (int i = 0; i < dimension; i++) {
+ embedding[i] = generateMockEmbeddingValue(inputText, i);
+ }
+
+ return embedding;
+ }
+
+ /**
+ * Generates a single embedding value for a specific dimension based on input text.
+ * Uses a deterministic function so tests are repeatable.
+ */
+ private float generateMockEmbeddingValue(String inputText, int dimension) {
+ // Create a deterministic value based on input text and dimension
+ int hash = (inputText.hashCode() + dimension * 31) % 10000;
+ return hash / 10000.0f; // Normalize to [0, 1) range
+ }
+}
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..97792e4387b22 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
@@ -10,6 +10,7 @@
import org.apache.lucene.util.SetOnce;
import org.elasticsearch.action.ActionListener;
import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.threadpool.ThreadPool;
import org.elasticsearch.xpack.esql.EsqlTestUtils;
import org.elasticsearch.xpack.esql.core.expression.Alias;
import org.elasticsearch.xpack.esql.core.expression.Expression;
@@ -17,13 +18,18 @@
import org.elasticsearch.xpack.esql.core.tree.Source;
import org.elasticsearch.xpack.esql.expression.function.scalar.string.Concat;
import org.elasticsearch.xpack.esql.expression.predicate.operator.arithmetic.Add;
+import org.elasticsearch.xpack.esql.optimizer.rules.logical.preoptimizer.LogicalPlanPreOptimizerRule;
import org.elasticsearch.xpack.esql.plan.logical.Eval;
import org.elasticsearch.xpack.esql.plan.logical.Filter;
import org.elasticsearch.xpack.esql.plan.logical.Limit;
import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
import org.elasticsearch.xpack.esql.plan.logical.Project;
+import org.junit.After;
+import org.junit.Before;
import java.util.List;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.as;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.fieldAttribute;
@@ -33,9 +39,21 @@
public class LogicalPlanPreOptimizerTests extends ESTestCase {
+ private ThreadPool threadPool;
+
+ @Before
+ public void setUpThreadPool() {
+ threadPool = createThreadPool();
+ }
+
+ @After
+ public void tearDownThreadPool() {
+ terminate(threadPool);
+ }
+
public void testPlanIsMarkedAsPreOptimized() throws Exception {
for (int round = 0; round < 100; round++) {
- // We want to make sure that the pre-optimizer woks for a wide range of plans
+ // We want to make sure that the pre-optimizer works for a wide range of plans
preOptimizedPlan(randomPlan());
}
}
@@ -52,6 +70,75 @@ public void testPreOptimizeFailsIfPlanIsNotAnalyzed() throws Exception {
});
}
+ public void testPreOptimizerRulesAreAppliedInOrder() throws Exception {
+ LogicalPlan plan = EsqlTestUtils.relation();
+ plan.setPreOptimized();
+
+ StringBuilder executionOrder = new StringBuilder();
+
+ // Create mock rules that track execution order
+ LogicalPlanPreOptimizerRule rule1 = createOrderTrackingRule("A", executionOrder);
+ LogicalPlanPreOptimizerRule rule2 = createOrderTrackingRule("B", executionOrder);
+ LogicalPlanPreOptimizerRule rule3 = createOrderTrackingRule("C", executionOrder);
+
+ LogicalPlanPreOptimizer preOptimizer = new LogicalPlanPreOptimizer(List.of(rule1, rule2, rule3));
+
+ SetOnce resultHolder = new SetOnce<>();
+
+ preOptimizer.preOptimize(plan, ActionListener.wrap(resultHolder::set, ESTestCase::fail));
+
+ assertBusy(() -> {
+ assertThat(resultHolder.get(), notNullValue());
+ // Rules should be applied in the order they were provided
+ assertThat(executionOrder.toString(), equalTo("ABC"));
+ assertThat(resultHolder.get().preOptimized(), equalTo(true));
+ });
+ }
+
+ public void testPreOptimizerWithEmptyRulesList() throws Exception {
+ LogicalPlan plan = EsqlTestUtils.relation();
+ plan.setPreOptimized();
+
+ LogicalPlanPreOptimizer preOptimizer = new LogicalPlanPreOptimizer(List.of());
+
+ SetOnce resultHolder = new SetOnce<>();
+
+ preOptimizer.preOptimize(plan, ActionListener.wrap(resultHolder::set, ESTestCase::fail));
+
+ assertBusy(() -> {
+ assertThat(resultHolder.get(), notNullValue());
+ assertThat(resultHolder.get().preOptimized(), equalTo(true));
+ // The plan should be the same as the original (no modifications)
+ assertThat(resultHolder.get(), equalTo(plan));
+ });
+ }
+
+ public void testPreOptimizerRuleFailurePropagatesError() throws Exception {
+ LogicalPlan plan = EsqlTestUtils.relation();
+ plan.setPreOptimized();
+
+ RuntimeException expectedError = new RuntimeException("Mock rule failure");
+
+ AtomicInteger ruleACounter = new AtomicInteger();
+ LogicalPlanPreOptimizerRule ruleA = createMockRule(ruleACounter);
+ LogicalPlanPreOptimizerRule ruleB = createFailingRule(expectedError);
+ AtomicInteger ruleCCounter = new AtomicInteger();
+ LogicalPlanPreOptimizerRule ruleC = createMockRule(ruleCCounter);
+
+ LogicalPlanPreOptimizer preOptimizer = new LogicalPlanPreOptimizer(List.of(ruleA, ruleB, ruleC));
+
+ SetOnce exceptionHolder = new SetOnce<>();
+
+ preOptimizer.preOptimize(plan, ActionListener.wrap(r -> fail("Should have failed"), exceptionHolder::set));
+
+ assertBusy(() -> {
+ assertThat(exceptionHolder.get(), notNullValue());
+ assertThat(exceptionHolder.get(), equalTo(expectedError));
+ assertThat(ruleACounter.get(), equalTo(1));
+ assertThat(ruleCCounter.get(), equalTo(0));
+ });
+ }
+
public LogicalPlan preOptimizedPlan(LogicalPlan plan) throws Exception {
// set plan as analyzed
plan.setPreOptimized();
@@ -72,7 +159,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);
}
@@ -107,4 +194,28 @@ private Expression randomCondition() {
return EsqlTestUtils.greaterThanOf(randomExpression(), randomExpression());
}
+
+ // Helper methods for creating mock rules
+
+ private LogicalPlanPreOptimizerRule createMockRule(AtomicInteger executionCounter) {
+ return (plan, listener) -> {
+ threadPool.schedule(() -> {
+ executionCounter.incrementAndGet();
+ listener.onResponse(plan); // Return the plan unchanged
+ }, randomTimeValue(1, 100, TimeUnit.MILLISECONDS), threadPool.executor(ThreadPool.Names.GENERIC));
+ };
+ }
+
+ private LogicalPlanPreOptimizerRule createOrderTrackingRule(String ruleId, StringBuilder executionOrder) {
+ return (plan, listener) -> {
+ threadPool.schedule(() -> {
+ executionOrder.append(ruleId);
+ listener.onResponse(plan); // Return the plan unchanged
+ }, randomTimeValue(1, 100, TimeUnit.MILLISECONDS), threadPool.executor(ThreadPool.Names.GENERIC));
+ };
+ }
+
+ private LogicalPlanPreOptimizerRule createFailingRule(Exception error) {
+ return (plan, listener) -> listener.onFailure(error);
+ }
}
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/preoptimizer/FoldInferenceFunctionsTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/preoptimizer/FoldInferenceFunctionsTests.java
new file mode 100644
index 0000000000000..50fcdab626501
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/optimizer/rules/logical/preoptimizer/FoldInferenceFunctionsTests.java
@@ -0,0 +1,155 @@
+/*
+ * 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.apache.lucene.util.BytesRef;
+import org.apache.lucene.util.SetOnce;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xpack.esql.core.expression.Alias;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.Literal;
+import org.elasticsearch.xpack.esql.core.tree.Source;
+import org.elasticsearch.xpack.esql.expression.function.inference.InferenceFunction;
+import org.elasticsearch.xpack.esql.expression.function.inference.TextEmbedding;
+import org.elasticsearch.xpack.esql.expression.function.vector.Knn;
+import org.elasticsearch.xpack.esql.inference.InferenceFunctionEvaluator;
+import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
+import org.elasticsearch.xpack.esql.plan.logical.Eval;
+import org.elasticsearch.xpack.esql.plan.logical.Filter;
+import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+
+import java.util.List;
+
+import static org.elasticsearch.xpack.esql.EsqlTestUtils.as;
+import static org.elasticsearch.xpack.esql.EsqlTestUtils.getFieldAttribute;
+import static org.elasticsearch.xpack.esql.EsqlTestUtils.of;
+import static org.elasticsearch.xpack.esql.core.type.DataType.DENSE_VECTOR;
+import static org.elasticsearch.xpack.esql.optimizer.LocalLogicalPlanOptimizerTests.relation;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.hasSize;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+
+public class FoldInferenceFunctionsTests extends ESTestCase {
+ /**
+ * Tests that the rule correctly evaluates TEXT_EMBEDDING functions in Eval nodes.
+ */
+ public void testEvalFunctionEmbedding() throws Exception {
+ for (int round = 0; round < 100; round++) {
+ // Setup: Create a plan with an Eval node containing a TEXT_EMBEDDING function
+ String inferenceId = randomUUID();
+ String query = String.join(" ", randomArray(1, String[]::new, () -> randomAlphaOfLength(randomIntBetween(1, 10))));
+ int dimensions = randomIntBetween(1, 2048);
+ String fieldName = randomIdentifier();
+
+ logger.info("query sent: {}", query);
+
+ EsRelation relation = relation();
+ Eval eval = new Eval(
+ Source.EMPTY,
+ relation,
+ List.of(new Alias(Source.EMPTY, fieldName, new TextEmbedding(Source.EMPTY, of(query), of(inferenceId))))
+ );
+ eval.setAnalyzed();
+
+ SetOnce resultHolder = new SetOnce<>();
+ createFoldInferenceFunctionRule(dimensions).apply(eval, ActionListener.wrap(resultHolder::set, ESTestCase::fail));
+
+ assertBusy(() -> {
+ Eval preOptimizedEval = as(resultHolder.get(), Eval.class);
+ 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(generateTestEmbedding(query, dimensions)));
+ });
+
+ }
+ }
+
+ /**
+ * Tests that the rule correctly evaluates TEXT_EMBEDDING functions as KNN query.
+ */
+ public void testKnnFunctionEmbedding() throws Exception {
+ for (int round = 0; round < 100; round++) {
+ // Setup: Create a plan with a Filter node containing a KNN predicate with a TEXT_EMBEDDING function
+ String inferenceId = randomUUID();
+ String query = String.join(" ", randomArray(1, String[]::new, () -> randomAlphaOfLength(randomIntBetween(1, 10))));
+ int dimensions = randomIntBetween(1, 2048);
+
+ EsRelation relation = relation();
+ Filter filter = new Filter(
+ Source.EMPTY,
+ relation,
+ new Knn(Source.EMPTY, getFieldAttribute("a"), new TextEmbedding(Source.EMPTY, of(query), of(inferenceId)), null)
+ );
+ Knn knn = as(filter.condition(), Knn.class);
+
+ SetOnce resultHolder = new SetOnce<>();
+ createFoldInferenceFunctionRule(dimensions).apply(filter, ActionListener.wrap(resultHolder::set, ESTestCase::fail));
+
+ assertBusy(() -> {
+ Filter preOptimizedFilter = as(resultHolder.get(), Filter.class);
+
+ 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(generateTestEmbedding(query, dimensions)));
+ });
+ }
+ }
+
+ @SuppressWarnings("unchecked")
+ private LogicalPlanPreOptimizerRule createFoldInferenceFunctionRule(int dimensions) {
+ InferenceFunctionEvaluator functionEvaluator = mock(InferenceFunctionEvaluator.class);
+
+ doAnswer((i) -> {
+ ActionListener l = i.getArgument(1, ActionListener.class);
+ InferenceFunction function = i.getArgument(0, InferenceFunction.class);
+ if (function instanceof TextEmbedding textEmbedding) {
+ BytesRef bytesRef = (BytesRef) as(textEmbedding.inputText(), Literal.class).value();
+ float[] embedding = generateTestEmbedding(bytesRef.utf8ToString(), dimensions);
+ l.onResponse(Literal.of(function, embedding));
+ } else {
+ fail("Unexpected function type: " + function.getClass().getSimpleName());
+ }
+ return null;
+ }).when(functionEvaluator).fold(any(InferenceFunction.class), any(ActionListener.class));
+ return new FoldInferenceFunctions(functionEvaluator);
+ }
+
+ /**
+ * Generates a deterministic mock embedding vector based on the input text.
+ * This ensures our tests are repeatable and verifiable.
+ */
+ private float[] generateTestEmbedding(String inputText, int dimensions) {
+ float[] embedding = new float[dimensions];
+
+ for (int i = 0; i < dimensions; i++) {
+ embedding[i] = generateMockFloatEmbeddingValue(inputText, i);
+ }
+
+ return embedding;
+ }
+
+ /**
+ * Generates a single embedding value for a specific dimension based on input text.
+ * Uses a deterministic function so tests are repeatable.
+ */
+ private float generateMockFloatEmbeddingValue(String inputText, int dimension) {
+ // Create a deterministic value based on input text and dimension
+ int hash = (inputText.hashCode() + dimension * 31) % 10000;
+ return hash / 10000.0f; // Normalize to [0, 1) range
+ }
+}
diff --git a/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/esql/60_usage.yml b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/esql/60_usage.yml
index 72cb6efe42efc..a1211b9e62ff3 100644
--- a/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/esql/60_usage.yml
+++ b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/esql/60_usage.yml
@@ -89,6 +89,7 @@ setup:
- set: {esql.functions.to_long: functions_to_long}
- set: {esql.functions.coalesce: functions_coalesce}
- set: {esql.functions.categorize: functions_categorize}
+ - set: {esql.functions.text_embedding: functions_text_embedding}
- do:
esql.query:
@@ -131,6 +132,7 @@ setup:
- gt: {esql.functions.to_long: $functions_to_long}
- match: {esql.functions.coalesce: $functions_coalesce}
- gt: {esql.functions.categorize: $functions_categorize}
+ - match: {esql.functions.text_embedding: $functions_text_embedding}
# Testing for the entire function set isn't feasible, so we just check that we return the correct count as an approximation.
# - length: {esql.functions: 179} # check the "sister" test below for a likely update to the same esql.functions length check
---