Optionally allow text similarity reranking to fail

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -177,8 +177,8 @@ static TransportVersion def(int id) {
     public static final TransportVersion COHERE_BIT_EMBEDDING_TYPE_SUPPORT_ADDED = def(9_001_0_00);
     public static final TransportVersion REMOVE_SNAPSHOT_FAILURES = def(9_002_0_00);
     public static final TransportVersion TRANSPORT_STATS_HANDLING_TIME_REQUIRED = def(9_003_0_00);
-
     public static final TransportVersion REMOVE_DESIRED_NODE_VERSION = def(9_004_0_00);
+    public static final TransportVersion LENIENT_RERANKERS = def(9_005_0_00);
 
     /*
      * STOP! READ THIS FIRST! No, really,

server/src/main/java/org/elasticsearch/action/search/RankFeaturePhase.java

@@ -26,6 +26,7 @@
 import org.elasticsearch.search.rank.feature.RankFeatureShardRequest;
 import org.elasticsearch.transport.Transport;
 
+import java.util.Arrays;
 import java.util.List;
 
 /**
@@ -181,6 +182,11 @@ private void onPhaseDone(
         RankFeaturePhaseRankCoordinatorContext rankFeaturePhaseRankCoordinatorContext,
         SearchPhaseController.ReducedQueryPhase reducedQueryPhase
     ) {
+        RankFeatureDoc[] docs = rankPhaseResults.getSuccessfulResults()
+            .flatMap(r -> Arrays.stream(r.rankFeatureResult().shardResult().rankFeatureDocs))
+            .filter(rfd -> rfd.featureData != null)
+            .toArray(RankFeatureDoc[]::new);
+
         ThreadedActionListener<RankFeatureDoc[]> rankResultListener = new ThreadedActionListener<>(
             context::execute,
             new ActionListener<>() {
@@ -196,21 +202,26 @@ public void onResponse(RankFeatureDoc[] docsWithUpdatedScores) {
 
                 @Override
                 public void onFailure(Exception e) {
-                    context.onPhaseFailure(NAME, "Computing updated ranks for results failed", e);
+                    if (rankFeaturePhaseRankCoordinatorContext.isLenient()) {
+                        // TODO: handle the exception somewhere
+                        logger.warn("Exception computing updated ranks. Continuing with existing ranks.", e);
+                        // use the existing docs as-is
+                        // AbstractThreadedActionListener forks onFailure to the same executor as onResponse,
+                        // so we can just call this direct
+                        onResponse(docs);
+                    } else {
+                        context.onPhaseFailure(NAME, "Computing updated ranks for results failed", e);
+                    }
                 }
             }
         );
-        rankFeaturePhaseRankCoordinatorContext.computeRankScoresForGlobalResults(
-            rankPhaseResults.getAtomicArray().asList().stream().map(SearchPhaseResult::rankFeatureResult).toList(),
-            rankResultListener
-        );
+        rankFeaturePhaseRankCoordinatorContext.computeRankScoresForGlobalResults(docs, rankResultListener);
     }
 
     private SearchPhaseController.ReducedQueryPhase newReducedQueryPhaseResults(
         SearchPhaseController.ReducedQueryPhase reducedQueryPhase,
         ScoreDoc[] scoreDocs
     ) {
-
         return new SearchPhaseController.ReducedQueryPhase(
             reducedQueryPhase.totalHits(),
             reducedQueryPhase.fetchHits(),

server/src/main/java/org/elasticsearch/search/rank/RankBuilder.java

@@ -11,6 +11,7 @@
 
 import org.apache.lucene.search.Explanation;
 import org.apache.lucene.search.Query;
+import org.elasticsearch.TransportVersions;
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -42,21 +43,32 @@
 public abstract class RankBuilder implements VersionedNamedWriteable, ToXContentObject {
 
     public static final ParseField RANK_WINDOW_SIZE_FIELD = new ParseField("rank_window_size");
+    public static final ParseField LENIENT_FIELD = new ParseField("lenient");
 
     public static final int DEFAULT_RANK_WINDOW_SIZE = SearchService.DEFAULT_SIZE;
 
     private final int rankWindowSize;
+    private final boolean lenient;
 
-    public RankBuilder(int rankWindowSize) {
+    public RankBuilder(int rankWindowSize, boolean lenient) {
         this.rankWindowSize = rankWindowSize;
+        this.lenient = lenient;
     }
 
     public RankBuilder(StreamInput in) throws IOException {
         rankWindowSize = in.readVInt();
+        if (in.getTransportVersion().onOrAfter(TransportVersions.LENIENT_RERANKERS)) {
+            lenient = in.readBoolean();
+        } else {
+            lenient = false;
+        }
     }
 
     public final void writeTo(StreamOutput out) throws IOException {
         out.writeVInt(rankWindowSize);
+        if (out.getTransportVersion().onOrAfter(TransportVersions.LENIENT_RERANKERS)) {
+            out.writeBoolean(lenient);
+        }
         doWriteTo(out);
     }
 
@@ -67,6 +79,9 @@ public final XContentBuilder toXContent(XContentBuilder builder, Params params)
         builder.startObject();
         builder.startObject(getWriteableName());
         builder.field(RANK_WINDOW_SIZE_FIELD.getPreferredName(), rankWindowSize);
+        if (lenient) {
+            builder.field(LENIENT_FIELD.getPreferredName(), lenient);
+        }
         doXContent(builder, params);
         builder.endObject();
         builder.endObject();
@@ -79,6 +94,10 @@ public int rankWindowSize() {
         return rankWindowSize;
     }
 
+    public boolean isLenient() {
+        return lenient;
+    }
+
     /**
      * Specify whether this rank builder is a compound builder or not. A compound builder is a rank builder that requires
      * two or more queries to be executed in order to generate the final result.

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

@@ -33,11 +33,17 @@ public abstract class RankFeaturePhaseRankCoordinatorContext {
     protected final int size;
     protected final int from;
     protected final int rankWindowSize;
+    protected final boolean lenient;
 
-    public RankFeaturePhaseRankCoordinatorContext(int size, int from, int rankWindowSize) {
+    public RankFeaturePhaseRankCoordinatorContext(int size, int from, int rankWindowSize, boolean lenient) {
         this.size = size < 0 ? DEFAULT_SIZE : size;
         this.from = from < 0 ? DEFAULT_FROM : from;
         this.rankWindowSize = rankWindowSize;
+        this.lenient = lenient;
+    }
+
+    public boolean isLenient() {
+        return lenient;
     }
 
     /**
@@ -51,9 +57,9 @@ public RankFeaturePhaseRankCoordinatorContext(int size, int from, int rankWindow
      * @param originalDocs documents to process
      */
     protected RankFeatureDoc[] preprocess(RankFeatureDoc[] originalDocs) {
-        return Arrays.stream(originalDocs)
-            .sorted(Comparator.comparing((RankFeatureDoc doc) -> doc.score).reversed())
-            .toArray(RankFeatureDoc[]::new);
+        RankFeatureDoc[] sorted = originalDocs.clone();
+        Arrays.sort(sorted, Comparator.comparing((RankFeatureDoc doc) -> doc.score).reversed());
+        return sorted;
     }
 
     /**
@@ -64,16 +70,10 @@ protected RankFeatureDoc[] preprocess(RankFeatureDoc[] originalDocs) {
      * Once all the scores have been computed, we sort the results, perform any pagination needed, and then call the `onFinish` consumer
      * with the final array of {@link ScoreDoc} results.
      *
-     * @param rankSearchResults a list of rank feature results from each shard
+     * @param featureDocs       an array of rank feature results from each shard
      * @param rankListener      a rankListener to handle the global ranking result
      */
-    public void computeRankScoresForGlobalResults(
-        List<RankFeatureResult> rankSearchResults,
-        ActionListener<RankFeatureDoc[]> rankListener
-    ) {
-        // extract feature data from each shard rank-feature phase result
-        RankFeatureDoc[] featureDocs = extractFeatureDocs(rankSearchResults);
-
+    public void computeRankScoresForGlobalResults(RankFeatureDoc[] featureDocs, ActionListener<RankFeatureDoc[]> rankListener) {
         // generate the final `topResults` results, and pass them to fetch phase through the `rankListener`
         computeScores(featureDocs, rankListener.delegateFailureAndWrap((listener, scores) -> {
             for (int i = 0; i < featureDocs.length; i++) {

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

@@ -775,7 +775,7 @@ public void sendExecuteRankFeature(
     }
 
     private RankFeaturePhaseRankCoordinatorContext defaultRankFeaturePhaseRankCoordinatorContext(int size, int from, int rankWindowSize) {
-        return new RankFeaturePhaseRankCoordinatorContext(size, from, rankWindowSize) {
+        return new RankFeaturePhaseRankCoordinatorContext(size, from, rankWindowSize, false) {
 
             @Override
             protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[]> scoreListener) {
@@ -785,16 +785,8 @@ protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[
             }
 
             @Override
-            public void computeRankScoresForGlobalResults(
-                List<RankFeatureResult> rankSearchResults,
-                ActionListener<RankFeatureDoc[]> rankListener
-            ) {
-                List<RankFeatureDoc> features = new ArrayList<>();
-                for (RankFeatureResult rankFeatureResult : rankSearchResults) {
-                    RankFeatureShardResult shardResult = rankFeatureResult.shardResult();
-                    features.addAll(Arrays.stream(shardResult.rankFeatureDocs).toList());
-                }
-                rankListener.onResponse(features.toArray(new RankFeatureDoc[0]));
+            public void computeRankScoresForGlobalResults(RankFeatureDoc[] featureDocs, ActionListener<RankFeatureDoc[]> rankListener) {
+                rankListener.onResponse(featureDocs);
             }
 
             @Override
@@ -875,7 +867,7 @@ private RankBuilder rankBuilder(
         RankFeaturePhaseRankShardContext rankFeaturePhaseRankShardContext,
         RankFeaturePhaseRankCoordinatorContext rankFeaturePhaseRankCoordinatorContext
     ) {
-        return new RankBuilder(rankWindowSize) {
+        return new RankBuilder(rankWindowSize, false) {
             @Override
             protected void doWriteTo(StreamOutput out) throws IOException {
                 // no-op

server/src/test/java/org/elasticsearch/search/SearchServiceSingleNodeTests.java

@@ -687,7 +687,7 @@ public RankFeaturePhaseRankCoordinatorContext buildRankFeaturePhaseCoordinatorCo
                                     int from,
                                     Client client
                                 ) {
-                                    return new RankFeaturePhaseRankCoordinatorContext(size, from, DEFAULT_RANK_WINDOW_SIZE) {
+                                    return new RankFeaturePhaseRankCoordinatorContext(size, from, DEFAULT_RANK_WINDOW_SIZE, false) {
                                         @Override
                                         protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[]> scoreListener) {
                                             float[] scores = new float[featureDocs.length];
@@ -831,7 +831,7 @@ public RankFeaturePhaseRankCoordinatorContext buildRankFeaturePhaseCoordinatorCo
                                 int from,
                                 Client client
                             ) {
-                                return new RankFeaturePhaseRankCoordinatorContext(size, from, DEFAULT_RANK_WINDOW_SIZE) {
+                                return new RankFeaturePhaseRankCoordinatorContext(size, from, DEFAULT_RANK_WINDOW_SIZE, false) {
                                     @Override
                                     protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[]> scoreListener) {
                                         throw new IllegalStateException("should have failed earlier");
@@ -947,7 +947,7 @@ public RankFeaturePhaseRankCoordinatorContext buildRankFeaturePhaseCoordinatorCo
                                     int from,
                                     Client client
                                 ) {
-                                    return new RankFeaturePhaseRankCoordinatorContext(size, from, DEFAULT_RANK_WINDOW_SIZE) {
+                                    return new RankFeaturePhaseRankCoordinatorContext(size, from, DEFAULT_RANK_WINDOW_SIZE, false) {
                                         @Override
                                         protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[]> scoreListener) {
                                             float[] scores = new float[featureDocs.length];
@@ -1075,7 +1075,7 @@ public RankFeaturePhaseRankCoordinatorContext buildRankFeaturePhaseCoordinatorCo
                                     int from,
                                     Client client
                                 ) {
-                                    return new RankFeaturePhaseRankCoordinatorContext(size, from, DEFAULT_RANK_WINDOW_SIZE) {
+                                    return new RankFeaturePhaseRankCoordinatorContext(size, from, DEFAULT_RANK_WINDOW_SIZE, false) {
                                         @Override
                                         protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[]> scoreListener) {
                                             float[] scores = new float[featureDocs.length];

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

@@ -117,7 +117,7 @@ public boolean isCancelled() {
     }
 
     private RankBuilder getRankBuilder(final String field) {
-        return new RankBuilder(DEFAULT_RANK_WINDOW_SIZE) {
+        return new RankBuilder(DEFAULT_RANK_WINDOW_SIZE, false) {
             @Override
             protected void doWriteTo(StreamOutput out) throws IOException {
                 // no-op
@@ -171,7 +171,7 @@ public RankShardResult buildRankFeatureShardResult(SearchHits hits, int shardId)
             // no work to be done on the coordinator node for the rank feature phase
             @Override
             public RankFeaturePhaseRankCoordinatorContext buildRankFeaturePhaseCoordinatorContext(int size, int from, Client client) {
-                return new RankFeaturePhaseRankCoordinatorContext(size, from, DEFAULT_RANK_WINDOW_SIZE) {
+                return new RankFeaturePhaseRankCoordinatorContext(size, from, DEFAULT_RANK_WINDOW_SIZE, false) {
                     @Override
                     protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[]> scoreListener) {
                         throw new AssertionError("not expected");

test/framework/src/main/java/org/elasticsearch/search/rank/TestRankBuilder.java

@@ -52,7 +52,7 @@ public static TestRankBuilder randomRankBuilder() {
     }
 
     public TestRankBuilder(int windowSize) {
-        super(windowSize);
+        super(windowSize, false);
     }
 
     public TestRankBuilder(StreamInput in) throws IOException {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/random/RandomRankBuilder.java

@@ -61,7 +61,7 @@ public class RandomRankBuilder extends RankBuilder {
     private final Integer seed;
 
     public RandomRankBuilder(int rankWindowSize, String field, Integer seed) {
-        super(rankWindowSize);
+        super(rankWindowSize, false);
 
         if (field == null || field.isEmpty()) {
             throw new IllegalArgumentException("field is required");

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/rank/random/RandomRankFeaturePhaseRankCoordinatorContext.java

@@ -11,8 +11,6 @@
 import org.elasticsearch.search.rank.context.RankFeaturePhaseRankCoordinatorContext;
 import org.elasticsearch.search.rank.feature.RankFeatureDoc;
 
-import java.util.Arrays;
-import java.util.Comparator;
 import java.util.Random;
 
 /**
@@ -24,7 +22,7 @@ public class RandomRankFeaturePhaseRankCoordinatorContext extends RankFeaturePha
     private final Integer seed;
 
     public RandomRankFeaturePhaseRankCoordinatorContext(int size, int from, int rankWindowSize, Integer seed) {
-        super(size, from, rankWindowSize);
+        super(size, from, rankWindowSize, false);
         this.seed = seed;
     }
 
@@ -40,16 +38,4 @@ protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[
         }
         scoreListener.onResponse(scores);
     }
-
-    /**
-     * Sorts documents by score descending.
-     * @param originalDocs documents to process
-     */
-    @Override
-    protected RankFeatureDoc[] preprocess(RankFeatureDoc[] originalDocs) {
-        return Arrays.stream(originalDocs)
-            .sorted(Comparator.comparing((RankFeatureDoc doc) -> doc.score).reversed())
-            .toArray(RankFeatureDoc[]::new);
-    }
-
 }

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

@@ -57,8 +57,9 @@ public class TextSimilarityRankBuilder extends RankBuilder {
         String field = (String) args[2];
         Integer rankWindowSize = args[3] == null ? DEFAULT_RANK_WINDOW_SIZE : (Integer) args[3];
         Float minScore = (Float) args[4];
+        boolean lenient = args[5] != null && (Boolean) args[5];
 
-        return new TextSimilarityRankBuilder(field, inferenceId, inferenceText, rankWindowSize, minScore);
+        return new TextSimilarityRankBuilder(field, inferenceId, inferenceText, rankWindowSize, minScore, lenient);
     });
 
     static {
@@ -67,15 +68,23 @@ public class TextSimilarityRankBuilder extends RankBuilder {
         PARSER.declareString(constructorArg(), FIELD_FIELD);
         PARSER.declareInt(optionalConstructorArg(), RANK_WINDOW_SIZE_FIELD);
         PARSER.declareFloat(optionalConstructorArg(), MIN_SCORE_FIELD);
+        PARSER.declareBoolean(optionalConstructorArg(), LENIENT_FIELD);
     }
 
     private final String inferenceId;
     private final String inferenceText;
     private final String field;
     private final Float minScore;
 
-    public TextSimilarityRankBuilder(String field, String inferenceId, String inferenceText, int rankWindowSize, Float minScore) {
-        super(rankWindowSize);
+    public TextSimilarityRankBuilder(
+        String field,
+        String inferenceId,
+        String inferenceText,
+        int rankWindowSize,
+        Float minScore,
+        boolean lenient
+    ) {
+        super(rankWindowSize, lenient);
         this.inferenceId = inferenceId;
         this.inferenceText = inferenceText;
         this.field = field;
@@ -103,7 +112,7 @@ public TransportVersion getMinimalSupportedVersion() {
 
     @Override
     public void doWriteTo(StreamOutput out) throws IOException {
-        // rankWindowSize serialization is handled by the parent class RankBuilder
+        // rankWindowSize & lenient serialization is handled by the parent class RankBuilder
         out.writeString(inferenceId);
         out.writeString(inferenceText);
         out.writeString(field);
@@ -112,7 +121,7 @@ public void doWriteTo(StreamOutput out) throws IOException {
 
     @Override
     public void doXContent(XContentBuilder builder, Params params) throws IOException {
-        // rankWindowSize serialization is handled by the parent class RankBuilder
+        // rankWindowSize & lenient serialization is handled by the parent class RankBuilder
         builder.field(INFERENCE_ID_FIELD.getPreferredName(), inferenceId);
         builder.field(INFERENCE_TEXT_FIELD.getPreferredName(), inferenceText);
         builder.field(FIELD_FIELD.getPreferredName(), field);
@@ -177,7 +186,8 @@ public RankFeaturePhaseRankCoordinatorContext buildRankFeaturePhaseCoordinatorCo
             client,
             inferenceId,
             inferenceText,
-            minScore
+            minScore,
+            isLenient()
         );
     }