Add rescore vector builder to KnnSearchBuilder

Author: carlosdelest

server/src/internalClusterTest/java/org/elasticsearch/search/nested/VectorNestedIT.java

@@ -69,7 +69,7 @@ public void testSimpleNested() throws Exception {
 
         assertResponse(
             prepareSearch("test").setKnnSearch(
-                List.of(new KnnSearchBuilder("nested.vector", new float[] { 1, 1, 1 }, 1, 1, null).innerHit(new InnerHitBuilder()))
+                List.of(new KnnSearchBuilder("nested.vector", new float[] { 1, 1, 1 }, 1, 1, null, null).innerHit(new InnerHitBuilder()))
             ).setAllowPartialSearchResults(false),
             response -> assertThat(response.getHits().getHits().length, greaterThan(0))
         );

server/src/internalClusterTest/java/org/elasticsearch/search/profile/dfs/DfsProfilerIT.java

@@ -19,6 +19,7 @@
 import org.elasticsearch.search.profile.query.CollectorResult;
 import org.elasticsearch.search.profile.query.QueryProfileShardResult;
 import org.elasticsearch.search.vectors.KnnSearchBuilder;
+import org.elasticsearch.search.vectors.RescoreVectorBuilder;
 import org.elasticsearch.test.ESIntegTestCase;
 import org.elasticsearch.xcontent.XContentFactory;
 
@@ -71,6 +72,7 @@ public void testProfileDfs() throws Exception {
                 new float[] { randomFloat(), randomFloat(), randomFloat() },
                 randomIntBetween(5, 10),
                 50,
+                randomBoolean() ? null : new RescoreVectorBuilder(randomFloatBetween(1.0f, 10.0f, false)),
                 randomBoolean() ? null : randomFloat()
             );
             if (randomBoolean()) {

server/src/internalClusterTest/java/org/elasticsearch/search/retriever/RetrieverTelemetryIT.java

@@ -84,7 +84,9 @@ public void testTelemetryForRetrievers() throws IOException {
 
         // search#1 - this will record 1 entry for "retriever" in `sections`, and 1 for "knn" under `retrievers`
         {
-            performSearch(new SearchSourceBuilder().retriever(new KnnRetrieverBuilder("vector", new float[] { 1.0f }, null, 10, 15, null)));
+            performSearch(
+                new SearchSourceBuilder().retriever(new KnnRetrieverBuilder("vector", new float[] { 1.0f }, null, 10, 15, null, null))
+            );
         }
 
         // search#2 - this will record 1 entry for "retriever" in `sections`, 1 for "standard" under `retrievers`, and 1 for "range" under
@@ -112,7 +114,9 @@ public void testTelemetryForRetrievers() throws IOException {
         // search#5 - t
         // his will record 1 entry for "knn" in `sections`
         {
-            performSearch(new SearchSourceBuilder().knnSearch(List.of(new KnnSearchBuilder("vector", new float[] { 1.0f }, 10, 15, null))));
+            performSearch(
+                new SearchSourceBuilder().knnSearch(List.of(new KnnSearchBuilder("vector", new float[] { 1.0f }, 10, 15, null, null)))
+            );
         }
 
         // search#6 - this will record 1 entry for "query" in `sections`, and 1 for "match_all" under `queries`

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

@@ -56,6 +56,7 @@ public class KnnSearchBuilder implements Writeable, ToXContentFragment, Rewritea
     public static final ParseField NAME_FIELD = AbstractQueryBuilder.NAME_FIELD;
     public static final ParseField BOOST_FIELD = AbstractQueryBuilder.BOOST_FIELD;
     public static final ParseField INNER_HITS_FIELD = new ParseField("inner_hits");
+    public static final ParseField RESCORE_FIELD = new ParseField("rescore");
 
     @SuppressWarnings("unchecked")
     private static final ConstructingObjectParser<KnnSearchBuilder.Builder, Void> PARSER = new ConstructingObjectParser<>("knn", args -> {
@@ -65,7 +66,8 @@ public class KnnSearchBuilder implements Writeable, ToXContentFragment, Rewritea
             .queryVectorBuilder((QueryVectorBuilder) args[4])
             .k((Integer) args[2])
             .numCandidates((Integer) args[3])
-            .similarity((Float) args[5]);
+            .similarity((Float) args[5])
+            .rescoreVectorBuilder((RescoreVectorBuilder) args[6]);
     });
 
     static {
@@ -78,13 +80,18 @@ public class KnnSearchBuilder implements Writeable, ToXContentFragment, Rewritea
         );
         PARSER.declareInt(optionalConstructorArg(), K_FIELD);
         PARSER.declareInt(optionalConstructorArg(), NUM_CANDS_FIELD);
-
         PARSER.declareNamedObject(
             optionalConstructorArg(),
             (p, c, n) -> p.namedObject(QueryVectorBuilder.class, n, c),
             QUERY_VECTOR_BUILDER_FIELD
         );
         PARSER.declareFloat(optionalConstructorArg(), VECTOR_SIMILARITY);
+        PARSER.declareField(
+            optionalConstructorArg(),
+            (p, c) -> RescoreVectorBuilder.fromXContent(p),
+            RESCORE_FIELD,
+            ObjectParser.ValueType.OBJECT_OR_NULL
+        );
         PARSER.declareFieldArray(
             KnnSearchBuilder.Builder::addFilterQueries,
             (p, c) -> AbstractQueryBuilder.parseTopLevelQuery(p),
@@ -116,6 +123,7 @@ public static KnnSearchBuilder.Builder fromXContent(XContentParser parser) throw
     String queryName;
     float boost = DEFAULT_BOOST;
     InnerHitBuilder innerHitBuilder;
+    final RescoreVectorBuilder rescoreVectorBuilder;
 
     /**
      * Defines a kNN search.
@@ -124,14 +132,23 @@ public static KnnSearchBuilder.Builder fromXContent(XContentParser parser) throw
      * @param queryVector the query vector
      * @param k           the final number of nearest neighbors to return as top hits
      * @param numCands    the number of nearest neighbor candidates to consider per shard
+     * @param rescoreVectorBuilder rescore vector information
      */
-    public KnnSearchBuilder(String field, float[] queryVector, int k, int numCands, Float similarity) {
+    public KnnSearchBuilder(
+        String field,
+        float[] queryVector,
+        int k,
+        int numCands,
+        RescoreVectorBuilder rescoreVectorBuilder,
+        Float similarity
+    ) {
         this(
             field,
             Objects.requireNonNull(VectorData.fromFloats(queryVector), format("[%s] cannot be null", QUERY_VECTOR_FIELD)),
             null,
             k,
             numCands,
+            rescoreVectorBuilder,
             similarity
         );
     }
@@ -144,8 +161,15 @@ public KnnSearchBuilder(String field, float[] queryVector, int k, int numCands,
      * @param k           the final number of nearest neighbors to return as top hits
      * @param numCands    the number of nearest neighbor candidates to consider per shard
      */
-    public KnnSearchBuilder(String field, VectorData queryVector, int k, int numCands, Float similarity) {
-        this(field, queryVector, null, k, numCands, similarity);
+    public KnnSearchBuilder(
+        String field,
+        VectorData queryVector,
+        int k,
+        int numCands,
+        RescoreVectorBuilder rescoreVectorBuilder,
+        Float similarity
+    ) {
+        this(field, queryVector, null, k, numCands, rescoreVectorBuilder, similarity);
     }
 
     /**
@@ -156,13 +180,21 @@ public KnnSearchBuilder(String field, VectorData queryVector, int k, int numCand
      * @param k                  the final number of nearest neighbors to return as top hits
      * @param numCands           the number of nearest neighbor candidates to consider per shard
      */
-    public KnnSearchBuilder(String field, QueryVectorBuilder queryVectorBuilder, int k, int numCands, Float similarity) {
+    public KnnSearchBuilder(
+        String field,
+        QueryVectorBuilder queryVectorBuilder,
+        int k,
+        int numCands,
+        RescoreVectorBuilder rescoreVectorBuilder,
+        Float similarity
+    ) {
         this(
             field,
             null,
             Objects.requireNonNull(queryVectorBuilder, format("[%s] cannot be null", QUERY_VECTOR_BUILDER_FIELD.getPreferredName())),
             k,
             numCands,
+            rescoreVectorBuilder,
             similarity
         );
     }
@@ -173,16 +205,30 @@ public KnnSearchBuilder(
         QueryVectorBuilder queryVectorBuilder,
         int k,
         int numCands,
+        RescoreVectorBuilder rescoreVectorBuilder,
         Float similarity
     ) {
-        this(field, queryVectorBuilder, queryVector, new ArrayList<>(), k, numCands, similarity, null, null, DEFAULT_BOOST);
+        this(
+            field,
+            queryVectorBuilder,
+            queryVector,
+            new ArrayList<>(),
+            k,
+            numCands,
+            rescoreVectorBuilder,
+            similarity,
+            null,
+            null,
+            DEFAULT_BOOST
+        );
     }
 
     private KnnSearchBuilder(
         String field,
         Supplier<float[]> querySupplier,
         Integer k,
         Integer numCands,
+        RescoreVectorBuilder rescoreVectorBuilder,
         List<QueryBuilder> filterQueries,
         Float similarity
     ) {
@@ -194,6 +240,7 @@ private KnnSearchBuilder(
         this.filterQueries = filterQueries;
         this.querySupplier = querySupplier;
         this.similarity = similarity;
+        this.rescoreVectorBuilder = rescoreVectorBuilder;
     }
 
     private KnnSearchBuilder(
@@ -203,6 +250,7 @@ private KnnSearchBuilder(
         List<QueryBuilder> filterQueries,
         int k,
         int numCandidates,
+        RescoreVectorBuilder rescoreVectorBuilder,
         Float similarity,
         InnerHitBuilder innerHitBuilder,
         String queryName,
@@ -242,6 +290,7 @@ private KnnSearchBuilder(
         this.queryVectorBuilder = queryVectorBuilder;
         this.k = k;
         this.numCands = numCandidates;
+        this.rescoreVectorBuilder = rescoreVectorBuilder;
         this.innerHitBuilder = innerHitBuilder;
         this.similarity = similarity;
         this.queryName = queryName;
@@ -280,6 +329,11 @@ public KnnSearchBuilder(StreamInput in) throws IOException {
         if (in.getTransportVersion().onOrAfter(V_8_11_X)) {
             this.innerHitBuilder = in.readOptionalWriteable(InnerHitBuilder::new);
         }
+        if (in.getTransportVersion().onOrAfter(TransportVersions.KNN_QUERY_RESCORE_OVERSAMPLE)) {
+            this.rescoreVectorBuilder = in.readOptional(RescoreVectorBuilder::new);
+        } else {
+            this.rescoreVectorBuilder = null;
+        }
     }
 
     public int k() {
@@ -290,6 +344,10 @@ public int getNumCands() {
         return numCands;
     }
 
+    public RescoreVectorBuilder getRescoreVectorBuilder() {
+        return rescoreVectorBuilder;
+    }
+
     public QueryVectorBuilder getQueryVectorBuilder() {
         return queryVectorBuilder;
     }
@@ -358,7 +416,7 @@ public KnnSearchBuilder rewrite(QueryRewriteContext ctx) throws IOException {
             if (querySupplier.get() == null) {
                 return this;
             }
-            return new KnnSearchBuilder(field, querySupplier.get(), k, numCands, similarity).boost(boost)
+            return new KnnSearchBuilder(field, querySupplier.get(), k, numCands, rescoreVectorBuilder, similarity).boost(boost)
                 .queryName(queryName)
                 .addFilterQueries(filterQueries)
                 .innerHit(innerHitBuilder);
@@ -381,7 +439,7 @@ public KnnSearchBuilder rewrite(QueryRewriteContext ctx) throws IOException {
                 }
                 ll.onResponse(null);
             })));
-            return new KnnSearchBuilder(field, toSet::get, k, numCands, filterQueries, similarity).boost(boost)
+            return new KnnSearchBuilder(field, toSet::get, k, numCands, rescoreVectorBuilder, filterQueries, similarity).boost(boost)
                 .queryName(queryName)
                 .innerHit(innerHitBuilder);
         }
@@ -395,7 +453,7 @@ public KnnSearchBuilder rewrite(QueryRewriteContext ctx) throws IOException {
             rewrittenQueries.add(rewrittenQuery);
         }
         if (changed) {
-            return new KnnSearchBuilder(field, queryVector, k, numCands, similarity).boost(boost)
+            return new KnnSearchBuilder(field, queryVector, k, numCands, rescoreVectorBuilder, similarity).boost(boost)
                 .queryName(queryName)
                 .addFilterQueries(rewrittenQueries)
                 .innerHit(innerHitBuilder);
@@ -407,7 +465,7 @@ public KnnVectorQueryBuilder toQueryBuilder() {
         if (queryVectorBuilder != null) {
             throw new IllegalArgumentException("missing rewrite");
         }
-        return new KnnVectorQueryBuilder(field, queryVector, null, numCands, null, similarity).boost(boost)
+        return new KnnVectorQueryBuilder(field, queryVector, null, numCands, rescoreVectorBuilder, similarity).boost(boost)
             .queryName(queryName)
             .addFilterQueries(filterQueries);
     }
@@ -423,6 +481,7 @@ public boolean equals(Object o) {
         KnnSearchBuilder that = (KnnSearchBuilder) o;
         return k == that.k
             && numCands == that.numCands
+            && Objects.equals(rescoreVectorBuilder, that.rescoreVectorBuilder)
             && Objects.equals(field, that.field)
             && Objects.equals(queryVector, that.queryVector)
             && Objects.equals(queryVectorBuilder, that.queryVectorBuilder)
@@ -442,6 +501,7 @@ public int hashCode() {
             numCands,
             querySupplier,
             queryVectorBuilder,
+            rescoreVectorBuilder,
             similarity,
             Objects.hashCode(queryVector),
             Objects.hashCode(filterQueries),
@@ -486,6 +546,11 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         if (queryName != null) {
             builder.field(NAME_FIELD.getPreferredName(), queryName);
         }
+        if (rescoreVectorBuilder != null) {
+            builder.startObject(RESCORE_FIELD.getPreferredName());
+            rescoreVectorBuilder.toXContent(builder, params);
+            builder.endObject();
+        }
 
         return builder;
     }
@@ -526,6 +591,9 @@ public void writeTo(StreamOutput out) throws IOException {
         if (out.getTransportVersion().onOrAfter(V_8_11_X)) {
             out.writeOptionalWriteable(innerHitBuilder);
         }
+        if (out.getTransportVersion().onOrAfter(TransportVersions.KNN_QUERY_RESCORE_OVERSAMPLE)) {
+            out.writeOptionalWriteable(rescoreVectorBuilder);
+        }
     }
 
     public static class Builder {
@@ -540,6 +608,7 @@ public static class Builder {
         private String queryName;
         private float boost = DEFAULT_BOOST;
         private InnerHitBuilder innerHitBuilder;
+        private RescoreVectorBuilder rescoreVectorBuilder;
 
         public Builder addFilterQueries(List<QueryBuilder> filterQueries) {
             Objects.requireNonNull(filterQueries);
@@ -592,6 +661,11 @@ public Builder similarity(Float similarity) {
             return this;
         }
 
+        public Builder rescoreVectorBuilder(RescoreVectorBuilder rescoreVectorBuilder) {
+            this.rescoreVectorBuilder = rescoreVectorBuilder;
+            return this;
+        }
+
         public KnnSearchBuilder build(int size) {
             int requestSize = size < 0 ? DEFAULT_SIZE : size;
             int adjustedK = k == null ? requestSize : k;
@@ -605,6 +679,7 @@ public KnnSearchBuilder build(int size) {
                 filterQueries,
                 adjustedK,
                 adjustedNumCandidates,
+                rescoreVectorBuilder,
                 similarity,
                 innerHitBuilder,
                 queryName,

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

@@ -344,8 +344,8 @@ public void testRewriteShardSearchRequestWithRank() {
         SearchSourceBuilder ssb = new SearchSourceBuilder().query(bm25)
             .knnSearch(
                 List.of(
-                    new KnnSearchBuilder("vector", new float[] { 0.0f }, 10, 100, null),
-                    new KnnSearchBuilder("vector2", new float[] { 0.0f }, 10, 100, null)
+                    new KnnSearchBuilder("vector", new float[] { 0.0f }, 10, 100, null, null),
+                    new KnnSearchBuilder("vector2", new float[] { 0.0f }, 10, 100, null, null)
                 )
             )
             .rankBuilder(new TestRankBuilder(100));