Add text embedding function definition.

Author: afoucret

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

@@ -1303,6 +1303,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

@@ -12,6 +12,7 @@
 import org.elasticsearch.xpack.esql.expression.function.UnsupportedAttribute;
 import org.elasticsearch.xpack.esql.expression.function.aggregate.AggregateWritables;
 import org.elasticsearch.xpack.esql.expression.function.fulltext.FullTextWritables;
+import org.elasticsearch.xpack.esql.expression.function.inference.InferenceWritables;
 import org.elasticsearch.xpack.esql.expression.function.scalar.ScalarFunctionWritables;
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.FromBase64;
 import org.elasticsearch.xpack.esql.expression.function.scalar.convert.ToAggregateMetricDouble;
@@ -118,6 +119,7 @@ public static List<NamedWriteableRegistry.Entry> getNamedWriteables() {
         entries.addAll(fullText());
         entries.addAll(unaryScalars());
         entries.addAll(vector());
+        entries.addAll(inference());
         return entries;
     }
 
@@ -260,4 +262,8 @@ private static List<NamedWriteableRegistry.Entry> fullText() {
     private static List<NamedWriteableRegistry.Entry> vector() {
         return VectorWritables.getNamedWritables();
     }
+
+    private static List<NamedWriteableRegistry.Entry> inference() {
+        return InferenceWritables.getNamedWritables();
+    }
 }

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

@@ -193,6 +193,7 @@
 import org.elasticsearch.xpack.esql.expression.function.vector.L1Norm;
 import org.elasticsearch.xpack.esql.expression.function.vector.L2Norm;
 import org.elasticsearch.xpack.esql.expression.function.vector.Magnitude;
+import org.elasticsearch.xpack.esql.expression.function.inference.TextEmbedding;
 import org.elasticsearch.xpack.esql.parser.ParsingException;
 import org.elasticsearch.xpack.esql.session.Configuration;
 
@@ -519,7 +520,8 @@ private static FunctionDefinition[][] snapshotFunctions() {
                 def(Magnitude.class, Magnitude::new, "v_magnitude"),
                 def(Hamming.class, Hamming::new, "v_hamming"),
                 def(UrlEncode.class, UrlEncode::new, "url_encode"),
-                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/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -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/expression/ExpressionWritables.java

@@ -0,0 +1,36 @@
+/*
+ * 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.NamedWriteableRegistry;
+import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * Defines the named writables for inference functions in ESQL.
+ */
+public final class InferenceWritables {
+
+    private InferenceWritables() {
+        // Utility class
+        throw new UnsupportedOperationException();
+    }
+
+    public static List<NamedWriteableRegistry.Entry> getNamedWritables() {
+        List<NamedWriteableRegistry.Entry> entries = new ArrayList<>();
+
+        if (EsqlCapabilities.Cap.TEXT_EMBEDDING_FUNCTION.isEnabled()) {
+            entries.add(TextEmbedding.ENTRY);
+        }
+
+        return Collections.unmodifiableList(entries);
+    }
+}

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java

@@ -0,0 +1,164 @@
+/*
+ * 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.NamedWriteableRegistry;
+import org.elasticsearch.common.io.stream.StreamInput;
+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.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 org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
+
+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.isString;
+
+/**
+ * TEXT_EMBEDDING function converts text to dense vector embeddings using an inference endpoint.
+ */
+public class TextEmbedding extends InferenceFunction<TextEmbedding> {
+
+    public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
+        Expression.class,
+        "TextEmbedding",
+        TextEmbedding::new
+    );
+
+    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
+    )
+    public TextEmbedding(
+        Source source,
+        @Param(name = "text", type = { "keyword", "text" }, description = "Text to embed") Expression inputText,
+        @Param(
+            name = InferenceFunction.INFERENCE_ID_PARAMETER_NAME,
+            type = { "keyword", "text" },
+            description = "Identifier of the inference endpoint"
+        ) Expression inferenceId
+    ) {
+        super(source, List.of(inputText, inferenceId));
+        this.inferenceId = inferenceId;
+        this.inputText = inputText;
+    }
+
+    private TextEmbedding(StreamInput in) throws IOException {
+        this(Source.readFrom((PlanStreamInput) in), in.readNamedWriteable(Expression.class), in.readNamedWriteable(Expression.class));
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        source().writeTo(out);
+        out.writeNamedWriteable(inputText);
+        out.writeNamedWriteable(inferenceId);
+    }
+
+    @Override
+    public String getWriteableName() {
+        return ENTRY.name;
+    }
+
+    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(isString(inputText, sourceText(), FIRST));
+
+        if (textResolution.unresolved()) {
+            return textResolution;
+        }
+
+        TypeResolution inferenceIdResolution = isNotNull(inferenceId, sourceText(), SECOND).and(isString(inferenceId, sourceText(), SECOND))
+            .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);
+    }
+}

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/inference/InferenceFunction.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<TestCaseSupplier> cases() {
+        return paramsToSuppliers(TextEmbeddingTests.parameters());
+    }
+
+    @Override
+    protected Expression build(Source source, List<Expression> args) {
+        return new TextEmbedding(source, args.get(0), args.get(1));
+    }
+
+    @Override
+    protected Matcher<String> expectedTypeErrorMatcher(List<Set<DataType>> validPerPosition, List<DataType> signature) {
+        return equalTo(typeErrorMessage(true, validPerPosition, signature, (v, p) -> "string"));
+    }
+
+    protected static String typeErrorMessage(
+        boolean includeOrdinal,
+        List<Set<DataType>> validPerPosition,
+        List<DataType> 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
+        );
+    }
+}