[ES|QL] TEXT_EMBEDDING function definition (#135059)

Author: afoucret

docs/reference/query-languages/esql/_snippets/functions/parameters/text_embedding.md

@@ -0,0 +1,15 @@
+placeholder 
+required_capability: text_embedding_function
+required_capability: not_existing_capability
+
+// tag::embedding-eval[]
+ROW input="Who is Victor Hugo?"
+| EVAL embedding = TEXT_EMBEDDING("Who is Victor Hugo?", "test_dense_inference")
+;
+// end::embedding-eval[]
+
+
+input:keyword       | embedding:dense_vector
+Who is Victor Hugo? | [56.0, 50.0, 48.0]
+;
+

docs/reference/query-languages/esql/images/functions/text_embedding.svg

@@ -1319,6 +1319,11 @@ public enum Cap {
          */
         KNN_FUNCTION_V5(Build.current().isSnapshot()),
 
+        /**
+         * Support for the {@code TEXT_EMBEDDING} function for generating dense vector embeddings.
+         */
+        TEXT_EMBEDDING_FUNCTION(Build.current().isSnapshot()),
+
         /**
          * Support for the LIKE operator with a list of wildcards.
          */

docs/reference/query-languages/esql/kibana/definition/functions/text_embedding.json

@@ -73,6 +73,7 @@
 import org.elasticsearch.xpack.esql.expression.function.aggregate.SummationMode;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.Values;
 import org.elasticsearch.xpack.esql.expression.function.grouping.GroupingFunction;
+import org.elasticsearch.xpack.esql.expression.function.inference.InferenceFunction;
 import org.elasticsearch.xpack.esql.expression.function.scalar.EsqlScalarFunction;
 import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Case;
 import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Greatest;
@@ -1419,7 +1420,8 @@ private static class ResolveInference extends ParameterizedRule<LogicalPlan, Log
 
         @Override
         public LogicalPlan apply(LogicalPlan plan, AnalyzerContext context) {
-            return plan.transformDown(InferencePlan.class, p -> resolveInferencePlan(p, context));
+            return plan.transformDown(InferencePlan.class, p -> resolveInferencePlan(p, context))
+                .transformExpressionsOnly(InferenceFunction.class, f -> resolveInferenceFunction(f, context));
         }
 
         private LogicalPlan resolveInferencePlan(InferencePlan<?> plan, AnalyzerContext context) {
@@ -1448,6 +1450,36 @@ private LogicalPlan resolveInferencePlan(InferencePlan<?> plan, AnalyzerContext
 
             return plan;
         }
+
+        private InferenceFunction<?> resolveInferenceFunction(InferenceFunction<?> inferenceFunction, AnalyzerContext context) {
+            if (inferenceFunction.inferenceId().resolved()
+                && inferenceFunction.inferenceId().foldable()
+                && DataType.isString(inferenceFunction.inferenceId().dataType())) {
+
+                String inferenceId = BytesRefs.toString(inferenceFunction.inferenceId().fold(FoldContext.small()));
+                ResolvedInference resolvedInference = context.inferenceResolution().getResolvedInference(inferenceId);
+
+                if (resolvedInference == null) {
+                    String error = context.inferenceResolution().getError(inferenceId);
+                    return inferenceFunction.withInferenceResolutionError(inferenceId, error);
+                }
+
+                if (resolvedInference.taskType() != inferenceFunction.taskType()) {
+                    String error = "cannot use inference endpoint ["
+                        + inferenceId
+                        + "] with task type ["
+                        + resolvedInference.taskType()
+                        + "] within a "
+                        + context.functionRegistry().snapshotRegistry().functionName(inferenceFunction.getClass())
+                        + " function. Only inference endpoints with the task type ["
+                        + inferenceFunction.taskType()
+                        + "] are supported.";
+                    return inferenceFunction.withInferenceResolutionError(inferenceId, error);
+                }
+            }
+
+            return inferenceFunction;
+        }
     }
 
     private static class AddImplicitLimit extends ParameterizedRule<LogicalPlan, LogicalPlan, AnalyzerContext> {

docs/reference/query-languages/esql/kibana/docs/functions/text_embedding.md

@@ -63,6 +63,7 @@
 import org.elasticsearch.xpack.esql.expression.function.grouping.Bucket;
 import org.elasticsearch.xpack.esql.expression.function.grouping.Categorize;
 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.conditional.Case;
 import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Greatest;
 import org.elasticsearch.xpack.esql.expression.function.scalar.conditional.Least;
@@ -543,7 +544,8 @@ private static FunctionDefinition[][] snapshotFunctions() {
                 def(Hamming.class, Hamming::new, "v_hamming"),
                 def(UrlEncode.class, UrlEncode::new, "url_encode"),
                 def(UrlEncodeComponent.class, UrlEncodeComponent::new, "url_encode_component"),
-                def(UrlDecode.class, UrlDecode::new, "url_decode") } };
+                def(UrlDecode.class, UrlDecode::new, "url_decode"),
+                def(TextEmbedding.class, bi(TextEmbedding::new), "text_embedding") } };
     }
 
     public EsqlFunctionRegistry snapshotRegistry() {

x-pack/plugin/esql/qa/testFixtures/src/main/resources/text-embedding.csv-spec

@@ -0,0 +1,41 @@
+/*
+ * 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.inference.TaskType;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.function.Function;
+import org.elasticsearch.xpack.esql.core.tree.Source;
+
+import java.util.List;
+
+/**
+ * Base class for ESQL functions that use inference endpoints (e.g., TEXT_EMBEDDING).
+ */
+public abstract class InferenceFunction<PlanType extends InferenceFunction<PlanType>> extends Function {
+
+    public static final String INFERENCE_ID_PARAMETER_NAME = "inference_id";
+
+    protected InferenceFunction(Source source, List<Expression> children) {
+        super(source, children);
+    }
+
+    /** The inference endpoint identifier expression. */
+    public abstract Expression inferenceId();
+
+    /** The task type required by this function (e.g., TEXT_EMBEDDING). */
+    public abstract TaskType taskType();
+
+    /** Returns a copy with inference resolution error for display to user. */
+    public abstract PlanType withInferenceResolutionError(String inferenceId, String error);
+
+    /** True if this function contains nested inference function calls. */
+    public boolean hasNestedInferenceFunction() {
+        return anyMatch(e -> e instanceof InferenceFunction && e != this);
+    }
+}

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -0,0 +1,157 @@
+/*
+ * 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.common.io.stream.StreamOutput;
+import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.UnresolvedAttribute;
+import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
+import org.elasticsearch.xpack.esql.core.tree.Source;
+import org.elasticsearch.xpack.esql.core.type.DataType;
+import org.elasticsearch.xpack.esql.expression.function.Example;
+import org.elasticsearch.xpack.esql.expression.function.FunctionAppliesTo;
+import org.elasticsearch.xpack.esql.expression.function.FunctionAppliesToLifecycle;
+import org.elasticsearch.xpack.esql.expression.function.FunctionInfo;
+import org.elasticsearch.xpack.esql.expression.function.Param;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Objects;
+
+import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.FIRST;
+import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.ParamOrdinal.SECOND;
+import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isFoldable;
+import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isNotNull;
+import static org.elasticsearch.xpack.esql.core.expression.TypeResolutions.isType;
+
+/**
+ * TEXT_EMBEDDING function converts text to dense vector embeddings using an inference endpoint.
+ */
+public class TextEmbedding extends InferenceFunction<TextEmbedding> {
+
+    private final Expression inferenceId;
+    private final Expression inputText;
+
+    @FunctionInfo(
+        returnType = "dense_vector",
+        description = "Generates dense vector embeddings for text using a specified inference endpoint.",
+        appliesTo = { @FunctionAppliesTo(lifeCycle = FunctionAppliesToLifecycle.DEVELOPMENT) },
+        preview = true,
+        examples = {
+            @Example(
+                description = "Generate text embeddings using the 'test_dense_inference' inference endpoint.",
+                file = "text-embedding",
+                tag = "embedding-eval"
+            ) }
+    )
+    public TextEmbedding(
+        Source source,
+        @Param(name = "text", type = { "keyword" }, description = "Text to generate embeddings from") Expression inputText,
+        @Param(
+            name = InferenceFunction.INFERENCE_ID_PARAMETER_NAME,
+            type = { "keyword" },
+            description = "Identifier of the inference endpoint"
+        ) Expression inferenceId
+    ) {
+        super(source, List.of(inputText, inferenceId));
+        this.inferenceId = inferenceId;
+        this.inputText = inputText;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        throw new UnsupportedOperationException("doesn't escape the node");
+    }
+
+    @Override
+    public String getWriteableName() {
+        throw new UnsupportedOperationException("doesn't escape the node");
+    }
+
+    public Expression inputText() {
+        return inputText;
+    }
+
+    @Override
+    public Expression inferenceId() {
+        return inferenceId;
+    }
+
+    @Override
+    public boolean foldable() {
+        return inferenceId.foldable() && inputText.foldable();
+    }
+
+    @Override
+    public DataType dataType() {
+        return DataType.DENSE_VECTOR;
+    }
+
+    @Override
+    protected TypeResolution resolveType() {
+        if (childrenResolved() == false) {
+            return new TypeResolution("Unresolved children");
+        }
+
+        TypeResolution textResolution = isNotNull(inputText, sourceText(), FIRST).and(isFoldable(inputText, sourceText(), FIRST))
+            .and(isType(inputText, DataType.KEYWORD::equals, sourceText(), FIRST, "string"));
+
+        if (textResolution.unresolved()) {
+            return textResolution;
+        }
+
+        TypeResolution inferenceIdResolution = isNotNull(inferenceId, sourceText(), SECOND).and(
+            isType(inferenceId, DataType.KEYWORD::equals, sourceText(), SECOND, "string")
+        ).and(isFoldable(inferenceId, sourceText(), SECOND));
+
+        if (inferenceIdResolution.unresolved()) {
+            return inferenceIdResolution;
+        }
+
+        return TypeResolution.TYPE_RESOLVED;
+    }
+
+    @Override
+    public TaskType taskType() {
+        return TaskType.TEXT_EMBEDDING;
+    }
+
+    @Override
+    public TextEmbedding withInferenceResolutionError(String inferenceId, String error) {
+        return new TextEmbedding(source(), inputText, new UnresolvedAttribute(inferenceId().source(), inferenceId, error));
+    }
+
+    @Override
+    public Expression replaceChildren(List<Expression> newChildren) {
+        return new TextEmbedding(source(), newChildren.get(0), newChildren.get(1));
+    }
+
+    @Override
+    protected NodeInfo<? extends Expression> info() {
+        return NodeInfo.create(this, TextEmbedding::new, inputText, inferenceId);
+    }
+
+    @Override
+    public String toString() {
+        return "TEXT_EMBEDDING(" + inputText + ", " + inferenceId + ")";
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (o == null || getClass() != o.getClass()) return false;
+        if (super.equals(o) == false) return false;
+        TextEmbedding textEmbedding = (TextEmbedding) o;
+        return Objects.equals(inferenceId, textEmbedding.inferenceId) && Objects.equals(inputText, textEmbedding.inputText);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(super.hashCode(), inferenceId, inputText);
+    }
+}