Better implementation of OutputBuilder.

Author: afoucret

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

@@ -112,9 +112,9 @@ public Page getOutput() {
                 outputBuilder.addInferenceResponse(response);
             }
             return outputBuilder.buildOutput();
-
-        } finally {
+        } catch (Exception e) {
             releaseFetchedOnAnyThread(ongoingInferenceResult);
+            throw e;
         }
     }
 

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

@@ -33,7 +33,6 @@ public CompletionOperatorOutputBuilder(BytesRefBlock.Builder outputBlockBuilder,
     @Override
     public void close() {
         Releasables.close(outputBlockBuilder);
-        releasePageOnAnyThread(inputPage);
     }
 
     /**
@@ -72,13 +71,13 @@ public void addInferenceResponse(InferenceAction.Response inferenceResponse) {
     }
 
     /**
-     * Builds the final output page by appending the completion output block to a shallow copy of the input page.
+     * Builds the final output page by appending the completion output block to the input page.
      */
     @Override
     public Page buildOutput() {
         Block outputBlock = outputBlockBuilder.build();
         assert outputBlock.getPositionCount() == inputPage.getPositionCount();
-        return inputPage.shallowCopy().appendBlock(outputBlock);
+        return inputPage.appendBlock(outputBlock);
     }
 
     private ChatCompletionResults inferenceResults(InferenceAction.Response inferenceResponse) {

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

@@ -7,7 +7,6 @@
 
 package org.elasticsearch.xpack.esql.inference.rerank;
 
-import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.DoubleBlock;
 import org.elasticsearch.compute.data.Page;
 import org.elasticsearch.core.Releasables;
@@ -18,6 +17,7 @@
 
 import java.util.Comparator;
 import java.util.Iterator;
+import java.util.stream.IntStream;
 
 /**
  * Builds the output page for the {@link RerankOperator} by adding
@@ -39,7 +39,6 @@ public RerankOperatorOutputBuilder(DoubleBlock.Builder scoreBlockBuilder, Page i
     @Override
     public void close() {
         Releasables.close(scoreBlockBuilder);
-        releasePageOnAnyThread(inputPage);
     }
 
     /**
@@ -49,21 +48,24 @@ public void close() {
     @Override
     public Page buildOutput() {
         int blockCount = Integer.max(inputPage.getBlockCount(), scoreChannel + 1);
-        Block[] blocks = new Block[blockCount];
+        Page outputPage = inputPage.appendBlock(scoreBlockBuilder.build());
+
+        if (scoreChannel == inputPage.getBlockCount()) {
+            // Just need to append the block at the end
+            // We can just return the output page we have just created
+            return outputPage;
+        }
 
         try {
-            for (int b = 0; b < blockCount; b++) {
-                if (b == scoreChannel) {
-                    blocks[b] = scoreBlockBuilder.build();
-                } else {
-                    blocks[b] = inputPage.getBlock(b);
-                    blocks[b].incRef();
-                }
-            }
-            return new Page(blocks);
-        } catch (Exception e) {
-            Releasables.close(blocks);
-            throw (e);
+            // We need to project the last column to the score channel.
+            int[] blockNapping = IntStream.range(0, inputPage.getBlockCount())
+                .map(channel -> channel == scoreChannel ? inputPage.getBlockCount() : channel)
+                .toArray();
+
+            return outputPage.projectBlocks(blockNapping);
+        } finally {
+            // Releasing the output page since projection is incrementing block references.
+            releasePageOnAnyThread(outputPage);
         }
     }
 

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

@@ -23,15 +23,15 @@
 
 public class CompletionOperatorOutputBuilderTests extends ComputeTestCase {
 
-    public void testBuildSmallOutput() {
+    public void testBuildSmallOutput() throws Exception {
         assertBuildOutput(between(1, 100));
     }
 
-    public void testBuildLargeOutput() {
+    public void testBuildLargeOutput() throws Exception {
         assertBuildOutput(between(10_000, 100_000));
     }
 
-    private void assertBuildOutput(int size) {
+    private void assertBuildOutput(int size) throws Exception {
         final Page inputPage = randomInputPage(size, between(1, 20));
         try (
             CompletionOperatorOutputBuilder outputBuilder = new CompletionOperatorOutputBuilder(
@@ -50,11 +50,9 @@ private void assertBuildOutput(int size) {
             assertOutputContent(outputPage.getBlock(outputPage.getBlockCount() - 1));
 
             outputPage.releaseBlocks();
-
-        } finally {
-            inputPage.releaseBlocks();
         }
 
+        allBreakersEmpty();
     }
 
     private void assertOutputContent(BytesRefBlock block) {

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

@@ -23,15 +23,15 @@
 
 public class RerankOperatorOutputBuilderTests extends ComputeTestCase {
 
-    public void testBuildSmallOutput() {
+    public void testBuildSmallOutput() throws Exception {
         assertBuildOutput(between(1, 100));
     }
 
-    public void testBuildLargeOutput() {
+    public void testBuildLargeOutput() throws Exception {
         assertBuildOutput(between(10_000, 100_000));
     }
 
-    private void assertBuildOutput(int size) {
+    private void assertBuildOutput(int size) throws Exception {
         final Page inputPage = randomInputPage(size, between(1, 20));
         final int scoreChannel = randomIntBetween(0, inputPage.getBlockCount());
         try (
@@ -61,9 +61,9 @@ private void assertBuildOutput(int size) {
                 outputPage.releaseBlocks();
             }
 
-        } finally {
-            inputPage.releaseBlocks();
         }
+
+        allBreakersEmpty();
     }
 
     private float relevanceScore(int position) {