More flexible output builder.

Author: afoucret

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/InferenceOperator.java

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.esql.inference;
 
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.AsyncOperator;
@@ -102,7 +103,7 @@ public Page getOutput() {
             return null;
         }
 
-        try (OutputBuilder outputBuilder = outputBuilder(ongoingInferenceResult.inputPage)) {
+        try (OutputBuilder<Page> outputBuilder = outputBuilder(ongoingInferenceResult.inputPage)) {
             for (InferenceAction.Response response : ongoingInferenceResult.responses) {
                 outputBuilder.addInferenceResponse(response);
             }
@@ -125,12 +126,12 @@ public Page getOutput() {
      *
      * @param input The corresponding input page used to generate the inference requests.
      */
-    protected abstract OutputBuilder outputBuilder(Page input);
+    protected abstract OutputBuilder<Page> outputBuilder(Page input);
 
     /**
-     * An interface for accumulating inference responses and constructing a result {@link Page}.
+     * An interface for accumulating inference responses and constructing a result (can be a {@link  Page} or a {@link Block}).
      */
-    public interface OutputBuilder extends Releasable {
+    public interface OutputBuilder<T> extends Releasable {
 
         /**
          * Adds an inference response to the output.
@@ -144,11 +145,11 @@ public interface OutputBuilder extends Releasable {
         void addInferenceResponse(InferenceAction.Response inferenceResponse);
 
         /**
-         * Builds the final output page from accumulated inference responses.
+         * Builds the final output from accumulated inference responses.
          *
-         * @return The constructed output page.
+         * @return The constructed output block.
          */
-        Page buildOutput();
+        T buildOutput();
 
         static <IR extends InferenceServiceResults> IR inferenceResults(InferenceAction.Response inferenceResponse, Class<IR> clazz) {
             InferenceServiceResults results = inferenceResponse.getResults();

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/completion/CompletionOperatorOutputBuilder.java

@@ -20,7 +20,7 @@
  * {@link CompletionOperatorOutputBuilder} builds the output page for {@link CompletionOperator} by converting {@link ChatCompletionResults}
  * into a {@link BytesRefBlock}.
  */
-public class CompletionOperatorOutputBuilder implements InferenceOperator.OutputBuilder {
+public class CompletionOperatorOutputBuilder implements InferenceOperator.OutputBuilder<Page> {
     private final Page inputPage;
     private final BytesRefBlock.Builder outputBlockBuilder;
     private final BytesRefBuilder bytesRefBuilder = new BytesRefBuilder();

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/rerank/RerankOperatorOutputBuilder.java

@@ -24,7 +24,7 @@
  *  * reranked relevance scores into the specified score channel of the input page.
  */
 
-public class RerankOperatorOutputBuilder implements InferenceOperator.OutputBuilder {
+public class RerankOperatorOutputBuilder implements InferenceOperator.OutputBuilder<Page> {
 
     private final Page inputPage;
     private final DoubleBlock.Builder scoreBlockBuilder;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/textembedding/TextEmbeddingOperatorOutputBuilder.java

@@ -21,7 +21,7 @@
  * {@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 {
+public class TextEmbeddingOperatorOutputBuilder implements InferenceOperator.OutputBuilder<Page> {
     private final Page inputPage;
     private final FloatBlock.Builder outputBlockBuilder;
 

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/inference/TextEmbeddingTests.java

@@ -50,7 +50,7 @@ public static Iterable<Object[]> parameters() {
                         List.of(inputTextDataType, inferenceIdDataType),
                         () -> new TestCaseSupplier.TestCase(
                             List.of(
-                                new TestCaseSupplier.TypedData(randomBytesReference(10).toBytesRef(), inputTextDataType, "inputText"),
+                                new TestCaseSupplier.TypedData(randomBytesReference(10).toBytesRef(), inputTextDataType, "text"),
                                 new TestCaseSupplier.TypedData(randomBytesReference(10).toBytesRef(), inferenceIdDataType, "inference_id")
                             ),
                             Matchers.blankOrNullString(),