+ * If the response is not of type {@link TextEmbeddingResults} an {@link IllegalStateException} is thrown.
+ *
+ *
+ * The responses must be added in the same order as the corresponding inference requests were generated.
+ * Failing to preserve order may lead to incorrect or misaligned output rows.
+ *
+ */
+ @Override
+ public void addInferenceResponse(InferenceAction.Response inferenceResponse) {
+ EmbeddingValueReader embeddingValueReader = EmbeddingValueReader.of(inferenceResponse, dimensions);
+ while (embeddingValueReader.hasNext()) {
+ writeEmbeddings(embeddingValueReader.next());
+ }
+ }
+
+ private void writeEmbeddings(float[] values) {
+ outputBlockBuilder.beginPositionEntry();
+ for (float value : values) {
+ outputBlockBuilder.appendFloat(value);
+ }
+ outputBlockBuilder.endPositionEntry();
+ }
+
+ private static class EmbeddingValueReader implements Iterator {
+ private final int dimensions;
+
+ private final Iterator> embeddingsIterator;
+
+ private EmbeddingValueReader(Iterator> embeddingsIterator, int dimensions) {
+ this.dimensions = dimensions;
+ this.embeddingsIterator = embeddingsIterator;
+ }
+
+ public boolean hasNext() {
+ return embeddingsIterator.hasNext();
+ }
+
+ public float[] next() {
+ EmbeddingResults.Embedding embedding = embeddingsIterator.next();
+ float[] values = switch (embedding) {
+ case TextEmbeddingFloatResults.Embedding textEmbeddingFloat -> textEmbeddingFloat.values();
+ case TextEmbeddingByteResults.Embedding textEmbeddingBytes -> toFloatArray(textEmbeddingBytes.values());
+ default -> throw new IllegalStateException("Unsupported embedding type [" + embedding.getClass() + "]");
+ };
+
+ assert values.length == dimensions : "Unexpected vector size: " + values.length;
+
+ return values;
+ }
+
+ private static float[] toFloatArray(byte[] bytes) {
+ float[] floatValues = new float[bytes.length];
+ IntStream.range(0, floatValues.length).forEach(i -> floatValues[i] = ((Byte) bytes[i]).floatValue());
+ return floatValues;
+ }
+
+ public static EmbeddingValueReader of(InferenceAction.Response inferenceResponse, int dimensions) {
+ TextEmbeddingResults inferenceResults = InferenceOperator.OutputBuilder.inferenceResults(
+ inferenceResponse,
+ TextEmbeddingResults.class
+ );
+ return new EmbeddingValueReader(inferenceResults.embeddings().iterator(), dimensions);
+ }
+ }
+}
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/embedding/DenseEmbeddingOperatorRequestIterator.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/embedding/DenseEmbeddingOperatorRequestIterator.java
new file mode 100644
index 0000000000000..e1f3a3c4ebec1
--- /dev/null
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/embedding/DenseEmbeddingOperatorRequestIterator.java
@@ -0,0 +1,83 @@
+/*
+ * 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.embedding;
+
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.common.lucene.BytesRefs;
+import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRequestIterator;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.NoSuchElementException;
+
+/**
+ * Iterator over input data blocks to create batched inference requests for the dense vector text embedding task.
+ *
+ * This iterator reads from a {@link BytesRefBlock} containing text to be embedded. It slices the input into batches
+ * of configurable size and converts each batch into an {@link InferenceAction.Request} with the task type {@link TaskType#TEXT_EMBEDDING}.
+ */
+public class DenseEmbeddingOperatorRequestIterator implements BulkInferenceRequestIterator {
+ private final BytesRefBlock inputBlock;
+ private final String inferenceId;
+ private final int batchSize;
+ private int remainingPositions;
+
+ public DenseEmbeddingOperatorRequestIterator(BytesRefBlock inputBlock, String inferenceId, int batchSize) {
+ this.inputBlock = inputBlock;
+ this.inferenceId = inferenceId;
+ this.batchSize = batchSize;
+ this.remainingPositions = inputBlock.getPositionCount();
+ }
+
+ @Override
+ public boolean hasNext() {
+ return remainingPositions > 0;
+ }
+
+ @Override
+ public InferenceAction.Request next() {
+ if (hasNext() == false) {
+ throw new NoSuchElementException();
+ }
+
+ final int inputSize = Math.min(remainingPositions, batchSize);
+ final List inputs = new ArrayList<>(inputSize);
+ BytesRef scratch = new BytesRef();
+
+ int startIndex = inputBlock.getPositionCount() - remainingPositions;
+ for (int i = 0; i < inputSize; i++) {
+ int pos = startIndex + i;
+ if (inputBlock.isNull(pos)) {
+ inputs.add("");
+ } else {
+ scratch = inputBlock.getBytesRef(inputBlock.getFirstValueIndex(pos), scratch);
+ inputs.add(BytesRefs.toString(scratch));
+ }
+ }
+
+ remainingPositions -= inputSize;
+ return inferenceRequest(inputs);
+ }
+
+ @Override
+ public int estimatedSize() {
+ return inputBlock.getPositionCount();
+ }
+
+ private InferenceAction.Request inferenceRequest(List inputs) {
+ return InferenceAction.Request.builder(inferenceId, TaskType.TEXT_EMBEDDING).setInput(inputs).build();
+ }
+
+ @Override
+ public void close() {
+
+ }
+}
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/PlanWritables.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/PlanWritables.java
index 2fe9f5182ae00..940d8adda546c 100644
--- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/PlanWritables.java
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/PlanWritables.java
@@ -26,6 +26,7 @@
import org.elasticsearch.xpack.esql.plan.logical.TopN;
import org.elasticsearch.xpack.esql.plan.logical.inference.Completion;
import org.elasticsearch.xpack.esql.plan.logical.inference.Rerank;
+import org.elasticsearch.xpack.esql.plan.logical.inference.embedding.DenseVectorEmbedding;
import org.elasticsearch.xpack.esql.plan.logical.join.InlineJoin;
import org.elasticsearch.xpack.esql.plan.logical.join.Join;
import org.elasticsearch.xpack.esql.plan.logical.local.EsqlProject;
@@ -55,6 +56,7 @@
import org.elasticsearch.xpack.esql.plan.physical.TopNExec;
import org.elasticsearch.xpack.esql.plan.physical.inference.CompletionExec;
import org.elasticsearch.xpack.esql.plan.physical.inference.RerankExec;
+import org.elasticsearch.xpack.esql.plan.physical.inference.embedding.DenseVectorEmbeddingExec;
import java.util.ArrayList;
import java.util.List;
@@ -72,6 +74,7 @@ public static List logical() {
return List.of(
Aggregate.ENTRY,
Completion.ENTRY,
+ DenseVectorEmbedding.ENTRY,
Dissect.ENTRY,
Enrich.ENTRY,
EsRelation.ENTRY,
@@ -99,6 +102,7 @@ public static List physical() {
return List.of(
AggregateExec.ENTRY,
CompletionExec.ENTRY,
+ DenseVectorEmbeddingExec.ENTRY,
DissectExec.ENTRY,
EnrichExec.ENTRY,
EsQueryExec.ENTRY,
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/inference/InferencePlan.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/inference/InferencePlan.java
index 620e8726865d6..1ad945dffa4a8 100644
--- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/inference/InferencePlan.java
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/inference/InferencePlan.java
@@ -8,6 +8,7 @@
package org.elasticsearch.xpack.esql.plan.logical.inference;
import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.inference.ModelConfigurations;
import org.elasticsearch.inference.TaskType;
import org.elasticsearch.xpack.esql.core.expression.Expression;
import org.elasticsearch.xpack.esql.core.expression.UnresolvedAttribute;
@@ -69,4 +70,9 @@ public int hashCode() {
public PlanType withInferenceResolutionError(String inferenceId, String error) {
return withInferenceId(new UnresolvedAttribute(inferenceId().source(), inferenceId, error));
}
+
+ @SuppressWarnings("unchecked")
+ public PlanType withModelConfigurations(ModelConfigurations modelConfig) {
+ return (PlanType) this;
+ }
}
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/inference/embedding/DenseVectorEmbedding.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/inference/embedding/DenseVectorEmbedding.java
new file mode 100644
index 0000000000000..c654bcff590be
--- /dev/null
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/inference/embedding/DenseVectorEmbedding.java
@@ -0,0 +1,210 @@
+/*
+ * 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.plan.logical.inference.embedding;
+
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.lucene.BytesRefs;
+import org.elasticsearch.inference.ModelConfigurations;
+import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.xpack.esql.capabilities.TelemetryAware;
+import org.elasticsearch.xpack.esql.core.capabilities.Unresolvable;
+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.expression.NameId;
+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.io.stream.PlanStreamInput;
+import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+import org.elasticsearch.xpack.esql.plan.logical.inference.InferencePlan;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Objects;
+
+import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
+
+public class DenseVectorEmbedding extends InferencePlan implements TelemetryAware {
+
+ public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
+ LogicalPlan.class,
+ "DenseVectorEmbedding",
+ DenseVectorEmbedding::new
+ );
+
+ private final Expression input;
+ private final Expression dimensions;
+ private final Attribute targetField;
+ private List lazyOutput;
+
+ public DenseVectorEmbedding(Source source, LogicalPlan child, Expression inferenceId, Expression input, Attribute targetField) {
+ this(source, child, inferenceId, new UnresolvedDimensions(inferenceId), input, targetField);
+ }
+
+ DenseVectorEmbedding(
+ Source source,
+ LogicalPlan child,
+ Expression inferenceId,
+ Expression dimensions,
+ Expression input,
+ Attribute targetField
+ ) {
+ super(source, child, inferenceId);
+ this.input = input;
+ this.targetField = targetField;
+ this.dimensions = dimensions;
+ }
+
+ public DenseVectorEmbedding(StreamInput in) throws IOException {
+ this(
+ Source.readFrom((PlanStreamInput) in),
+ in.readNamedWriteable(LogicalPlan.class),
+ in.readNamedWriteable(Expression.class),
+ in.readNamedWriteable(Expression.class),
+ in.readNamedWriteable(Expression.class),
+ in.readNamedWriteable(Attribute.class)
+ );
+ }
+
+ @Override
+ public void writeTo(StreamOutput out) throws IOException {
+ source().writeTo(out);
+ out.writeNamedWriteable(child());
+ out.writeNamedWriteable(inferenceId());
+ out.writeNamedWriteable(dimensions);
+ out.writeNamedWriteable(input);
+ out.writeNamedWriteable(targetField);
+ }
+
+ public Expression input() {
+ return input;
+ }
+
+ public Attribute embeddingField() {
+ return targetField;
+ }
+
+ @Override
+ public TaskType taskType() {
+ return TaskType.TEXT_EMBEDDING;
+ }
+
+ public Expression dimensions() {
+ return dimensions;
+ }
+
+ @Override
+ public String getWriteableName() {
+ return ENTRY.name;
+ }
+
+ @Override
+ public List output() {
+ if (lazyOutput == null) {
+ lazyOutput = mergeOutputAttributes(List.of(targetField), child().output());
+ }
+ return lazyOutput;
+ }
+
+ @Override
+ public List generatedAttributes() {
+ return List.of(targetField);
+ }
+
+ @Override
+ public DenseVectorEmbedding withGeneratedNames(List newNames) {
+ checkNumberOfNewNames(newNames);
+ return new DenseVectorEmbedding(source(), child(), inferenceId(), dimensions, input, this.renameTargetField(newNames.get(0)));
+ }
+
+ private Attribute renameTargetField(String newName) {
+ if (newName.equals(targetField.name())) {
+ return targetField;
+ }
+
+ return targetField.withName(newName).withId(new NameId());
+ }
+
+ @Override
+ public boolean expressionsResolved() {
+ return super.expressionsResolved() && input.resolved() && targetField.resolved() && dimensions.resolved();
+ }
+
+ @Override
+ public DenseVectorEmbedding withInferenceId(Expression newInferenceId) {
+ return new DenseVectorEmbedding(source(), child(), newInferenceId, dimensions, input, targetField);
+ }
+
+ public DenseVectorEmbedding withDimensions(Expression newDimensions) {
+ return new DenseVectorEmbedding(source(), child(), inferenceId(), newDimensions, input, targetField);
+ }
+
+ public DenseVectorEmbedding withTargetField(Attribute targetField) {
+ return new DenseVectorEmbedding(source(), child(), inferenceId(), dimensions, input, targetField);
+ }
+
+ @Override
+ public DenseVectorEmbedding withModelConfigurations(ModelConfigurations modelConfig) {
+ boolean hasChanged = false;
+ Expression newDimensions = dimensions;
+
+ if (dimensions.resolved() == false
+ && modelConfig.getServiceSettings() != null
+ && modelConfig.getServiceSettings().dimensions() > 0) {
+ hasChanged = true;
+ newDimensions = new Literal(Source.EMPTY, modelConfig.getServiceSettings().dimensions(), DataType.INTEGER);
+ }
+
+ return hasChanged ? withDimensions(newDimensions) : this;
+ }
+
+ @Override
+ public DenseVectorEmbedding replaceChild(LogicalPlan newChild) {
+ return new DenseVectorEmbedding(source(), newChild, inferenceId(), dimensions, input, targetField);
+ }
+
+ @Override
+ protected NodeInfo info() {
+ return NodeInfo.create(this, DenseVectorEmbedding::new, child(), inferenceId(), dimensions, input, targetField);
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ if (super.equals(o) == false) return false;
+ DenseVectorEmbedding that = (DenseVectorEmbedding) o;
+ return Objects.equals(input, that.input)
+ && Objects.equals(dimensions, that.dimensions)
+ && Objects.equals(targetField, that.targetField);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(super.hashCode(), input, targetField, dimensions);
+ }
+
+ private static class UnresolvedDimensions extends Literal implements Unresolvable {
+
+ private final String inferenceId;
+
+ private UnresolvedDimensions(Expression inferenceId) {
+ super(Source.EMPTY, null, DataType.NULL);
+ this.inferenceId = BytesRefs.toString(inferenceId.fold(FoldContext.small()));
+ }
+
+ @Override
+ public String unresolvedMessage() {
+ return "Dimensions cannot be resolved for inference endpoint[" + inferenceId + "]";
+ }
+ }
+}
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/inference/embedding/DenseVectorEmbeddingExec.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/inference/embedding/DenseVectorEmbeddingExec.java
new file mode 100644
index 0000000000000..4eae30daeef8e
--- /dev/null
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/physical/inference/embedding/DenseVectorEmbeddingExec.java
@@ -0,0 +1,130 @@
+/*
+ * 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.plan.physical.inference.embedding;
+
+import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.xpack.esql.core.expression.Attribute;
+import org.elasticsearch.xpack.esql.core.expression.AttributeSet;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
+import org.elasticsearch.xpack.esql.core.tree.Source;
+import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
+import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
+import org.elasticsearch.xpack.esql.plan.physical.UnaryExec;
+import org.elasticsearch.xpack.esql.plan.physical.inference.InferenceExec;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Objects;
+
+import static org.elasticsearch.xpack.esql.expression.NamedExpressions.mergeOutputAttributes;
+
+public class DenseVectorEmbeddingExec extends InferenceExec {
+
+ public static final NamedWriteableRegistry.Entry ENTRY = new NamedWriteableRegistry.Entry(
+ PhysicalPlan.class,
+ "DenseVectorEmbeddingExec",
+ DenseVectorEmbeddingExec::new
+ );
+
+ private final Expression input;
+ private final Expression dimensions;
+ private final Attribute targetField;
+ private List lazyOutput;
+
+ public DenseVectorEmbeddingExec(
+ Source source,
+ PhysicalPlan child,
+ Expression inferenceId,
+ Expression dimensions,
+ Expression input,
+ Attribute targetField
+ ) {
+ super(source, child, inferenceId);
+ this.input = input;
+ this.dimensions = dimensions;
+ this.targetField = targetField;
+ }
+
+ public DenseVectorEmbeddingExec(StreamInput in) throws IOException {
+ this(
+ Source.readFrom((PlanStreamInput) in),
+ in.readNamedWriteable(PhysicalPlan.class),
+ in.readNamedWriteable(Expression.class),
+ in.readNamedWriteable(Expression.class),
+ in.readNamedWriteable(Expression.class),
+ in.readNamedWriteable(Attribute.class)
+ );
+ }
+
+ public Expression input() {
+ return input;
+ }
+
+ public Attribute targetField() {
+ return targetField;
+ }
+
+ @Override
+ public String getWriteableName() {
+ return ENTRY.name;
+ }
+
+ @Override
+ public void writeTo(StreamOutput out) throws IOException {
+ super.writeTo(out);
+ out.writeNamedWriteable(dimensions);
+ out.writeNamedWriteable(input);
+ out.writeNamedWriteable(targetField);
+ }
+
+ public Expression dimensions() {
+ return dimensions;
+ }
+
+ @Override
+ protected NodeInfo info() {
+ return NodeInfo.create(this, DenseVectorEmbeddingExec::new, child(), inferenceId(), input, dimensions, targetField);
+ }
+
+ @Override
+ public UnaryExec replaceChild(PhysicalPlan newChild) {
+ return new DenseVectorEmbeddingExec(source(), newChild, inferenceId(), input, dimensions, targetField);
+ }
+
+ @Override
+ public List output() {
+ if (lazyOutput == null) {
+ lazyOutput = mergeOutputAttributes(List.of(targetField), child().output());
+ }
+ return lazyOutput;
+ }
+
+ @Override
+ protected AttributeSet computeReferences() {
+ return input.references();
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) return true;
+ if (o == null || getClass() != o.getClass()) return false;
+ if (super.equals(o) == false) return false;
+ DenseVectorEmbeddingExec that = (DenseVectorEmbeddingExec) o;
+ return Objects.equals(input, that.input)
+ && Objects.equals(dimensions, that.dimensions)
+ && Objects.equals(targetField, that.targetField);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(super.hashCode(), input, dimensions, targetField);
+ }
+}
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java
index a92d2f439a0ea..333eab10b217a 100644
--- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/LocalExecutionPlanner.java
@@ -88,6 +88,7 @@
import org.elasticsearch.xpack.esql.inference.InferenceRunner;
import org.elasticsearch.xpack.esql.inference.XContentRowEncoder;
import org.elasticsearch.xpack.esql.inference.completion.CompletionOperator;
+import org.elasticsearch.xpack.esql.inference.embedding.DenseEmbeddingOperator;
import org.elasticsearch.xpack.esql.inference.rerank.RerankOperator;
import org.elasticsearch.xpack.esql.plan.logical.Fork;
import org.elasticsearch.xpack.esql.plan.physical.AggregateExec;
@@ -119,6 +120,7 @@
import org.elasticsearch.xpack.esql.plan.physical.TopNExec;
import org.elasticsearch.xpack.esql.plan.physical.inference.CompletionExec;
import org.elasticsearch.xpack.esql.plan.physical.inference.RerankExec;
+import org.elasticsearch.xpack.esql.plan.physical.inference.embedding.DenseVectorEmbeddingExec;
import org.elasticsearch.xpack.esql.planner.EsPhysicalOperationProviders.ShardContext;
import org.elasticsearch.xpack.esql.plugin.QueryPragmas;
import org.elasticsearch.xpack.esql.score.ScoreMapper;
@@ -266,6 +268,8 @@ private PhysicalOperation plan(PhysicalPlan node, LocalExecutionPlannerContext c
return planChangePoint(changePoint, context);
} else if (node instanceof CompletionExec completion) {
return planCompletion(completion, context);
+ } else if (node instanceof DenseVectorEmbeddingExec embedding) {
+ return planDenseVectorEmbedding(embedding, context);
} else if (node instanceof SampleExec Sample) {
return planSample(Sample, context);
}
@@ -319,6 +323,31 @@ private PhysicalOperation planCompletion(CompletionExec completion, LocalExecuti
return source.with(new CompletionOperator.Factory(inferenceRunner, inferenceId, promptEvaluatorFactory), outputLayout);
}
+ private PhysicalOperation planDenseVectorEmbedding(DenseVectorEmbeddingExec embedding, LocalExecutionPlannerContext context) {
+ PhysicalOperation source = plan(embedding.child(), context);
+ String inferenceId = BytesRefs.toString(embedding.inferenceId().fold(context.foldCtx()));
+
+ int dimensions;
+ if (embedding.dimensions() instanceof Literal literal) {
+ Object val = literal.value() instanceof BytesRef br ? BytesRefs.toString(br) : literal.value();
+ dimensions = stringToInt(val.toString());
+ } else {
+ throw new EsqlIllegalArgumentException("dimensions only supported with literal values");
+ }
+
+ Layout outputLayout = source.layout.builder().append(embedding.targetField()).build();
+ EvalOperator.ExpressionEvaluator.Factory inputEvaluatorFactory = EvalMapper.toEvaluator(
+ context.foldCtx(),
+ embedding.input(),
+ source.layout
+ );
+
+ return source.with(
+ new DenseEmbeddingOperator.Factory(inferenceRunner, inferenceId, dimensions, inputEvaluatorFactory),
+ outputLayout
+ );
+ }
+
private PhysicalOperation planRrfScoreEvalExec(RrfScoreEvalExec rrf, LocalExecutionPlannerContext context) {
PhysicalOperation source = plan(rrf.child(), context);
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/mapper/MapperUtils.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/mapper/MapperUtils.java
index 4851de1616844..884f50b71a549 100644
--- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/mapper/MapperUtils.java
+++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/mapper/MapperUtils.java
@@ -29,6 +29,7 @@
import org.elasticsearch.xpack.esql.plan.logical.UnaryPlan;
import org.elasticsearch.xpack.esql.plan.logical.inference.Completion;
import org.elasticsearch.xpack.esql.plan.logical.inference.Rerank;
+import org.elasticsearch.xpack.esql.plan.logical.inference.embedding.DenseVectorEmbedding;
import org.elasticsearch.xpack.esql.plan.logical.local.LocalRelation;
import org.elasticsearch.xpack.esql.plan.logical.show.ShowInfo;
import org.elasticsearch.xpack.esql.plan.physical.AggregateExec;
@@ -47,6 +48,7 @@
import org.elasticsearch.xpack.esql.plan.physical.TimeSeriesAggregateExec;
import org.elasticsearch.xpack.esql.plan.physical.inference.CompletionExec;
import org.elasticsearch.xpack.esql.plan.physical.inference.RerankExec;
+import org.elasticsearch.xpack.esql.plan.physical.inference.embedding.DenseVectorEmbeddingExec;
import org.elasticsearch.xpack.esql.planner.AbstractPhysicalOperationProviders;
import java.util.List;
@@ -106,6 +108,17 @@ static PhysicalPlan mapUnary(UnaryPlan p, PhysicalPlan child) {
return new CompletionExec(completion.source(), child, completion.inferenceId(), completion.prompt(), completion.targetField());
}
+ if (p instanceof DenseVectorEmbedding embedding) {
+ return new DenseVectorEmbeddingExec(
+ embedding.source(),
+ child,
+ embedding.inferenceId(),
+ embedding.dimensions(),
+ embedding.input(),
+ embedding.embeddingField()
+ );
+ }
+
if (p instanceof Enrich enrich) {
return new EnrichExec(
enrich.source(),
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 4ff65f59bbd72..920201897199b 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
@@ -85,7 +85,6 @@
import org.elasticsearch.xpack.esql.plan.logical.TopN;
import org.elasticsearch.xpack.esql.plan.logical.UnresolvedRelation;
import org.elasticsearch.xpack.esql.plan.logical.inference.Completion;
-import org.elasticsearch.xpack.esql.plan.logical.inference.InferencePlan;
import org.elasticsearch.xpack.esql.plan.logical.join.InlineJoin;
import org.elasticsearch.xpack.esql.plan.logical.join.JoinTypes;
import org.elasticsearch.xpack.esql.plan.logical.join.LookupJoin;
@@ -372,7 +371,7 @@ public void analyzedPlan(
l -> enrichPolicyResolver.resolvePolicies(unresolvedPolicies, executionInfo, l)
)
.andThen((l, enrichResolution) -> resolveFieldNames(parsed, enrichResolution, l))
- .andThen((l, preAnalysisResult) -> resolveInferences(preAnalysis.inferencePlans, preAnalysisResult, l));
+ .andThen((l, preAnalysisResult) -> resolveInferences(preAnalysis.inferenceIds, preAnalysisResult, l));
// first resolve the lookup indices, then the main indices
for (var index : preAnalysis.lookupIndices) {
listener = listener.andThen((l, preAnalysisResult) -> { preAnalyzeLookupIndex(index, preAnalysisResult, l); });
@@ -588,12 +587,8 @@ private static void resolveFieldNames(LogicalPlan parsed, EnrichResolution enric
}
}
- private void resolveInferences(
- List> inferencePlans,
- PreAnalysisResult preAnalysisResult,
- ActionListener l
- ) {
- inferenceRunner.resolveInferenceIds(inferencePlans, l.map(preAnalysisResult::withInferenceResolution));
+ private void resolveInferences(Set inferenceIds, PreAnalysisResult preAnalysisResult, ActionListener l) {
+ inferenceRunner.resolveInferenceIds(inferenceIds, l.map(preAnalysisResult::withInferenceResolution));
}
static PreAnalysisResult fieldNames(LogicalPlan parsed, Set enrichPolicyMatchFields, PreAnalysisResult result) {
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 5e6c37545a396..a016efa137956 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
@@ -8,6 +8,7 @@
package org.elasticsearch.xpack.esql.analysis;
import org.elasticsearch.index.IndexMode;
+import org.elasticsearch.inference.ModelConfigurations;
import org.elasticsearch.inference.TaskType;
import org.elasticsearch.xpack.core.enrich.EnrichPolicy;
import org.elasticsearch.xpack.esql.EsqlTestUtils;
@@ -39,6 +40,8 @@
import static org.elasticsearch.xpack.esql.EsqlTestUtils.TEST_VERIFIER;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.configuration;
import static org.elasticsearch.xpack.esql.EsqlTestUtils.emptyInferenceResolution;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
public final class AnalyzerTestUtils {
@@ -189,12 +192,18 @@ public static EnrichResolution defaultEnrichResolution() {
public static InferenceResolution defaultInferenceResolution() {
return InferenceResolution.builder()
- .withResolvedInference(new ResolvedInference("reranking-inference-id", TaskType.RERANK))
- .withResolvedInference(new ResolvedInference("completion-inference-id", TaskType.COMPLETION))
+ .withResolvedInference(mockedResolvedInference("reranking-inference-id", TaskType.RERANK))
+ .withResolvedInference(mockedResolvedInference("completion-inference-id", TaskType.COMPLETION))
.withError("error-inference-id", "error with inference resolution")
.build();
}
+ private static ResolvedInference mockedResolvedInference(String inferenceId, TaskType taskType) {
+ ModelConfigurations modelConfigurations = mock(ModelConfigurations.class);
+ when(modelConfigurations.getTaskType()).thenReturn(taskType);
+ return new ResolvedInference(inferenceId, modelConfigurations);
+ }
+
public static void loadEnrichPolicyResolution(
EnrichResolution enrich,
String policyType,
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java
index 00f20b9376a6f..4e256a1563af2 100644
--- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java
@@ -798,6 +798,7 @@ public static void testFunctionInfo() {
Set returnTypes = Arrays.stream(description.returnType())
.filter(t -> DataType.UNDER_CONSTRUCTION.containsKey(DataType.fromNameOrAlias(t)) == false)
.collect(Collectors.toCollection(TreeSet::new));
+
assertEquals(returnFromSignature, returnTypes);
}
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/DenseVectorEmbeddingErrorTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/DenseVectorEmbeddingErrorTests.java
new file mode 100644
index 0000000000000..4158b2cc9380e
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/DenseVectorEmbeddingErrorTests.java
@@ -0,0 +1,79 @@
+/*
+ * 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 java.util.stream.Stream;
+
+import static org.elasticsearch.common.logging.LoggerMessageFormat.format;
+import static org.hamcrest.Matchers.equalTo;
+
+public class DenseVectorEmbeddingErrorTests extends ErrorsForCasesWithoutExamplesTestCase {
+
+ @Before
+ public void checkCapability() {
+ assumeTrue("DENSE_VECTOR_EMBEDDING_FUNCTION is not enabled", EsqlCapabilities.Cap.DENSE_VECTOR_EMBEDDING_FUNCTION.isEnabled());
+ }
+
+ @Override
+ protected List cases() {
+ return paramsToSuppliers(DenseVectorEmbeddingFunctionTests.parameters());
+ }
+
+ @Override
+ protected Stream> testCandidates(List cases, Set> valid) {
+ // Don't test null, as it is not allowed but the expected message is not a type error - so we check it separately in VerifierTests
+ return super.testCandidates(cases, valid).filter(sig -> false == sig.contains(DataType.NULL));
+ }
+
+ @Override
+ protected Expression build(Source source, List args) {
+ return new DenseVectorEmbeddingFunction(source, args.get(0), args.get(1));
+ }
+
+ @Override
+ protected Matcher expectedTypeErrorMatcher(List> validPerPosition, List signature) {
+ return equalTo(errorMessageString(validPerPosition, signature, (v, p) -> "string"));
+ }
+
+ private static String errorMessageString(
+ List> validPerPosition,
+ List signature,
+ AbstractFunctionTestCase.PositionalErrorMessageSupplier positionalErrorMessageSupplier
+ ) {
+ for (int i = 0; i < signature.size(); i++) {
+ if (validPerPosition.get(i).contains(signature.get(i)) == false) {
+ // Map expressions have different error messages
+ if (i == signature.size() - 1) {
+ return format(
+ null,
+ "{} argument of [{}] must be a map expression, received []",
+ TypeResolutions.ParamOrdinal.fromIndex(i).name().toLowerCase(Locale.ROOT),
+ sourceForSignature(signature)
+ );
+ }
+ break;
+ }
+ }
+
+ return typeErrorMessage(true, validPerPosition, signature, positionalErrorMessageSupplier);
+ }
+}
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/DenseVectorEmbeddingFunctionTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/DenseVectorEmbeddingFunctionTests.java
new file mode 100644
index 0000000000000..c493cf7177043
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/DenseVectorEmbeddingFunctionTests.java
@@ -0,0 +1,103 @@
+/*
+ * 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.apache.lucene.util.BytesRef;
+import org.elasticsearch.xpack.core.security.authc.support.mapper.expressiondsl.FieldExpression;
+import org.elasticsearch.xpack.esql.action.EsqlCapabilities;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.Literal;
+import org.elasticsearch.xpack.esql.core.expression.MapExpression;
+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 java.util.stream.Stream;
+
+import static org.elasticsearch.xpack.esql.core.tree.Source.EMPTY;
+import static org.elasticsearch.xpack.esql.core.type.DataType.DENSE_VECTOR;
+import static org.elasticsearch.xpack.esql.core.type.DataType.UNSUPPORTED;
+import static org.hamcrest.Matchers.equalTo;
+
+@FunctionName("text_dense_vector_embedding")
+public class DenseVectorEmbeddingFunctionTests extends AbstractFunctionTestCase {
+ @Before
+ public void checkCapability() {
+ assumeTrue("DENSE_VECTOR_EMBEDDING_FUNCTION is not enabled", EsqlCapabilities.Cap.DENSE_VECTOR_EMBEDDING_FUNCTION.isEnabled());
+ }
+
+ public DenseVectorEmbeddingFunctionTests(@Name("TestCase") Supplier testCaseSupplier) {
+ this.testCase = testCaseSupplier.get();
+ }
+
+ @ParametersFactory
+ public static Iterable parameters() {
+ List> paramList = List.of(List.of());
+ List>> paramSuppliers = List.of(
+ () -> Stream.concat(
+ textTestCases("inputText", new BytesRef("input text value")).stream(),
+ textTestCases("inputText", new FieldExpression("foo", List.of(new FieldExpression.FieldValue("foo")))).stream()
+ ).toList(),
+ DenseVectorEmbeddingFunctionTests::optionsTestCases
+ );
+
+ for (Supplier> paramSupplier : paramSuppliers) {
+ List> newParams = new ArrayList<>();
+ for (List params : paramList) {
+ for (TestCaseSupplier.TypedData value : paramSupplier.get()) {
+ List combination = new ArrayList<>(params);
+ combination.add(value);
+ newParams.add(combination);
+ }
+ }
+ paramList = newParams;
+ }
+
+ return parameterSuppliersFromTypedData(paramList.stream().map(args -> {
+ List dataTypes = args.stream().map(TestCaseSupplier.TypedData::type).toList();
+ return new TestCaseSupplier(
+ "TextDenseVectorEmbedding[" + dataTypes + "]",
+ dataTypes,
+ () -> new TestCaseSupplier.TestCase(args, Matchers.blankOrNullString(), DENSE_VECTOR, equalTo(true))
+ );
+ }).toList());
+ }
+
+ @Override
+ protected Expression build(Source source, List args) {
+ return new DenseVectorEmbeddingFunction(source, args.get(0), args.get(1));
+ }
+
+ private static List textTestCases(String name, Object value) {
+ return DataType.stringTypes().stream().map(dataType -> new TestCaseSupplier.TypedData(value, dataType, name)).toList();
+ }
+
+ private static List optionsTestCases() {
+ Literal inferenceIdOptionName = Literal.keyword(EMPTY, "inference_id");
+ return DataType.stringTypes()
+ .stream()
+ .map(
+ dataType -> new TestCaseSupplier.TypedData(
+ new MapExpression(EMPTY, List.of(inferenceIdOptionName, new Literal(EMPTY, new BytesRef("inferenceId"), dataType))),
+ UNSUPPORTED,
+ "options"
+ )
+ )
+ .toList();
+ }
+}
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/DenseVectorSerializationTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/DenseVectorSerializationTests.java
new file mode 100644
index 0000000000000..57661cf9a3ae0
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/DenseVectorSerializationTests.java
@@ -0,0 +1,39 @@
+/*
+ * 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.util.ArrayList;
+import java.util.List;
+
+public class DenseVectorSerializationTests extends AbstractExpressionSerializationTests {
+
+ @Before
+ public void checkCapability() {
+ assumeTrue("DENSE_VECTOR_EMBEDDING_FUNCTION is not enabled", EsqlCapabilities.Cap.DENSE_VECTOR_EMBEDDING_FUNCTION.isEnabled());
+ }
+
+ protected DenseVectorEmbeddingFunction createTestInstance() {
+ Source source = randomSource();
+ Expression inputText = randomChild();
+ Expression options = randomChild();
+ return new DenseVectorEmbeddingFunction(source, inputText, options);
+ }
+
+ @Override
+ protected DenseVectorEmbeddingFunction mutateInstance(DenseVectorEmbeddingFunction instance) {
+ List newChildren = new ArrayList<>(instance.children());
+ newChildren.set(randomInt(newChildren.size() - 1), randomChild());
+ return instance.replaceChildren(newChildren);
+ }
+}
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceOperatorTestCase.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceOperatorTestCase.java
index c49e301968aa0..9a000778d4003 100644
--- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceOperatorTestCase.java
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceOperatorTestCase.java
@@ -93,7 +93,7 @@ protected Page createPage(int positionOffset, int length) {
if (randomInt() % 100 == 0) {
builder.appendNull();
} else {
- builder.appendBytesRef(new BytesRef(randomAlphaOfLength(10)));
+ builder.appendBytesRef(new BytesRef(randomAlphaOfLength(randomInt(10))));
}
}
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceRunnerTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceRunnerTests.java
index ef7b3984bd532..b63a4e11e60c0 100644
--- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceRunnerTests.java
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceRunnerTests.java
@@ -26,11 +26,15 @@
import org.elasticsearch.xpack.esql.core.tree.Source;
import org.elasticsearch.xpack.esql.plan.logical.inference.InferencePlan;
import org.elasticsearch.xpack.esql.plugin.EsqlPlugin;
+import org.hamcrest.FeatureMatcher;
+import org.hamcrest.Matcher;
import org.junit.After;
import org.junit.Before;
import java.util.List;
+import java.util.Set;
+import static org.hamcrest.Matchers.allOf;
import static org.hamcrest.Matchers.contains;
import static org.hamcrest.Matchers.empty;
import static org.hamcrest.Matchers.equalTo;
@@ -65,31 +69,30 @@ public void shutdownThreadPool() {
public void testResolveInferenceIds() throws Exception {
InferenceRunner inferenceRunner = new InferenceRunner(mockClient(), threadPool);
- List> inferencePlans = List.of(mockInferencePlan("rerank-plan"));
+ Set inferenceIds = Set.of("rerank-plan");
SetOnce inferenceResolutionSetOnce = new SetOnce<>();
- inferenceRunner.resolveInferenceIds(inferencePlans, ActionListener.wrap(inferenceResolutionSetOnce::set, e -> {
+ inferenceRunner.resolveInferenceIds(inferenceIds, ActionListener.wrap(inferenceResolutionSetOnce::set, e -> {
throw new RuntimeException(e);
}));
assertBusy(() -> {
InferenceResolution inferenceResolution = inferenceResolutionSetOnce.get();
assertNotNull(inferenceResolution);
- assertThat(inferenceResolution.resolvedInferences(), contains(new ResolvedInference("rerank-plan", TaskType.RERANK)));
+ assertThat(
+ inferenceResolution.resolvedInferences(),
+ contains(allOf(inferenceId(equalTo("rerank-plan")), taskType(equalTo(TaskType.RERANK))))
+ );
assertThat(inferenceResolution.hasError(), equalTo(false));
});
}
public void testResolveMultipleInferenceIds() throws Exception {
InferenceRunner inferenceRunner = new InferenceRunner(mockClient(), threadPool);
- List> inferencePlans = List.of(
- mockInferencePlan("rerank-plan"),
- mockInferencePlan("rerank-plan"),
- mockInferencePlan("completion-plan")
- );
+ Set inferenceIds = Set.of("rerank-plan", "completion-plan");
SetOnce inferenceResolutionSetOnce = new SetOnce<>();
- inferenceRunner.resolveInferenceIds(inferencePlans, ActionListener.wrap(inferenceResolutionSetOnce::set, e -> {
+ inferenceRunner.resolveInferenceIds(inferenceIds, ActionListener.wrap(inferenceResolutionSetOnce::set, e -> {
throw new RuntimeException(e);
}));
@@ -100,8 +103,8 @@ public void testResolveMultipleInferenceIds() throws Exception {
assertThat(
inferenceResolution.resolvedInferences(),
contains(
- new ResolvedInference("rerank-plan", TaskType.RERANK),
- new ResolvedInference("completion-plan", TaskType.COMPLETION)
+ allOf(inferenceId(equalTo("rerank-plan")), taskType(equalTo(TaskType.RERANK))),
+ allOf(inferenceId(equalTo("completion-plan")), taskType(equalTo(TaskType.COMPLETION)))
)
);
assertThat(inferenceResolution.hasError(), equalTo(false));
@@ -110,11 +113,11 @@ public void testResolveMultipleInferenceIds() throws Exception {
public void testResolveMissingInferenceIds() throws Exception {
InferenceRunner inferenceRunner = new InferenceRunner(mockClient(), threadPool);
- List> inferencePlans = List.of(mockInferencePlan("missing-plan"));
+ Set inferenceIds = Set.of("missing-plan");
SetOnce inferenceResolutionSetOnce = new SetOnce<>();
- inferenceRunner.resolveInferenceIds(inferencePlans, ActionListener.wrap(inferenceResolutionSetOnce::set, e -> {
+ inferenceRunner.resolveInferenceIds(inferenceIds, ActionListener.wrap(inferenceResolutionSetOnce::set, e -> {
throw new RuntimeException(e);
}));
@@ -184,4 +187,22 @@ private static InferencePlan mockInferencePlan(String inferenceId) {
when(plan.inferenceId()).thenReturn(Literal.keyword(Source.EMPTY, inferenceId));
return plan;
}
+
+ private FeatureMatcher inferenceId(Matcher matcher) {
+ return new FeatureMatcher<>(matcher, "inference id", "inferenceId") {
+ @Override
+ protected String featureValueOf(ResolvedInference resolution) {
+ return resolution.inferenceId();
+ }
+ };
+ }
+
+ private FeatureMatcher taskType(Matcher matcher) {
+ return new FeatureMatcher<>(matcher, "task type", "taskType") {
+ @Override
+ protected TaskType featureValueOf(ResolvedInference resolution) {
+ return resolution.taskType();
+ }
+ };
+ }
}
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/ResolvedInferenceTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/ResolvedInferenceTests.java
deleted file mode 100644
index b4dfd87224a3a..0000000000000
--- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/ResolvedInferenceTests.java
+++ /dev/null
@@ -1,41 +0,0 @@
-/*
- * 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.TransportVersion;
-import org.elasticsearch.inference.TaskType;
-import org.elasticsearch.test.AbstractWireTestCase;
-import org.elasticsearch.test.ESTestCase;
-
-import java.io.IOException;
-
-public class ResolvedInferenceTests extends AbstractWireTestCase {
-
- @Override
- protected ResolvedInference createTestInstance() {
- return new ResolvedInference(randomIdentifier(), randomTaskType());
- }
-
- @Override
- protected ResolvedInference mutateInstance(ResolvedInference instance) throws IOException {
- if (randomBoolean()) {
- return new ResolvedInference(randomValueOtherThan(instance.inferenceId(), ESTestCase::randomIdentifier), instance.taskType());
- }
-
- return new ResolvedInference(instance.inferenceId(), randomValueOtherThan(instance.taskType(), this::randomTaskType));
- }
-
- @Override
- protected ResolvedInference copyInstance(ResolvedInference instance, TransportVersion version) throws IOException {
- return copyInstance(instance, getNamedWriteableRegistry(), (out, v) -> v.writeTo(out), in -> new ResolvedInference(in), version);
- }
-
- private TaskType randomTaskType() {
- return randomFrom(TaskType.values());
- }
-}
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/embedding/DenseEmbeddingOperatorOutputBuilderTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/embedding/DenseEmbeddingOperatorOutputBuilderTests.java
new file mode 100644
index 0000000000000..13a31b3f1a74d
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/embedding/DenseEmbeddingOperatorOutputBuilderTests.java
@@ -0,0 +1,159 @@
+/*
+ * 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.embedding;
+
+import com.carrotsearch.randomizedtesting.annotations.Name;
+import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
+
+import org.elasticsearch.compute.data.Block;
+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.results.TextEmbeddingBitResults;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingByteResults;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingFloatResults;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingResults;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class DenseEmbeddingOperatorOutputBuilderTests extends ComputeTestCase {
+
+ private static final List DIMENSIONS = List.of(1, 32, 128, 512, 2048, 5096);
+ private static final List INPUT_SIZES = List.of(10, 100, 1_000, 10_000);
+ private static final List BATCH_SIZES = List.of(1, 10, 100, 1000);
+ private static final List>> EMBEDDING_TYPES = List.of(
+ TextEmbeddingBitResults.class,
+ TextEmbeddingByteResults.class,
+ TextEmbeddingFloatResults.class
+ );
+
+ private final int dimensions;
+ private final int inputPageSize;
+ private final int batchSize;
+ private final Class> embeddingType;
+
+ @ParametersFactory
+ public static Iterable parameters() {
+ List params = new ArrayList<>();
+ params.add(new Object[] {});
+
+ for (List axis : List.of(DIMENSIONS, INPUT_SIZES, BATCH_SIZES, EMBEDDING_TYPES)) {
+
+ List newParams = new ArrayList<>();
+ for (Object[] combination : params) {
+ for (Object element : axis) {
+ Object[] newCombination = new Object[combination.length + 1];
+ System.arraycopy(combination, 0, newCombination, 0, combination.length);
+ newCombination[newCombination.length - 1] = element;
+ newParams.add(newCombination);
+ }
+ }
+ params = newParams;
+ }
+
+ return params;
+ }
+
+ public DenseEmbeddingOperatorOutputBuilderTests(
+ @Name("dimensions") int dimensions,
+ @Name("inputPageSize") int inputPageSize,
+ @Name("batchSize") int batchSize,
+ @Name("embeddingType") Class embeddingType
+ ) {
+ this.dimensions = dimensions;
+ this.inputPageSize = inputPageSize;
+ this.batchSize = batchSize;
+ this.embeddingType = embeddingType;
+ }
+
+ public void testOutput() {
+ final Page inputPage = randomInputPage(inputPageSize, between(1, 20));
+ try (
+ DenseEmbeddingOperatorOutputBuilder outputBuilder = new DenseEmbeddingOperatorOutputBuilder(
+ blockFactory().newFloatBlockBuilder(inputPageSize * dimensions),
+ inputPage,
+ dimensions
+ )
+ ) {
+ for (int currentPos = 0; currentPos < inputPage.getPositionCount(); currentPos += batchSize) {
+ outputBuilder.addInferenceResponse(
+ DenseEmbeddingUtils.inferenceResponse(
+ currentPos,
+ Math.min(inputPage.getPositionCount() - currentPos, batchSize),
+ embeddingType,
+ dimensions
+ )
+ );
+ }
+
+ final Page outputPage = outputBuilder.buildOutput();
+ try {
+ assertThat(outputPage.getPositionCount(), equalTo(inputPage.getPositionCount()));
+ assertThat(outputPage.getBlockCount(), equalTo(inputPage.getBlockCount() + 1));
+ assertOutputContent(outputPage.getBlock(inputPage.getBlockCount()));
+ } finally {
+ outputPage.releaseBlocks();
+ }
+
+ } finally {
+ inputPage.releaseBlocks();
+ }
+ }
+
+ private void assertOutputContent(FloatBlock block) {
+
+ for (int currentPos = 0; currentPos < block.getPositionCount(); currentPos++) {
+ assertThat(block.isNull(currentPos), equalTo(false));
+ assertThat(block.getValueCount(currentPos), equalTo(dimensions));
+ float[] expectedEmbeddingValues = new float[dimensions];
+ if (embeddingType.equals(TextEmbeddingByteResults.class) || embeddingType.equals(TextEmbeddingBitResults.class)) {
+ byte[] bytesValues = embeddingType.equals(TextEmbeddingBitResults.class)
+ ? DenseEmbeddingUtils.toBitArray(currentPos, dimensions)
+ : DenseEmbeddingUtils.toByteArray(currentPos, dimensions);
+ assertThat(bytesValues.length, equalTo(dimensions));
+ for (int i = 0; i < bytesValues.length; i++) {
+ expectedEmbeddingValues[i] = bytesValues[i];
+ }
+ } else if (embeddingType.equals(TextEmbeddingFloatResults.class)) {
+ expectedEmbeddingValues = DenseEmbeddingUtils.toFloatArray(currentPos, dimensions);
+ }
+
+ for (int valueIndex = 0; valueIndex < block.getValueCount(currentPos); valueIndex++) {
+ assertThat(block.getFloat(block.getFirstValueIndex(currentPos) + valueIndex), equalTo(expectedEmbeddingValues[valueIndex]));
+ }
+ }
+ }
+
+ 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.randomElementType(),
+ 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/embedding/DenseEmbeddingOperatorRequestIteratorTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/embedding/DenseEmbeddingOperatorRequestIteratorTests.java
new file mode 100644
index 0000000000000..7570dabb40cad
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/embedding/DenseEmbeddingOperatorRequestIteratorTests.java
@@ -0,0 +1,65 @@
+/*
+ * 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.embedding;
+
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.compute.test.ComputeTestCase;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+
+import java.util.List;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class DenseEmbeddingOperatorRequestIteratorTests extends ComputeTestCase {
+
+ public void testIterateSmallInput() {
+ assertIterate(between(1, 100), randomIntBetween(1, 1_000));
+ }
+
+ public void testIterateLargeInput() {
+ assertIterate(between(10_000, 100_000), randomIntBetween(1, 1_000));
+ }
+
+ private void assertIterate(int size, int batchSize) {
+ final String inferenceId = randomIdentifier();
+
+ try (
+ BytesRefBlock inputBlock = randomInputBlock(size);
+ DenseEmbeddingOperatorRequestIterator requestIterator = new DenseEmbeddingOperatorRequestIterator(
+ inputBlock,
+ inferenceId,
+ batchSize
+ )
+ ) {
+ BytesRef scratch = new BytesRef();
+
+ for (int currentPos = 0; requestIterator.hasNext();) {
+ InferenceAction.Request request = requestIterator.next();
+
+ assertThat(request.getInferenceEntityId(), equalTo(inferenceId));
+ List inputs = request.getInput();
+ for (String input : inputs) {
+ scratch = inputBlock.getBytesRef(inputBlock.getFirstValueIndex(currentPos), scratch);
+ assertThat(input, equalTo(scratch.utf8ToString()));
+ currentPos++;
+ }
+ }
+ }
+ }
+
+ 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();
+ }
+ }
+}
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/embedding/DenseEmbeddingOperatorTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/embedding/DenseEmbeddingOperatorTests.java
new file mode 100644
index 0000000000000..d0eafc4154835
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/embedding/DenseEmbeddingOperatorTests.java
@@ -0,0 +1,162 @@
+/*
+ * 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.embedding;
+
+import com.carrotsearch.randomizedtesting.annotations.Name;
+import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
+
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.common.lucene.BytesRefs;
+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.TextEmbeddingBitResults;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingByteResults;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingFloatResults;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingResults;
+import org.elasticsearch.xpack.esql.inference.InferenceOperatorTestCase;
+import org.hamcrest.Matcher;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.elasticsearch.xpack.esql.inference.embedding.DenseEmbeddingUtils.toBitArray;
+import static org.elasticsearch.xpack.esql.inference.embedding.DenseEmbeddingUtils.toByteArray;
+import static org.elasticsearch.xpack.esql.inference.embedding.DenseEmbeddingUtils.toFloatArray;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.hasSize;
+
+public class DenseEmbeddingOperatorTests extends InferenceOperatorTestCase> {
+
+ private static final String SIMPLE_INFERENCE_ID = "test_dense_embedding";
+
+ private static final List DIMENSIONS = List.of(1, 32, 128, 512, 2048, 5096);
+ private static final List>> EMBEDDING_TYPES = List.of(
+ TextEmbeddingBitResults.class,
+ TextEmbeddingByteResults.class,
+ TextEmbeddingFloatResults.class
+ );
+
+ @ParametersFactory
+ public static Iterable parameters() {
+ List params = new ArrayList<>();
+ params.add(new Object[] {});
+
+ for (List axis : List.of(DIMENSIONS, EMBEDDING_TYPES)) {
+
+ List newParams = new ArrayList<>();
+ for (Object[] combination : params) {
+ for (Object element : axis) {
+ Object[] newCombination = new Object[combination.length + 1];
+ System.arraycopy(combination, 0, newCombination, 0, combination.length);
+ newCombination[newCombination.length - 1] = element;
+ newParams.add(newCombination);
+ }
+ }
+ params = newParams;
+ }
+
+ return params;
+ }
+
+ private final int dimensions;
+ private final Class> embeddingType;
+
+ public DenseEmbeddingOperatorTests(
+ @Name("dimensions") int dimensions,
+ @Name("embeddingType") Class> embeddingType
+ ) {
+ this.dimensions = dimensions;
+ this.embeddingType = embeddingType;
+ }
+
+ @Override
+ protected Operator.OperatorFactory simple(SimpleOptions options) {
+ return new DenseEmbeddingOperator.Factory(mockedSimpleInferenceRunner(), SIMPLE_INFERENCE_ID, dimensions, evaluatorFactory(0));
+ }
+
+ @Override
+ protected void assertSimpleOutput(List inputPages, List resultPages) {
+ assertThat(inputPages, hasSize(resultPages.size()));
+
+ for (int pageId = 0; pageId < inputPages.size(); pageId++) {
+ Page inputPage = inputPages.get(pageId);
+ Page resultPage = resultPages.get(pageId);
+
+ assertThat(resultPage.getPositionCount(), equalTo(inputPage.getPositionCount()));
+ assertThat(resultPage.getBlockCount(), equalTo(inputPage.getBlockCount() + 1));
+
+ for (int channel = 0; channel < inputPage.getBlockCount(); channel++) {
+ Block inputBlock = inputPage.getBlock(channel);
+ Block resultBlock = resultPage.getBlock(channel);
+
+ assertThat(resultBlock.getPositionCount(), equalTo(resultPage.getPositionCount()));
+ assertThat(resultBlock.elementType(), equalTo(inputBlock.elementType()));
+
+ if (channel == 0) {
+ assertEmbeddingContent((BytesRefBlock) inputBlock, resultPage.getBlock(inputPage.getBlockCount()));
+ }
+ }
+ }
+ }
+
+ private void assertEmbeddingContent(BytesRefBlock inputBlock, FloatBlock embeddingsBlock) {
+ BytesRef scratch = new BytesRef();
+ for (int position = 0; position < inputBlock.getPositionCount(); position++) {
+ String textContent = "";
+ if (inputBlock.isNull(position) == false) {
+ scratch = inputBlock.getBytesRef(inputBlock.getFirstValueIndex(position), scratch);
+ textContent = BytesRefs.toString(scratch);
+ }
+ float[] expectedEmbedding = new float[dimensions];
+ if (embeddingType.equals(TextEmbeddingFloatResults.class)) {
+ expectedEmbedding = toFloatArray(textContent.length(), dimensions);
+ }
+ if (embeddingType.equals(TextEmbeddingByteResults.class) || embeddingType.equals(TextEmbeddingBitResults.class)) {
+ byte[] bytesValues = embeddingType.equals(TextEmbeddingBitResults.class)
+ ? toBitArray(textContent.length(), dimensions)
+ : toByteArray(textContent.length(), dimensions);
+ assertThat(bytesValues.length, equalTo(dimensions));
+ for (int i = 0; i < bytesValues.length; i++) {
+ expectedEmbedding[i] = bytesValues[i];
+ }
+ }
+
+ assertThat(embeddingsBlock.isNull(position), equalTo(false));
+ assertThat(embeddingsBlock.getValueCount(position), equalTo(dimensions));
+ for (int valueIndex = 0; valueIndex < dimensions; valueIndex++) {
+ assertThat(
+ embeddingsBlock.getFloat(embeddingsBlock.getFirstValueIndex(position) + valueIndex),
+ equalTo(expectedEmbedding[valueIndex])
+ );
+ }
+ }
+ }
+
+ @Override
+ protected Matcher expectedDescriptionOfSimple() {
+ return expectedToStringOfSimple();
+ }
+
+ @Override
+ protected Matcher expectedToStringOfSimple() {
+ return equalTo("DenseEmbeddingOperator[inference_id=[" + SIMPLE_INFERENCE_ID + "]]");
+ }
+
+ @Override
+ protected TextEmbeddingResults mockInferenceResult(InferenceAction.Request request) {
+ return DenseEmbeddingUtils.inferenceResults(
+ request.getInput().stream().mapToInt(String::length).toArray(),
+ embeddingType,
+ dimensions
+ );
+ }
+}
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/embedding/DenseEmbeddingUtils.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/embedding/DenseEmbeddingUtils.java
new file mode 100644
index 0000000000000..193d9975e7281
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/embedding/DenseEmbeddingUtils.java
@@ -0,0 +1,88 @@
+/*
+ * 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.embedding;
+
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingBitResults;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingByteResults;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingFloatResults;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingResults;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.IntStream;
+
+public class DenseEmbeddingUtils {
+
+ public static byte[] toBitArray(int position, int length) {
+ byte[] bits = new byte[length];
+ for (int i = 0; i < 32 && i < length; i++) {
+ bits[length - 1 - i] = (byte) ((position >> i) & 1);
+ }
+ return bits;
+ }
+
+ public static byte[] toByteArray(int position, int length) {
+ byte[] digits = new byte[length];
+ for (int i = length - 1; i >= 0; i--) {
+ digits[i] = (byte) (position % 10);
+ position /= 10;
+ }
+
+ return digits;
+ }
+
+ public static float[] toFloatArray(int position, int length) {
+ float[] floats = new float[length];
+ floats[0] = position;
+ return floats;
+ }
+
+ public static InferenceAction.Response inferenceResponse(
+ int startPosition,
+ int size,
+ Class> embeddingType,
+ int dimensions
+ ) {
+ return new InferenceAction.Response(
+ inferenceResults(IntStream.range(startPosition, startPosition + size).toArray(), embeddingType, dimensions)
+ );
+ }
+
+ public static TextEmbeddingResults inferenceResults(
+ int[] values,
+ Class> embeddingType,
+ int dimensions
+ ) {
+ if (embeddingType.equals(TextEmbeddingBitResults.class)) {
+ List embeddings = new ArrayList<>();
+ for (int i = 0; i < values.length; i++) {
+ embeddings.add(new TextEmbeddingByteResults.Embedding(toBitArray(values[i], dimensions)));
+ }
+ return new TextEmbeddingBitResults(embeddings);
+ }
+
+ if (embeddingType.equals(TextEmbeddingByteResults.class)) {
+ List embeddings = new ArrayList<>();
+ for (int i = 0; i < values.length; i++) {
+ embeddings.add(new TextEmbeddingByteResults.Embedding(toByteArray(values[i], dimensions)));
+ }
+ return new TextEmbeddingByteResults(embeddings);
+ }
+
+ if (embeddingType.equals(TextEmbeddingFloatResults.class)) {
+ List embeddings = new ArrayList<>();
+ for (int i = 0; i < values.length; i++) {
+ embeddings.add(new TextEmbeddingFloatResults.Embedding(toFloatArray(values[i], dimensions)));
+ }
+ return new TextEmbeddingFloatResults(embeddings);
+ }
+
+ throw new AssertionError("Unexpected Embedding type [" + embeddingType + "]");
+ }
+}
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/logical/inference/embedding/DenseVectorEmbeddingSerializationTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/logical/inference/embedding/DenseVectorEmbeddingSerializationTests.java
new file mode 100644
index 0000000000000..f82a090d0801a
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/logical/inference/embedding/DenseVectorEmbeddingSerializationTests.java
@@ -0,0 +1,73 @@
+/*
+ * 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.plan.logical.inference.embedding;
+
+import org.elasticsearch.xpack.esql.core.expression.Attribute;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.Literal;
+import org.elasticsearch.xpack.esql.core.type.DataType;
+import org.elasticsearch.xpack.esql.expression.function.ReferenceAttributeTests;
+import org.elasticsearch.xpack.esql.plan.logical.AbstractLogicalPlanSerializationTests;
+import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+
+import java.io.IOException;
+
+import static org.elasticsearch.xpack.esql.core.tree.Source.EMPTY;
+
+public class DenseVectorEmbeddingSerializationTests extends AbstractLogicalPlanSerializationTests {
+
+ @Override
+ protected DenseVectorEmbedding createTestInstance() {
+ return new DenseVectorEmbedding(
+ randomSource(),
+ randomChild(0),
+ randomInferenceId(),
+ randomDimensions(),
+ randomInput(),
+ randomTargetField()
+ );
+ }
+
+ @Override
+ protected DenseVectorEmbedding mutateInstance(DenseVectorEmbedding instance) throws IOException {
+ LogicalPlan child = instance.child();
+ Expression inferenceId = instance.inferenceId();
+ Expression dimensions = instance.dimensions();
+ Expression input = instance.input();
+ Attribute targetField = instance.embeddingField();
+
+ switch (between(0, 4)) {
+ case 0 -> child = randomValueOtherThan(child, () -> randomChild(0));
+ case 1 -> inferenceId = randomValueOtherThan(inferenceId, this::randomInferenceId);
+ case 2 -> dimensions = randomValueOtherThan(instance.dimensions(), this::randomDimensions);
+ case 3 -> input = randomValueOtherThan(input, this::randomInput);
+ case 4 -> targetField = randomValueOtherThan(targetField, this::randomTargetField);
+ }
+ return new DenseVectorEmbedding(instance.source(), child, inferenceId, dimensions, input, targetField);
+ }
+
+ private Literal randomInferenceId() {
+ return Literal.keyword(EMPTY, randomIdentifier());
+ }
+
+ private Expression randomInput() {
+ return randomBoolean() ? Literal.keyword(EMPTY, randomIdentifier()) : randomAttribute();
+ }
+
+ private Attribute randomTargetField() {
+ return ReferenceAttributeTests.randomReferenceAttribute(randomBoolean());
+ }
+
+ private Attribute randomAttribute() {
+ return ReferenceAttributeTests.randomReferenceAttribute(randomBoolean());
+ }
+
+ private Expression randomDimensions() {
+ return new Literal(EMPTY, randomInt(), DataType.INTEGER);
+ }
+}
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/physical/inference/embedding/DenseVectorEmbeddingExecTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/physical/inference/embedding/DenseVectorEmbeddingExecTests.java
new file mode 100644
index 0000000000000..c079860869920
--- /dev/null
+++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/physical/inference/embedding/DenseVectorEmbeddingExecTests.java
@@ -0,0 +1,73 @@
+/*
+ * 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.plan.physical.inference.embedding;
+
+import org.elasticsearch.xpack.esql.core.expression.Attribute;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
+import org.elasticsearch.xpack.esql.core.expression.Literal;
+import org.elasticsearch.xpack.esql.core.type.DataType;
+import org.elasticsearch.xpack.esql.expression.function.ReferenceAttributeTests;
+import org.elasticsearch.xpack.esql.plan.physical.AbstractPhysicalPlanSerializationTests;
+import org.elasticsearch.xpack.esql.plan.physical.PhysicalPlan;
+
+import java.io.IOException;
+
+import static org.elasticsearch.xpack.esql.core.tree.Source.EMPTY;
+
+public class DenseVectorEmbeddingExecTests extends AbstractPhysicalPlanSerializationTests {
+
+ @Override
+ protected DenseVectorEmbeddingExec createTestInstance() {
+ return new DenseVectorEmbeddingExec(
+ randomSource(),
+ randomChild(0),
+ randomInferenceId(),
+ randomDimensions(),
+ randomInput(),
+ randomTargetField()
+ );
+ }
+
+ @Override
+ protected DenseVectorEmbeddingExec mutateInstance(DenseVectorEmbeddingExec instance) throws IOException {
+ PhysicalPlan child = instance.child();
+ Expression inferenceId = instance.inferenceId();
+ Expression dimensions = instance.dimensions();
+ Expression input = instance.input();
+ Attribute targetField = instance.targetField();
+
+ switch (between(0, 4)) {
+ case 0 -> child = randomValueOtherThan(child, () -> randomChild(0));
+ case 1 -> inferenceId = randomValueOtherThan(inferenceId, this::randomInferenceId);
+ case 2 -> dimensions = randomValueOtherThan(input, this::randomAttribute);
+ case 3 -> input = randomValueOtherThan(input, this::randomInput);
+ case 4 -> targetField = randomValueOtherThan(targetField, this::randomTargetField);
+ }
+ return new DenseVectorEmbeddingExec(instance.source(), child, inferenceId, dimensions, input, targetField);
+ }
+
+ private Literal randomInferenceId() {
+ return Literal.keyword(EMPTY, randomIdentifier());
+ }
+
+ private Expression randomInput() {
+ return randomBoolean() ? Literal.keyword(EMPTY, randomIdentifier()) : randomAttribute();
+ }
+
+ private Attribute randomTargetField() {
+ return ReferenceAttributeTests.randomReferenceAttribute(randomBoolean());
+ }
+
+ private Attribute randomAttribute() {
+ return ReferenceAttributeTests.randomReferenceAttribute(randomBoolean());
+ }
+
+ private Expression randomDimensions() {
+ return new Literal(EMPTY, randomInt(), DataType.INTEGER);
+ }
+}