Minor improvements

Author: afoucret

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

@@ -56,7 +56,7 @@ public void addInput(Page input) {
             super.addInput(input.appendBlock(promptEvaluator.eval(input)));
         } catch (Exception e) {
             releasePageOnAnyThread(input);
-            throw (e);
+            throw e;
         }
     }
 

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

@@ -7,6 +7,8 @@
 
 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.compute.operator.DriverContext;
 import org.elasticsearch.compute.operator.EvalOperator.ExpressionEvaluator;
@@ -54,10 +56,11 @@ public RerankOperator(
     @Override
     public void addInput(Page input) {
         try {
-            super.addInput(input.appendBlock(rowEncoder.eval(input)));
+            Block inputBlock = rowEncoder.eval(input);
+            super.addInput(input.appendBlock(inputBlock));
         } catch (Exception e) {
             releasePageOnAnyThread(input);
-            throw (e);
+            throw e;
         }
     }
 
@@ -85,8 +88,9 @@ protected RerankOperatorRequestIterator requests(Page inputPage) {
      */
     @Override
     protected RerankOperatorOutputBuilder outputBuilder(Page input) {
+        DoubleBlock.Builder outputBlockBuilder = blockFactory().newDoubleBlockBuilder(input.getPositionCount());
         return new RerankOperatorOutputBuilder(
-            blockFactory().newDoubleBlockBuilder(input.getPositionCount()),
+            outputBlockBuilder,
             input.projectBlocks(IntStream.range(0, input.getBlockCount() - 1).toArray()),
             scoreChannel
         );

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

@@ -18,6 +18,7 @@
 import org.elasticsearch.common.util.concurrent.EsExecutors;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
+import org.elasticsearch.compute.data.BlockUtils;
 import org.elasticsearch.compute.data.BooleanBlock;
 import org.elasticsearch.compute.data.BytesRefBlock;
 import org.elasticsearch.compute.data.DoubleBlock;
@@ -78,18 +79,18 @@ protected ThreadPool threadPool() {
 
     @Override
     protected SourceOperator simpleInput(BlockFactory blockFactory, int size) {
-        int minSize = Integer.min(1, size / 10);
-        return new AbstractBlockSourceOperator(blockFactory, between(minSize, 8 * 1024)) {
+        return new AbstractBlockSourceOperator(blockFactory, 8 * 1024) {
             @Override
             protected int remaining() {
                 return size - currentPosition;
             }
 
             @Override
             protected Page createPage(int positionOffset, int length) {
+                length = Integer.min(length, remaining());
                 try (var builder = blockFactory.newBytesRefVectorBuilder(length)) {
                     for (int i = 0; i < length; i++) {
-                        builder.appendBytesRef(new BytesRef(randomAlphaOfLength(1000)));
+                        builder.appendBytesRef(new BytesRef(randomAlphaOfLength(10)));
                     }
                     currentPosition += length;
                     return new Page(builder.build().asBlock());
@@ -195,9 +196,7 @@ protected EvalOperator.ExpressionEvaluator.Factory evaluatorFactory(int channel)
         return context -> new EvalOperator.ExpressionEvaluator() {
             @Override
             public Block eval(Page page) {
-                Block b = page.getBlock(channel);
-                b.incRef();
-                return b;
+                return BlockUtils.deepCopyOf(page.getBlock(channel), blockFactory());
             }
 
             @Override

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

@@ -11,19 +11,14 @@
 import org.elasticsearch.compute.data.BytesRefBlock;
 import org.elasticsearch.compute.data.DoubleBlock;
 import org.elasticsearch.compute.data.Page;
-import org.elasticsearch.compute.operator.EvalOperator;
 import org.elasticsearch.compute.operator.Operator;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
 import org.elasticsearch.xpack.core.inference.results.RankedDocsResults;
-import org.elasticsearch.xpack.esql.action.ColumnInfoImpl;
-import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.inference.InferenceOperatorTestCase;
-import org.elasticsearch.xpack.esql.inference.XContentRowEncoder;
 import org.hamcrest.Matcher;
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
 
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasSize;
@@ -35,18 +30,7 @@ public class RerankOperatorTests extends InferenceOperatorTestCase<RankedDocsRes
 
     @Override
     protected Operator.OperatorFactory simple(SimpleOptions options) {
-        Map<ColumnInfoImpl, EvalOperator.ExpressionEvaluator.Factory> fieldEvaluators = Map.of(
-            new ColumnInfoImpl(randomIdentifier(), DataType.TEXT, List.of()),
-            evaluatorFactory(0)
-        );
-
-        return new RerankOperator.Factory(
-            mockedSimpleInferenceRunner(),
-            SIMPLE_INFERENCE_ID,
-            SIMPLE_QUERY,
-            XContentRowEncoder.yamlRowEncoderFactory(fieldEvaluators),
-            1
-        );
+        return new RerankOperator.Factory(mockedSimpleInferenceRunner(), SIMPLE_INFERENCE_ID, SIMPLE_QUERY, evaluatorFactory(0), 1);
     }
 
     @Override