Implement comprehensive top N parameter handling for text similarity reranker

docs/changelog/142039.yaml

@@ -0,0 +1,5 @@
+area: Ranking
+issues: []
+pr: 142039
+summary: Implement comprehensive top N parameter handling for text similarity reranker
+type: bug

server/src/main/java/org/elasticsearch/inference/TopNProvider.java

@@ -0,0 +1,14 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.inference;
+
+public interface TopNProvider {
+    Integer getTopN();
+}

x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestRerankingServiceExtension.java

@@ -31,6 +31,7 @@
 import org.elasticsearch.inference.SettingsConfiguration;
 import org.elasticsearch.inference.TaskSettings;
 import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.inference.TopNProvider;
 import org.elasticsearch.inference.UnifiedCompletionRequest;
 import org.elasticsearch.inference.configuration.SettingsConfigurationFieldType;
 import org.elasticsearch.rest.RestStatus;
@@ -40,7 +41,6 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
-import java.util.Comparator;
 import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.List;
@@ -194,7 +194,14 @@ private RankedDocsResults makeResults(List<String> input, TestRerankingServiceEx
                 for (int i = 0; i < totalResults; i++) {
                     results.add(new RankedDocsResults.RankedDoc(i, Float.parseFloat(input.get(i)), input.get(i)));
                 }
-                return new RankedDocsResults(results.stream().sorted(Comparator.reverseOrder()).toList());
+
+                // RankedDoc's compareTo implementation already sorts by score descending, so we don't need to reverse the sort order
+                var sortedResultsStream = results.stream().sorted();
+                if (taskSettings.topN != null) {
+                    sortedResultsStream = sortedResultsStream.limit(taskSettings.topN);
+                }
+
+                return new RankedDocsResults(sortedResultsStream.toList());
             } catch (NumberFormatException ex) {
                 return makeResultFromTextInput(input, taskSettings);
             }
@@ -216,6 +223,10 @@ private RankedDocsResults makeResultFromTextInput(List<String> input, TestRerank
             }
             // Ensure result are sorted by descending score
             results.sort((a, b) -> -Float.compare(a.relevanceScore(), b.relevanceScore()));
+            if (taskSettings.topN != null && taskSettings.topN < results.size()) {
+                results = results.subList(0, taskSettings.topN);
+            }
+
             return new RankedDocsResults(results);
         }
 
@@ -257,9 +268,14 @@ public static InferenceServiceConfiguration get() {
         }
     }
 
-    public record TestTaskSettings(boolean shouldFailValidation, boolean useTextLength, float minScore, float resultDiff)
-        implements
-            TaskSettings {
+    public record TestTaskSettings(
+        boolean shouldFailValidation,
+        boolean useTextLength,
+        float minScore,
+        float resultDiff,
+        Integer topN,
+        boolean hideTopN
+    ) implements TaskSettings, TopNProvider {
 
         static final String NAME = "test_reranking_task_settings";
 
@@ -268,6 +284,8 @@ public static TestTaskSettings fromMap(Map<String, Object> map) {
             boolean useTextLength = false;
             float minScore = random.nextFloat(-1f, 1f);
             float resultDiff = 0.2f;
+            Integer topN = null;
+            boolean hideTopN = false;
 
             if (map.containsKey("should_fail_validation")) {
                 shouldFailValidation = Boolean.parseBoolean(map.remove("should_fail_validation").toString());
@@ -285,11 +303,19 @@ public static TestTaskSettings fromMap(Map<String, Object> map) {
                 resultDiff = Float.parseFloat(map.remove("result_diff").toString());
             }
 
-            return new TestTaskSettings(shouldFailValidation, useTextLength, minScore, resultDiff);
+            if (map.containsKey("top_n")) {
+                topN = Integer.parseInt(map.remove("top_n").toString());
+            }
+
+            if (map.containsKey("hide_top_n")) {
+                hideTopN = Boolean.parseBoolean(map.remove("hide_top_n").toString());
+            }
+
+            return new TestTaskSettings(shouldFailValidation, useTextLength, minScore, resultDiff, topN, hideTopN);
         }
 
         public TestTaskSettings(StreamInput in) throws IOException {
-            this(in.readBoolean(), in.readBoolean(), in.readFloat(), in.readFloat());
+            this(in.readBoolean(), in.readBoolean(), in.readFloat(), in.readFloat(), in.readOptionalInt(), in.readBoolean());
         }
 
         @Override
@@ -303,18 +329,30 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeBoolean(useTextLength);
             out.writeFloat(minScore);
             out.writeFloat(resultDiff);
+            out.writeOptionalInt(topN);
+            out.writeBoolean(hideTopN);
         }
 
         @Override
         public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
             builder.startObject();
+            builder.field("should_fail_validation", shouldFailValidation);
             builder.field("use_text_length", useTextLength);
             builder.field("min_score", minScore);
             builder.field("result_diff", resultDiff);
+            if (topN != null) {
+                builder.field("top_n", topN);
+            }
+            builder.field("hide_top_n", hideTopN);
             builder.endObject();
             return builder;
         }
 
+        @Override
+        public Integer getTopN() {
+            return hideTopN ? null : topN;
+        }
+
         @Override
         public String getWriteableName() {
             return NAME;
@@ -332,7 +370,9 @@ public TaskSettings updatedTaskSettings(Map<String, Object> newSettingsMap) {
                 newSettingsMap.containsKey("should_fail_validation") ? newSettingsObject.shouldFailValidation() : shouldFailValidation,
                 newSettingsMap.containsKey("use_text_length") ? newSettingsObject.useTextLength() : useTextLength,
                 newSettingsMap.containsKey("min_score") ? newSettingsObject.minScore() : minScore,
-                newSettingsMap.containsKey("result_diff") ? newSettingsObject.resultDiff() : resultDiff
+                newSettingsMap.containsKey("result_diff") ? newSettingsObject.resultDiff() : resultDiff,
+                newSettingsMap.containsKey("top_n") ? newSettingsObject.topN() : topN,
+                newSettingsMap.containsKey("hide_top_n") ? newSettingsObject.hideTopN() : hideTopN
             );
         }
     }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferenceFeatures.java

@@ -69,6 +69,9 @@ public class InferenceFeatures implements FeatureSpecification {
     public static final NodeFeature TEXT_SIMILARITY_RERANKER_INFERENCE_ID_CHUNKING = new NodeFeature(
         "text_similarity_reranker_inference_id_chunking"
     );
+    public static final NodeFeature TEXT_SIMILARITY_RERANKER_COMPREHENSIVE_TOP_N_HANDLING = new NodeFeature(
+        "text_similarity_reranker.comprehensive_top_n_handling"
+    );
     public static final NodeFeature INFERENCE_AUTH_POLLER_PERSISTENT_TASK = new NodeFeature("inference.auth_poller.persistent_task");
     public static final NodeFeature INFERENCE_CCM_ENABLEMENT_SERVICE = new NodeFeature("inference.ccm.enablement_service");
 
@@ -133,7 +136,8 @@ public Set<NodeFeature> getTestFeatures() {
                 SEARCH_USAGE_EXTENDED_DATA,
                 TEXT_SIMILARITY_RANK_DOC_EXPLAIN_CHUNKS,
                 RETRIEVER_RESULT_DIVERSIFICATION_USES_QUERY_VECTOR_BUILDER,
-                TEXT_SIMILARITY_RERANKER_INFERENCE_ID_CHUNKING
+                TEXT_SIMILARITY_RERANKER_INFERENCE_ID_CHUNKING,
+                TEXT_SIMILARITY_RERANKER_COMPREHENSIVE_TOP_N_HANDLING
             )
         );
         testFeatures.addAll(getFeatures());

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

@@ -196,10 +196,7 @@ public RankFeaturePhaseRankCoordinatorContext buildRankFeaturePhaseCoordinatorCo
             inferenceId,
             inferenceText,
             minScore,
-            failuresAllowed,
-            chunkScorerConfig != null
-                ? new ChunkScorerConfig(chunkScorerConfig.size, inferenceText, chunkScorerConfig.chunkingSettings())
-                : null
+            failuresAllowed
         );
     }