Move input block building outside the async portion.

Author: afoucret

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

@@ -72,10 +72,7 @@ public Page getOutput() {
     @Override
     protected void performAsync(Page input, ActionListener<OngoingInference> listener) {
         try {
-            BulkInferenceRequestIterator requests = requests(input);
-            listener = ActionListener.releaseBefore(requests, listener);
-
-            bulkInferenceExecutor.execute(requests, listener.map(responses -> new OngoingInference(input, responses)));
+            bulkInferenceExecutor.execute(requests(input), listener.map(responses -> new OngoingInference(input, responses)));
         } catch (Exception e) {
             listener.onFailure(e);
         }

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

@@ -28,7 +28,6 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
-import java.util.Objects;
 import java.util.stream.Collectors;
 
 /**
@@ -109,7 +108,6 @@ public BytesRefBlock eval(Page page) {
         } catch (IOException e) {
             throw new UncheckedIOException(e);
         } finally {
-            Arrays.stream(fieldValueBlocks).filter(Objects::nonNull).forEach(Block::allowPassingToDifferentDriver);
             Releasables.closeExpectNoException(fieldValueBlocks);
         }
     }

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

@@ -7,11 +7,10 @@
 
 package org.elasticsearch.xpack.esql.inference.bulk;
 
-import org.elasticsearch.core.Releasable;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
 
 import java.util.Iterator;
 
-public interface BulkInferenceRequestIterator extends Iterator<InferenceAction.Request>, Releasable {
+public interface BulkInferenceRequestIterator extends Iterator<InferenceAction.Request> {
 
 }

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

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.esql.inference.completion;
 
-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;
@@ -54,6 +53,16 @@ public CompletionOperator(
         this.promptEvaluator = promptEvaluator;
     }
 
+    @Override
+    public void addInput(Page input) {
+        try {
+            super.addInput(input.appendBlock(promptEvaluator.eval(input)));
+        } catch (Exception e) {
+            releasePageOnAnyThread(input);
+            throw e;
+        }
+    }
+
     @Override
     public Class<ChatCompletionResults> inferenceResultsClass() {
         return ChatCompletionResults.class;
@@ -71,7 +80,7 @@ public String toString() {
 
     @Override
     protected BulkInferenceRequestIterator requests(Page inputPage) {
-        return new CompletionOperatorRequestIterator((BytesRefBlock) promptEvaluator.eval(inputPage), inferenceId());
+        return new CompletionOperatorRequestIterator(inputPage, inferenceId());
     }
 
     @Override

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

@@ -16,6 +16,7 @@
 import org.elasticsearch.xpack.esql.inference.InferenceOperator;
 
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.stream.IntStream;
 
 public class CompletionOperatorOutputBuilder implements InferenceOperator.OutputBuilder<ChatCompletionResults> {
     private final Page inputPage;
@@ -52,6 +53,6 @@ public void addInferenceResults(ChatCompletionResults completionResults) {
     public Page buildOutput() {
         Block outputBlock = outputBlockBuilder.build();
         assert outputBlock.getPositionCount() == inputPage.getPositionCount();
-        return inputPage.shallowCopy().appendBlock(outputBlock);
+        return inputPage.projectBlocks(IntStream.range(0, inputPage.getBlockCount() - 1).toArray()).appendBlock(outputBlock);
     }
 }

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

@@ -9,8 +9,7 @@
 
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.compute.data.BytesRefBlock;
-import org.elasticsearch.core.Releasable;
-import org.elasticsearch.core.Releasables;
+import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
 import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRequestIterator;
@@ -25,17 +24,14 @@ public class CompletionOperatorRequestIterator implements BulkInferenceRequestIt
     private final int size;
     private int currentPos = 0;
 
-    public CompletionOperatorRequestIterator(BytesRefBlock promptBlock, String inferenceId) {
+    public CompletionOperatorRequestIterator(Page inputPage, String inferenceId) {
+        assert inputPage.getBlockCount() > 0;
+        BytesRefBlock promptBlock = inputPage.getBlock(inputPage.getBlockCount() - 1);
         this.promptReader = new PromptReader(promptBlock);
         this.size = promptBlock.getPositionCount();
         this.inferenceId = inferenceId;
     }
 
-    @Override
-    public void close() {
-        Releasables.close(promptReader);
-    }
-
     @Override
     public boolean hasNext() {
         return currentPos < size;
@@ -53,7 +49,7 @@ private InferenceAction.Request inferenceRequest(String prompt) {
         return InferenceAction.Request.builder(inferenceId, TaskType.COMPLETION).setInput(List.of(prompt)).build();
     }
 
-    private static class PromptReader implements Releasable {
+    private static class PromptReader {
         private final BytesRefBlock promptBlock;
         private BytesRef readBuffer = new BytesRef();
         private StringBuilder strBuilder = new StringBuilder();
@@ -76,11 +72,5 @@ public String readPrompt(int pos) {
 
             return strBuilder.toString();
         }
-
-        @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;
@@ -69,6 +68,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
     public Class<RankedDocsResults> inferenceResultsClass() {
         return RankedDocsResults.class;
@@ -86,7 +95,7 @@ public String toString() {
 
     @Override
     protected RerankOperatorRequestIterator requests(Page inputPage) {
-        return new RerankOperatorRequestIterator((BytesRefBlock) rowEncoder.eval(inputPage), inferenceId(), queryText, batchSize);
+        return new RerankOperatorRequestIterator(inputPage, inferenceId(), queryText, batchSize);
     }
 
     @Override

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

@@ -35,7 +35,7 @@ public void close() {
 
     @Override
     public Page buildOutput() {
-        int blockCount = Integer.max(inputPage.getBlockCount(), scoreChannel + 1);
+        int blockCount = Integer.max(inputPage.getBlockCount() - 1, scoreChannel + 1);
         Block[] blocks = new Block[blockCount];
 
         try {

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

@@ -10,7 +10,7 @@
 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.compute.data.Page;
 import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
 import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRequestIterator;
@@ -26,8 +26,9 @@ public class RerankOperatorRequestIterator implements BulkInferenceRequestIterat
     private final int batchSize;
     private int remainingPositions;
 
-    public RerankOperatorRequestIterator(BytesRefBlock inputBlock, String inferenceId, String queryText, int batchSize) {
-        this.inputBlock = inputBlock;
+    public RerankOperatorRequestIterator(Page inputPage, String inferenceId, String queryText, int batchSize) {
+        assert inputPage.getBlockCount() > 0;
+        this.inputBlock = inputPage.getBlock(inputPage.getBlockCount() - 1);
         this.inferenceId = inferenceId;
         this.queryText = queryText;
         this.batchSize = batchSize;
@@ -64,12 +65,6 @@ public InferenceAction.Request next() {
         return inferenceRequest(inputs);
     }
 
-    @Override
-    public void close() {
-        inputBlock.allowPassingToDifferentDriver();
-        Releasables.close(inputBlock);
-    }
-
     private InferenceAction.Request inferenceRequest(List<String> inputs) {
         return InferenceAction.Request.builder(inferenceId, TaskType.RERANK).setInput(inputs).setQuery(queryText).build();
     }

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

@@ -79,7 +79,8 @@ protected ThreadPool threadPool() {
 
     @Override
     protected SourceOperator simpleInput(BlockFactory blockFactory, int size) {
-        return new AbstractBlockSourceOperator(blockFactory, 8 * 1024) {
+        final int minPageSize = Math.max(1, size / 100);
+        return new AbstractBlockSourceOperator(blockFactory, between(minPageSize, size)) {
             @Override
             protected int remaining() {
                 return size - currentPosition;