Adding more tests for RerankOperatorOutputBuilder and RerankOperatorRequestIterator

Author: afoucret

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

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

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

@@ -0,0 +1,113 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.inference.bulk;
+
+import org.elasticsearch.compute.data.Block;
+import org.elasticsearch.compute.data.DoubleBlock;
+import org.elasticsearch.compute.data.Page;
+import org.elasticsearch.compute.test.ComputeTestCase;
+import org.elasticsearch.compute.test.RandomBlock;
+import org.elasticsearch.core.Releasables;
+import org.elasticsearch.logging.LogManager;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+import org.elasticsearch.xpack.core.inference.results.RankedDocsResults;
+import org.elasticsearch.xpack.esql.inference.rerank.RerankOperatorOutputBuilder;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class RerankOperatorOutputBuilderTests extends ComputeTestCase {
+
+    public void testBuildSmallOutput() {
+        assertBuildOutput(between(1, 100));
+    }
+
+    public void testBuildLargeOutput() {
+        assertBuildOutput(between(10_000, 100_000));
+    }
+
+    private void assertBuildOutput(int size) {
+        final Page inputPage = randomInputPage(size, between(1, 20));
+        final int scoreChannel = randomIntBetween(0, inputPage.getBlockCount());
+        try (
+            RerankOperatorOutputBuilder outputBuilder = new RerankOperatorOutputBuilder(
+                blockFactory().newDoubleBlockBuilder(size),
+                inputPage,
+                scoreChannel
+            )
+        ) {
+            int batchSize = randomIntBetween(1, size);
+            for (int currentPos = 0; currentPos < inputPage.getPositionCount();) {
+                List<RankedDocsResults.RankedDoc> rankedDocs = new ArrayList<>();
+                for (int rankedDocIndex = 0; rankedDocIndex < batchSize && currentPos < inputPage.getPositionCount(); rankedDocIndex++) {
+                    rankedDocs.add(new RankedDocsResults.RankedDoc(rankedDocIndex, relevanceScore(currentPos), randomIdentifier()));
+                    currentPos++;
+                }
+
+                outputBuilder.addInferenceResponse(new InferenceAction.Response(new RankedDocsResults(rankedDocs)));
+            }
+
+            final Page outputPage = outputBuilder.buildOutput();
+            try {
+                assertThat(outputPage.getPositionCount(), equalTo(inputPage.getPositionCount()));
+                LogManager.getLogger(RerankOperatorOutputBuilderTests.class)
+                    .info(
+                        "{} , {}, {}, {}",
+                        scoreChannel,
+                        inputPage.getBlockCount(),
+                        outputPage.getBlockCount(),
+                        Math.max(scoreChannel + 1, inputPage.getBlockCount())
+                    );
+                assertThat(outputPage.getBlockCount(), equalTo(Integer.max(scoreChannel + 1, inputPage.getBlockCount())));
+                assertOutputContent(outputPage.getBlock(scoreChannel));
+            } finally {
+                outputPage.releaseBlocks();
+            }
+
+        } finally {
+            inputPage.releaseBlocks();
+        }
+    }
+
+    private float relevanceScore(int position) {
+        return (float) 1 / (1 + position);
+    }
+
+    private void assertOutputContent(DoubleBlock block) {
+        for (int currentPos = 0; currentPos < block.getPositionCount(); currentPos++) {
+            assertThat(block.isNull(currentPos), equalTo(false));
+            assertThat(block.getValueCount(currentPos), equalTo(1));
+            assertThat(block.getDouble(block.getFirstValueIndex(currentPos)), equalTo((double) relevanceScore(currentPos)));
+        }
+    }
+
+    private Page randomInputPage(int positionCount, int columnCount) {
+        final Block[] blocks = new Block[columnCount];
+        try {
+            for (int i = 0; i < columnCount; i++) {
+                blocks[i] = RandomBlock.randomBlock(
+                    blockFactory(),
+                    RandomBlock.randomElementType(),
+                    positionCount,
+                    randomBoolean(),
+                    0,
+                    0,
+                    randomInt(10),
+                    randomInt(10)
+                ).block();
+            }
+
+            return new Page(blocks);
+        } catch (Exception e) {
+            Releasables.close(blocks);
+            throw (e);
+        }
+    }
+}

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

@@ -0,0 +1,64 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.inference.bulk;
+
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.compute.data.BytesRefBlock;
+import org.elasticsearch.compute.test.ComputeTestCase;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+import org.elasticsearch.xpack.esql.inference.rerank.RerankOperatorRequestIterator;
+
+import java.util.List;
+
+import static org.hamcrest.Matchers.equalTo;
+
+public class RerankOperatorRequestIteratorTests extends ComputeTestCase {
+
+    public void testIterateSmallInput() {
+        assertIterate(between(1, 100), randomIntBetween(1, 1_000));
+    }
+
+    public void testIterateLargeInput() {
+        assertIterate(between(10_000, 100_000), randomIntBetween(1, 1_000));
+    }
+
+    private void assertIterate(int size, int batchSize) {
+        final BytesRefBlock inputBlock = randomInputBlock(size);
+        final String inferenceId = randomIdentifier();
+        final String queryText = randomIdentifier();
+
+        try (
+            RerankOperatorRequestIterator requestIterator = new RerankOperatorRequestIterator(inputBlock, inferenceId, queryText, batchSize)
+        ) {
+            BytesRef scratch = new BytesRef();
+
+            for (int currentPos = 0; requestIterator.hasNext();) {
+                InferenceAction.Request request = requestIterator.next();
+
+                assertThat(request.getInferenceEntityId(), equalTo(inferenceId));
+                assertThat(request.getQuery(), equalTo(queryText));
+                List<String> inputs = request.getInput();
+                for (String input : inputs) {
+                    scratch = inputBlock.getBytesRef(inputBlock.getFirstValueIndex(currentPos), scratch);
+                    assertThat(input, equalTo(scratch.utf8ToString()));
+                    currentPos++;
+                }
+            }
+        }
+    }
+
+    private BytesRefBlock randomInputBlock(int size) {
+        try (BytesRefBlock.Builder blockBuilder = blockFactory().newBytesRefBlockBuilder(size)) {
+            for (int i = 0; i < size; i++) {
+                blockBuilder.appendBytesRef(new BytesRef(randomAlphaOfLength(10)));
+            }
+
+            return blockBuilder.build();
+        }
+    }
+}

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

@@ -32,7 +32,7 @@ public void testBuildLargeOutput() {
     }
 
     private void assertBuildOutput(int size) {
-        Page inputPage = randomInputPage(size, between(1, 20));
+        final Page inputPage = randomInputPage(size, between(1, 20));
         try (
             CompletionOperatorOutputBuilder outputBuilder = new CompletionOperatorOutputBuilder(
                 blockFactory().newBytesRefBlockBuilder(size),
@@ -44,7 +44,7 @@ private void assertBuildOutput(int size) {
                 outputBuilder.addInferenceResponse(new InferenceAction.Response(new ChatCompletionResults(results)));
             }
 
-            Page outputPage = outputBuilder.buildOutput();
+            final Page outputPage = outputBuilder.buildOutput();
             assertThat(outputPage.getPositionCount(), equalTo(inputPage.getPositionCount()));
             assertThat(outputPage.getBlockCount(), equalTo(inputPage.getBlockCount() + 1));
             assertOutputContent(outputPage.getBlock(outputPage.getBlockCount() - 1));
@@ -68,7 +68,7 @@ private void assertOutputContent(BytesRefBlock block) {
     }
 
     private Page randomInputPage(int positionCount, int columnCount) {
-        Block[] blocks = new Block[columnCount];
+        final Block[] blocks = new Block[columnCount];
         try {
             for (int i = 0; i < columnCount; i++) {
                 blocks[i] = RandomBlock.randomBlock(