Improved input management.

Author: afoucret

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

@@ -106,13 +106,8 @@ public Page getOutput() {
         }
 
         try (OutputBuilder outputBuilder = outputBuilder(ongoingInferenceResult.inputPage)) {
-            assert ongoingInferenceResult.inputPage.getPositionCount() == ongoingInferenceResult.responses.size();
             for (InferenceAction.Response response : ongoingInferenceResult.responses) {
-                try {
-                    outputBuilder.addInferenceResponse(response);
-                } catch (IllegalArgumentException e) {
-                    throw new IllegalStateException("Invalid inference response", e);
-                }
+                outputBuilder.addInferenceResponse(response);
             }
             return outputBuilder.buildOutput();
 
@@ -168,6 +163,10 @@ static <IR extends InferenceServiceResults> IR inferenceResults(InferenceAction.
                 format("Inference result has wrong type. Got [{}] while expecting [{}]", results.getClass().getName(), clazz.getName())
             );
         }
+
+        default void releasePageOnAnyThread(Page page) {
+            InferenceOperator.releasePageOnAnyThread(page);
+        }
     }
 
     /**

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

@@ -19,6 +19,8 @@
 import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceExecutionConfig;
 import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRequestIterator;
 
+import java.util.stream.IntStream;
+
 /**
  * {@link CompletionOperator} is an {@link InferenceOperator} that performs inference using prompt-based model (e.g., text completion).
  * It evaluates a prompt expression for each input row, constructs inference requests, and emits the model responses as output.
@@ -48,14 +50,25 @@ public String toString() {
         return "CompletionOperator[inference_id=[" + inferenceId() + "]]";
     }
 
+    @Override
+    public void addInput(Page input) {
+        try {
+            super.addInput(input.appendBlock(promptEvaluator.eval(input)));
+        } catch (Exception e) {
+            releasePageOnAnyThread(input);
+            throw(e);
+        }
+    }
+
     /**
      * Constructs the completion inference requests iterator for the given input page by evaluating the prompt expression.
      *
      * @param inputPage The input data page.
      */
     @Override
     protected BulkInferenceRequestIterator requests(Page inputPage) {
-        return new CompletionOperatorRequestIterator((BytesRefBlock) promptEvaluator.eval(inputPage), inferenceId());
+        int inputBlockChannel = inputPage.getBlockCount() - 1;
+        return new CompletionOperatorRequestIterator(inputPage.getBlock(inputBlockChannel), inferenceId());
     }
 
     /**
@@ -65,7 +78,8 @@ protected BulkInferenceRequestIterator requests(Page inputPage) {
      */
     @Override
     protected CompletionOperatorOutputBuilder outputBuilder(Page input) {
-        return new CompletionOperatorOutputBuilder(blockFactory().newBytesRefBlockBuilder(input.getPositionCount()), input);
+        BytesRefBlock.Builder outputBlockBuilder = blockFactory().newBytesRefBlockBuilder(input.getPositionCount());
+        return new CompletionOperatorOutputBuilder(outputBlockBuilder, input.projectBlocks(IntStream.range(0, input.getBlockCount() - 1).toArray()));
     }
 
     /**

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

@@ -33,6 +33,7 @@ public CompletionOperatorOutputBuilder(BytesRefBlock.Builder outputBlockBuilder,
     @Override
     public void close() {
         Releasables.close(outputBlockBuilder);
+        releasePageOnAnyThread(inputPage);
     }
 
     /**
@@ -63,7 +64,6 @@ public void addInferenceResponse(InferenceAction.Response inferenceResponse) {
             bytesRefBuilder.clear();
         }
         outputBlockBuilder.endPositionEntry();
-
     }
 
     /**

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

@@ -115,8 +115,7 @@ public int estimatedSize() {
 
         @Override
         public void close() {
-            promptBlock.allowPassingToDifferentDriver();
-            Releasables.close(promptBlock);
+
         }
     }
 }

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

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.esql.inference.rerank;
 
-import org.elasticsearch.compute.data.BytesRefBlock;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.compute.operator.DriverContext;
 import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
@@ -18,6 +17,8 @@
 import org.elasticsearch.xpack.esql.inference.InferenceRunner;
 import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceExecutionConfig;
 
+import java.util.stream.IntStream;
+
 /**
  * {@link RerankOperator} is an inference operator that compute scores for rows using a reranking model.
  */
@@ -50,6 +51,16 @@ public RerankOperator(
         this.scoreChannel = scoreChannel;
     }
 
+    @Override
+    public void addInput(Page input) {
+        try {
+            super.addInput(input.appendBlock(rowEncoder.eval(input)));
+        } catch (Exception e) {
+            releasePageOnAnyThread(input);
+            throw(e);
+        }
+    }
+
     @Override
     protected void doClose() {
         Releasables.close(rowEncoder);
@@ -65,15 +76,20 @@ public String toString() {
      */
     @Override
     protected RerankOperatorRequestIterator requests(Page inputPage) {
-        return new RerankOperatorRequestIterator((BytesRefBlock) rowEncoder.eval(inputPage), inferenceId(), queryText, batchSize);
+        int inputBlockChannel = inputPage.getBlockCount() - 1;
+        return new RerankOperatorRequestIterator(inputPage.getBlock(inputBlockChannel), inferenceId(), queryText, batchSize);
     }
 
     /**
      * Returns the output builder responsible for collecting inference responses and building the output page.
      */
     @Override
     protected RerankOperatorOutputBuilder outputBuilder(Page input) {
-        return new RerankOperatorOutputBuilder(blockFactory().newDoubleBlockBuilder(input.getPositionCount()), input, scoreChannel);
+        return new RerankOperatorOutputBuilder(
+            blockFactory().newDoubleBlockBuilder(input.getPositionCount()),
+            input.projectBlocks(IntStream.range(0, input.getBlockCount() - 1).toArray()),
+            scoreChannel
+        );
     }
 
     /**

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

@@ -17,6 +17,7 @@
 import org.elasticsearch.xpack.esql.inference.InferenceOperator;
 
 import java.util.Comparator;
+import java.util.Iterator;
 
 /**
  * Builds the output page for the {@link RerankOperator} by adding
@@ -38,6 +39,7 @@ public RerankOperatorOutputBuilder(DoubleBlock.Builder scoreBlockBuilder, Page i
     @Override
     public void close() {
         Releasables.close(scoreBlockBuilder);
+        releasePageOnAnyThread(inputPage);
     }
 
     /**
@@ -77,11 +79,10 @@ public Page buildOutput() {
      */
     @Override
     public void addInferenceResponse(InferenceAction.Response inferenceResponse) {
-        inferenceResults(inferenceResponse).getRankedDocs()
-            .stream()
-            .sorted(Comparator.comparingInt(RankedDocsResults.RankedDoc::index))
-            .mapToDouble(RankedDocsResults.RankedDoc::relevanceScore)
-            .forEach(scoreBlockBuilder::appendDouble);
+        Iterator<RankedDocsResults.RankedDoc> sortedRankedDocIterator = inferenceResults(inferenceResponse).getRankedDocs().stream().sorted(Comparator.comparingInt(RankedDocsResults.RankedDoc::index)).iterator();
+        while (sortedRankedDocIterator.hasNext()) {
+            scoreBlockBuilder.appendDouble(sortedRankedDocIterator.next().relevanceScore());
+        }
     }
 
     private RankedDocsResults inferenceResults(InferenceAction.Response inferenceResponse) {

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

@@ -10,7 +10,6 @@
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.common.lucene.BytesRefs;
 import org.elasticsearch.compute.data.BytesRefBlock;
-import org.elasticsearch.core.Releasables;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
 import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRequestIterator;
@@ -81,7 +80,6 @@ private InferenceAction.Request inferenceRequest(List<String> inputs) {
 
     @Override
     public void close() {
-        inputBlock.allowPassingToDifferentDriver();
-        Releasables.close(inputBlock);
+
     }
 }

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

@@ -78,7 +78,8 @@ protected ThreadPool threadPool() {
 
     @Override
     protected SourceOperator simpleInput(BlockFactory blockFactory, int size) {
-        return new AbstractBlockSourceOperator(blockFactory, 8 * 1024) {
+        int minSize = Integer.min(1, size / 10);
+        return new AbstractBlockSourceOperator(blockFactory, between(minSize, 8 * 1024)) {
             @Override
             protected int remaining() {
                 return size - currentPosition;

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/bulk/RerankOperatorRequestIteratorTests.java

@@ -28,11 +28,11 @@ public void testIterateLargeInput() {
     }
 
     private void assertIterate(int size, int batchSize) {
-        final BytesRefBlock inputBlock = randomInputBlock(size);
         final String inferenceId = randomIdentifier();
         final String queryText = randomIdentifier();
 
         try (
+            BytesRefBlock inputBlock = randomInputBlock(size);
             RerankOperatorRequestIterator requestIterator = new RerankOperatorRequestIterator(inputBlock, inferenceId, queryText, batchSize)
         ) {
             BytesRef scratch = new BytesRef();

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/completion/CompletionOperatorRequestIteratorTests.java

@@ -25,10 +25,9 @@ public void testIterateLargeInput() {
     }
 
     private void assertIterate(int size) {
-        final BytesRefBlock inputBlock = randomInputBlock(size);
         final String inferenceId = randomIdentifier();
 
-        try (CompletionOperatorRequestIterator requestIterator = new CompletionOperatorRequestIterator(inputBlock, inferenceId)) {
+        try (BytesRefBlock inputBlock = randomInputBlock(size); CompletionOperatorRequestIterator requestIterator = new CompletionOperatorRequestIterator(inputBlock, inferenceId)) {
             BytesRef scratch = new BytesRef();
 
             for (int currentPos = 0; requestIterator.hasNext(); currentPos++) {