Remove docIndices from RankFeatureDoc

Author: kderusso

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

@@ -29,7 +29,6 @@ public class RankFeatureDoc extends RankDoc {
 
     // TODO: update to support more than 1 fields; and not restrict to string data
     public List<String> featureData;
-    public List<Integer> docIndices;
 
     public RankFeatureDoc(int doc, float score, int shardIndex) {
         super(doc, score, shardIndex);
@@ -39,7 +38,6 @@ public RankFeatureDoc(StreamInput in) throws IOException {
         super(in);
         if (in.getTransportVersion().onOrAfter(TransportVersions.RERANK_SNIPPETS)) {
             featureData = in.readOptionalStringCollectionAsList();
-            docIndices = in.readOptionalCollectionAsList(StreamInput::readVInt);
         } else {
             String featureDataString = in.readOptionalString();
             featureData = featureDataString == null ? null : List.of(featureDataString);
@@ -55,15 +53,14 @@ public void featureData(List<String> featureData) {
         this.featureData = featureData;
     }
 
-    public void docIndices(List<Integer> docIndices) {
-        this.docIndices = docIndices;
+    public void featureData(String featureData) {
+        this.featureData = List.of(featureData);
     }
 
     @Override
     protected void doWriteTo(StreamOutput out) throws IOException {
         if (out.getTransportVersion().onOrAfter(TransportVersions.RERANK_SNIPPETS)) {
             out.writeOptionalStringCollection(featureData);
-            out.writeOptionalCollection(docIndices, StreamOutput::writeVInt);
         } else {
             out.writeOptionalString(featureData.get(0));
         }
@@ -72,12 +69,12 @@ protected void doWriteTo(StreamOutput out) throws IOException {
     @Override
     protected boolean doEquals(RankDoc rd) {
         RankFeatureDoc other = (RankFeatureDoc) rd;
-        return Objects.equals(this.featureData, other.featureData) && Objects.equals(this.docIndices, other.docIndices);
+        return Objects.equals(this.featureData, other.featureData);
     }
 
     @Override
     protected int doHashCode() {
-        return Objects.hash(featureData, docIndices);
+        return Objects.hash(featureData);
     }
 
     @Override
@@ -88,6 +85,5 @@ public String getWriteableName() {
     @Override
     protected void doToXContent(XContentBuilder builder, Params params) throws IOException {
         builder.array("featureData", featureData);
-        builder.array("docIndices", docIndices);
     }
 }

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

@@ -22,7 +22,6 @@
 import org.elasticsearch.search.rank.feature.RerankSnippetInput;
 import org.elasticsearch.xcontent.Text;
 
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
@@ -50,27 +49,20 @@ public RerankingRankFeaturePhaseRankShardContext(String field, RerankSnippetInpu
     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];
                 DocumentField docField = hit.field(field);
                 if (docField != null && snippets == null) {
-                    rankFeatureDocs[i].featureData(List.of(docField.getValue().toString()));
+                    rankFeatureDocs[i].featureData(docField.getValue().toString());
                 }
                 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).toList();
-                        List<Integer> docIndices = new ArrayList<>();
-                        for (String s : snippets) {
-                            docIndices.add(docIndex);
-                        }
                         rankFeatureDocs[i].featureData(snippets);
-                        rankFeatureDocs[i].docIndices(docIndices);
                     }
                 }
-                docIndex++;
             }
             return new RankFeatureShardResult(rankFeatureDocs);
         } catch (Exception ex) {

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

@@ -23,7 +23,6 @@
 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;
 
@@ -192,13 +191,13 @@ private float[] extractScoresFromRankedDocs(List<RankedDocsResults.RankedDoc> ra
     }
 
     private float[] extractScoresFromRankedSnippets(List<RankedDocsResults.RankedDoc> rankedDocs, RankFeatureDoc[] featureDocs) {
-        int[] docMappings = Arrays.stream(featureDocs).flatMapToInt(f -> f.docIndices.stream().mapToInt(Integer::intValue)).toArray();
-
+        int numSnippets = rankedDocs.size() / featureDocs.length;
         float[] scores = new float[featureDocs.length];
         boolean[] hasScore = new boolean[featureDocs.length];
 
         for (int i = 0; i < rankedDocs.size(); i++) {
-            int docId = docMappings[i];
+            // TODO this naively assumes that we always get the requested number of snippets per ranked document
+            int docId = i / numSnippets;
             float score = rankedDocs.get(i).relevanceScore();
 
             if (hasScore[docId] == false) {