Optionally allow text similarity reranking to fail (#121784) (#128250)

Backport #121784 to 8.19

Author: thecoop

docs/changelog/121784.yaml

@@ -0,0 +1,5 @@
+pr: 121784
+summary: Optionally allow text similarity reranking to fail
+area: Search
+type: enhancement
+issues: []

server/src/internalClusterTest/java/org/elasticsearch/search/rank/FieldBasedRerankerIT.java

@@ -205,7 +205,7 @@ public RankShardResult buildRankFeatureShardResult(SearchHits hits, int shardId)
 
         @Override
         public RankFeaturePhaseRankCoordinatorContext buildRankFeaturePhaseCoordinatorContext(int size, int from, Client client) {
-            return new RankFeaturePhaseRankCoordinatorContext(size, from, rankWindowSize()) {
+            return new RankFeaturePhaseRankCoordinatorContext(size, from, rankWindowSize(), false) {
                 @Override
                 protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[]> scoreListener) {
                     float[] scores = new float[featureDocs.length];
@@ -346,7 +346,7 @@ public RankShardResult buildRankFeatureShardResult(SearchHits hits, int shardId)
         @Override
         public RankFeaturePhaseRankCoordinatorContext buildRankFeaturePhaseCoordinatorContext(int size, int from, Client client) {
             if (this.throwingRankBuilderType == ThrowingRankBuilderType.THROWING_RANK_FEATURE_PHASE_COORDINATOR_CONTEXT)
-                return new RankFeaturePhaseRankCoordinatorContext(size, from, rankWindowSize()) {
+                return new RankFeaturePhaseRankCoordinatorContext(size, from, rankWindowSize(), false) {
                     @Override
                     protected void computeScores(RankFeatureDoc[] featureDocs, ActionListener<float[]> scoreListener) {
                         throw new UnsupportedOperationException("rfc - simulated failure");

server/src/internalClusterTest/java/org/elasticsearch/search/rank/MockedRequestActionBasedRerankerIT.java

@@ -249,7 +249,7 @@ public static class TestRerankingRankFeaturePhaseRankCoordinatorContext extends
             String inferenceText,
             float minScore
         ) {
-            super(size, from, windowSize);
+            super(size, from, windowSize, false);
             this.client = client;
             this.inferenceId = inferenceId;
             this.inferenceText = inferenceText;

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

@@ -223,6 +223,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion FIELD_CAPS_ADD_CLUSTER_ALIAS = def(8_841_0_32);
     public static final TransportVersion INCLUDE_INDEX_MODE_IN_GET_DATA_STREAM_BACKPORT_8_19 = def(8_841_0_33);
     public static final TransportVersion ESQL_HASH_OPERATOR_STATUS_OUTPUT_TIME_8_19 = def(8_841_0_34);
+    public static final TransportVersion RERANKER_FAILURES_ALLOWED_8_19 = def(8_841_0_35);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -20,12 +20,14 @@
 import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.dfs.AggregatedDfs;
 import org.elasticsearch.search.internal.ShardSearchContextId;
+import org.elasticsearch.search.rank.RankDoc;
 import org.elasticsearch.search.rank.context.RankFeaturePhaseRankCoordinatorContext;
 import org.elasticsearch.search.rank.feature.RankFeatureDoc;
 import org.elasticsearch.search.rank.feature.RankFeatureResult;
 import org.elasticsearch.search.rank.feature.RankFeatureShardRequest;
 import org.elasticsearch.transport.Transport;
 
+import java.util.Arrays;
 import java.util.List;
 
 /**
@@ -186,7 +188,7 @@ private void onPhaseDone(
             new ActionListener<>() {
                 @Override
                 public void onResponse(RankFeatureDoc[] docsWithUpdatedScores) {
-                    RankFeatureDoc[] topResults = rankFeaturePhaseRankCoordinatorContext.rankAndPaginate(docsWithUpdatedScores);
+                    RankDoc[] topResults = rankFeaturePhaseRankCoordinatorContext.rankAndPaginate(docsWithUpdatedScores, true);
                     SearchPhaseController.ReducedQueryPhase reducedRankFeaturePhase = newReducedQueryPhaseResults(
                         reducedQueryPhase,
                         topResults
@@ -196,12 +198,36 @@ public void onResponse(RankFeatureDoc[] docsWithUpdatedScores) {
 
                 @Override
                 public void onFailure(Exception e) {
-                    context.onPhaseFailure(NAME, "Computing updated ranks for results failed", e);
+                    if (rankFeaturePhaseRankCoordinatorContext.failuresAllowed()) {
+                        // TODO: handle the exception somewhere
+                        // don't want to log the entire stack trace, it's not helpful here
+                        logger.warn("Exception computing updated ranks, continuing with existing ranks: {}", e.toString());
+                        // use the existing score docs as-is
+                        // downstream things expect every doc to have a score, so we need to infer a score here
+                        // if the doc doesn't otherwise have one. We can use the rank to infer a possible score instead (1/rank).
+                        ScoreDoc[] inputDocs = reducedQueryPhase.sortedTopDocs().scoreDocs();
+                        RankFeatureDoc[] rankDocs = new RankFeatureDoc[inputDocs.length];
+                        for (int i = 0; i < inputDocs.length; i++) {
+                            ScoreDoc doc = inputDocs[i];
+                            rankDocs[i] = new RankFeatureDoc(doc.doc, Float.isNaN(doc.score) ? 1f / (i + 1) : doc.score, doc.shardIndex);
+                        }
+                        RankDoc[] topResults = rankFeaturePhaseRankCoordinatorContext.rankAndPaginate(rankDocs, false);
+                        SearchPhaseController.ReducedQueryPhase reducedRankFeaturePhase = newReducedQueryPhaseResults(
+                            reducedQueryPhase,
+                            topResults
+                        );
+                        moveToNextPhase(rankPhaseResults, reducedRankFeaturePhase);
+                    } else {
+                        context.onPhaseFailure(NAME, "Computing updated ranks for results failed", e);
+                    }
                 }
             }
         );
         rankFeaturePhaseRankCoordinatorContext.computeRankScoresForGlobalResults(
-            rankPhaseResults.getAtomicArray().asList().stream().map(SearchPhaseResult::rankFeatureResult).toList(),
+            rankPhaseResults.getSuccessfulResults()
+                .flatMap(r -> Arrays.stream(r.rankFeatureResult().shardResult().rankFeatureDocs))
+                .filter(rfd -> rfd.featureData != null)
+                .toArray(RankFeatureDoc[]::new),
             rankResultListener
         );
     }
@@ -210,7 +236,6 @@ 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

@@ -117,9 +117,8 @@ public final boolean equals(Object obj) {
         if (obj == null || getClass() != obj.getClass()) {
             return false;
         }
-        @SuppressWarnings("unchecked")
         RankBuilder other = (RankBuilder) obj;
-        return Objects.equals(rankWindowSize, other.rankWindowSize()) && doEquals(other);
+        return rankWindowSize == other.rankWindowSize && doEquals(other);
     }
 
     protected abstract boolean doEquals(RankBuilder other);

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

@@ -12,13 +12,9 @@
 import org.apache.lucene.search.ScoreDoc;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.search.rank.feature.RankFeatureDoc;
-import org.elasticsearch.search.rank.feature.RankFeatureResult;
-import org.elasticsearch.search.rank.feature.RankFeatureShardResult;
 
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Comparator;
-import java.util.List;
 
 import static org.elasticsearch.search.SearchService.DEFAULT_FROM;
 import static org.elasticsearch.search.SearchService.DEFAULT_SIZE;
@@ -33,11 +29,17 @@ public abstract class RankFeaturePhaseRankCoordinatorContext {
     protected final int size;
     protected final int from;
     protected final int rankWindowSize;
+    protected final boolean failuresAllowed;
 
-    public RankFeaturePhaseRankCoordinatorContext(int size, int from, int rankWindowSize) {
+    public RankFeaturePhaseRankCoordinatorContext(int size, int from, int rankWindowSize, boolean failuresAllowed) {
         this.size = size < 0 ? DEFAULT_SIZE : size;
         this.from = from < 0 ? DEFAULT_FROM : from;
         this.rankWindowSize = rankWindowSize;
+        this.failuresAllowed = failuresAllowed;
+    }
+
+    public boolean failuresAllowed() {
+        return failuresAllowed;
     }
 
     /**
@@ -48,12 +50,13 @@ public RankFeaturePhaseRankCoordinatorContext(int size, int from, int rankWindow
 
     /**
      * Preprocesses the provided documents: sorts them by score descending.
-     * @param originalDocs documents to process
+     *
+     * @param originalDocs   documents to process
+     * @param rerankedScores {@code true} if the document scores have been reranked
      */
-    protected RankFeatureDoc[] preprocess(RankFeatureDoc[] originalDocs) {
-        return Arrays.stream(originalDocs)
-            .sorted(Comparator.comparing((RankFeatureDoc doc) -> doc.score).reversed())
-            .toArray(RankFeatureDoc[]::new);
+    protected RankFeatureDoc[] preprocess(RankFeatureDoc[] originalDocs, boolean rerankedScores) {
+        Arrays.sort(originalDocs, Comparator.comparing((RankFeatureDoc doc) -> doc.score).reversed());
+        return originalDocs;
     }
 
     /**
@@ -64,16 +67,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++) {
@@ -86,28 +83,17 @@ public void computeRankScoresForGlobalResults(
     /**
      * Ranks the provided {@link RankFeatureDoc} array and paginates the results based on the `from` and `size` parameters. Filters out
      * documents that have a relevance score less than min_score.
+     *
      * @param rankFeatureDocs documents to process
+     * @param rerankedScores {@code true} if the document scores have been reranked
      */
-    public RankFeatureDoc[] rankAndPaginate(RankFeatureDoc[] rankFeatureDocs) {
-        RankFeatureDoc[] sortedDocs = preprocess(rankFeatureDocs);
+    public RankFeatureDoc[] rankAndPaginate(RankFeatureDoc[] rankFeatureDocs, boolean rerankedScores) {
+        RankFeatureDoc[] sortedDocs = preprocess(rankFeatureDocs, rerankedScores);
         RankFeatureDoc[] topResults = new RankFeatureDoc[Math.max(0, Math.min(size, sortedDocs.length - from))];
         for (int rank = 0; rank < topResults.length; ++rank) {
             topResults[rank] = sortedDocs[from + rank];
             topResults[rank].rank = from + rank + 1;
         }
         return topResults;
     }
-
-    private RankFeatureDoc[] extractFeatureDocs(List<RankFeatureResult> rankSearchResults) {
-        List<RankFeatureDoc> docFeatures = new ArrayList<>();
-        for (RankFeatureResult rankFeatureResult : rankSearchResults) {
-            RankFeatureShardResult shardResult = rankFeatureResult.shardResult();
-            for (RankFeatureDoc rankFeatureDoc : shardResult.rankFeatureDocs) {
-                if (rankFeatureDoc.featureData != null) {
-                    docFeatures.add(rankFeatureDoc);
-                }
-            }
-        }
-        return docFeatures.toArray(new RankFeatureDoc[0]);
-    }
 }

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,20 +785,12 @@ 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
-            public RankFeatureDoc[] rankAndPaginate(RankFeatureDoc[] rankFeatureDocs) {
+            public RankFeatureDoc[] rankAndPaginate(RankFeatureDoc[] rankFeatureDocs, boolean rerankedScores) {
                 Arrays.sort(rankFeatureDocs, Comparator.comparing((RankFeatureDoc doc) -> doc.score).reversed());
                 RankFeatureDoc[] topResults = new RankFeatureDoc[Math.max(0, Math.min(size, rankFeatureDocs.length - from))];
                 // perform pagination

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

@@ -693,7 +693,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];
@@ -837,7 +837,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");
@@ -953,7 +953,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];
@@ -1081,7 +1081,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

@@ -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");