Add text embedding output builder.

Author: afoucret

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/results/TextEmbeddingByteResults.java

@@ -155,7 +155,7 @@ public String toString() {
             return Strings.toString(this);
         }
 
-        float[] toFloatArray() {
+        public float[] toFloatArray() {
             float[] floatArray = new float[values.length];
             for (int i = 0; i < values.length; i++) {
                 floatArray[i] = ((Byte) values[i]).floatValue();

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/textembedding/TextEmbeddingOperatorOutputBuilder.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.inference.textembedding;
+
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.FloatBlock;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.core.Releasables;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingByteResults;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingFloatResults;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingResults;
+import org.elasticsearch.xpack.esql.inference.InferenceOperator;
+
+/**
+ * {@link TextEmbeddingOperatorOutputBuilder} builds the output page for text embedding by converting
+ * {@link TextEmbeddingResults} into a {@link FloatBlock} containing dense vector embeddings.
+ */
+public class TextEmbeddingOperatorOutputBuilder implements InferenceOperator.OutputBuilder {
+    private final Page inputPage;
+    private final FloatBlock.Builder outputBlockBuilder;
+
+    public TextEmbeddingOperatorOutputBuilder(FloatBlock.Builder outputBlockBuilder, Page inputPage) {
+        this.inputPage = inputPage;
+        this.outputBlockBuilder = outputBlockBuilder;
+    }
+
+    @Override
+    public void close() {
+        Releasables.close(outputBlockBuilder);
+    }
+
+    /**
+     * Adds an inference response to the output builder.
+     *
+     * <p>
+     * If the response is null or not of type {@link TextEmbeddingResults} an {@link IllegalStateException} is thrown.
+     * Else, the embedding vector is added to the output block as a multi-value position.
+     * </p>
+     *
+     * <p>
+     * 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.
+     * </p>
+     */
+    @Override
+    public void addInferenceResponse(InferenceAction.Response inferenceResponse) {
+        if (inferenceResponse == null) {
+            outputBlockBuilder.appendNull();
+            return;
+        }
+
+        TextEmbeddingResults<?> embeddingResults = inferenceResults(inferenceResponse);
+
+        var embeddings = embeddingResults.embeddings();
+        if (embeddings.isEmpty()) {
+            outputBlockBuilder.appendNull();
+            return;
+        }
+
+        float[] embeddingArray = getEmbeddingAsFloatArray(embeddingResults);
+
+        outputBlockBuilder.beginPositionEntry();
+        for (float component : embeddingArray) {
+            outputBlockBuilder.appendFloat(component);
+        }
+        outputBlockBuilder.endPositionEntry();
+    }
+
+    /**
+     * Builds the final output page by appending the embedding output block to the input page.
+     */
+    @Override
+    public Page buildOutput() {
+        Block outputBlock = outputBlockBuilder.build();
+        assert outputBlock.getPositionCount() == inputPage.getPositionCount();
+        return inputPage.appendBlock(outputBlock);
+    }
+
+    private TextEmbeddingResults<?> inferenceResults(InferenceAction.Response inferenceResponse) {
+        return InferenceOperator.OutputBuilder.inferenceResults(inferenceResponse, TextEmbeddingResults.class);
+    }
+
+    /**
+     * Extracts the embedding as a float array from the embedding result.
+     */
+    private float[] getEmbeddingAsFloatArray(TextEmbeddingResults<?> embedding) {
+        return switch (embedding.embeddings().get(0)) {
+            case TextEmbeddingFloatResults.Embedding floatEmbedding -> floatEmbedding.values();
+            case TextEmbeddingByteResults.Embedding byteEmbedding -> byteEmbedding.toFloatArray();
+            default -> throw new IllegalArgumentException(
+                "Unsupported embedding type: "
+                    + embedding.embeddings().get(0).getClass().getName()
+                    + ". Expected TextEmbeddingFloatResults.Embedding or TextEmbeddingByteResults.Embedding."
+            );
+        };
+    }
+}

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/textembedding/TextEmbeddingOperatorOutputBuilderTests.java

@@ -0,0 +1,247 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.inference.textembedding;
+
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.ElementType;
+import org.elasticsearch.compute.data.FloatBlock;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.test.ComputeTestCase;
+import org.elasticsearch.compute.test.RandomBlock;
+import org.elasticsearch.core.Releasables;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingByteResults;
+import org.elasticsearch.xpack.core.inference.results.TextEmbeddingFloatResults;
+
+import java.util.List;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class TextEmbeddingOperatorOutputBuilderTests extends ComputeTestCase {
+
+    public void testBuildSmallOutputWithFloatEmbeddings() throws Exception {
+        assertBuildOutputWithFloatEmbeddings(between(1, 100));
+    }
+
+    public void testBuildLargeOutputWithFloatEmbeddings() throws Exception {
+        assertBuildOutputWithFloatEmbeddings(between(1_000, 10_000));
+    }
+
+    public void testBuildSmallOutputWithByteEmbeddings() throws Exception {
+        assertBuildOutputWithByteEmbeddings(between(1, 100));
+    }
+
+    public void testBuildLargeOutputWithByteEmbeddings() throws Exception {
+        assertBuildOutputWithByteEmbeddings(between(1_000, 10_000));
+    }
+
+    public void testHandleNullResponses() throws Exception {
+        final int size = between(10, 100);
+        final Page inputPage = randomInputPage(size, between(1, 5));
+
+        try (
+            TextEmbeddingOperatorOutputBuilder outputBuilder = new TextEmbeddingOperatorOutputBuilder(
+                blockFactory().newFloatBlockBuilder(size),
+                inputPage
+            )
+        ) {
+            // Add some null responses
+            for (int currentPos = 0; currentPos < inputPage.getPositionCount(); currentPos++) {
+                if (randomBoolean()) {
+                    outputBuilder.addInferenceResponse(null);
+                } else {
+                    float[] embedding = randomFloatEmbedding(randomIntBetween(50, 200));
+                    outputBuilder.addInferenceResponse(createFloatEmbeddingResponse(embedding));
+                }
+            }
+
+            final Page outputPage = outputBuilder.buildOutput();
+            assertThat(outputPage.getPositionCount(), equalTo(inputPage.getPositionCount()));
+            assertThat(outputPage.getBlockCount(), equalTo(inputPage.getBlockCount() + 1));
+
+            FloatBlock outputBlock = (FloatBlock) outputPage.getBlock(outputPage.getBlockCount() - 1);
+            assertThat(outputBlock.getPositionCount(), equalTo(size));
+
+            outputPage.releaseBlocks();
+        }
+
+        allBreakersEmpty();
+    }
+
+    public void testHandleEmptyEmbeddings() throws Exception {
+        final int size = between(5, 50);
+        final Page inputPage = randomInputPage(size, between(1, 3));
+
+        try (
+            TextEmbeddingOperatorOutputBuilder outputBuilder = new TextEmbeddingOperatorOutputBuilder(
+                blockFactory().newFloatBlockBuilder(size),
+                inputPage
+            )
+        ) {
+            // Add responses with empty embeddings
+            for (int currentPos = 0; currentPos < inputPage.getPositionCount(); currentPos++) {
+                outputBuilder.addInferenceResponse(createEmptyFloatEmbeddingResponse());
+            }
+
+            final Page outputPage = outputBuilder.buildOutput();
+            FloatBlock outputBlock = (FloatBlock) outputPage.getBlock(outputPage.getBlockCount() - 1);
+
+            // All positions should be null due to empty embeddings
+            for (int pos = 0; pos < outputBlock.getPositionCount(); pos++) {
+                assertThat(outputBlock.isNull(pos), equalTo(true));
+            }
+
+            outputPage.releaseBlocks();
+        }
+
+        allBreakersEmpty();
+    }
+
+    private void assertBuildOutputWithFloatEmbeddings(int size) throws Exception {
+        final Page inputPage = randomInputPage(size, between(1, 10));
+        final int embeddingDim = randomIntBetween(50, 1536); // Common embedding dimensions
+        final float[][] expectedEmbeddings = new float[size][];
+
+        try (
+            TextEmbeddingOperatorOutputBuilder outputBuilder = new TextEmbeddingOperatorOutputBuilder(
+                blockFactory().newFloatBlockBuilder(size),
+                inputPage
+            )
+        ) {
+            for (int currentPos = 0; currentPos < inputPage.getPositionCount(); currentPos++) {
+                float[] embedding = randomFloatEmbedding(embeddingDim);
+                expectedEmbeddings[currentPos] = embedding;
+                outputBuilder.addInferenceResponse(createFloatEmbeddingResponse(embedding));
+            }
+
+            final Page outputPage = outputBuilder.buildOutput();
+            assertThat(outputPage.getPositionCount(), equalTo(inputPage.getPositionCount()));
+            assertThat(outputPage.getBlockCount(), equalTo(inputPage.getBlockCount() + 1));
+
+            assertFloatEmbeddingContent((FloatBlock) outputPage.getBlock(outputPage.getBlockCount() - 1), expectedEmbeddings);
+
+            outputPage.releaseBlocks();
+        }
+
+        allBreakersEmpty();
+    }
+
+    private void assertBuildOutputWithByteEmbeddings(int size) throws Exception {
+        final Page inputPage = randomInputPage(size, between(1, 10));
+        final int embeddingDim = randomIntBetween(50, 1536);
+        final byte[][] expectedByteEmbeddings = new byte[size][];
+
+        try (
+            TextEmbeddingOperatorOutputBuilder outputBuilder = new TextEmbeddingOperatorOutputBuilder(
+                blockFactory().newFloatBlockBuilder(size),
+                inputPage
+            )
+        ) {
+            for (int currentPos = 0; currentPos < inputPage.getPositionCount(); currentPos++) {
+                byte[] embedding = randomByteEmbedding(embeddingDim);
+                expectedByteEmbeddings[currentPos] = embedding;
+                outputBuilder.addInferenceResponse(createByteEmbeddingResponse(embedding));
+            }
+
+            final Page outputPage = outputBuilder.buildOutput();
+            assertThat(outputPage.getPositionCount(), equalTo(inputPage.getPositionCount()));
+            assertThat(outputPage.getBlockCount(), equalTo(inputPage.getBlockCount() + 1));
+
+            assertByteEmbeddingContent((FloatBlock) outputPage.getBlock(outputPage.getBlockCount() - 1), expectedByteEmbeddings);
+
+            outputPage.releaseBlocks();
+        }
+
+        allBreakersEmpty();
+    }
+
+    private void assertFloatEmbeddingContent(FloatBlock block, float[][] expectedEmbeddings) {
+        for (int currentPos = 0; currentPos < block.getPositionCount(); currentPos++) {
+            assertThat(block.isNull(currentPos), equalTo(false));
+            assertThat(block.getValueCount(currentPos), equalTo(expectedEmbeddings[currentPos].length));
+
+            int firstValueIndex = block.getFirstValueIndex(currentPos);
+            for (int i = 0; i < expectedEmbeddings[currentPos].length; i++) {
+                float actualValue = block.getFloat(firstValueIndex + i);
+                float expectedValue = expectedEmbeddings[currentPos][i];
+                assertThat(actualValue, equalTo(expectedValue));
+            }
+        }
+    }
+
+    private void assertByteEmbeddingContent(FloatBlock block, byte[][] expectedByteEmbeddings) {
+        for (int currentPos = 0; currentPos < block.getPositionCount(); currentPos++) {
+            assertThat(block.isNull(currentPos), equalTo(false));
+            assertThat(block.getValueCount(currentPos), equalTo(expectedByteEmbeddings[currentPos].length));
+
+            int firstValueIndex = block.getFirstValueIndex(currentPos);
+            for (int i = 0; i < expectedByteEmbeddings[currentPos].length; i++) {
+                float actualValue = block.getFloat(firstValueIndex + i);
+                // Convert byte to float the same way as TextEmbeddingByteResults.Embedding.toFloatArray()
+                float expectedValue = expectedByteEmbeddings[currentPos][i];
+                assertThat(actualValue, equalTo(expectedValue));
+            }
+        }
+    }
+
+    private float[] randomFloatEmbedding(int dimension) {
+        float[] embedding = new float[dimension];
+        for (int i = 0; i < dimension; i++) {
+            embedding[i] = randomFloat();
+        }
+        return embedding;
+    }
+
+    private byte[] randomByteEmbedding(int dimension) {
+        byte[] embedding = new byte[dimension];
+        for (int i = 0; i < dimension; i++) {
+            embedding[i] = randomByte();
+        }
+        return embedding;
+    }
+
+    private static InferenceAction.Response createFloatEmbeddingResponse(float[] embedding) {
+        var embeddingResult = new TextEmbeddingFloatResults.Embedding(embedding);
+        var textEmbeddingResults = new TextEmbeddingFloatResults(List.of(embeddingResult));
+        return new InferenceAction.Response(textEmbeddingResults);
+    }
+
+    private static InferenceAction.Response createByteEmbeddingResponse(byte[] embedding) {
+        var embeddingResult = new TextEmbeddingByteResults.Embedding(embedding);
+        var textEmbeddingResults = new TextEmbeddingByteResults(List.of(embeddingResult));
+        return new InferenceAction.Response(textEmbeddingResults);
+    }
+
+    private static InferenceAction.Response createEmptyFloatEmbeddingResponse() {
+        var textEmbeddingResults = new TextEmbeddingFloatResults(List.of());
+        return new InferenceAction.Response(textEmbeddingResults);
+    }
+
+    private Page randomInputPage(int positionCount, int columnCount) {
+        final Block[] blocks = new Block[columnCount];
+        try {
+            for (int i = 0; i < columnCount; i++) {
+                blocks[i] = RandomBlock.randomBlock(
+                    blockFactory(),
+                    RandomBlock.randomElementExcluding(List.of(ElementType.AGGREGATE_METRIC_DOUBLE)),
+                    positionCount,
+                    randomBoolean(),
+                    0,
+                    0,
+                    randomInt(10),
+                    randomInt(10)
+                ).block();
+            }
+
+            return new Page(blocks);
+        } catch (Exception e) {
+            Releasables.close(blocks);
+            throw (e);
+        }
+    }
+}