Use oversample to modify k instead of num_candidates for rescoring

carlosdelest · carlosdelest · commit 0f956b4c81ff · 2025-01-09T18:07:17.000+01:00
diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.retrievers/20_knn_retriever.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.retrievers/20_knn_retriever.yml
@@ -100,7 +100,7 @@ setup:
             k: 3
             num_candidates: 3
             rescore_vector:
-              num_candidates_factor: 1.5
+              oversample: 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/210_knn_search_profile.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/210_knn_search_profile.yml
@@ -106,9 +106,9 @@ setup:
             k: 3
             num_candidates: 3
             "rescore_vector":
-              "num_candidates_factor": 2.0
+              "oversample": 2.0
 
-  # We expect the knn search ops + rescoring num_cnaidates (for rescoring) per shard
+  # We expect the knn search ops + rescoring k * oversample (for rescoring) per shard
   - match: { profile.shards.0.dfs.knn.0.vector_operations_count: 6 }
 
   # Search with similarity to check number of operations are propagated correctly
@@ -131,7 +131,7 @@ setup:
             num_candidates: 3
             similarity: 100000
             "rescore_vector":
-              "num_candidates_factor": 2.0
+              "oversample": 2.0
 
-  # We expect the knn search ops + rescoring num_cnaidates (for rescoring) per shard
+  # We expect the knn search ops + rescoring k * oversample (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:
-              num_candidates_factor: 1.5
+              oversample: 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
@@ -140,7 +140,7 @@ setup:
             k: 3
             num_candidates: 3
             rescore_vector:
-              num_candidates_factor: 1.5
+              oversample: 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:
-                num_candidates_factor: 1.5
+                oversample: 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:
-              num_candidates_factor: 1.5
+              oversample: 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
@@ -139,7 +139,7 @@ setup:
             k: 3
             num_candidates: 3
             rescore_vector:
-              num_candidates_factor: 1.5
+              oversample: 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:
-              num_candidates_factor: 1.5
+              oversample: 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:
-              num_candidates_factor: 1.5
+              oversample: 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:
-              num_candidates_factor: 1.5
+              oversample: 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:
-              num_candidates_factor: 1.5
+              oversample: 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:
-              num_candidates_factor: 1.5
+              oversample: 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:
-              num_candidates_factor: 1.5
+              oversample: 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
@@ -123,7 +123,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 NUM_CANDS_OVERSAMPLE_LIMIT = 10_000; // Max oversample allowed for k and num_candidates
+    public static final int OVERSAMPLE_LIMIT = 10_000; // Max oversample allowed
 
     private static DenseVectorFieldMapper toType(FieldMapper in) {
         return (DenseVectorFieldMapper) in;
@@ -2019,7 +2019,7 @@ public Query createKnnQuery(
             VectorData queryVector,
             int k,
             int numCands,
-            Float numCandsFactor,
+            Float oversample,
             Query filter,
             Float similarityThreshold,
             BitSetProducer parentFilter
@@ -2035,7 +2035,7 @@ public Query createKnnQuery(
                     queryVector.asFloatVector(),
                     k,
                     numCands,
-                    numCandsFactor,
+                    oversample,
                     filter,
                     similarityThreshold,
                     parentFilter
@@ -2045,7 +2045,11 @@ public Query createKnnQuery(
         }
 
         private boolean needsRescore(Float rescoreOversample) {
-            return rescoreOversample != null && (indexOptions != null && indexOptions.type != null && indexOptions.type.isQuantized());
+            return rescoreOversample != null && isQuantized();
+        }
+
+        private boolean isQuantized() {
+            return indexOptions != null && indexOptions.type != null && indexOptions.type.isQuantized();
         }
 
         private Query createKnnBitQuery(
@@ -2101,7 +2105,7 @@ private Query createKnnFloatQuery(
             float[] queryVector,
             int k,
             int numCands,
-            Float numCandsFactor,
+            Float oversample,
             Query filter,
             Float similarityThreshold,
             BitSetProducer parentFilter
@@ -2122,18 +2126,17 @@ && isNotUnitVector(squaredMagnitude)) {
                 }
             }
 
-            Integer adjustedK = k;
-            int adjustedNumCands = numCands;
-            if (needsRescore(numCandsFactor)) {
-                // Get all candidates, get top k as part of rescoring
-                adjustedK = null;
-                // numCands * numCandsFactor <= NUM_CANDS_OVERSAMPLE_LIMIT. Adjust otherwise.
-                adjustedNumCands = Math.min((int) Math.ceil(numCands * numCandsFactor), NUM_CANDS_OVERSAMPLE_LIMIT);
+            int adjustedK = k;
+            boolean rescore = needsRescore(oversample);
+            if (rescore) {
+                // Will get k * oversample for rescoring, and get the top k
+                adjustedK = Math.min((int) Math.ceil(k * oversample), OVERSAMPLE_LIMIT);
+                numCands = Math.max(adjustedK, numCands);
             }
             Query knnQuery = parentFilter != null
-                ? new ESDiversifyingChildrenFloatKnnVectorQuery(name(), queryVector, filter, adjustedK, adjustedNumCands, parentFilter)
-                : new ESKnnFloatVectorQuery(name(), queryVector, adjustedK, adjustedNumCands, filter);
-            if (needsRescore(numCandsFactor)) {
+                ? new ESDiversifyingChildrenFloatKnnVectorQuery(name(), queryVector, filter, adjustedK, numCands, parentFilter)
+                : new ESKnnFloatVectorQuery(name(), queryVector, adjustedK, numCands, filter);
+            if (rescore) {
                 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
@@ -522,7 +522,7 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException {
 
         DenseVectorFieldType vectorFieldType = (DenseVectorFieldType) fieldType;
         String parentPath = context.nestedLookup().getNestedParent(fieldName);
-        Float numCandidatesFactor = rescoreVectorBuilder() == null ? null : rescoreVectorBuilder.numCandidatesFactor();
+        Float oversample = rescoreVectorBuilder() == null ? null : rescoreVectorBuilder.overesample();
 
         BitSetProducer parentBitSet = null;
         if (parentPath != null) {
@@ -557,15 +557,7 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException {
             }
         }
 
-        return vectorFieldType.createKnnQuery(
-            queryVector,
-            k,
-            adjustedNumCands,
-            numCandidatesFactor,
-            filterQuery,
-            vectorSimilarity,
-            parentBitSet
-        );
+        return vectorFieldType.createKnnQuery(queryVector, k, adjustedNumCands, oversample, filterQuery, vectorSimilarity, parentBitSet);
     }
 
     @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,24 +23,24 @@
 
 public class RescoreVectorBuilder implements Writeable, ToXContentObject {
 
-    public static final ParseField NUM_CANDIDATES_FACTOR_FIELD = new ParseField("num_candidates_factor");
+    public static final ParseField OVERSAMPLE_FIELD = new ParseField("oversample");
     public static final float MIN_OVERSAMPLE = 1.0F;
     private static final ConstructingObjectParser<RescoreVectorBuilder, Void> PARSER = new ConstructingObjectParser<>(
         "rescore_vector",
         args -> new RescoreVectorBuilder((Float) args[0])
     );
 
     static {
-        PARSER.declareFloat(ConstructingObjectParser.constructorArg(), NUM_CANDIDATES_FACTOR_FIELD);
+        PARSER.declareFloat(ConstructingObjectParser.constructorArg(), OVERSAMPLE_FIELD);
     }
 
     // Oversample is required as of now as it is the only field in the rescore vector
     private final float numCandidatesFactor;
 
     public RescoreVectorBuilder(float numCandidatesFactor) {
-        Objects.requireNonNull(numCandidatesFactor, "[" + NUM_CANDIDATES_FACTOR_FIELD.getPreferredName() + "] must be set");
+        Objects.requireNonNull(numCandidatesFactor, "[" + OVERSAMPLE_FIELD.getPreferredName() + "] must be set");
         if (numCandidatesFactor < MIN_OVERSAMPLE) {
-            throw new IllegalArgumentException("[" + NUM_CANDIDATES_FACTOR_FIELD.getPreferredName() + "] must be >= " + MIN_OVERSAMPLE);
+            throw new IllegalArgumentException("[" + OVERSAMPLE_FIELD.getPreferredName() + "] must be >= " + MIN_OVERSAMPLE);
         }
         this.numCandidatesFactor = numCandidatesFactor;
     }
@@ -57,7 +57,7 @@ public void writeTo(StreamOutput out) throws IOException {
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();
-        builder.field(NUM_CANDIDATES_FACTOR_FIELD.getPreferredName(), numCandidatesFactor);
+        builder.field(OVERSAMPLE_FIELD.getPreferredName(), numCandidatesFactor);
         builder.endObject();
         return builder;
     }
@@ -79,7 +79,7 @@ public int hashCode() {
         return Objects.hashCode(numCandidatesFactor);
     }
 
-    public float numCandidatesFactor() {
+    public float overesample() {
         return numCandidatesFactor;
     }
 }
diff --git a/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java b/server/src/test/java/org/elasticsearch/index/mapper/vectors/DenseVectorFieldTypeTests.java
@@ -36,6 +36,7 @@
 import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.BBQ_MIN_DIMS;
 import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.ElementType.BYTE;
 import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.ElementType.FLOAT;
+import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.OVERSAMPLE_LIMIT;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.instanceOf;
@@ -456,20 +457,18 @@ public void testRescoreOversampleModifiesNumCandidates() {
         );
 
         // Total results is k, internal k is multiplied by oversample
-        checkRescoreQueryParameters(fieldType, 10, 200, randomInt(), 2.5F, null, 500, 10);
+        checkRescoreQueryParameters(fieldType, 10, 200, 2.5F, 25, 200, 10);
         // If numCands < k, update numCands to k
-        checkRescoreQueryParameters(fieldType, 10, 20, randomInt(), 2.5F, null, 50, 10);
-        // Oversampling limits for num candidates
-        checkRescoreQueryParameters(fieldType, 1000, 1000, randomInt(), 11.0F, null, 10000, 1000);
-        checkRescoreQueryParameters(fieldType, 5000, 7500, randomInt(), 2.5F, null, 10000, 5000);
+        checkRescoreQueryParameters(fieldType, 10, 20, 2.5F, 25, 25, 10);
+        // Oversampling limits for k
+        checkRescoreQueryParameters(fieldType, 1000, 1000, 11.0F, OVERSAMPLE_LIMIT, OVERSAMPLE_LIMIT, 1000);
     }
 
     private static void checkRescoreQueryParameters(
         DenseVectorFieldType fieldType,
         int k,
         int candidates,
-        int requestSize,
-        float numCandsFactor,
+        float oversample,
         Integer expectedK,
         int expectedCandidates,
         int expectedResults
@@ -478,7 +477,7 @@ private static void checkRescoreQueryParameters(
             VectorData.fromFloats(new float[] { 1, 4, 10 }),
             k,
             candidates,
-            numCandsFactor,
+            oversample,
             null,
             null,
             null
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.NUM_CANDS_OVERSAMPLE_LIMIT;
+import static org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper.OVERSAMPLE_LIMIT;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasSize;
@@ -194,10 +194,10 @@ protected void doAssertLuceneQuery(KnnVectorQueryBuilder queryBuilder, Query que
         Integer k = queryBuilder.k();
         Integer numCands = queryBuilder.numCands();
         if (queryBuilder.rescoreVectorBuilder() != null && isQuantizedElementType()) {
-            Float numCandsFactor = queryBuilder.rescoreVectorBuilder().numCandidatesFactor();
+            Float numCandsFactor = queryBuilder.rescoreVectorBuilder().overesample();
             int minCands = k == null ? 1 : k;
             numCands = Math.max(minCands, (int) Math.ceil(numCands * numCandsFactor));
-            numCands = Math.min(numCands, NUM_CANDS_OVERSAMPLE_LIMIT);
+            numCands = Math.min(numCands, OVERSAMPLE_LIMIT);
         }
 
         Query knnVectorQueryBuilt = switch (elementType()) {
diff --git a/server/src/test/java/org/elasticsearch/search/vectors/KnnSearchBuilderTests.java b/server/src/test/java/org/elasticsearch/search/vectors/KnnSearchBuilderTests.java
@@ -259,7 +259,7 @@ public void testInvalidRescoreVectorBuilder() {
             IllegalArgumentException.class,
             () -> new KnnSearchBuilder("field", randomVector(3), 10, 100, new RescoreVectorBuilder(0.99F), null)
         );
-        assertThat(e.getMessage(), containsString("[num_candidates_factor] must be >= 1.0"));
+        assertThat(e.getMessage(), containsString("[oversample] must be >= 1.0"));
     }
 
     public void testRewrite() throws Exception {
