Fix test compilation errors

Author: kderusso

server/src/main/java/org/elasticsearch/search/rank/context/RankFeaturePhaseRankCoordinatorContext.java

@@ -11,6 +11,7 @@
 
 import org.apache.lucene.search.ScoreDoc;
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.search.rank.feature.RankFeatureDoc;
 import org.elasticsearch.search.rank.feature.RerankSnippetInput;
 
@@ -33,12 +34,16 @@ public abstract class RankFeaturePhaseRankCoordinatorContext {
     protected final boolean failuresAllowed;
     protected final RerankSnippetInput snippets;
 
+    public RankFeaturePhaseRankCoordinatorContext(int size, int from, int rankWindowSize, boolean failuresAllowed) {
+        this(size, from, rankWindowSize, failuresAllowed, null);
+    }
+
     public RankFeaturePhaseRankCoordinatorContext(
         int size,
         int from,
         int rankWindowSize,
         boolean failuresAllowed,
-        RerankSnippetInput snippets
+        @Nullable RerankSnippetInput snippets
     ) {
         this.size = size < 0 ? DEFAULT_SIZE : size;
         this.from = from < 0 ? DEFAULT_FROM : from;
@@ -55,7 +60,12 @@ public RerankSnippetInput snippets() {
         return snippets;
     }
 
-    public abstract Integer tokenSizeLimit();
+    /**
+     * @return If snippets are requested, this should be overridden with the token size limit of the associated model.
+     */
+    public Integer tokenSizeLimit() {
+        return 0;
+    }
 
     /**
      * Computes the updated scores for a list of features (i.e. document-based data). We also pass along an ActionListener

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

@@ -53,10 +53,6 @@ public void featureData(List<String> featureData) {
         this.featureData = featureData;
     }
 
-    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)) {

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

@@ -41,15 +41,24 @@ public class RankFeatureShardRequest extends AbstractTransportRequest implements
     private final int[] docIds;
 
     private final RerankSnippetInput snippets;
-    private final int tokenSizeLimit;
+    private final Integer tokenSizeLimit;
+
+    public RankFeatureShardRequest(
+        OriginalIndices originalIndices,
+        ShardSearchContextId contextId,
+        ShardSearchRequest shardSearchRequest,
+        List<Integer> docIds
+    ) {
+        this(originalIndices, contextId, shardSearchRequest, docIds, null, null);
+    }
 
     public RankFeatureShardRequest(
         OriginalIndices originalIndices,
         ShardSearchContextId contextId,
         ShardSearchRequest shardSearchRequest,
         List<Integer> docIds,
         @Nullable RerankSnippetInput snippets,
-        int tokenSizeLimit
+        Integer tokenSizeLimit
     ) {
         this.originalIndices = originalIndices;
         this.shardSearchRequest = shardSearchRequest;
@@ -67,10 +76,10 @@ public RankFeatureShardRequest(StreamInput in) throws IOException {
         contextId = in.readOptionalWriteable(ShardSearchContextId::new);
         if (in.getTransportVersion().onOrAfter(TransportVersions.RERANK_SNIPPETS)) {
             snippets = in.readOptionalWriteable(RerankSnippetInput::new);
-            this.tokenSizeLimit = in.readVInt();
+            this.tokenSizeLimit = in.readOptionalInt();
         } else {
             snippets = null;
-            this.tokenSizeLimit = 0;
+            this.tokenSizeLimit = null;
         }
     }
 
@@ -83,7 +92,7 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeOptionalWriteable(contextId);
         if (out.getTransportVersion().onOrAfter(TransportVersions.RERANK_SNIPPETS)) {
             out.writeOptionalWriteable(snippets);
-            out.writeVInt(tokenSizeLimit);
+            out.writeOptionalInt(tokenSizeLimit);
         }
     }
 
@@ -119,7 +128,7 @@ public RerankSnippetInput snippets() {
         return snippets;
     }
 
-    public int getTokenSizeLimit() {
+    public Integer getTokenSizeLimit() {
         return tokenSizeLimit;
     }
 

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

@@ -54,7 +54,7 @@ public RankShardResult buildRankFeatureShardResult(SearchHits hits, int shardId)
                 SearchHit hit = hits.getHits()[i];
                 DocumentField docField = hit.field(field);
                 if (docField != null && snippets == null) {
-                    rankFeatureDocs[i].featureData(docField.getValue().toString());
+                    rankFeatureDocs[i].featureData(List.of(docField.getValue().toString()));
                 }
                 Map<String, HighlightField> highlightFields = hit.getHighlightFields();
                 if (highlightFields != null) {

server/src/test/java/org/elasticsearch/action/search/RankFeaturePhaseTests.java

@@ -813,7 +813,7 @@ public RankShardResult buildRankFeatureShardResult(SearchHits hits, int shardId)
                     SearchHit hit = hits.getHits()[i];
                     rankFeatureDocs[i] = new RankFeatureDoc(hit.docId(), hit.getScore(), shardId);
                     rankFeatureDocs[i].score += 100f;
-                    rankFeatureDocs[i].featureData("ranked_" + hit.docId());
+                    rankFeatureDocs[i].featureData(List.of("ranked_" + hit.docId()));
                     rankFeatureDocs[i].rank = i + 1;
                 }
                 return new RankFeatureShardResult(rankFeatureDocs);

server/src/test/java/org/elasticsearch/search/rank/RankFeatureShardPhaseTests.java

@@ -160,7 +160,7 @@ public RankShardResult buildRankFeatureShardResult(SearchHits hits, int shardId)
                         for (int i = 0; i < hits.getHits().length; i++) {
                             SearchHit hit = hits.getHits()[i];
                             rankFeatureDocs[i] = new RankFeatureDoc(hit.docId(), hit.getScore(), shardId);
-                            rankFeatureDocs[i].featureData(hit.getFields().get(field).getValue());
+                            rankFeatureDocs[i].featureData(List.of(hit.getFields().get(field).getValue()));
                             rankFeatureDocs[i].rank = i + 1;
                         }
                         return new RankFeatureShardResult(rankFeatureDocs);

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

@@ -13,6 +13,8 @@
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xpack.core.inference.action.GetInferenceModelAction;
 
+import java.util.List;
+
 import static org.elasticsearch.action.support.ActionTestUtils.assertNoFailureListener;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.argThat;
@@ -32,16 +34,17 @@ public class TextSimilarityRankFeaturePhaseRankCoordinatorContextTests extends E
         "my-inference-id",
         "some query",
         0.0f,
-        false
+        false,
+        null
     );
 
     public void testComputeScores() {
         RankFeatureDoc featureDoc1 = new RankFeatureDoc(0, 1.0f, 0);
-        featureDoc1.featureData("text 1");
+        featureDoc1.featureData(List.of("text 1"));
         RankFeatureDoc featureDoc2 = new RankFeatureDoc(1, 3.0f, 1);
-        featureDoc2.featureData("text 2");
+        featureDoc2.featureData(List.of("text 2"));
         RankFeatureDoc featureDoc3 = new RankFeatureDoc(2, 2.0f, 0);
-        featureDoc3.featureData("text 3");
+        featureDoc3.featureData(List.of("text 3"));
         RankFeatureDoc[] featureDocs = new RankFeatureDoc[] { featureDoc1, featureDoc2, featureDoc3 };
 
         subject.computeScores(featureDocs, assertNoFailureListener(f -> assertArrayEquals(new float[] { 1.0f, 3.0f, 2.0f }, f, 0.0f)));

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

@@ -32,7 +32,7 @@ public class TextSimilarityRankMultiNodeTests extends AbstractRerankerIT {
 
     @Override
     protected RankBuilder getRankBuilder(int rankWindowSize, String rankFeatureField) {
-        return new TextSimilarityRankBuilder(rankFeatureField, inferenceId, inferenceText, rankWindowSize, minScore, false);
+        return new TextSimilarityRankBuilder(rankFeatureField, inferenceId, inferenceText, rankWindowSize, minScore, false, null);
     }
 
     @Override
@@ -53,7 +53,8 @@ protected RankBuilder getThrowingRankBuilder(
             inferenceText,
             minScore,
             failuresAllowed,
-            type.name()
+            type.name(),
+            null
         );
     }
 

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

@@ -58,7 +58,8 @@ public static TextSimilarityRankRetrieverBuilder createRandomTextSimilarityRankR
             randomAlphaOfLength(20),
             randomAlphaOfLength(50),
             randomIntBetween(100, 10000),
-            randomBoolean()
+            randomBoolean(),
+            null
         );
     }
 

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

@@ -139,7 +139,8 @@ public void testTelemetryForRRFRetriever() throws IOException {
                         "some_inference_text",
                         "some_field",
                         10,
-                        false
+                        false,
+                        null
                     )
                 )
             );