Skip to content

Commit 3103cae

Browse files
usamoiusamoi
committed
perf: disable SAQ for rabitq8 (#434)
Signed-off-by: usamoi <usamoi@outlook.com>
1 parent 578859a commit 3103cae

4 files changed

Lines changed: 171 additions & 127 deletions

File tree

‎crates/rabitq/src/extended.rs‎

Lines changed: 3 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -112,15 +112,14 @@ pub fn ugly_code<const BITS: usize>(vector: &[f32]) -> Code {
112112
f32::vector_mul_scalar_inplace(&mut vector, 1.0 / dis_u_2.sqrt());
113113
vector
114114
};
115-
let (scale, delta) = {
115+
let scale = {
116116
let mut o = normalized_vector.clone();
117117
f32::vector_abs_inplace(&mut o);
118118
ugly_find_scale::<BITS>(&o)
119119
};
120120
let mut code = Vec::with_capacity(n as _);
121121
for i in 0..n {
122122
let v = scale * normalized_vector[i];
123-
let v = v + v.signum() * delta[i] as f32;
124123
let c = v.floor().clamp(min as f32, max as f32) as i32;
125124
code.push((c + (1 << (BITS - 1))) as _);
126125
}
@@ -251,55 +250,10 @@ fn find_scale<const B: usize>(o: &[f32]) -> f32 {
251250
x_m as f32 + f32::EPSILON
252251
}
253252

254-
fn ugly_find_scale<const B: usize>(o: &[f32]) -> (f32, Vec<i32>) {
253+
fn ugly_find_scale<const B: usize>(o: &[f32]) -> f32 {
255254
assert!((1..=8).contains(&B));
256255

257-
let dim = o.len();
258-
259-
let mut code = Vec::<u8>::with_capacity(dim);
260-
let mut numerator_m = 0.0f64;
261-
let mut sqr_denominator_m = 0.0f64;
262-
263-
let scale = (1 << (B - 1)) as f32 / f32::reduce_min_max_of_x(o).1;
264-
for i in 0..dim {
265-
code.push((o[i] as f64 * scale as f64) as u8);
266-
let value = code[i] as f64 + 0.5;
267-
numerator_m += value * o[i] as f64;
268-
sqr_denominator_m += value * value;
269-
}
270-
let mut y_m = numerator_m / sqr_denominator_m.sqrt();
271-
272-
let mut delta = vec![0_i32; dim];
273-
for _ in 0..8 {
274-
for i in 0..dim {
275-
if code[i] < (1 << (B - 1)) - 1 {
276-
let numerator = numerator_m + o[i] as f64;
277-
let sqr_denominator = sqr_denominator_m + 2.0 * (code[i] as f64 + 1.0);
278-
let y = numerator / sqr_denominator.sqrt();
279-
if y > y_m {
280-
y_m = y;
281-
numerator_m = numerator;
282-
sqr_denominator_m = sqr_denominator;
283-
code[i] += 1;
284-
delta[i] += 1;
285-
}
286-
}
287-
if code[i] > 0 {
288-
let numerator = numerator_m - o[i] as f64;
289-
let sqr_denominator = sqr_denominator_m - 2.0 * code[i] as f64;
290-
let y = numerator / sqr_denominator.sqrt();
291-
if y > y_m {
292-
y_m = y;
293-
numerator_m = numerator;
294-
sqr_denominator_m = sqr_denominator;
295-
code[i] -= 1;
296-
delta[i] -= 1;
297-
}
298-
}
299-
}
300-
}
301-
302-
(scale, delta)
256+
(1 << (B - 1)) as f32 / f32::reduce_min_max_of_x(o).1
303257
}
304258

305259
pub fn pack_code<const BITS: usize>(input: &[u8]) -> [Vec<u64>; BITS] {

‎tests/vchordg/index_vector_rabitq.slt‎

Lines changed: 85 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -20,16 +20,16 @@ CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabi
2020
query I
2121
SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10;
2222
----
23-
155
2423
1608
25-
1643
24+
155
2625
174
2726
818
2827
1603
2928
1629
3029
60
3130
218
3231
1080
32+
1568
3333

3434
statement ok
3535
DROP INDEX ti;
@@ -40,16 +40,16 @@ CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabi
4040
query I
4141
SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10;
4242
----
43-
155
4443
1608
45-
1643
44+
155
4645
174
4746
818
4847
1603
4948
1629
5049
60
5150
218
5251
1080
52+
1639
5353

5454
statement ok
5555
DROP INDEX ti;
@@ -60,16 +60,16 @@ CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val)::rabitq8(64)) rabi
6060
query I
6161
SELECT index FROM t ORDER BY quantize_to_rabitq8(val) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10;
6262
----
63-
155
6463
1608
65-
1643
64+
155
6665
174
6766
818
6867
1603
6968
1629
7069
60
7170
218
7271
1080
72+
1568
7373

7474
statement ok
7575
DROP INDEX ti;
@@ -78,17 +78,18 @@ statement ok
7878
CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_l2_ops);
7979

8080
query I
81-
SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 9;
81+
SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <-> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10;
8282
----
83-
155
8483
1608
85-
1643
84+
155
8685
174
87-
818
8886
1603
87+
818
8988
1629
9089
60
9190
218
91+
1080
92+
79
9293

9394
statement ok
9495
DROP INDEX ti;
@@ -97,17 +98,18 @@ statement ok
9798
CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_ip_ops);
9899

99100
query I
100-
SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 9;
101+
SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <#> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10;
101102
----
102-
155
103103
1608
104-
1643
104+
155
105105
174
106106
818
107107
1603
108108
1629
109109
60
110110
218
111+
1080
112+
79
111113

112114
statement ok
113115
DROP INDEX ti;
@@ -116,17 +118,18 @@ statement ok
116118
CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq8(val::halfvec)::rabitq8(64)) rabitq8_cosine_ops);
117119

118120
query I
119-
SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 9;
121+
SELECT index FROM t ORDER BY quantize_to_rabitq8(val::halfvec) <=> quantize_to_rabitq8(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10;
120122
----
121-
155
122123
1608
123-
1643
124+
155
124125
174
125126
818
126127
1603
127128
1629
128129
60
129130
218
131+
1080
132+
79
130133

131134
statement ok
132135
DROP INDEX ti;
@@ -135,11 +138,18 @@ statement ok
135138
CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_l2_ops);
136139

137140
query I
138-
SELECT (SUM(index) = 6162 OR SUM(index) = 6289)::int FROM (
139-
SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10
140-
) AS s(index);
141+
SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10;
141142
----
142-
1
143+
155
144+
1643
145+
174
146+
818
147+
1207
148+
1608
149+
1940
150+
331
151+
1568
152+
1839
143153

144154
statement ok
145155
DROP INDEX ti;
@@ -148,11 +158,18 @@ statement ok
148158
CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_ip_ops);
149159

150160
query I
151-
SELECT (SUM(index) = 6162 OR SUM(index) = 6289)::int FROM (
152-
SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10
153-
) AS s(index);
161+
SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10;
154162
----
155-
1
163+
155
164+
1643
165+
174
166+
818
167+
1207
168+
1608
169+
1940
170+
331
171+
1568
172+
1839
156173

157174
statement ok
158175
DROP INDEX ti;
@@ -161,11 +178,18 @@ statement ok
161178
CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val)::rabitq4(64)) rabitq4_cosine_ops);
162179

163180
query I
164-
SELECT (SUM(index) = 6162 OR SUM(index) = 6289)::int FROM (
165-
SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10
166-
) AS s(index);
181+
SELECT index FROM t ORDER BY quantize_to_rabitq4(val) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::vector) LIMIT 10;
167182
----
168-
1
183+
155
184+
1643
185+
174
186+
818
187+
1207
188+
1608
189+
1940
190+
331
191+
1568
192+
1839
169193

170194
statement ok
171195
DROP INDEX ti;
@@ -174,11 +198,18 @@ statement ok
174198
CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_l2_ops);
175199

176200
query I
177-
SELECT SUM(index) FROM (
178-
SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10
179-
) AS s(index);
201+
SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <-> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10;
180202
----
181-
7290
203+
155
204+
1643
205+
174
206+
818
207+
1207
208+
1608
209+
1940
210+
331
211+
1839
212+
1568
182213

183214
statement ok
184215
DROP INDEX ti;
@@ -187,11 +218,18 @@ statement ok
187218
CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_ip_ops);
188219

189220
query I
190-
SELECT SUM(index) FROM (
191-
SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10
192-
) AS s(index);
221+
SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <#> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10;
193222
----
194-
7290
223+
155
224+
1643
225+
174
226+
818
227+
1207
228+
1608
229+
1940
230+
331
231+
1568
232+
1839
195233

196234
statement ok
197235
DROP INDEX ti;
@@ -200,11 +238,18 @@ statement ok
200238
CREATE INDEX ti ON t USING vchordg ((quantize_to_rabitq4(val::halfvec)::rabitq4(64)) rabitq4_cosine_ops);
201239

202240
query I
203-
SELECT SUM(index) FROM (
204-
SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10
205-
) AS s(index);
241+
SELECT index FROM t ORDER BY quantize_to_rabitq4(val::halfvec) <=> quantize_to_rabitq4(array_cat(ARRAY[0.6, 0.8], ARRAY(SELECT 0.0 FROM generate_series(1, 62)))::halfvec) LIMIT 10;
206242
----
207-
7290
243+
155
244+
1643
245+
174
246+
818
247+
1207
248+
1608
249+
1940
250+
331
251+
1568
252+
1839
208253

209254
statement ok
210255
DROP INDEX ti;

‎tests/vchordrq/index_multivector_rabitq.slt‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -74,13 +74,13 @@ LIMIT 10;
7474
1207
7575
919
7676
1821
77-
1076
7877
1639
78+
537
7979
174
80-
1125
8180
79
82-
1163
83-
537
81+
194
82+
1920
83+
239
8484

8585
statement ok
8686
DROP INDEX pg_temp.ti;

0 commit comments

Comments
 (0)