1use crate::error::{Result, RuvectorError};
16use crate::index::VectorIndex;
17use crate::types::{
18 DistanceMetric, HnswConfig, SearchPolicy, SearchQuantization, SearchResult, VectorId,
19};
20use dashmap::DashMap;
21use hnsw_rs::prelude::*;
22use parking_lot::RwLock;
23use ruvector_turboquant::{bits1, score, Metric, Turbo4Codec, Turbo4Query};
24use std::sync::atomic::{AtomicU64, Ordering};
25use std::sync::Arc;
26
27struct EscalationParams {
31 margin_threshold: f32,
32 ef_mult: usize,
33 max_rounds: usize,
34}
35
36impl EscalationParams {
37 fn for_policy(policy: SearchPolicy) -> Self {
46 match policy {
47 SearchPolicy::Quality => Self {
48 margin_threshold: 0.05,
49 ef_mult: 3,
50 max_rounds: 2,
51 },
52 SearchPolicy::Balanced | SearchPolicy::MaxCompression => Self {
53 margin_threshold: 0.0,
54 ef_mult: 1,
55 max_rounds: 0,
56 },
57 }
58 }
59}
60
61struct Turbo4DistanceFn {
73 metric: Metric,
74 dim: usize,
75 code_len: usize,
76 query_len: usize,
77 combined_len: usize,
79 bits_query_len: usize,
80 bits_code_len: usize,
81}
82
83impl Distance<u8> for Turbo4DistanceFn {
84 #[inline(always)]
85 fn eval(&self, a: &[u8], b: &[u8]) -> f32 {
86 if a.len() == self.combined_len && b.len() == self.combined_len {
89 score::symmetric_distance(
90 self.metric,
91 &a[self.bits_code_len..],
92 &b[self.bits_code_len..],
93 self.dim,
94 )
95 } else if a.len() == self.bits_query_len && b.len() == self.combined_len {
96 bits1::query_blob_distance(self.metric, a, &b[..self.bits_code_len], self.dim)
97 } else if b.len() == self.bits_query_len && a.len() == self.combined_len {
98 bits1::query_blob_distance(self.metric, b, &a[..self.bits_code_len], self.dim)
99 } else if a.len() == self.code_len && b.len() == self.code_len {
100 score::symmetric_distance(self.metric, a, b, self.dim)
101 } else if a.len() == self.query_len && b.len() == self.code_len {
102 score::asymmetric_distance(self.metric, a, b, self.dim)
103 } else if b.len() == self.query_len && a.len() == self.code_len {
104 score::asymmetric_distance(self.metric, b, a, self.dim)
105 } else {
106 debug_assert!(
107 false,
108 "unrecognized Turbo4 blob lengths {}/{}",
109 a.len(),
110 b.len()
111 );
112 f32::MAX
113 }
114 }
115}
116
117fn to_turbo_metric(metric: DistanceMetric) -> Result<Metric> {
118 match metric {
119 DistanceMetric::Euclidean => Ok(Metric::Euclidean),
120 DistanceMetric::Cosine => Ok(Metric::Cosine),
121 DistanceMetric::DotProduct => Ok(Metric::DotProduct),
122 DistanceMetric::Manhattan => Err(RuvectorError::InvalidParameter(
123 "Turbo4 quantization does not support the Manhattan metric \
124 (it does not decompose over the level dot product); use \
125 Euclidean, Cosine, or DotProduct, or disable quantization"
126 .into(),
127 )),
128 }
129}
130
131struct Turbo4Inner {
132 hnsw: Hnsw<'static, u8, Turbo4DistanceFn>,
133 codes: DashMap<VectorId, Vec<u8>>,
136 id_to_idx: DashMap<VectorId, usize>,
137 idx_to_id: DashMap<usize, VectorId>,
138 next_idx: usize,
139}
140
141#[derive(bincode::Encode, bincode::Decode)]
147struct Turbo4State {
148 codes: Vec<(String, Vec<u8>)>,
149 id_to_idx: Vec<(String, usize)>,
150 idx_to_id: Vec<(usize, String)>,
151 next_idx: usize,
152 m: usize,
153 ef_construction: usize,
154 ef_search: usize,
155 max_elements: usize,
156 dimensions: usize,
157 metric: u8,
159 rotation_seed: u64,
160 rescore_multiplier: usize,
161 policy: u8,
163 search_quant: u8,
165}
166
167fn policy_to_u8(p: SearchPolicy) -> u8 {
168 match p {
169 SearchPolicy::Quality => 0,
170 SearchPolicy::Balanced => 1,
171 SearchPolicy::MaxCompression => 2,
172 }
173}
174
175fn policy_from_u8(v: u8) -> SearchPolicy {
176 match v {
177 0 => SearchPolicy::Quality,
178 2 => SearchPolicy::MaxCompression,
179 _ => SearchPolicy::Balanced,
180 }
181}
182
183pub struct Turbo4HnswIndex {
186 inner: Arc<RwLock<Turbo4Inner>>,
187 codec: Arc<Turbo4Codec>,
188 config: HnswConfig,
189 metric: Metric,
190 dimensions: usize,
191 rotation_seed: u64,
192 rescore_multiplier: usize,
193 policy: SearchPolicy,
194 search_quantization: SearchQuantization,
195 t4_off: usize,
198 escalation: EscalationParams,
199 queries: AtomicU64,
201 escalated: AtomicU64,
202}
203
204impl Turbo4HnswIndex {
205 pub fn new(
207 dimensions: usize,
208 metric: DistanceMetric,
209 config: HnswConfig,
210 rotation_seed: u64,
211 rescore_multiplier: usize,
212 policy: SearchPolicy,
213 search_quantization: SearchQuantization,
214 ) -> Result<Self> {
215 let metric = to_turbo_metric(metric)?;
216 let codec = Turbo4Codec::new(dimensions, rotation_seed).map_err(|e| {
217 RuvectorError::InvalidParameter(format!("Turbo4 codec init failed: {e}"))
218 })?;
219 let cascade = search_quantization == SearchQuantization::RaBitQ1;
220 let bits_code_len = bits1::code1_len(dimensions);
221 let distance_fn = Turbo4DistanceFn {
226 metric,
227 dim: dimensions,
228 code_len: codec.code_len(),
229 query_len: codec.query_len(),
230 combined_len: if cascade {
231 bits_code_len + codec.code_len()
232 } else {
233 0
234 },
235 bits_query_len: if cascade {
236 bits1::query1_len(dimensions)
237 } else {
238 0
239 },
240 bits_code_len,
241 };
242 let hnsw = Hnsw::<u8, Turbo4DistanceFn>::new(
243 config.m,
244 config.max_elements,
245 dimensions,
246 config.ef_construction,
247 distance_fn,
248 );
249 Ok(Self {
250 inner: Arc::new(RwLock::new(Turbo4Inner {
251 hnsw,
252 codes: DashMap::new(),
253 id_to_idx: DashMap::new(),
254 idx_to_id: DashMap::new(),
255 next_idx: 0,
256 })),
257 codec: Arc::new(codec),
258 config,
259 metric,
260 dimensions,
261 rotation_seed,
262 rescore_multiplier: rescore_multiplier.max(1),
263 policy,
264 search_quantization,
265 t4_off: if cascade { bits_code_len } else { 0 },
266 escalation: EscalationParams::for_policy(policy),
267 queries: AtomicU64::new(0),
268 escalated: AtomicU64::new(0),
269 })
270 }
271
272 fn encode_blob(&self, vector: &[f32]) -> Result<Vec<u8>> {
274 if self.search_quantization == SearchQuantization::RaBitQ1 {
275 let (turbo4, bits) = self
276 .codec
277 .encode_dual(vector)
278 .map_err(|e| RuvectorError::InvalidInput(e.to_string()))?;
279 let mut blob = bits;
280 blob.extend_from_slice(&turbo4);
281 Ok(blob)
282 } else {
283 self.codec
284 .encode(vector)
285 .map_err(|e| RuvectorError::InvalidInput(e.to_string()))
286 }
287 }
288
289 pub fn serialize(&self) -> Result<Vec<u8>> {
292 let inner = self.inner.read();
293 let metric_u8 = match self.metric {
294 Metric::Euclidean => 0u8,
295 Metric::Cosine => 1,
296 Metric::DotProduct => 2,
297 };
298 let state = Turbo4State {
299 codes: inner
300 .codes
301 .iter()
302 .map(|e| (e.key().clone(), e.value().clone()))
303 .collect(),
304 id_to_idx: inner
305 .id_to_idx
306 .iter()
307 .map(|e| (e.key().clone(), *e.value()))
308 .collect(),
309 idx_to_id: inner
310 .idx_to_id
311 .iter()
312 .map(|e| (*e.key(), e.value().clone()))
313 .collect(),
314 next_idx: inner.next_idx,
315 m: self.config.m,
316 ef_construction: self.config.ef_construction,
317 ef_search: self.config.ef_search,
318 max_elements: self.config.max_elements,
319 dimensions: self.dimensions,
320 metric: metric_u8,
321 rotation_seed: self.rotation_seed,
322 rescore_multiplier: self.rescore_multiplier,
323 policy: policy_to_u8(self.policy),
324 search_quant: match self.search_quantization {
325 SearchQuantization::Turbo4Direct => 0,
326 SearchQuantization::RaBitQ1 => 1,
327 },
328 };
329 bincode::encode_to_vec(&state, bincode::config::standard()).map_err(|e| {
330 RuvectorError::SerializationError(format!("Failed to serialize Turbo4 index: {e}"))
331 })
332 }
333
334 pub fn deserialize(bytes: &[u8]) -> Result<Self> {
338 let (state, _): (Turbo4State, usize) =
339 bincode::decode_from_slice(bytes, bincode::config::standard()).map_err(|e| {
340 RuvectorError::SerializationError(format!(
341 "Failed to deserialize Turbo4 index: {e}"
342 ))
343 })?;
344 let metric = match state.metric {
345 0 => DistanceMetric::Euclidean,
346 1 => DistanceMetric::Cosine,
347 2 => DistanceMetric::DotProduct,
348 other => {
349 return Err(RuvectorError::SerializationError(format!(
350 "unknown Turbo4 metric tag {other}"
351 )))
352 }
353 };
354 let config = HnswConfig {
355 m: state.m,
356 ef_construction: state.ef_construction,
357 ef_search: state.ef_search,
358 max_elements: state.max_elements,
359 };
360 let index = Self::new(
361 state.dimensions,
362 metric,
363 config,
364 state.rotation_seed,
365 state.rescore_multiplier,
366 policy_from_u8(state.policy),
367 if state.search_quant == 1 {
368 SearchQuantization::RaBitQ1
369 } else {
370 SearchQuantization::Turbo4Direct
371 },
372 )?;
373 {
374 let mut inner = index.inner.write();
375 let codes_lookup: std::collections::HashMap<&str, &Vec<u8>> =
376 state.codes.iter().map(|(id, c)| (id.as_str(), c)).collect();
377 let mut sorted: Vec<(usize, &String)> =
378 state.idx_to_id.iter().map(|(i, id)| (*i, id)).collect();
379 sorted.sort_unstable_by_key(|(i, _)| *i);
380 for (idx, id) in &sorted {
381 if let Some(blob) = codes_lookup.get(id.as_str()) {
382 inner.hnsw.insert_data(blob, *idx);
383 }
384 }
385 for (id, blob) in state.codes {
386 inner.codes.insert(id, blob);
387 }
388 for (id, idx) in state.id_to_idx {
389 inner.id_to_idx.insert(id, idx);
390 }
391 for (idx, id) in state.idx_to_id {
392 inner.idx_to_id.insert(idx, id);
393 }
394 inner.next_idx = state.next_idx;
395 }
396 Ok(index)
397 }
398
399 pub fn code_len(&self) -> usize {
401 self.codec.code_len()
402 }
403
404 pub fn adaptive_stats(&self) -> (u64, u64) {
407 (
408 self.queries.load(Ordering::Relaxed),
409 self.escalated.load(Ordering::Relaxed),
410 )
411 }
412
413 fn traverse_and_rescore(
417 &self,
418 inner: &Turbo4Inner,
419 traversal_blob: &[u8],
420 prepared: &Turbo4Query,
421 fetch: usize,
422 ef: usize,
423 ) -> Vec<SearchResult> {
424 let neighbors = inner.hnsw.search(traversal_blob, fetch, ef);
425 let mut results: Vec<SearchResult> = neighbors
426 .into_iter()
427 .filter_map(|n| {
428 let id = inner.idx_to_id.get(&n.d_id)?.clone();
429 let code = inner.codes.get(&id)?;
430 let dist = score::rescore(
431 self.metric,
432 prepared,
433 &code.value()[self.t4_off..],
434 self.dimensions,
435 );
436 Some(SearchResult {
437 id,
438 score: dist,
439 vector: None,
440 metadata: None,
441 })
442 })
443 .collect();
444 results.sort_unstable_by(|a, b| {
445 a.score
446 .partial_cmp(&b.score)
447 .unwrap_or(std::cmp::Ordering::Equal)
448 });
449 results
450 }
451
452 fn boundary_margin(results: &[SearchResult], k: usize) -> f32 {
455 if results.len() <= k || k == 0 {
456 return f32::INFINITY;
457 }
458 let kept = results[k - 1].score;
459 let dropped = results[k].score;
460 (dropped - kept) / kept.abs().max(1e-6)
461 }
462
463 pub fn search_with_ef(
465 &self,
466 query: &[f32],
467 k: usize,
468 ef_search: usize,
469 ) -> Result<Vec<SearchResult>> {
470 if query.len() != self.dimensions {
471 return Err(RuvectorError::DimensionMismatch {
472 expected: self.dimensions,
473 actual: query.len(),
474 });
475 }
476 if k == 0 {
477 return Ok(vec![]);
478 }
479 let prepared = self
480 .codec
481 .encode_query(query)
482 .map_err(|e| RuvectorError::InvalidInput(e.to_string()))?;
483 let traversal_blob: Vec<u8> = if self.search_quantization == SearchQuantization::RaBitQ1 {
486 bits1::Bits1Query::new(&prepared.rotated)
487 .map_err(|e| RuvectorError::InvalidInput(e.to_string()))?
488 .to_blob()
489 } else {
490 prepared.blob.clone()
491 };
492
493 let inner = self.inner.read();
494 if inner.codes.is_empty() {
496 return Ok(vec![]);
497 }
498 self.queries.fetch_add(1, Ordering::Relaxed);
499
500 let mut fetch = k.saturating_mul(self.rescore_multiplier);
502 let mut ef = ef_search.max(fetch);
503 let mut results = self.traverse_and_rescore(&inner, &traversal_blob, &prepared, fetch, ef);
504
505 let mut round = 0;
510 let mut did_escalate = false;
511 while round < self.escalation.max_rounds
512 && Self::boundary_margin(&results, k) < self.escalation.margin_threshold
513 {
514 did_escalate = true;
515 fetch = fetch.saturating_mul(2);
516 ef = ef.saturating_mul(self.escalation.ef_mult);
517 let wider = self.traverse_and_rescore(&inner, &traversal_blob, &prepared, fetch, ef);
518
519 let stable = wider.len() >= k
520 && results.len() >= k
521 && wider[..k]
522 .iter()
523 .zip(&results[..k])
524 .all(|(a, b)| a.id == b.id);
525 results = wider;
526 round += 1;
527 if stable {
528 break;
529 }
530 }
531 if did_escalate {
532 self.escalated.fetch_add(1, Ordering::Relaxed);
533 }
534
535 results.truncate(k);
536 Ok(results)
537 }
538}
539
540impl VectorIndex for Turbo4HnswIndex {
541 fn add(&mut self, id: VectorId, vector: Vec<f32>) -> Result<()> {
542 if vector.len() != self.dimensions {
543 return Err(RuvectorError::DimensionMismatch {
544 expected: self.dimensions,
545 actual: vector.len(),
546 });
547 }
548 let blob = self.encode_blob(&vector)?;
549 drop(vector); let mut inner = self.inner.write();
552 let idx = inner.next_idx;
553 inner.next_idx += 1;
554 inner.hnsw.insert_data(&blob, idx);
555 inner.codes.insert(id.clone(), blob);
556 inner.id_to_idx.insert(id.clone(), idx);
557 inner.idx_to_id.insert(idx, id);
558 Ok(())
559 }
560
561 fn add_batch(&mut self, entries: Vec<(VectorId, Vec<f32>)>) -> Result<()> {
562 for (_, vector) in &entries {
563 if vector.len() != self.dimensions {
564 return Err(RuvectorError::DimensionMismatch {
565 expected: self.dimensions,
566 actual: vector.len(),
567 });
568 }
569 }
570 let mut encoded = Vec::with_capacity(entries.len());
572 for (id, vector) in entries {
573 let blob = self.encode_blob(&vector)?;
574 encoded.push((id, blob));
575 }
576
577 let mut inner = self.inner.write();
578 let base = inner.next_idx;
579 inner.next_idx += encoded.len();
580
581 const PARALLEL_THRESHOLD: usize = 10_000;
584 if encoded.len() >= PARALLEL_THRESHOLD {
585 let datas: Vec<(&[u8], usize)> = encoded
586 .iter()
587 .enumerate()
588 .map(|(i, (_, blob))| (blob.as_slice(), base + i))
589 .collect();
590 inner.hnsw.parallel_insert_slice(&datas);
591 } else {
592 for (i, (_, blob)) in encoded.iter().enumerate() {
593 inner.hnsw.insert_data(blob, base + i);
594 }
595 }
596
597 for (i, (id, blob)) in encoded.into_iter().enumerate() {
598 inner.codes.insert(id.clone(), blob);
599 inner.id_to_idx.insert(id.clone(), base + i);
600 inner.idx_to_id.insert(base + i, id);
601 }
602 Ok(())
603 }
604
605 fn search(&self, query: &[f32], k: usize) -> Result<Vec<SearchResult>> {
606 self.search_with_ef(query, k, self.config.ef_search)
607 }
608
609 fn remove(&mut self, id: &VectorId) -> Result<bool> {
610 let inner = self.inner.write();
611 let removed = inner.codes.remove(id).is_some();
614 if removed {
615 if let Some((_, idx)) = inner.id_to_idx.remove(id) {
616 inner.idx_to_id.remove(&idx);
617 }
618 }
619 Ok(removed)
620 }
621
622 fn len(&self) -> usize {
623 self.inner.read().codes.len()
624 }
625}
626
627#[cfg(test)]
628mod tests {
629 use super::*;
630
631 fn gauss_vecs(count: usize, dim: usize, seed: u64) -> Vec<Vec<f32>> {
633 let mut state = seed;
634 let mut next = move || {
635 state = state.wrapping_add(0x9E3779B97F4A7C15);
636 let mut z = state;
637 z = (z ^ (z >> 30)).wrapping_mul(0xBF58476D1CE4E5B9);
638 z = (z ^ (z >> 27)).wrapping_mul(0x94D049BB133111EB);
639 z ^ (z >> 31)
640 };
641 (0..count)
642 .map(|_| {
643 let mut v = Vec::with_capacity(dim);
644 while v.len() < dim {
645 let u1 = (next() >> 11) as f64 / (1u64 << 53) as f64;
646 let u2 = (next() >> 11) as f64 / (1u64 << 53) as f64;
647 let r = (-2.0 * u1.max(1e-12).ln()).sqrt();
648 let (s, c) = (2.0 * std::f64::consts::PI * u2).sin_cos();
649 v.push((r * c) as f32);
650 if v.len() < dim {
651 v.push((r * s) as f32);
652 }
653 }
654 v
655 })
656 .collect()
657 }
658
659 fn l2(a: &[f32], b: &[f32]) -> f32 {
660 a.iter()
661 .zip(b)
662 .map(|(x, y)| (x - y) * (x - y))
663 .sum::<f32>()
664 .sqrt()
665 }
666
667 #[test]
668 fn manhattan_is_rejected() {
669 let err = Turbo4HnswIndex::new(
670 64,
671 DistanceMetric::Manhattan,
672 HnswConfig::default(),
673 42,
674 4,
675 SearchPolicy::Balanced,
676 SearchQuantization::Turbo4Direct,
677 );
678 assert!(err.is_err());
679 }
680
681 fn clustered_vecs(
684 n: usize,
685 dim: usize,
686 cluster_size: usize,
687 spread: f32,
688 seed: u64,
689 ) -> Vec<Vec<f32>> {
690 let n_centers = n.div_ceil(cluster_size);
691 let centers = gauss_vecs(n_centers, dim, seed);
692 let noise = gauss_vecs(n, dim, seed ^ 0xABCD_EF01);
693 (0..n)
694 .map(|i| {
695 centers[i / cluster_size]
696 .iter()
697 .zip(&noise[i])
698 .map(|(c, e)| c + spread * e)
699 .collect()
700 })
701 .collect()
702 }
703
704 #[test]
713 fn recall_close_to_f32_hnsw_baseline() -> Result<()> {
714 use crate::index::hnsw::HnswIndex;
715
716 let dim = 128;
717 let n = 500;
718 let config = HnswConfig {
719 m: 16,
720 ef_construction: 200,
721 ef_search: 100,
722 max_elements: 1000,
723 };
724 let vectors = clustered_vecs(n, dim, 10, 0.35, 7);
725 let entries: Vec<_> = vectors
726 .iter()
727 .enumerate()
728 .map(|(i, v)| (format!("v{i}"), v.clone()))
729 .collect();
730
731 let mut t4 = Turbo4HnswIndex::new(
732 dim,
733 DistanceMetric::Euclidean,
734 config.clone(),
735 42,
736 8,
737 SearchPolicy::Balanced,
738 SearchQuantization::Turbo4Direct,
739 )?;
740 t4.add_batch(entries.clone())?;
741 assert_eq!(t4.len(), n);
742 let mut f32_ix = HnswIndex::new(dim, DistanceMetric::Euclidean, config)?;
743 f32_ix.add_batch(entries)?;
744
745 let qnoise = gauss_vecs(20, dim, 998877);
748 let queries: Vec<Vec<f32>> = (0..20)
749 .map(|j| {
750 vectors[(j * 25) % n]
751 .iter()
752 .zip(&qnoise[j])
753 .map(|(v, e)| v + 0.2 * e)
754 .collect()
755 })
756 .collect();
757 let (mut t4_hits, mut f32_hits, mut total) = (0usize, 0usize, 0usize);
758 for q in &queries {
759 let mut truth: Vec<(usize, f32)> = vectors
760 .iter()
761 .enumerate()
762 .map(|(i, v)| (i, l2(q, v)))
763 .collect();
764 truth.sort_by(|a, b| a.1.total_cmp(&b.1));
765 let top10: std::collections::HashSet<String> =
766 truth[..10].iter().map(|(i, _)| format!("v{i}")).collect();
767
768 let t4_res = t4.search(q, 10)?;
769 assert_eq!(t4_res.len(), 10);
770 t4_hits += t4_res.iter().filter(|r| top10.contains(&r.id)).count();
771 f32_hits += f32_ix
772 .search(q, 10)?
773 .iter()
774 .filter(|r| top10.contains(&r.id))
775 .count();
776 total += 10;
777 }
778 let t4_recall = t4_hits as f32 / total as f32;
779 let f32_recall = f32_hits as f32 / total as f32;
780 assert!(
783 t4_recall >= f32_recall - 0.03,
784 "Turbo4 recall@10 {t4_recall} vs f32 baseline {f32_recall}: loss above 3pp \
785 on clustered data ({t4_hits}/{f32_hits}/{total})"
786 );
787 assert!(
788 t4_recall >= 0.90,
789 "Turbo4 absolute recall@10 {t4_recall} below floor"
790 );
791 Ok(())
792 }
793
794 #[test]
799 fn gaussian_worst_case_recall_floor() -> Result<()> {
800 let dim = 128;
801 let n = 500;
802 let config = HnswConfig {
803 m: 16,
804 ef_construction: 200,
805 ef_search: 100,
806 max_elements: 1000,
807 };
808 let vectors = gauss_vecs(n, dim, 7);
809 let entries: Vec<_> = vectors
810 .iter()
811 .enumerate()
812 .map(|(i, v)| (format!("v{i}"), v.clone()))
813 .collect();
814 let mut t4 = Turbo4HnswIndex::new(
815 dim,
816 DistanceMetric::Euclidean,
817 config,
818 42,
819 8,
820 SearchPolicy::Balanced,
821 SearchQuantization::Turbo4Direct,
822 )?;
823 t4.add_batch(entries)?;
824
825 let queries = gauss_vecs(20, dim, 12345);
826 let (mut hits, mut total) = (0usize, 0usize);
827 for q in &queries {
828 let mut truth: Vec<(usize, f32)> = vectors
829 .iter()
830 .enumerate()
831 .map(|(i, v)| (i, l2(q, v)))
832 .collect();
833 truth.sort_by(|a, b| a.1.total_cmp(&b.1));
834 let top10: std::collections::HashSet<String> =
835 truth[..10].iter().map(|(i, _)| format!("v{i}")).collect();
836 hits += t4
837 .search(q, 10)?
838 .iter()
839 .filter(|r| top10.contains(&r.id))
840 .count();
841 total += 10;
842 }
843 let recall = hits as f32 / total as f32;
844 assert!(
845 recall >= 0.75,
846 "Turbo4 worst-case Gaussian recall@10 {recall} below floor ({hits}/{total})"
847 );
848 Ok(())
849 }
850
851 #[test]
852 fn self_query_returns_self_first() -> Result<()> {
853 let dim = 64;
854 let config = HnswConfig::default();
855 let mut index = Turbo4HnswIndex::new(
856 dim,
857 DistanceMetric::Cosine,
858 config,
859 42,
860 4,
861 SearchPolicy::Balanced,
862 SearchQuantization::Turbo4Direct,
863 )?;
864 let vectors = gauss_vecs(50, dim, 3);
865 for (i, v) in vectors.iter().enumerate() {
866 index.add(format!("v{i}"), v.clone())?;
867 }
868 let results = index.search(&vectors[7], 5)?;
869 assert_eq!(results[0].id, "v7");
870 Ok(())
871 }
872
873 #[test]
877 fn adaptive_escalation_follows_policy() -> Result<()> {
878 let dim = 128;
879 let n = 400;
880 let config = HnswConfig {
881 m: 16,
882 ef_construction: 200,
883 ef_search: 60,
884 max_elements: 1000,
885 };
886 let vectors = gauss_vecs(n, dim, 21);
888 let entries: Vec<_> = vectors
889 .iter()
890 .enumerate()
891 .map(|(i, v)| (format!("v{i}"), v.clone()))
892 .collect();
893
894 let mut fixed = Turbo4HnswIndex::new(
895 dim,
896 DistanceMetric::Euclidean,
897 config.clone(),
898 42,
899 2,
900 SearchPolicy::MaxCompression,
901 SearchQuantization::Turbo4Direct,
902 )?;
903 fixed.add_batch(entries.clone())?;
904 let mut adaptive = Turbo4HnswIndex::new(
905 dim,
906 DistanceMetric::Euclidean,
907 config,
908 42,
909 2,
910 SearchPolicy::Quality,
911 SearchQuantization::Turbo4Direct,
912 )?;
913 adaptive.add_batch(entries)?;
914
915 let queries = gauss_vecs(25, dim, 4242);
916 let (mut fixed_hits, mut adaptive_hits, mut total) = (0usize, 0usize, 0usize);
917 for q in &queries {
918 let mut truth: Vec<(usize, f32)> = vectors
919 .iter()
920 .enumerate()
921 .map(|(i, v)| (i, l2(q, v)))
922 .collect();
923 truth.sort_by(|a, b| a.1.total_cmp(&b.1));
924 let top10: std::collections::HashSet<String> =
925 truth[..10].iter().map(|(i, _)| format!("v{i}")).collect();
926 fixed_hits += fixed
927 .search(q, 10)?
928 .iter()
929 .filter(|r| top10.contains(&r.id))
930 .count();
931 adaptive_hits += adaptive
932 .search(q, 10)?
933 .iter()
934 .filter(|r| top10.contains(&r.id))
935 .count();
936 total += 10;
937 }
938
939 let (fq, fe) = fixed.adaptive_stats();
940 assert_eq!(fq, 25);
941 assert_eq!(fe, 0, "MaxCompression must never escalate");
942
943 let (aq, ae) = adaptive.adaptive_stats();
944 assert_eq!(aq, 25);
945 assert!(
946 ae > 0,
947 "Quality policy should escalate on concentrated Gaussian queries"
948 );
949 assert!(
954 adaptive_hits + 5 >= fixed_hits,
955 "escalation must not lose recall: adaptive {adaptive_hits} vs fixed {fixed_hits} of {total}"
956 );
957 Ok(())
958 }
959
960 #[test]
964 fn easy_queries_do_not_escalate() -> Result<()> {
965 let dim = 64;
966 let config = HnswConfig::default();
967 let mut index = Turbo4HnswIndex::new(
968 dim,
969 DistanceMetric::Euclidean,
970 config,
971 42,
972 4,
973 SearchPolicy::Balanced,
974 SearchQuantization::Turbo4Direct,
975 )?;
976 let vectors = clustered_vecs(200, dim, 5, 0.05, 33);
980 for (i, v) in vectors.iter().enumerate() {
981 index.add(format!("v{i}"), v.iter().map(|x| x * 10.0).collect())?;
982 }
983 for j in 0..10 {
984 let q: Vec<f32> = vectors[j * 20].iter().map(|x| x * 10.0).collect();
985 index.search(&q, 5)?;
986 }
987 let (queries, escalated) = index.adaptive_stats();
988 assert_eq!(queries, 10);
989 assert!(
990 escalated <= 2,
991 "well-separated queries escalated {escalated}/10 times"
992 );
993 Ok(())
994 }
995
996 #[test]
1000 fn serialization_roundtrip() -> Result<()> {
1001 let dim = 64;
1002 let mut index = Turbo4HnswIndex::new(
1003 dim,
1004 DistanceMetric::Euclidean,
1005 HnswConfig::default(),
1006 42,
1007 4,
1008 SearchPolicy::Balanced,
1009 SearchQuantization::Turbo4Direct,
1010 )?;
1011 let vectors = clustered_vecs(120, dim, 6, 0.2, 55);
1012 for (i, v) in vectors.iter().enumerate() {
1013 index.add(format!("v{i}"), v.clone())?;
1014 }
1015 let bytes = index.serialize()?;
1016 let restored = Turbo4HnswIndex::deserialize(&bytes)?;
1017 assert_eq!(restored.len(), 120);
1018 assert_eq!(restored.code_len(), index.code_len());
1019
1020 for probe in [3usize, 47, 99] {
1021 let orig = index.search(&vectors[probe], 3)?;
1022 let rest = restored.search(&vectors[probe], 3)?;
1023 assert_eq!(rest[0].id, format!("v{probe}"), "self-query after restore");
1024 assert!(
1025 (orig[0].score - rest[0].score).abs() < 1e-6,
1026 "rescored self distance must be identical: {} vs {}",
1027 orig[0].score,
1028 rest[0].score
1029 );
1030 }
1031 Ok(())
1032 }
1033
1034 #[test]
1039 fn cascade_recall_close_to_direct_mode() -> Result<()> {
1040 let dim = 128;
1041 let n = 500;
1042 let config = HnswConfig {
1043 m: 16,
1044 ef_construction: 200,
1045 ef_search: 100,
1046 max_elements: 1000,
1047 };
1048 let vectors = clustered_vecs(n, dim, 10, 0.35, 7);
1049 let entries: Vec<_> = vectors
1050 .iter()
1051 .enumerate()
1052 .map(|(i, v)| (format!("v{i}"), v.clone()))
1053 .collect();
1054
1055 let mut direct = Turbo4HnswIndex::new(
1056 dim,
1057 DistanceMetric::Euclidean,
1058 config.clone(),
1059 42,
1060 8,
1061 SearchPolicy::Balanced,
1062 SearchQuantization::Turbo4Direct,
1063 )?;
1064 direct.add_batch(entries.clone())?;
1065 let mut cascade = Turbo4HnswIndex::new(
1066 dim,
1067 DistanceMetric::Euclidean,
1068 config,
1069 42,
1070 8,
1071 SearchPolicy::Balanced,
1072 SearchQuantization::RaBitQ1,
1073 )?;
1074 cascade.add_batch(entries)?;
1075
1076 let qnoise = gauss_vecs(20, dim, 998877);
1077 let queries: Vec<Vec<f32>> = (0..20)
1078 .map(|j| {
1079 vectors[(j * 25) % n]
1080 .iter()
1081 .zip(&qnoise[j])
1082 .map(|(v, e)| v + 0.2 * e)
1083 .collect()
1084 })
1085 .collect();
1086
1087 let (mut d_hits, mut c_hits, mut total) = (0usize, 0usize, 0usize);
1088 for q in &queries {
1089 let mut truth: Vec<(usize, f32)> = vectors
1090 .iter()
1091 .enumerate()
1092 .map(|(i, v)| (i, l2(q, v)))
1093 .collect();
1094 truth.sort_by(|a, b| a.1.total_cmp(&b.1));
1095 let top10: std::collections::HashSet<String> =
1096 truth[..10].iter().map(|(i, _)| format!("v{i}")).collect();
1097 d_hits += direct
1098 .search(q, 10)?
1099 .iter()
1100 .filter(|r| top10.contains(&r.id))
1101 .count();
1102 c_hits += cascade
1103 .search(q, 10)?
1104 .iter()
1105 .filter(|r| top10.contains(&r.id))
1106 .count();
1107 total += 10;
1108 }
1109 let (d_recall, c_recall) = (d_hits as f32 / total as f32, c_hits as f32 / total as f32);
1110 assert!(
1111 c_recall >= d_recall - 0.05,
1112 "cascade recall@10 {c_recall} vs direct {d_recall} ({c_hits}/{d_hits}/{total})"
1113 );
1114 assert!(c_recall >= 0.85, "cascade absolute recall {c_recall}");
1115 Ok(())
1116 }
1117
1118 #[test]
1121 fn cascade_serialization_roundtrip() -> Result<()> {
1122 let dim = 64;
1123 let mut index = Turbo4HnswIndex::new(
1124 dim,
1125 DistanceMetric::Cosine,
1126 HnswConfig::default(),
1127 42,
1128 4,
1129 SearchPolicy::Balanced,
1130 SearchQuantization::RaBitQ1,
1131 )?;
1132 let vectors = clustered_vecs(100, dim, 5, 0.2, 77);
1133 for (i, v) in vectors.iter().enumerate() {
1134 index.add(format!("v{i}"), v.clone())?;
1135 }
1136 let restored = Turbo4HnswIndex::deserialize(&index.serialize()?)?;
1137 assert_eq!(restored.len(), 100);
1138 for probe in [2usize, 55, 98] {
1139 let orig = index.search(&vectors[probe], 3)?;
1140 let rest = restored.search(&vectors[probe], 3)?;
1141 assert_eq!(rest[0].id, format!("v{probe}"));
1142 assert!((orig[0].score - rest[0].score).abs() < 1e-6);
1143 }
1144 Ok(())
1145 }
1146
1147 #[test]
1148 fn empty_index_is_safe() -> Result<()> {
1149 let index = Turbo4HnswIndex::new(
1150 64,
1151 DistanceMetric::Euclidean,
1152 HnswConfig::default(),
1153 42,
1154 4,
1155 SearchPolicy::Balanced,
1156 SearchQuantization::Turbo4Direct,
1157 )?;
1158 assert!(index.search(&vec![0.5; 64], 5)?.is_empty());
1159 Ok(())
1160 }
1161}