Support reranking based on max score of multiple snippets per document

Author: kderusso

server/src/main/java/org/elasticsearch/search/rank/feature/RankFeatureDoc.java

@@ -30,6 +30,7 @@ public class RankFeatureDoc extends RankDoc {
     // TODO: update to support more than 1 fields; and not restrict to string data
     public String featureData;
     public List<String> snippets;
+    public List<Integer> docIndices;
 
     public RankFeatureDoc(int doc, float score, int shardIndex) {
         super(doc, score, shardIndex);
@@ -40,6 +41,7 @@ public RankFeatureDoc(StreamInput in) throws IOException {
         featureData = in.readOptionalString();
         if (in.getTransportVersion().onOrAfter(TransportVersions.RERANK_SNIPPETS)) {
             snippets = in.readOptionalStringCollectionAsList();
+            docIndices = in.readOptionalCollectionAsList(StreamInput::readVInt);
         }
     }
 
@@ -56,23 +58,30 @@ public void snippets(List<String> snippets) {
         this.snippets = snippets;
     }
 
+    public void docIndices(List<Integer> docIndices) {
+        this.docIndices = docIndices;
+    }
+
     @Override
     protected void doWriteTo(StreamOutput out) throws IOException {
         out.writeOptionalString(featureData);
         if (out.getTransportVersion().onOrAfter(TransportVersions.RERANK_SNIPPETS)) {
             out.writeOptionalStringCollection(snippets);
+            out.writeOptionalCollection(docIndices, StreamOutput::writeVInt);
         }
     }
 
     @Override
     protected boolean doEquals(RankDoc rd) {
         RankFeatureDoc other = (RankFeatureDoc) rd;
-        return Objects.equals(this.featureData, other.featureData) && Objects.equals(this.snippets, other.snippets);
+        return Objects.equals(this.featureData, other.featureData)
+            && Objects.equals(this.snippets, other.snippets)
+            && Objects.equals(this.docIndices, other.docIndices);
     }
 
     @Override
     protected int doHashCode() {
-        return Objects.hash(featureData, snippets);
+        return Objects.hash(featureData, snippets, docIndices);
     }
 
     @Override
@@ -84,5 +93,6 @@ public String getWriteableName() {
     protected void doToXContent(XContentBuilder builder, Params params) throws IOException {
         builder.field("featureData", featureData);
         builder.array("snippets", snippets);
+        builder.array("docIndices", docIndices);
     }
 }

server/src/main/java/org/elasticsearch/search/rank/rerank/RerankingRankFeaturePhaseRankShardContext.java

@@ -21,10 +21,10 @@
 import org.elasticsearch.search.rank.feature.RankFeatureDoc;
 import org.elasticsearch.search.rank.feature.RankFeatureShardResult;
 
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
-import java.util.stream.Collectors;
 
 /**
  * The {@code ReRankingRankFeaturePhaseRankShardContext} is handles the {@code SearchHits} generated from the {@code RankFeatureShardPhase}
@@ -43,6 +43,7 @@ public RerankingRankFeaturePhaseRankShardContext(String field) {
     public RankShardResult buildRankFeatureShardResult(SearchHits hits, int shardId) {
         try {
             RankFeatureDoc[] rankFeatureDocs = new RankFeatureDoc[hits.getHits().length];
+            int docIndex = 0;
             for (int i = 0; i < hits.getHits().length; i++) {
                 rankFeatureDocs[i] = new RankFeatureDoc(hits.getHits()[i].docId(), hits.getHits()[i].getScore(), shardId);
                 SearchHit hit = hits.getHits()[i];
@@ -53,12 +54,16 @@ public RankShardResult buildRankFeatureShardResult(SearchHits hits, int shardId)
                 Map<String, HighlightField> highlightFields = hit.getHighlightFields();
                 if (highlightFields != null) {
                     if (highlightFields.containsKey(field)) {
-                        List<String> snippets = Arrays.stream(highlightFields.get(field).fragments())
-                            .map(Text::string)
-                            .collect(Collectors.toList());
+                        List<String> snippets = Arrays.stream(highlightFields.get(field).fragments()).map(Text::string).toList();
+                        List<Integer> docIndices = new ArrayList<>();
+                        for (String snippet : snippets) {
+                            docIndices.add(docIndex);
+                        }
                         rankFeatureDocs[i].snippets(snippets);
+                        rankFeatureDocs[i].docIndices(docIndices);
                     }
                 }
+                docIndex++;
             }
             return new RankFeatureShardResult(rankFeatureDocs);
         } catch (Exception ex) {

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/inference/action/InferenceAction.java

@@ -107,7 +107,7 @@ public Request(
             String query,
             Boolean returnDocuments,
             Integer topN,
-            List<String> input, // I think we need to add some metadata to the strings here and return this with each response
+            List<String> input,
             Map<String, Object> taskSettings,
             InputType inputType,
             TimeValue inferenceTimeout,

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

@@ -36,6 +36,7 @@
 import static org.elasticsearch.xpack.inference.rank.textsimilarity.TextSimilarityRankRetrieverBuilder.INFERENCE_ID_FIELD;
 import static org.elasticsearch.xpack.inference.rank.textsimilarity.TextSimilarityRankRetrieverBuilder.INFERENCE_TEXT_FIELD;
 import static org.elasticsearch.xpack.inference.rank.textsimilarity.TextSimilarityRankRetrieverBuilder.MIN_SCORE_FIELD;
+import static org.elasticsearch.xpack.inference.rank.textsimilarity.TextSimilarityRankRetrieverBuilder.SNIPPETS_FIELD;
 
 /**
  * A {@code RankBuilder} that enables ranking with text similarity model inference. Supports parameters for configuring the inference call.
@@ -133,7 +134,7 @@ public void doXContent(XContentBuilder builder, Params params) throws IOExceptio
             builder.field(FAILURES_ALLOWED_FIELD.getPreferredName(), true);
         }
         if (snippets != null) {
-
+            builder.field(SNIPPETS_FIELD.getPreferredName(), snippets);
         }
     }
 

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

@@ -22,6 +22,7 @@
 import org.elasticsearch.xpack.inference.services.googlevertexai.rerank.GoogleVertexAiRerankTaskSettings;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 
@@ -57,30 +58,32 @@ public TextSimilarityRankFeaturePhaseRankCoordinatorContext(
     @Override
     protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[]> scoreListener) {
 
-        // Reconcile the input strings with the documents that they belong to. Input size 6.
-        // Let's say we have 6 snippets that we reranked from 2 documents (3 snippets each)
-
         // 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> inferenceListener = scoreListener.delegateFailureAndWrap((l, r) -> {
             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();
+            final float[] scores;
+            if (featureDocs.length > 0 && featureDocs[0].snippets != null) {
+                scores = extractScoresFromRankedSnippets(rankedDocs, featureDocs);
+            } else {
+                scores = extractScoresFromRankedDocs(rankedDocs);
+            }
 
-            if (rankedDocs.size() != featureDocs.length) {
+            // Ensure we get exactly as many final scores as the number of docs we passed, otherwise we may return incorrect results
+            if (scores.length != featureDocs.length) {
                 l.onFailure(
                     new IllegalStateException(
                         "Reranker input document count and returned score count mismatch: ["
                             + featureDocs.length
                             + "] vs ["
-                            + rankedDocs.size()
+                            + scores.length
                             + "]"
                     )
                 );
             } else {
-                float[] scores = extractScoresFromRankedDocs(rankedDocs); // Return is size 2
                 l.onResponse(scores);
             }
         });
@@ -119,8 +122,7 @@ protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[
                 List<String> inferenceInputs = new ArrayList<>();
                 for (RankFeatureDoc featureDoc : featureDocs) {
                     if (featureDoc.snippets != null && featureDoc.snippets.isEmpty() == false) {
-                        // TODO support reranking multiple snippets
-                        inferenceInputs.add(featureDoc.snippets.get(0));
+                        inferenceInputs.addAll(featureDoc.snippets);
                     } else {
                         inferenceInputs.add(featureDoc.featureData);
                     }
@@ -181,7 +183,33 @@ private float[] extractScoresFromRankedDocs(List<RankedDocsResults.RankedDoc> ra
         for (RankedDocsResults.RankedDoc rankedDoc : rankedDocs) {
             scores[rankedDoc.index()] = rankedDoc.relevanceScore();
         }
-        return scores; // Return a float of size 2 (max score index per doc)
+        return scores;
+    }
+
+    private float[] extractScoresFromRankedSnippets(List<RankedDocsResults.RankedDoc> rankedDocs, RankFeatureDoc[] featureDocs) {
+        int[] docMappings = Arrays.stream(featureDocs).flatMapToInt(f -> f.docIndices.stream().mapToInt(Integer::intValue)).toArray();
+
+        float[] scores = new float[featureDocs.length];
+        boolean[] hasScore = new boolean[featureDocs.length];
+
+        for (int i = 0; i < rankedDocs.size(); i++) {
+            int docId = docMappings[i];
+            float score = rankedDocs.get(i).relevanceScore();
+
+            if (hasScore[docId] == false) {
+                scores[docId] = score;
+                hasScore[docId] = true;
+            } else {
+                scores[docId] = Math.max(scores[docId], score);
+            }
+        }
+
+        float[] result = new float[featureDocs.length];
+        for (int i = 0; i < featureDocs.length; i++) {
+            result[i] = hasScore[i] ? normalizeScore(scores[i]) : 0f;
+        }
+
+        return result;
     }
 
     private static float normalizeScore(float score) {