Transforming rank rrf to the corresponding retriever (#115026)

Author: pmpailis

server/src/main/java/org/elasticsearch/rest/action/search/SearchCapabilities.java

@@ -24,10 +24,13 @@ private SearchCapabilities() {}
     private static final String BIT_DENSE_VECTOR_SYNTHETIC_SOURCE_CAPABILITY = "bit_dense_vector_synthetic_source";
     /** Support Byte and Float with Bit dot product. */
     private static final String BYTE_FLOAT_BIT_DOT_PRODUCT_CAPABILITY = "byte_float_bit_dot_product";
+    /** Support transforming rank rrf queries to the corresponding rrf retriever. */
+    private static final String TRANSFORM_RANK_RRF_TO_RETRIEVER = "transform_rank_rrf_to_retriever";
 
     public static final Set<String> CAPABILITIES = Set.of(
         RANGE_REGEX_INTERVAL_QUERY_CAPABILITY,
         BIT_DENSE_VECTOR_SYNTHETIC_SOURCE_CAPABILITY,
-        BYTE_FLOAT_BIT_DOT_PRODUCT_CAPABILITY
+        BYTE_FLOAT_BIT_DOT_PRODUCT_CAPABILITY,
+        TRANSFORM_RANK_RRF_TO_RETRIEVER
     );
 }

server/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java

@@ -1638,6 +1638,18 @@ private SearchSourceBuilder parseXContent(
         }
 
         knnSearch = knnBuilders.stream().map(knnBuilder -> knnBuilder.build(size())).collect(Collectors.toList());
+        if (rankBuilder != null) {
+            if (retrieverBuilder != null) {
+                throw new IllegalArgumentException("Cannot specify both [rank] and [retriever].");
+            }
+            RetrieverBuilder transformedRetriever = rankBuilder.toRetriever(this, clusterSupportsFeature);
+            if (transformedRetriever != null) {
+                this.retriever(transformedRetriever);
+                rankBuilder = null;
+                subSearchSourceBuilders.clear();
+                knnSearch.clear();
+            }
+        }
         searchUsageConsumer.accept(searchUsage);
         return this;
     }

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

@@ -16,18 +16,24 @@
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.io.stream.VersionedNamedWriteable;
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.core.UpdateForV10;
+import org.elasticsearch.features.NodeFeature;
 import org.elasticsearch.search.SearchService;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.rank.context.QueryPhaseRankCoordinatorContext;
 import org.elasticsearch.search.rank.context.QueryPhaseRankShardContext;
 import org.elasticsearch.search.rank.context.RankFeaturePhaseRankCoordinatorContext;
 import org.elasticsearch.search.rank.context.RankFeaturePhaseRankShardContext;
+import org.elasticsearch.search.retriever.RetrieverBuilder;
 import org.elasticsearch.xcontent.ParseField;
 import org.elasticsearch.xcontent.ToXContentObject;
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
 import java.util.List;
 import java.util.Objects;
+import java.util.function.Predicate;
 
 /**
  * {@code RankBuilder} is used as a base class to manage input, parsing, and subsequent generation of appropriate contexts
@@ -109,6 +115,16 @@ public int rankWindowSize() {
      */
     public abstract RankFeaturePhaseRankCoordinatorContext buildRankFeaturePhaseCoordinatorContext(int size, int from, Client client);
 
+    /**
+     * Transforms the specific rank builder (as parsed through SearchSourceBuilder) to the corresponding retriever.
+     * This is used to ensure smooth deprecation of `rank` and `sub_searches` and move towards the retriever framework
+     */
+    @UpdateForV10(owner = UpdateForV10.Owner.SEARCH_RELEVANCE) // remove for 10.0 once we remove support for the rank parameter in SearchAPI
+    @Nullable
+    public RetrieverBuilder toRetriever(SearchSourceBuilder searchSourceBuilder, Predicate<NodeFeature> clusterSupportsFeature) {
+        return null;
+    }
+
     @Override
     public final boolean equals(Object obj) {
         if (this == obj) {

server/src/main/java/org/elasticsearch/search/vectors/KnnSearchBuilder.java

@@ -286,6 +286,10 @@ public int k() {
         return k;
     }
 
+    public int getNumCands() {
+        return numCands;
+    }
+
     public QueryVectorBuilder getQueryVectorBuilder() {
         return queryVectorBuilder;
     }

x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRankBuilder.java

@@ -14,23 +14,32 @@
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.features.NodeFeature;
 import org.elasticsearch.license.LicenseUtils;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
 import org.elasticsearch.search.rank.RankBuilder;
 import org.elasticsearch.search.rank.RankDoc;
 import org.elasticsearch.search.rank.context.QueryPhaseRankCoordinatorContext;
 import org.elasticsearch.search.rank.context.QueryPhaseRankShardContext;
 import org.elasticsearch.search.rank.context.RankFeaturePhaseRankCoordinatorContext;
 import org.elasticsearch.search.rank.context.RankFeaturePhaseRankShardContext;
+import org.elasticsearch.search.retriever.CompoundRetrieverBuilder;
+import org.elasticsearch.search.retriever.KnnRetrieverBuilder;
+import org.elasticsearch.search.retriever.RetrieverBuilder;
+import org.elasticsearch.search.retriever.StandardRetrieverBuilder;
+import org.elasticsearch.search.vectors.KnnSearchBuilder;
 import org.elasticsearch.xcontent.ConstructingObjectParser;
 import org.elasticsearch.xcontent.ParseField;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xpack.core.XPackPlugin;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
+import java.util.function.Predicate;
 
 import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg;
 
@@ -183,6 +192,37 @@ public RankFeaturePhaseRankCoordinatorContext buildRankFeaturePhaseCoordinatorCo
         return null;
     }
 
+    @Override
+    public RetrieverBuilder toRetriever(SearchSourceBuilder source, Predicate<NodeFeature> clusterSupportsFeature) {
+        if (false == clusterSupportsFeature.test(RRFRetrieverBuilder.RRF_RETRIEVER_COMPOSITION_SUPPORTED)) {
+            return null;
+        }
+        int totalQueries = source.subSearches().size() + source.knnSearch().size();
+        if (totalQueries < 2) {
+            throw new IllegalArgumentException("[rrf] requires at least 2 sub-queries to be defined");
+        }
+        List<CompoundRetrieverBuilder.RetrieverSource> retrieverSources = new ArrayList<>(totalQueries);
+        for (int i = 0; i < source.subSearches().size(); i++) {
+            RetrieverBuilder standardRetriever = new StandardRetrieverBuilder(source.subSearches().get(i).getQueryBuilder());
+            standardRetriever.retrieverName(source.subSearches().get(i).getQueryBuilder().queryName());
+            retrieverSources.add(new CompoundRetrieverBuilder.RetrieverSource(standardRetriever, null));
+        }
+        for (int i = 0; i < source.knnSearch().size(); i++) {
+            KnnSearchBuilder knnSearchBuilder = source.knnSearch().get(i);
+            RetrieverBuilder knnRetriever = new KnnRetrieverBuilder(
+                knnSearchBuilder.getField(),
+                knnSearchBuilder.getQueryVector().asFloatVector(),
+                knnSearchBuilder.getQueryVectorBuilder(),
+                knnSearchBuilder.k(),
+                knnSearchBuilder.getNumCands(),
+                knnSearchBuilder.getSimilarity()
+            );
+            knnRetriever.retrieverName(knnSearchBuilder.queryName());
+            retrieverSources.add(new CompoundRetrieverBuilder.RetrieverSource(knnRetriever, null));
+        }
+        return new RRFRetrieverBuilder(retrieverSources, rankWindowSize(), rankConstant());
+    }
+
     @Override
     protected boolean doEquals(RankBuilder other) {
         return Objects.equals(rankConstant, ((RRFRankBuilder) other).rankConstant);

x-pack/plugin/rank-rrf/src/yamlRestTest/resources/rest-api-spec/test/rrf/100_rank_rrf.yml

@@ -1,7 +1,11 @@
 setup:
   - requires:
-      cluster_features: "gte_v8.8.0"
-      reason: 'rank added in 8.8'
+      capabilities:
+        - method: POST
+          path: /_search
+          capabilities: [ transform_rank_rrf_to_retriever ]
+      test_runner_features: capabilities
+      reason: "Support for transforming deprecated rank_rrf queries to the corresponding rrf retriever is required"
   - skip:
       features: "warnings"
 
@@ -212,7 +216,7 @@ setup:
 "RRF rank should fail if size > rank_window_size":
 
   - do:
-      catch: "/\\[rank\\] requires \\[rank_window_size: 2\\] be greater than or equal to \\[size: 10\\]/"
+      catch: "/\\[rrf\\] requires \\[rank_window_size: 2\\] be greater than or equal to \\[size: 10\\]/"
       search:
         index: test
         body:
@@ -284,3 +288,22 @@ setup:
               rank_window_size: 10
               rank_constant: 0.3
           size: 10
+
+---
+"RRF rank should fail if we specify both rank and retriever":
+  - do:
+      catch: "/Cannot specify both \\[rank\\] and \\[retriever\\]./"
+      search:
+        index: test
+        body:
+          track_total_hits: true
+          fields: [ "text", "keyword" ]
+          retriever:
+            standard:
+              query:
+                match_all: {}
+          rank:
+            rrf:
+              rank_window_size: 10
+              rank_constant: 10
+          size: 10