Use 'num_candidates_factor' parameter and update num_candidates instead of k

carlosdelest · carlosdelest · commit 9907aadec687 · 2024-12-09T17:46:30.000+01:00
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/210_knn_search_profile.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/210_knn_search_profile.yml
@@ -105,7 +105,7 @@ setup:
             k: 3
             num_candidates: 3
             "rescore_vector":
-              "oversample": 2.0
+              "num_candidates_factor": 2.0
 
   # We expect the knn search ops + rescoring num_cnaidates (for rescoring) per shard
   - match: { profile.shards.0.dfs.knn.0.vector_operations_count: 6 }
@@ -123,7 +123,7 @@ setup:
             num_candidates: 3
             similarity: 100000.0
             "rescore_vector":
-              "oversample": 2.0
+              "num_candidates_factor": 2.0
 
   # We expect the knn search ops + rescoring num_cnaidates (for rescoring) per shard
   - match: { profile.shards.0.dfs.knn.0.vector_operations_count: 6 }
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/40_knn_search.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/40_knn_search.yml
@@ -589,7 +589,7 @@ setup:
             k: 3
             num_candidates: 3
             rescore_vector:
-              oversample: 1.5
+              num_candidates_factor: 1.5
 
   # Compare scores as hit IDs may change depending on how things are distributed
   - match: { hits.total: 3 }
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_bbq_hnsw.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_bbq_hnsw.yml
@@ -113,7 +113,7 @@ setup:
             k: 3
             num_candidates: 3
             rescore_vector:
-              oversample: 1.5
+              num_candidates_factor: 1.5
 
   # Get rescoring scores - hit ordering may change depending on how things are distributed
   - match: { hits.total: 3 }
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_byte_quantized.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_byte_quantized.yml
@@ -398,7 +398,7 @@ setup:
               field: vector
               query_vector: [0.5, 111.3, -13.0, 14.8, -156.0]
               rescore_vector:
-                oversample: 1.5
+                num_candidates_factor: 1.5
 
   # Get rescoring scores - hit ordering may change depending on how things are distributed
   - match: { hits.total: 3 }
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_half_byte_quantized.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_half_byte_quantized.yml
@@ -575,7 +575,7 @@ setup:
             k: 3
             num_candidates: 3
             rescore_vector:
-              oversample: 1.5
+              num_candidates_factor: 1.5
 
   # Get rescoring scores - hit ordering may change depending on how things are distributed
   - match: { hits.total: 3 }
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_bbq_flat.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_bbq_flat.yml
@@ -112,7 +112,7 @@ setup:
             k: 3
             num_candidates: 3
             rescore_vector:
-              oversample: 1.5
+              num_candidates_factor: 1.5
 
   # Get rescoring scores - hit ordering may change depending on how things are distributed
   - match: { hits.total: 3 }
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_flat.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_flat.yml
@@ -304,7 +304,7 @@ setup:
             k: 3
             num_candidates: 3
             rescore_vector:
-              oversample: 1.5
+              num_candidates_factor: 1.5
 
   # Compare scores as hit IDs may change depending on how things are distributed
   - match: { hits.total: 3 }
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_int4_flat.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_int4_flat.yml
@@ -371,7 +371,7 @@ setup:
             k: 3
             num_candidates: 3
             rescore_vector:
-              oversample: 1.5
+              num_candidates_factor: 1.5
 
   # Get rescoring scores - hit ordering may change depending on how things are distributed
   - match: { hits.total: 3 }
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_int8_flat.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/42_knn_search_int8_flat.yml
@@ -288,7 +288,7 @@ setup:
             k: 3
             num_candidates: 3
             rescore_vector:
-              oversample: 1.5
+              num_candidates_factor: 1.5
 
   # Get rescoring scores - hit ordering may change depending on how things are distributed
   - match: { hits.total: 3 }
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/45_knn_search_bit.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/45_knn_search_bit.yml
@@ -454,7 +454,7 @@ setup:
             k: 3
             num_candidates: 3
             rescore_vector:
-              oversample: 1.5
+              num_candidates_factor: 1.5
 
   # Compare scores as hit IDs may change depending on how things are distributed
   - match: { hits.total: 3 }
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/45_knn_search_bit_flat.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/45_knn_search_bit_flat.yml
@@ -270,7 +270,7 @@ setup:
             k: 3
             num_candidates: 3
             rescore_vector:
-              oversample: 1.5
+              num_candidates_factor: 1.5
 
   # Compare scores as hit IDs may change depending on how things are distributed
   - match: { hits.total: 3 }
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/45_knn_search_byte.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/45_knn_search_byte.yml
@@ -303,7 +303,7 @@ setup:
             k: 3
             num_candidates: 3
             rescore_vector:
-              oversample: 1.5
+              num_candidates_factor: 1.5
 
   # Compare scores as hit IDs may change depending on how things are distributed
   - match: { hits.total: 3 }
diff --git a/server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java b/server/src/main/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldMapper.java
@@ -124,7 +124,7 @@ public static boolean isNotUnitVector(float magnitude) {
 
     public static short MIN_DIMS_FOR_DYNAMIC_FLOAT_MAPPING = 128; // minimum number of dims for floats to be dynamically mapped to vector
     public static final int MAGNITUDE_BYTES = 4;
-    public static final int OVERSAMPLE_LIMIT = 10_000; // Max oversample allowed for k and num_candidates
+    public static final int NUM_CANDS_OVERSAMPLE_LIMIT = 10_000; // Max oversample allowed for k and num_candidates
 
     private static DenseVectorFieldMapper toType(FieldMapper in) {
         return (DenseVectorFieldMapper) in;
@@ -2008,7 +2008,7 @@ public Query createKnnQuery(
             VectorData queryVector,
             Integer k,
             int numCands,
-            Float rescoreOversample,
+            Float numCandsFactor,
             Query filter,
             Float similarityThreshold,
             BitSetProducer parentFilter
@@ -2024,7 +2024,7 @@ public Query createKnnQuery(
                     queryVector.asFloatVector(),
                     k,
                     numCands,
-                    rescoreOversample,
+                    numCandsFactor,
                     filter,
                     similarityThreshold,
                     parentFilter
@@ -2090,7 +2090,7 @@ private Query createKnnFloatQuery(
             float[] queryVector,
             Integer k,
             int numCands,
-            Float rescoreOversample,
+            Float numCandsFactor,
             Query filter,
             Float similarityThreshold,
             BitSetProducer parentFilter
@@ -2111,23 +2111,25 @@ && isNotUnitVector(squaredMagnitude)) {
                 }
             }
 
-            Integer adjustedK = k;
             int adjustedNumCands = numCands;
-            if (needsRescore(rescoreOversample) && adjustedK != null) {
-                adjustedK = Math.min(OVERSAMPLE_LIMIT, (int) Math.ceil(k * rescoreOversample));
-                adjustedNumCands = Math.max(adjustedK, numCands);
+            if (needsRescore(numCandsFactor)) {
+                // We shouldn't have less than k candidates (or 1 in case k is not set) to rescore
+                int minCands = k == null ? 1 : k;
+                // k <= numCands * numCandsFactor <= NUM_CANDS_OVERSAMPLE_LIMIT. Adjust otherwise.
+                adjustedNumCands = Math.max(minCands, (int) Math.ceil(numCands * numCandsFactor));
+                adjustedNumCands = Math.min(adjustedNumCands, NUM_CANDS_OVERSAMPLE_LIMIT);
             }
             Query knnQuery = parentFilter != null
-                ? new ESDiversifyingChildrenFloatKnnVectorQuery(name(), queryVector, filter, adjustedK, adjustedNumCands, parentFilter)
-                : new ESKnnFloatVectorQuery(name(), queryVector, adjustedK, adjustedNumCands, filter);
+                ? new ESDiversifyingChildrenFloatKnnVectorQuery(name(), queryVector, filter, k, adjustedNumCands, parentFilter)
+                : new ESKnnFloatVectorQuery(name(), queryVector, k, adjustedNumCands, filter);
             if (similarityThreshold != null) {
                 knnQuery = new VectorSimilarityQuery(
                     knnQuery,
                     similarityThreshold,
                     similarity.score(similarityThreshold, elementType, dims)
                 );
             }
-            if (needsRescore(rescoreOversample)) {
+            if (needsRescore(numCandsFactor)) {
                 knnQuery = new RescoreKnnVectorQuery(
                     name(),
                     queryVector,
diff --git a/server/src/main/java/org/elasticsearch/search/vectors/KnnVectorQueryBuilder.java b/server/src/main/java/org/elasticsearch/search/vectors/KnnVectorQueryBuilder.java
@@ -526,7 +526,7 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException {
 
         DenseVectorFieldType vectorFieldType = (DenseVectorFieldType) fieldType;
         String parentPath = context.nestedLookup().getNestedParent(fieldName);
-        Float rescoreOversample = rescoreVectorBuilder() == null ? null : rescoreVectorBuilder.oversample();
+        Float numCandidatesFactor = rescoreVectorBuilder() == null ? null : rescoreVectorBuilder.numCandidatesFactor();
 
         if (parentPath != null) {
             final BitSetProducer parentBitSet;
@@ -563,13 +563,13 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException {
                 queryVector,
                 k,
                 adjustedNumCands,
-                rescoreOversample,
+                numCandidatesFactor,
                 filterQuery,
                 vectorSimilarity,
                 parentBitSet
             );
         }
-        return vectorFieldType.createKnnQuery(queryVector, k, adjustedNumCands, rescoreOversample, filterQuery, vectorSimilarity, null);
+        return vectorFieldType.createKnnQuery(queryVector, k, adjustedNumCands, numCandidatesFactor, filterQuery, vectorSimilarity, null);
     }
 
     @Override
diff --git a/server/src/main/java/org/elasticsearch/search/vectors/RescoreVectorBuilder.java b/server/src/main/java/org/elasticsearch/search/vectors/RescoreVectorBuilder.java
@@ -23,41 +23,41 @@
 
 public class RescoreVectorBuilder implements Writeable, ToXContentObject {
 
-    public static final ParseField OVERSAMPLE_FIELD = new ParseField("oversample");
-    public static final float MIN_OVERSAMPLE = 1.0F;
+    public static final ParseField NUM_CANDIDATES_FACTOR_FIELD = new ParseField("num_candidates_factor");
+    public static final float MIN_OVERSAMPLE = 0.0F;
     private static final ConstructingObjectParser<RescoreVectorBuilder, Void> PARSER = new ConstructingObjectParser<>(
         "rescore_vector",
         args -> new RescoreVectorBuilder((Float) args[0])
     );
 
     static {
-        PARSER.declareFloat(ConstructingObjectParser.constructorArg(), OVERSAMPLE_FIELD);
+        PARSER.declareFloat(ConstructingObjectParser.constructorArg(), NUM_CANDIDATES_FACTOR_FIELD);
     }
 
     // Oversample is required as of now as it is the only field in the rescore vector
-    private final float oversample;
+    private final float numCandidatesFactor;
 
-    public RescoreVectorBuilder(Float oversample) {
-        Objects.requireNonNull(oversample, "[" + OVERSAMPLE_FIELD.getPreferredName() + "] must be set");
-        if (oversample <= MIN_OVERSAMPLE) {
-            throw new IllegalArgumentException("[" + OVERSAMPLE_FIELD.getPreferredName() + "] must be > " + MIN_OVERSAMPLE);
+    public RescoreVectorBuilder(Float numCandidatesFactor) {
+        Objects.requireNonNull(numCandidatesFactor, "[" + NUM_CANDIDATES_FACTOR_FIELD.getPreferredName() + "] must be set");
+        if (numCandidatesFactor <= MIN_OVERSAMPLE) {
+            throw new IllegalArgumentException("[" + NUM_CANDIDATES_FACTOR_FIELD.getPreferredName() + "] must be > " + MIN_OVERSAMPLE);
         }
-        this.oversample = oversample;
+        this.numCandidatesFactor = numCandidatesFactor;
     }
 
     public RescoreVectorBuilder(StreamInput in) throws IOException {
-        this.oversample = in.readFloat();
+        this.numCandidatesFactor = in.readFloat();
     }
 
     @Override
     public void writeTo(StreamOutput out) throws IOException {
-        out.writeFloat(oversample);
+        out.writeFloat(numCandidatesFactor);
     }
 
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();
-        builder.field(OVERSAMPLE_FIELD.getPreferredName(), oversample);
+        builder.field(NUM_CANDIDATES_FACTOR_FIELD.getPreferredName(), numCandidatesFactor);
         builder.endObject();
         return builder;
     }
@@ -71,15 +71,15 @@ public boolean equals(Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         RescoreVectorBuilder that = (RescoreVectorBuilder) o;
-        return Objects.equals(oversample, that.oversample);
+        return Objects.equals(numCandidatesFactor, that.numCandidatesFactor);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hashCode(oversample);
+        return Objects.hashCode(numCandidatesFactor);
     }
 
-    public Float oversample() {
-        return oversample;
+    public Float numCandidatesFactor() {
+        return numCandidatesFactor;
     }
 }
diff --git a/server/src/test/java/org/elasticsearch/search/vectors/AbstractKnnVectorQueryBuilderTestCase.java b/server/src/test/java/org/elasticsearch/search/vectors/AbstractKnnVectorQueryBuilderTestCase.java
@@ -46,7 +46,7 @@
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
-import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.OVERSAMPLE_LIMIT;
+import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.NUM_CANDS_OVERSAMPLE_LIMIT;
 import static org.elasticsearch.search.SearchService.DEFAULT_SIZE;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
@@ -200,9 +200,10 @@ protected void doAssertLuceneQuery(KnnVectorQueryBuilder queryBuilder, Query que
         Integer k = queryBuilder.k();
         Integer numCands = queryBuilder.numCands();
         if (queryBuilder.rescoreVectorBuilder() != null && isQuantizedElementType()) {
-            Float rescoreOversample = queryBuilder.rescoreVectorBuilder().oversample();
-            k = k == null ? null : Integer.valueOf(Math.min(OVERSAMPLE_LIMIT, (int) Math.ceil(k * rescoreOversample)));
-            numCands = numCands == null ? null : Math.max(k == null ? 0 : k, numCands);
+            Float numCandsFactor = queryBuilder.rescoreVectorBuilder().numCandidatesFactor();
+            int minCands = k == null ? 1 : k;
+            numCands = Math.max(minCands, (int) Math.ceil(numCands * numCandsFactor));
+            numCands = Math.min(numCands, NUM_CANDS_OVERSAMPLE_LIMIT);
         }
 
         Query knnVectorQueryBuilt = switch (elementType()) {
