Skip to content

Commit d40813d

Browse files
carlosdelestcarlosdelest
authored
[8.x] kNN vector rescoring for quantized vectors (elastic#116663) (elastic#118418)
* kNN vector rescoring for quantized vectors (elastic#116663) (cherry picked from commit 5996772) # Conflicts: # server/src/main/java/org/elasticsearch/search/vectors/KnnSearchBuilder.java # x-pack/plugin/rank-rrf/src/main/java/org/elasticsearch/xpack/rank/rrf/RRFRankBuilder.java * FloatVectorValues have a different interface in this Lucene version
1 parent 20f7919 commit d40813d

File tree

66 files changed

+2363
-296
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

66 files changed

+2363
-296
lines changed

‎docs/changelog/116663.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
pr: 116663
2+
summary: KNN vector rescoring for quantized vectors
3+
area: Vector Search
4+
type: feature
5+
issues: []

‎modules/percolator/src/internalClusterTest/java/org/elasticsearch/percolator/PercolatorQuerySearchIT.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1359,7 +1359,7 @@ public void testKnnQueryNotSupportedInPercolator() throws IOException {
13591359
""");
13601360
indicesAdmin().prepareCreate("index1").setMapping(mappings).get();
13611361
ensureGreen();
1362-
QueryBuilder knnVectorQueryBuilder = new KnnVectorQueryBuilder("my_vector", new float[] { 1, 1, 1, 1, 1 }, 10, 10, null);
1362+
QueryBuilder knnVectorQueryBuilder = new KnnVectorQueryBuilder("my_vector", new float[] { 1, 1, 1, 1, 1 }, 10, 10, null, null);
13631363

13641364
IndexRequestBuilder indexRequestBuilder = prepareIndex("index1").setId("knn_query1")
13651365
.setSource(jsonBuilder().startObject().field("my_query", knnVectorQueryBuilder).endObject());

‎rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.retrievers/20_knn_retriever.yml‎

Lines changed: 57 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ setup:
1818
dims: 5
1919
index: true
2020
index_options:
21-
type: hnsw
21+
type: int8_hnsw
2222
similarity: l2_norm
2323

2424
- do:
@@ -73,3 +73,59 @@ setup:
7373
- match: {hits.total.value: 1}
7474
- match: {hits.hits.0._id: "3"}
7575
- match: {hits.hits.0.fields.name.0: "rabbit.jpg"}
76+
77+
---
78+
"Vector rescoring has no effect for non-quantized vectors and provides same results as non-rescored knn":
79+
- requires:
80+
reason: 'Quantized vector rescoring is required'
81+
test_runner_features: [capabilities]
82+
capabilities:
83+
- method: GET
84+
path: /_search
85+
capabilities: [knn_quantized_vector_rescore]
86+
- skip:
87+
features: "headers"
88+
89+
# Rescore
90+
- do:
91+
headers:
92+
Content-Type: application/json
93+
search:
94+
rest_total_hits_as_int: true
95+
index: index1
96+
body:
97+
knn:
98+
field: vector
99+
query_vector: [2, 2, 2, 2, 3]
100+
k: 3
101+
num_candidates: 3
102+
rescore_vector:
103+
num_candidates_factor: 1.5
104+
105+
# Get rescoring scores - hit ordering may change depending on how things are distributed
106+
- match: { hits.total: 3 }
107+
- set: { hits.hits.0._score: rescore_score0 }
108+
- set: { hits.hits.1._score: rescore_score1 }
109+
- set: { hits.hits.2._score: rescore_score2 }
110+
111+
# Exact knn via script score
112+
- do:
113+
headers:
114+
Content-Type: application/json
115+
search:
116+
rest_total_hits_as_int: true
117+
index: index1
118+
body:
119+
query:
120+
script_score:
121+
query: {match_all: {} }
122+
script:
123+
source: "1.0 / (1.0 + Math.pow(l2norm(params.query_vector, 'vector'), 2.0))"
124+
params:
125+
query_vector: [2, 2, 2, 2, 3]
126+
127+
# Compare scores as hit IDs may change depending on how things are distributed
128+
- match: { hits.total: 3 }
129+
- match: { hits.hits.0._score: $rescore_score0 }
130+
- match: { hits.hits.1._score: $rescore_score1 }
131+
- match: { hits.hits.2._score: $rescore_score2 }

‎rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/210_knn_search_profile.yml‎

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
setup:
2+
- requires:
3+
reason: 'Quantized vector rescoring is required'
4+
test_runner_features: [ capabilities ]
5+
capabilities:
6+
- method: GET
7+
path: /_search
8+
capabilities: [ knn_quantized_vector_rescore ]
9+
- skip:
10+
features: "headers"
11+
12+
- do:
13+
indices.create:
14+
index: bbq_hnsw
15+
body:
16+
settings:
17+
index:
18+
number_of_shards: 1
19+
mappings:
20+
properties:
21+
vector:
22+
type: dense_vector
23+
dims: 64
24+
index: true
25+
similarity: max_inner_product
26+
index_options:
27+
type: bbq_hnsw
28+
29+
- do:
30+
index:
31+
index: bbq_hnsw
32+
id: "1"
33+
body:
34+
vector: [0.077, 0.32 , -0.205, 0.63 , 0.032, 0.201, 0.167, -0.313,
35+
0.176, 0.531, -0.375, 0.334, -0.046, 0.078, -0.349, 0.272,
36+
0.307, -0.083, 0.504, 0.255, -0.404, 0.289, -0.226, -0.132,
37+
-0.216, 0.49 , 0.039, 0.507, -0.307, 0.107, 0.09 , -0.265,
38+
-0.285, 0.336, -0.272, 0.369, -0.282, 0.086, -0.132, 0.475,
39+
-0.224, 0.203, 0.439, 0.064, 0.246, -0.396, 0.297, 0.242,
40+
-0.028, 0.321, -0.022, -0.009, -0.001 , 0.031, -0.533, 0.45,
41+
-0.683, 1.331, 0.194, -0.157, -0.1 , -0.279, -0.098, -0.176]
42+
# Flush in order to provoke a merge later
43+
- do:
44+
indices.flush:
45+
index: bbq_hnsw
46+
47+
- do:
48+
index:
49+
index: bbq_hnsw
50+
id: "2"
51+
body:
52+
vector: [0.196, 0.514, 0.039, 0.555, -0.042, 0.242, 0.463, -0.348,
53+
-0.08 , 0.442, -0.067, -0.05 , -0.001, 0.298, -0.377, 0.048,
54+
0.307, 0.159, 0.278, 0.119, -0.057, 0.333, -0.289, -0.438,
55+
-0.014, 0.361, -0.169, 0.292, -0.229, 0.123, 0.031, -0.138,
56+
-0.139, 0.315, -0.216, 0.322, -0.445, -0.059, 0.071, 0.429,
57+
-0.602, -0.142, 0.11 , 0.192, 0.259, -0.241, 0.181, -0.166,
58+
0.082, 0.107, -0.05 , 0.155, 0.011, 0.161, -0.486, 0.569,
59+
-0.489, 0.901, 0.208, 0.011, -0.209, -0.153, -0.27 , -0.013]
60+
# Flush in order to provoke a merge later
61+
- do:
62+
indices.flush:
63+
index: bbq_hnsw
64+
65+
- do:
66+
index:
67+
index: bbq_hnsw
68+
id: "3"
69+
body:
70+
name: rabbit.jpg
71+
vector: [0.139, 0.178, -0.117, 0.399, 0.014, -0.139, 0.347, -0.33 ,
72+
0.139, 0.34 , -0.052, -0.052, -0.249, 0.327, -0.288, 0.049,
73+
0.464, 0.338, 0.516, 0.247, -0.104, 0.259, -0.209, -0.246,
74+
-0.11 , 0.323, 0.091, 0.442, -0.254, 0.195, -0.109, -0.058,
75+
-0.279, 0.402, -0.107, 0.308, -0.273, 0.019, 0.082, 0.399,
76+
-0.658, -0.03 , 0.276, 0.041, 0.187, -0.331, 0.165, 0.017,
77+
0.171, -0.203, -0.198, 0.115, -0.007, 0.337, -0.444, 0.615,
78+
-0.657, 1.285, 0.2 , -0.062, 0.038, 0.089, -0.068, -0.058]
79+
# Flush in order to provoke a merge later
80+
- do:
81+
indices.flush:
82+
index: bbq_hnsw
83+
84+
- do:
85+
indices.forcemerge:
86+
index: bbq_hnsw
87+
max_num_segments: 1
88+
---
89+
"Profile rescored knn search":
90+
91+
- do:
92+
search:
93+
index: bbq_hnsw
94+
body:
95+
profile: true
96+
knn:
97+
field: vector
98+
query_vector: [0.128, 0.067, -0.08 , 0.395, -0.11 , -0.259, 0.473, -0.393,
99+
0.292, 0.571, -0.491, 0.444, -0.288, 0.198, -0.343, 0.015,
100+
0.232, 0.088, 0.228, 0.151, -0.136, 0.236, -0.273, -0.259,
101+
-0.217, 0.359, -0.207, 0.352, -0.142, 0.192, -0.061, -0.17 ,
102+
-0.343, 0.189, -0.221, 0.32 , -0.301, -0.1 , 0.005, 0.232,
103+
-0.344, 0.136, 0.252, 0.157, -0.13 , -0.244, 0.193, -0.034,
104+
-0.12 , -0.193, -0.102, 0.252, -0.185, -0.167, -0.575, 0.582,
105+
-0.426, 0.983, 0.212, 0.204, 0.03 , -0.276, -0.425, -0.158]
106+
k: 3
107+
num_candidates: 3
108+
"rescore_vector":
109+
"num_candidates_factor": 2.0
110+
111+
# We expect the knn search ops + rescoring num_cnaidates (for rescoring) per shard
112+
- match: { profile.shards.0.dfs.knn.0.vector_operations_count: 6 }
113+
114+
# Search with similarity to check number of operations are propagated correctly
115+
- do:
116+
search:
117+
index: bbq_hnsw
118+
body:
119+
profile: true
120+
knn:
121+
field: vector
122+
query_vector: [0.128, 0.067, -0.08 , 0.395, -0.11 , -0.259, 0.473, -0.393,
123+
0.292, 0.571, -0.491, 0.444, -0.288, 0.198, -0.343, 0.015,
124+
0.232, 0.088, 0.228, 0.151, -0.136, 0.236, -0.273, -0.259,
125+
-0.217, 0.359, -0.207, 0.352, -0.142, 0.192, -0.061, -0.17 ,
126+
-0.343, 0.189, -0.221, 0.32 , -0.301, -0.1 , 0.005, 0.232,
127+
-0.344, 0.136, 0.252, 0.157, -0.13 , -0.244, 0.193, -0.034,
128+
-0.12 , -0.193, -0.102, 0.252, -0.185, -0.167, -0.575, 0.582,
129+
-0.426, 0.983, 0.212, 0.204, 0.03 , -0.276, -0.425, -0.158]
130+
k: 3
131+
num_candidates: 3
132+
similarity: 100000
133+
"rescore_vector":
134+
"num_candidates_factor": 2.0
135+
136+
# We expect the knn search ops + rescoring num_cnaidates (for rescoring) per shard
137+
- match: { profile.shards.0.dfs.knn.0.vector_operations_count: 6 }

‎rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/40_knn_search.yml‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -541,3 +541,58 @@ setup:
541541
num_candidates: 3
542542

543543
- match: { hits.total.value: 0 }
544+
---
545+
"Vector rescoring has no effect for non-quantized vectors and provides same results as non-rescored knn":
546+
- requires:
547+
reason: 'Quantized vector rescoring is required'
548+
test_runner_features: [capabilities]
549+
capabilities:
550+
- method: GET
551+
path: /_search
552+
capabilities: [knn_quantized_vector_rescore]
553+
- skip:
554+
features: "headers"
555+
556+
# Non-rescored knn
557+
- do:
558+
headers:
559+
Content-Type: application/json
560+
search:
561+
rest_total_hits_as_int: true
562+
index: test
563+
body:
564+
fields: [ "name" ]
565+
knn:
566+
field: vector
567+
query_vector: [-0.5, 90.0, -10, 14.8, -156.0]
568+
k: 3
569+
num_candidates: 3
570+
571+
# Get scores - hit ordering may change depending on how things are distributed
572+
- match: { hits.total: 3 }
573+
- set: { hits.hits.0._score: knn_score0 }
574+
- set: { hits.hits.1._score: knn_score1 }
575+
- set: { hits.hits.2._score: knn_score2 }
576+
577+
# Rescored knn
578+
- do:
579+
headers:
580+
Content-Type: application/json
581+
search:
582+
rest_total_hits_as_int: true
583+
index: test
584+
body:
585+
fields: [ "name" ]
586+
knn:
587+
field: vector
588+
query_vector: [-0.5, 90.0, -10, 14.8, -156.0]
589+
k: 3
590+
num_candidates: 3
591+
rescore_vector:
592+
num_candidates_factor: 1.5
593+
594+
# Compare scores as hit IDs may change depending on how things are distributed
595+
- match: { hits.total: 3 }
596+
- match: { hits.hits.0._score: $knn_score0 }
597+
- match: { hits.hits.1._score: $knn_score1 }
598+
- match: { hits.hits.2._score: $knn_score2 }

‎rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/41_knn_search_bbq_hnsw.yml‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,75 @@ setup:
108108
- match: { hits.hits.1._id: "3" }
109109
- match: { hits.hits.2._id: "2" }
110110
---
111+
"Vector rescoring has same scoring as exact search for kNN section":
112+
- requires:
113+
reason: 'Quantized vector rescoring is required'
114+
test_runner_features: [capabilities]
115+
capabilities:
116+
- method: GET
117+
path: /_search
118+
capabilities: [knn_quantized_vector_rescore]
119+
- skip:
120+
features: "headers"
121+
122+
# Rescore
123+
- do:
124+
headers:
125+
Content-Type: application/json
126+
search:
127+
rest_total_hits_as_int: true
128+
index: bbq_hnsw
129+
body:
130+
knn:
131+
field: vector
132+
query_vector: [0.128, 0.067, -0.08 , 0.395, -0.11 , -0.259, 0.473, -0.393,
133+
0.292, 0.571, -0.491, 0.444, -0.288, 0.198, -0.343, 0.015,
134+
0.232, 0.088, 0.228, 0.151, -0.136, 0.236, -0.273, -0.259,
135+
-0.217, 0.359, -0.207, 0.352, -0.142, 0.192, -0.061, -0.17 ,
136+
-0.343, 0.189, -0.221, 0.32 , -0.301, -0.1 , 0.005, 0.232,
137+
-0.344, 0.136, 0.252, 0.157, -0.13 , -0.244, 0.193, -0.034,
138+
-0.12 , -0.193, -0.102, 0.252, -0.185, -0.167, -0.575, 0.582,
139+
-0.426, 0.983, 0.212, 0.204, 0.03 , -0.276, -0.425, -0.158]
140+
k: 3
141+
num_candidates: 3
142+
rescore_vector:
143+
num_candidates_factor: 1.5
144+
145+
# Get rescoring scores - hit ordering may change depending on how things are distributed
146+
- match: { hits.total: 3 }
147+
- set: { hits.hits.0._score: rescore_score0 }
148+
- set: { hits.hits.1._score: rescore_score1 }
149+
- set: { hits.hits.2._score: rescore_score2 }
150+
151+
# Exact knn via script score
152+
- do:
153+
headers:
154+
Content-Type: application/json
155+
search:
156+
rest_total_hits_as_int: true
157+
body:
158+
query:
159+
script_score:
160+
query: {match_all: {} }
161+
script:
162+
source: "double similarity = dotProduct(params.query_vector, 'vector'); return similarity < 0 ? 1 / (1 + -1 * similarity) : similarity + 1"
163+
params:
164+
query_vector: [0.128, 0.067, -0.08 , 0.395, -0.11 , -0.259, 0.473, -0.393,
165+
0.292, 0.571, -0.491, 0.444, -0.288, 0.198, -0.343, 0.015,
166+
0.232, 0.088, 0.228, 0.151, -0.136, 0.236, -0.273, -0.259,
167+
-0.217, 0.359, -0.207, 0.352, -0.142, 0.192, -0.061, -0.17 ,
168+
-0.343, 0.189, -0.221, 0.32 , -0.301, -0.1 , 0.005, 0.232,
169+
-0.344, 0.136, 0.252, 0.157, -0.13 , -0.244, 0.193, -0.034,
170+
-0.12 , -0.193, -0.102, 0.252, -0.185, -0.167, -0.575, 0.582,
171+
-0.426, 0.983, 0.212, 0.204, 0.03 , -0.276, -0.425, -0.158]
172+
173+
# Compare scores as hit IDs may change depending on how things are distributed
174+
- match: { hits.total: 3 }
175+
- match: { hits.hits.0._score: $rescore_score0 }
176+
- match: { hits.hits.1._score: $rescore_score1 }
177+
- match: { hits.hits.2._score: $rescore_score2 }
178+
179+
---
111180
"Test bad quantization parameters":
112181
- do:
113182
catch: bad_request

0 commit comments

Comments
 (0)