Vector rescoring oversamples k instead of num_candidates (elastic#119835)

carlosdelest · carlosdelest · commit ff562f657ca8 · 2025-01-10T20:34:58.000Z
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
@@ -82,7 +82,7 @@ setup:
       capabilities:
         - method: GET
           path: /_search
-          capabilities: [knn_quantized_vector_rescore]
+          capabilities: [knn_quantized_vector_rescore_oversample]
   - skip:
       features: "headers"
 
@@ -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
@@ -549,7 +549,7 @@ setup:
       capabilities:
         - method: GET
           path: /_search
-          capabilities: [knn_quantized_vector_rescore]
+          capabilities: [knn_quantized_vector_rescore_oversample]
   - skip:
       features: "headers"
 
@@ -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
@@ -115,7 +115,7 @@ setup:
       capabilities:
         - method: GET
           path: /_search
-          capabilities: [knn_quantized_vector_rescore]
+          capabilities: [knn_quantized_vector_rescore_oversample]
   - skip:
       features: "headers"
 
@@ -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
@@ -378,7 +378,7 @@ setup:
       capabilities:
         - method: GET
           path: /_search
-          capabilities: [knn_quantized_vector_rescore]
+          capabilities: [knn_quantized_vector_rescore_oversample]
   - skip:
       features: "headers"
 
@@ -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
@@ -556,7 +556,7 @@ setup:
       capabilities:
         - method: GET
           path: /_search
-          capabilities: [knn_quantized_vector_rescore]
+          capabilities: [knn_quantized_vector_rescore_oversample]
   - skip:
       features: "headers"
 
@@ -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
@@ -114,7 +114,7 @@ setup:
       capabilities:
         - method: GET
           path: /_search
-          capabilities: [knn_quantized_vector_rescore]
+          capabilities: [knn_quantized_vector_rescore_oversample]
   - skip:
       features: "headers"
 
@@ -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
@@ -264,7 +264,7 @@ setup:
       capabilities:
         - method: GET
           path: /_search
-          capabilities: [knn_quantized_vector_rescore]
+          capabilities: [knn_quantized_vector_rescore_oversample]
   - skip:
       features: "headers"
 
@@ -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
@@ -352,7 +352,7 @@ setup:
       capabilities:
         - method: GET
           path: /_search
-          capabilities: [knn_quantized_vector_rescore]
+          capabilities: [knn_quantized_vector_rescore_oversample]
   - skip:
       features: "headers"
 
@@ -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
@@ -269,7 +269,7 @@ setup:
       capabilities:
         - method: GET
           path: /_search
-          capabilities: [knn_quantized_vector_rescore]
+          capabilities: [knn_quantized_vector_rescore_oversample]
   - skip:
       features: "headers"
 
@@ -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
@@ -414,7 +414,7 @@ setup:
       capabilities:
         - method: GET
           path: /_search
-          capabilities: [knn_quantized_vector_rescore]
+          capabilities: [knn_quantized_vector_rescore_oversample]
   - skip:
       features: "headers"
 
@@ -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
@@ -230,7 +230,7 @@ setup:
       capabilities:
         - method: GET
           path: /_search
-          capabilities: [knn_quantized_vector_rescore]
+          capabilities: [knn_quantized_vector_rescore_oversample]
   - skip:
       features: "headers"
 
@@ -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
@@ -263,7 +263,7 @@ setup:
       capabilities:
         - method: GET
           path: /_search
-          capabilities: [knn_quantized_vector_rescore]
+          capabilities: [knn_quantized_vector_rescore_oversample]
   - skip:
       features: "headers"
 
@@ -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/rest/action/search/SearchCapabilities.java b/server/src/main/java/org/elasticsearch/rest/action/search/SearchCapabilities.java
@@ -38,7 +38,7 @@ private SearchCapabilities() {}
 
     private static final String RANDOM_SAMPLER_WITH_SCORED_SUBAGGS = "random_sampler_with_scored_subaggs";
     private static final String OPTIMIZED_SCALAR_QUANTIZATION_BBQ = "optimized_scalar_quantization_bbq";
-    private static final String KNN_QUANTIZED_VECTOR_RESCORE = "knn_quantized_vector_rescore";
+    private static final String KNN_QUANTIZED_VECTOR_RESCORE_OVERSAMPLE = "knn_quantized_vector_rescore_oversample";
 
     public static final Set<String> CAPABILITIES;
     static {
@@ -50,7 +50,7 @@ private SearchCapabilities() {}
         capabilities.add(NESTED_RETRIEVER_INNER_HITS_SUPPORT);
         capabilities.add(RANDOM_SAMPLER_WITH_SCORED_SUBAGGS);
         capabilities.add(OPTIMIZED_SCALAR_QUANTIZATION_BBQ);
-        capabilities.add(KNN_QUANTIZED_VECTOR_RESCORE);
+        capabilities.add(KNN_QUANTIZED_VECTOR_RESCORE_OVERSAMPLE);
         capabilities.add(MOVING_FN_RIGHT_MATH);
         capabilities.add(K_DEFAULT_TO_SIZE);
         capabilities.add(KQL_QUERY_SUPPORTED);
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.oversample();
 
         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
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
diff --git a/server/src/test/java/org/elasticsearch/search/vectors/AbstractKnnVectorQueryBuilderTestCase.java b/server/src/test/java/org/elasticsearch/search/vectors/AbstractKnnVectorQueryBuilderTestCase.java
diff --git a/server/src/test/java/org/elasticsearch/search/vectors/KnnSearchBuilderTests.java b/server/src/test/java/org/elasticsearch/search/vectors/KnnSearchBuilderTests.java
