Fix snippet calculation

Author: kderusso

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

@@ -23,6 +23,7 @@
 import org.elasticsearch.xpack.inference.services.huggingface.rerank.HuggingFaceRerankTaskSettings;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 
@@ -210,22 +211,23 @@ private float[] extractScoresFromRankedDocs(List<RankedDocsResults.RankedDoc> ra
     }
 
     private float[] extractScoresFromRankedSnippets(List<RankedDocsResults.RankedDoc> rankedDocs, RankFeatureDoc[] featureDocs) {
-        int numSnippets = rankedDocs.size() / featureDocs.length;
         float[] scores = new float[featureDocs.length];
         boolean[] hasScore = new boolean[featureDocs.length];
 
+        // We need to correlate the index/doc values of each RankedDoc in correlation with its associated RankFeatureDoc.
+        int[] rankedDocToFeatureDoc = Arrays.stream(featureDocs)
+            .flatMapToInt(
+                doc -> java.util.stream.IntStream.generate(() -> Arrays.asList(featureDocs).indexOf(doc)).limit(doc.featureData.size())
+            )
+            .limit(rankedDocs.size())
+            .toArray();
+
         for (int i = 0; i < rankedDocs.size(); i++) {
-            // TODO this naively assumes that we always get the requested number of snippets per ranked document
             RankedDocsResults.RankedDoc rankedDoc = rankedDocs.get(i);
-            int docId = rankedDoc.index() / numSnippets;
+            int docId = rankedDocToFeatureDoc[rankedDoc.index()];
             float score = rankedDoc.relevanceScore();
-
-            if (hasScore[docId] == false) {
-                scores[docId] = score;
-                hasScore[docId] = true;
-            } else {
-                scores[docId] = Math.max(scores[docId], score);
-            }
+            scores[docId] = hasScore[docId] == false ? score : Math.max(scores[docId], score);
+            hasScore[docId] = true;
         }
 
         float[] result = new float[featureDocs.length];