Fix score count validation in reranker response

docs/changelog/111212.yaml

@@ -0,0 +1,6 @@
+pr: 111212
+summary: Fix score count validation in reranker response
+area: Ranking
+type: bug
+issues:
+ - 111202

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankFeaturePhaseRankCoordinatorContext.java

@@ -54,12 +54,24 @@ protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[
         // Wrap the provided rankListener to an ActionListener that would handle the response from the inference service
         // and then pass the results
         final ActionListener<InferenceAction.Response> actionListener = scoreListener.delegateFailureAndWrap((l, r) -> {
-            float[] scores = extractScoresFromResponse(r);
-            if (scores.length != featureDocs.length) {
+            InferenceServiceResults results = r.getResults();
+            assert results instanceof RankedDocsResults;
+
+            // Ensure we get exactly as many scores as the number of docs we passed, otherwise we may return incorrect results
+            List<RankedDocsResults.RankedDoc> rankedDocs = ((RankedDocsResults) results).getRankedDocs();
+            if (rankedDocs.size() != featureDocs.length) {
                 l.onFailure(
-                    new IllegalStateException("Document and score count mismatch: [" + featureDocs.length + "] vs [" + scores.length + "]")
+                    new IllegalStateException(
+                        "Document and score count mismatch: ["
+                            + featureDocs.length
+                            + "] vs ["
+                            + rankedDocs.size()
+                            + "]. Check your rerank inference endpoint configuration and ensure it returns rank_window_size scores for "
+                            + "rank_window_size input documents."
+                    )
                 );
             } else {
+                float[] scores = extractScoresFromRankedDocs(rankedDocs);
                 l.onResponse(scores);
             }
         });
@@ -85,11 +97,7 @@ protected InferenceAction.Request generateRequest(List<String> docFeatures) {
         );
     }
 
-    private float[] extractScoresFromResponse(InferenceAction.Response response) {
-        InferenceServiceResults results = response.getResults();
-        assert results instanceof RankedDocsResults;
-
-        List<RankedDocsResults.RankedDoc> rankedDocs = ((RankedDocsResults) results).getRankedDocs();
+    private float[] extractScoresFromRankedDocs(List<RankedDocsResults.RankedDoc> rankedDocs) {
         float[] scores = new float[rankedDocs.size()];
         for (RankedDocsResults.RankedDoc rankedDoc : rankedDocs) {
             scores[rankedDoc.index()] = rankedDoc.relevanceScore();

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityRankTests.java

@@ -37,6 +37,20 @@ public class TextSimilarityRankTests extends ESSingleNodeTestCase {
      */
     public static class InvalidInferenceResultCountProvidingTextSimilarityRankBuilder extends TextSimilarityRankBuilder {
 
+        private boolean hasInvalidDocumentIndices = false;
+
+        public InvalidInferenceResultCountProvidingTextSimilarityRankBuilder(
+            String field,
+            String inferenceId,
+            String inferenceText,
+            int rankWindowSize,
+            Float minScore,
+            boolean hasInvalidDocumentIndices
+        ) {
+            this(field, inferenceId, inferenceText, rankWindowSize, minScore);
+            this.hasInvalidDocumentIndices = hasInvalidDocumentIndices;
+        }
+
         public InvalidInferenceResultCountProvidingTextSimilarityRankBuilder(
             String field,
             String inferenceId,
@@ -65,7 +79,7 @@ protected InferenceAction.Request generateRequest(List<String> docFeatures) {
                         inferenceId,
                         inferenceText,
                         docFeatures,
-                        Map.of("invalidInferenceResultCount", true),
+                        Map.of("invalidInferenceResultCount", true, "invalidDocumentIndices", hasInvalidDocumentIndices),
                         InputType.SEARCH,
                         InferenceAction.Request.DEFAULT_TIMEOUT
                     );
@@ -151,11 +165,12 @@ public void testRerankInferenceFailure() {
         );
     }
 
-    public void testRerankInferenceResultMismatch() {
+    public void testRerankInferenceResultCountMismatch() {
         ElasticsearchAssertions.assertFailures(
             // Execute search with text similarity reranking
             client.prepareSearch()
                 .setRankBuilder(
+                    // Simulate reranker returning different number of results from input
                     new InvalidInferenceResultCountProvidingTextSimilarityRankBuilder("text", "my-rerank-model", "my query", 100, 1.5f)
                 )
                 .setQuery(QueryBuilders.matchAllQuery()),
@@ -164,6 +179,27 @@ public void testRerankInferenceResultMismatch() {
         );
     }
 
+    public void testRerankInvalidDocumentIndices() {
+        ElasticsearchAssertions.assertFailures(
+            // Execute search with text similarity reranking
+            client.prepareSearch()
+                .setRankBuilder(
+                    // Simulate reranker returning different number of results from input, also invalid document indices in results
+                    new InvalidInferenceResultCountProvidingTextSimilarityRankBuilder(
+                        "text",
+                        "my-rerank-model",
+                        "my query",
+                        100,
+                        1.5f,
+                        true
+                    )
+                )
+                .setQuery(QueryBuilders.matchAllQuery()),
+            RestStatus.INTERNAL_SERVER_ERROR,
+            containsString("Failed to execute phase [rank-feature], Computing updated ranks for results failed")
+        );
+    }
+
     private static void assertHitHasRankScoreAndText(SearchHit hit, int expectedRank, float expectedScore, String expectedText) {
         assertEquals(expectedRank, hit.getRank());
         assertEquals(expectedScore, hit.getScore(), 0.0f);

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/rank/textsimilarity/TextSimilarityTestPlugin.java

@@ -115,9 +115,10 @@ public <Request extends ActionRequest, Response extends ActionResponse> void app
             }
 
             assert request instanceof InferenceAction.Request;
-            boolean shouldThrow = (boolean) ((InferenceAction.Request) request).getTaskSettings().getOrDefault("throwing", false);
-            boolean hasInvalidInferenceResultCount = (boolean) ((InferenceAction.Request) request).getTaskSettings()
-                .getOrDefault("invalidInferenceResultCount", false);
+            Map<String, Object> taskSettings = ((InferenceAction.Request) request).getTaskSettings();
+            boolean shouldThrow = (boolean) taskSettings.getOrDefault("throwing", false);
+            boolean hasInvalidInferenceResultCount = (boolean) taskSettings.getOrDefault("invalidInferenceResultCount", false);
+            boolean hasInvalidDocumentIndices = (boolean) taskSettings.getOrDefault("invalidDocumentIndices", false);
 
             if (shouldThrow) {
                 listener.onFailure(new UnsupportedOperationException("simulated failure"));
@@ -126,7 +127,13 @@ public <Request extends ActionRequest, Response extends ActionResponse> void app
                 List<String> inputs = ((InferenceAction.Request) request).getInput();
                 int resultCount = hasInvalidInferenceResultCount ? inputs.size() - 1 : inputs.size();
                 for (int i = 0; i < resultCount; i++) {
-                    rankedDocsResults.add(new RankedDocsResults.RankedDoc(i, Float.parseFloat(inputs.get(i)), inputs.get(i)));
+                    rankedDocsResults.add(
+                        new RankedDocsResults.RankedDoc(
+                            hasInvalidDocumentIndices ? i * 2 : i,
+                            Float.parseFloat(inputs.get(i)),
+                            inputs.get(i)
+                        )
+                    );
                 }
                 ActionResponse response = new InferenceAction.Response(new RankedDocsResults(rankedDocsResults));
                 listener.onResponse((Response) response);