1mod adjacency;
25mod build_kernels;
26mod fastscan;
27pub mod filter;
28pub(crate) mod kernels;
29mod projection;
30mod rerank;
31#[cfg(not(target_arch = "wasm32"))]
32mod update;
33mod vertex;
34mod vnni;
35#[cfg(not(target_arch = "wasm32"))]
36pub use update::{
37 Consolidation, ConsolidationMode, ConsolidationPolicy, UpdatableVGraph, UpdateError,
38};
39
40#[cfg(feature = "experimental")]
41pub mod diagnostics;
42
43#[cfg(feature = "experimental")]
44pub mod projection_diagnostic;
45
46#[cfg(feature = "experimental")]
47pub mod order_diagnostic;
48
49#[cfg(feature = "experimental")]
50pub mod ivf;
51
52use std::cell::RefCell;
53
54use rayon::prelude::*;
55
56use crate::core::error::WarmError;
57use adjacency::Adjacency;
58
59use kernels::{
60 as_f32, as_f32_mut, as_u32, as_u32_mut, dot, dot_i8_query, dot_i8_query_x4, QueryCode,
61 Rotation, SplitMix,
62};
63#[cfg(feature = "experimental")]
64pub use vertex::QueryProfile;
65pub use vertex::{
66 huge_page_status, HugePageStatus, Residency, SearchKnobs, VGraph, VGraphMemoryBreakdown,
67 VMetric, VSearchParams, VSearcher,
68};
69
70pub const DEFAULT_DEGREE: usize = 32;
72const NONE: u32 = u32::MAX;
73
74fn invalid(details: String) -> WarmError {
76 WarmError::InvalidParam { details }
77}
78const UPPER_DEGREE: usize = 16;
79const LAYER_RATIO: usize = 16;
80const MIN_LAYER: usize = 64;
81const BUILD_CODES_MIN_DIM: usize = 128;
85const FIRST_PASS_EF: usize = 200;
91
92#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
94pub enum Metric {
95 #[default]
97 Angular,
98 Euclidean,
103}
104
105impl Metric {
106 fn code(self) -> u64 {
107 match self {
108 Self::Angular => 0,
109 Self::Euclidean => 1,
110 }
111 }
112}
113
114#[derive(Clone, Debug)]
115pub struct BuildParams {
116 pub metric: Metric,
118 pub build_ef: usize,
121 pub degree: usize,
123 pub alpha: f32,
125 pub code_bits: usize,
128 pub seed: u64,
129 pub quantized_search: bool,
134}
135
136impl Default for BuildParams {
137 fn default() -> Self {
138 Self {
139 metric: Metric::Angular,
140 build_ef: 128,
141 degree: DEFAULT_DEGREE,
142 alpha: 1.0,
143 code_bits: 0,
144 seed: 0x5EED,
145 quantized_search: false,
146 }
147 }
148}
149
150#[derive(Clone, Copy, Debug)]
151pub struct SearchParams {
152 pub ef: usize,
154 pub max_exact: usize,
156}
157
158pub struct QGraph {
159 n: usize,
160 dim: usize,
161 degree: usize,
162 words: usize,
163 rotation: Rotation,
164 stride: usize,
165 codes_off: usize,
166 factors_off: usize,
167 ids_off: usize,
168 blocks: Vec<u64>,
169 entry: u32,
170 upper_ids: Vec<u32>,
172 layers: Vec<Vec<Vec<u32>>>,
174}
175
176#[derive(Clone, Copy)]
177struct Candidate {
178 dist: f32,
179 id: u32,
180 done: bool,
181}
182
183#[cfg(feature = "experimental-build-search")]
184#[inline(always)]
185fn insert_build_candidate(
186 pool: &mut Vec<Candidate>,
187 candidate: Candidate,
188 limit: usize,
189) -> Option<usize> {
190 debug_assert!(limit > 0 && pool.len() <= limit);
191 if pool.len() == limit && candidate.dist >= pool[limit - 1].dist {
192 return None;
193 }
194 let pos = pool.partition_point(|c| c.dist <= candidate.dist);
195 if pool.len() == limit {
196 let _ = pool.pop();
198 }
199 pool.insert(pos, candidate);
200 Some(pos)
201}
202
203struct Visited {
206 marks: Vec<u8>,
207 epoch: u8,
208}
209
210impl Visited {
211 fn new(n: usize) -> Self {
212 Self {
213 marks: vec![0; n],
214 epoch: 0,
215 }
216 }
217
218 fn next(&mut self) {
219 self.epoch = self.epoch.wrapping_add(1);
220 if self.epoch == 0 {
221 self.marks.fill(0);
222 self.epoch = 1;
223 }
224 }
225
226 fn insert(&mut self, id: u32) -> bool {
228 let m = &mut self.marks[id as usize];
229 let fresh = *m != self.epoch;
230 *m = self.epoch;
231 fresh
232 }
233}
234
235#[derive(Default)]
238struct BuildVisited {
239 bits: Vec<u64>,
240 touched: Vec<u32>,
241 pool: Vec<Candidate>,
242 fresh: Vec<u32>,
243 dists: Vec<f32>,
244}
245
246impl BuildVisited {
247 fn reset(&mut self, n: usize) {
248 for word in self.touched.drain(..) {
250 self.bits[word as usize] = 0;
251 }
252 self.bits.resize(n.div_ceil(64), 0);
253 }
254
255 fn insert(&mut self, id: u32) -> bool {
257 let (w, b) = (id as usize / 64, 1u64 << (id % 64));
258 let old = self.bits[w];
259 if old & b != 0 {
260 return false;
261 }
262 if old == 0 {
263 self.touched.push(w as u32);
264 }
265 self.bits[w] = old | b;
266 true
267 }
268}
269
270#[cfg(feature = "experimental")]
275#[doc(hidden)]
276pub fn set_build_phase_hook(hook: fn(&str)) -> bool {
277 prof::HOOK.set(hook).is_ok()
278}
279
280mod prof {
283 use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
284 use std::sync::OnceLock;
285 use std::time::Instant;
286
287 pub(super) static GREEDY: AtomicU64 = AtomicU64::new(0);
288 pub(super) static PRUNE: AtomicU64 = AtomicU64::new(0);
289 pub(super) static REV: AtomicU64 = AtomicU64::new(0);
290 pub(super) static FRESH: AtomicU64 = AtomicU64::new(0);
291 pub(super) static EXACT: AtomicU64 = AtomicU64::new(0);
292 pub(super) static HOOK: OnceLock<fn(&str)> = OnceLock::new();
294
295 pub(super) fn on() -> bool {
296 static ON: OnceLock<bool> = OnceLock::new();
297 *ON.get_or_init(|| {
298 HOOK.get().is_some()
299 || crate::compatibility::var_os("WARM_QGRAPH_BUILD_PROFILE").is_some()
300 })
301 }
302
303 pub(super) fn start() -> Option<Instant> {
304 on().then(Instant::now)
305 }
306
307 pub(super) fn add(c: &AtomicU64, t: Option<Instant>) {
308 if let Some(t) = t {
309 c.fetch_add(t.elapsed().as_nanos() as u64, Relaxed);
310 }
311 }
312
313 pub(super) fn count(c: &AtomicU64, n: u64) {
314 if on() {
315 c.fetch_add(n, Relaxed);
316 }
317 }
318
319 pub(super) fn report(phase: &str, t: Option<Instant>) {
320 let Some(t) = t else { return };
321 if let Some(hook) = HOOK.get() {
322 hook(phase);
323 }
324 let secs = |c: &AtomicU64| c.swap(0, Relaxed) as f64 * 1e-9;
325 eprintln!(
326 "PHASE {phase} wall={:.2}s greedy_cpu={:.1}s prune_cpu={:.1}s rev_wall={:.2}s fresh={} exact={}",
327 t.elapsed().as_secs_f64(),
328 secs(&GREEDY),
329 secs(&PRUNE),
330 secs(&REV),
331 FRESH.swap(0, Relaxed),
332 EXACT.swap(0, Relaxed),
333 );
334 }
335
336 pub(super) fn graph_hash(h: impl std::hash::Hash) {
337 if on() {
338 use std::hash::Hasher;
339 let mut s = std::collections::hash_map::DefaultHasher::new();
340 h.hash(&mut s);
341 eprintln!("GRAPH_HASH {:016x}", s.finish());
342 }
343 }
344}
345#[derive(Default)]
347struct InsertScratch {
348 cands: Vec<(f32, u32)>,
349 d: Vec<f32>,
350}
351
352#[derive(Default)]
354struct MergeScratch {
355 d: Vec<f32>,
356 cands: Vec<(f32, u32, bool)>,
357 open: Vec<u32>,
358 out: Vec<u32>,
359}
360
361#[derive(Default)]
363struct FillScratch {
364 ids: Vec<u32>,
365 skip: Vec<u32>,
366 bounds: Vec<(f32, f32)>,
367 upper: Vec<f32>,
368 d: Vec<f32>,
369 near: Vec<(f32, u32)>,
370}
371
372thread_local! {
373 static BUILD_VISITED: RefCell<BuildVisited> = RefCell::new(BuildVisited::default());
374 static INSERT_SCRATCH: RefCell<InsertScratch> = RefCell::default();
375 static FILL_SCRATCH: RefCell<FillScratch> = RefCell::default();
376}
377
378struct Rows<'a> {
379 data: &'a [f32],
380 dim: usize,
381 norms: Option<&'a [f32]>,
383 codes: Option<&'a BuildCodes>,
386 quantized: bool,
387}
388
389struct BuildCodes {
394 width: usize,
395 q: Vec<i8>,
396 meta: Vec<CodeMeta>,
398 sums: Vec<i32>,
401 euclidean: bool,
403 rel_slack: f32,
405}
406
407#[derive(Clone, Copy, Default)]
409#[repr(C, align(16))]
410struct CodeMeta {
411 scale: f32,
412 err: f32,
413 n2: f32,
414 n: f32,
415}
416
417impl BuildCodes {
418 fn new(data: &[f32], dim: usize, norms: Option<&[f32]>) -> Self {
419 let n = data.len() / dim;
420 let width = dim.next_multiple_of(16);
421 let mut q = vec![0i8; n * width];
422 let mut meta = vec![CodeMeta::default(); n];
423 q.par_chunks_mut(width)
424 .zip(meta.par_iter_mut())
425 .zip(data.par_chunks(dim))
426 .for_each(|((q, m), x)| {
427 let (s, e) = (&mut m.scale, &mut m.err);
428 let max = x.iter().fold(0f32, |m, v| m.max(v.abs()));
429 *s = if max > 0.0 { max / 127.0 } else { 0.0 };
430 let mut e2 = 0f64;
431 for (o, &v) in q.iter_mut().zip(x) {
432 let c = if *s > 0.0 {
433 (v / *s).round().clamp(-127.0, 127.0)
434 } else {
435 0.0
436 };
437 *o = c as i8;
438 let r = f64::from(v) - f64::from(*s) * f64::from(c);
439 e2 += r * r;
440 }
441 *e = (e2.sqrt() * (1.0 + 1e-6)) as f32 + 1e-7;
443 });
444 if let Some(norms) = norms {
445 for (m, &n2) in meta.iter_mut().zip(norms) {
446 (m.n2, m.n) = (n2, n2.sqrt());
447 }
448 }
449 let sums = q.par_chunks(width).map(vnni::query_sum).collect();
450 Self {
451 width,
452 q,
453 meta,
454 sums,
455 euclidean: norms.is_some(),
456 rel_slack: 1e-4f32.max((dim as f32 + 4.0) * 2.0 * f32::EPSILON),
459 }
460 }
461
462 fn code(&self, i: u32) -> &[i8] {
463 let i = i as usize;
464 &self.q[i * self.width..(i + 1) * self.width]
465 }
466
467 #[inline]
469 fn prefetch(&self, i: u32) {
470 build_kernels::prefetch_code(self.code(i));
471 build_kernels::prefetch_meta(&self.meta[i as usize]);
472 }
473
474 #[inline]
479 fn d2_lower(&self, p: u32, u: u32) -> f32 {
480 self.d2_bounds(p, u).0
481 }
482
483 #[inline]
485 fn raw(&self, p: u32, u: u32) -> i32 {
486 dot_i8_query(self.code(p), self.code(u), self.sums[p as usize])
487 }
488
489 #[inline]
491 fn d2_est(&self, p: u32, u: u32) -> f32 {
492 self.est_from_raw(p, u, self.raw(p, u))
493 }
494
495 #[inline]
498 fn d2_est4(&self, p: u32, u: [u32; 4]) -> [f32; 4] {
499 let raw = dot_i8_query_x4(self.code(p), u.map(|u| self.code(u)), self.sums[p as usize]);
500 [0, 1, 2, 3].map(|i| self.est_from_raw(p, u[i], raw[i]))
501 }
502
503 #[inline]
504 fn est_from_raw(&self, p: u32, u: u32, raw: i32) -> f32 {
505 let raw = raw as f32;
506 let (mp, mu) = (&self.meta[p as usize], &self.meta[u as usize]);
507 if self.euclidean {
508 (mp.n2 + mu.n2 - 2.0 * mp.scale * mu.scale * raw).max(0.0)
509 } else {
510 (2.0 - 2.0 * mp.scale * mu.scale * raw).max(0.0)
511 }
512 }
513
514 #[inline]
521 fn d2_bounds(&self, p: u32, u: u32) -> (f32, f32) {
522 const SLACK: f32 = 1e-4;
523 let (mp, mu) = (&self.meta[p as usize], &self.meta[u as usize]);
524 let raw = self.raw(p, u) as f32;
525 let ep = mp.err;
526 let est = mp.scale * mu.scale * raw;
527 if !self.euclidean {
528 let margin = ep + (1.0 + ep) * mu.err + SLACK;
529 return (2.0 - 2.0 * (est + margin), 2.0 - 2.0 * (est - margin));
530 }
531 let ((np2, np), (nu2, nu)) = ((mp.n2, mp.n), (mu.n2, mu.n));
532 let margin = ep * nu + (np + ep) * mu.err;
533 let slack = self.rel_slack * (np2 + nu2 + 2.0 * np * nu);
534 let base = np2 + nu2;
535 (
536 base - 2.0 * (est + margin) - slack,
537 base - 2.0 * (est - margin) + slack,
538 )
539 }
540}
541
542impl Rows<'_> {
543 fn row(&self, i: u32) -> &[f32] {
544 let i = i as usize;
545 &self.data[i * self.dim..(i + 1) * self.dim]
546 }
547
548 #[inline]
550 fn prefetch(&self, u: u32) {
551 build_kernels::prefetch_row(self.row(u));
552 if let Some(n) = self.norms {
553 build_kernels::prefetch_meta(&n[u as usize]);
554 }
555 }
556
557 #[inline]
560 fn dist(&self, a: u32, b: u32, dot: f32) -> f32 {
561 match self.norms {
562 None => (2.0 - 2.0 * dot).max(0.0),
563 Some(n) => (n[a as usize] + n[b as usize] - 2.0 * dot).max(0.0),
564 }
565 }
566
567 fn d2(&self, a: u32, b: u32) -> f32 {
569 self.dist(a, b, dot(self.row(a), self.row(b)))
570 }
571
572 fn d2_many(&self, p: u32, ids: &[u32], out: &mut Vec<f32>) {
574 out.clear();
575 let target = self.row(p);
576 let (quads, rest) = ids.as_chunks::<4>();
577 for q in quads {
578 let d = build_kernels::dot4(target, q.map(|u| self.row(u)));
579 out.extend(d.iter().zip(q).map(|(&x, &u)| self.dist(p, u, x)));
580 }
581 out.extend(rest.iter().map(|&u| self.d2(p, u)));
582 }
583
584 fn d2_stream(&self, p: u32, ids: &[u32], out: &mut Vec<f32>) {
586 const AHEAD: usize = 16;
587 out.clear();
588 let target = self.row(p);
589 for &u in ids.iter().take(AHEAD) {
590 self.prefetch(u);
591 }
592 let (quads, rest) = ids.as_chunks::<4>();
593 for (i, q) in quads.iter().enumerate() {
594 for &u in ids.iter().skip(AHEAD + 4 * i).take(4) {
595 self.prefetch(u);
596 }
597 let d = build_kernels::dot4(target, q.map(|u| self.row(u)));
598 out.extend(d.iter().zip(q).map(|(&x, &u)| self.dist(p, u, x)));
599 }
600 out.extend(rest.iter().map(|&u| self.d2(p, u)));
601 }
602
603 fn occluded(&self, c: u32, kept: &[u32], dc: f32, a2: f32) -> bool {
605 let rc = self.row(c);
606 if let Some(codes) = self.codes {
607 let mut open = [0u32; 4];
610 let mut n_open = 0;
611 for &s in kept {
612 let (lo, hi) = codes.d2_bounds(c, s);
613 if a2 * hi <= dc {
614 return true;
615 }
616 if a2 * lo > dc {
617 continue;
618 }
619 open[n_open] = s;
620 n_open += 1;
621 if n_open == 4 {
622 n_open = 0;
623 let d = build_kernels::dot4(rc, open.map(|u| self.row(u)));
624 if d.iter()
625 .zip(&open)
626 .any(|(&x, &s)| a2 * self.dist(c, s, x) <= dc)
627 {
628 return true;
629 }
630 }
631 }
632 return open[..n_open]
633 .iter()
634 .any(|&s| a2 * self.dist(s, c, dot(self.row(s), rc)) <= dc);
635 }
636 let (quads, rest) = kept.as_chunks::<4>();
637 for q in quads {
638 let d = build_kernels::dot4(rc, q.map(|u| self.row(u)));
639 if d.iter()
640 .zip(q)
641 .any(|(&x, &s)| a2 * self.dist(c, s, x) <= dc)
642 {
643 return true;
644 }
645 }
646 rest.iter()
647 .any(|&s| a2 * self.dist(s, c, dot(self.row(s), rc)) <= dc)
648 }
649
650 fn greedy(
652 &self,
653 graph: &Adjacency,
654 start: u32,
655 p: u32,
656 l: usize,
657 visited: &mut BuildVisited,
658 ) -> Vec<(f32, u32)> {
659 let mut expanded = Vec::new();
660 self.greedy_into(graph, start, p, l, visited, &mut expanded);
661 expanded
662 }
663
664 fn greedy_into(
666 &self,
667 graph: &Adjacency,
668 start: u32,
669 p: u32,
670 l: usize,
671 visited: &mut BuildVisited,
672 expanded: &mut Vec<(f32, u32)>,
673 ) {
674 let (mut n_fresh, mut n_exact) = (0u64, 0u64);
675 visited.reset(graph.len());
676 visited.insert(start);
677 let mut pool = std::mem::take(&mut visited.pool);
679 let mut fresh = std::mem::take(&mut visited.fresh);
680 let mut dists = std::mem::take(&mut visited.dists);
681 pool.clear();
682 pool.push(Candidate {
683 dist: self.d2(p, start),
684 id: start,
685 done: false,
686 });
687 expanded.clear();
688 let mut cur = 0;
689 while cur < pool.len() {
690 pool[cur].done = true;
691 let v = pool[cur].id;
692 expanded.push((pool[cur].dist, v));
693 fresh.clear();
695 match self.codes.filter(|_| self.quantized || pool.len() >= l) {
698 Some(codes) if self.quantized => {
699 for &u in graph.get(v as usize) {
700 if visited.insert(u) {
701 codes.prefetch(u);
702 fresh.push(u);
703 }
704 }
705 n_fresh += fresh.len() as u64;
706 }
707 Some(codes) => {
708 for &u in graph.get(v as usize) {
709 if visited.insert(u) {
710 codes.prefetch(u);
711 fresh.push(u);
712 }
713 }
714 let bound = pool[l - 1].dist;
715 n_fresh += fresh.len() as u64;
716 fresh.retain(|&u| codes.d2_lower(p, u) < bound);
717 for &u in &fresh {
718 build_kernels::prefetch_row(self.row(u));
719 }
720 }
721 None => {
722 for &u in graph.get(v as usize) {
723 if visited.insert(u) {
724 build_kernels::prefetch_row(self.row(u));
725 fresh.push(u);
726 }
727 }
728 n_fresh += fresh.len() as u64;
729 }
730 }
731 match self.codes.filter(|_| self.quantized) {
732 Some(codes) => {
733 dists.clear();
734 let (quads, rest) = fresh.as_chunks::<4>();
735 for &q in quads {
736 dists.extend(codes.d2_est4(p, q));
737 }
738 dists.extend(rest.iter().map(|&u| codes.d2_est(p, u)));
739 }
740 None => {
741 n_exact += fresh.len() as u64;
742 self.d2_many(p, &fresh, &mut dists);
743 }
744 }
745 let mut best = usize::MAX;
746 for (&u, &du) in fresh.iter().zip(&dists) {
747 #[cfg(feature = "experimental-build-search")]
748 let Some(pos) = insert_build_candidate(
749 &mut pool,
750 Candidate {
751 dist: du,
752 id: u,
753 done: false,
754 },
755 l,
756 ) else {
757 continue;
758 };
759 #[cfg(not(feature = "experimental-build-search"))]
760 let pos = {
761 if pool.len() >= l && du >= pool[l - 1].dist {
762 continue;
763 }
764 let pos = pool.partition_point(|c| c.dist <= du);
765 pool.insert(
766 pos,
767 Candidate {
768 dist: du,
769 id: u,
770 done: false,
771 },
772 );
773 pool.truncate(l);
774 pos
775 };
776 best = best.min(pos);
777 build_kernels::prefetch_meta(graph.len_word(u as usize));
779 }
780 cur = best.min(cur + 1);
781 while cur < pool.len() && pool[cur].done {
782 cur += 1;
783 }
784 if let Some(c) = pool.get(cur) {
785 build_kernels::prefetch_ids(graph.get(c.id as usize));
786 }
787 }
788 if self.quantized {
789 fresh.clear();
791 fresh.extend(expanded.iter().map(|e| e.1));
792 self.d2_stream(p, &fresh, &mut dists);
793 n_exact += fresh.len() as u64;
794 for (e, &d) in expanded.iter_mut().zip(&dists) {
795 e.0 = d;
796 }
797 }
798 (visited.pool, visited.fresh, visited.dists) = (pool, fresh, dists);
799 prof::count(&prof::FRESH, n_fresh);
800 prof::count(&prof::EXACT, n_exact);
801 }
802
803 fn prune(&self, p: u32, mut cands: Vec<(f32, u32)>, alpha: f32, r: usize) -> Vec<u32> {
805 let mut out = Vec::with_capacity(r);
806 self.prune_into(p, &mut cands, alpha, r, &mut out);
807 out
808 }
809
810 fn prune_into(
812 &self,
813 p: u32,
814 cands: &mut Vec<(f32, u32)>,
815 alpha: f32,
816 r: usize,
817 out: &mut Vec<u32>,
818 ) {
819 cands.retain(|c| c.1 != p);
820 cands.sort_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1)));
821 cands.dedup_by_key(|c| c.1);
822 let a2 = alpha * alpha;
823 out.clear();
824 for &(dc, c) in cands.iter() {
825 if out.len() == r {
826 break;
827 }
828 if !self.occluded(c, out, dc, a2) {
829 out.push(c);
830 }
831 }
832 }
833
834 #[cfg(test)]
839 fn prune_merged(&self, t: u32, list: &[u32], closed: usize, alpha: f32, r: usize) -> Vec<u32> {
840 let mut scratch = MergeScratch::default();
841 self.prune_merged_into(t, list, closed, alpha, r, &mut scratch);
842 scratch.out
843 }
844
845 fn prune_merged_into(
847 &self,
848 t: u32,
849 list: &[u32],
850 closed: usize,
851 alpha: f32,
852 r: usize,
853 scratch: &mut MergeScratch,
854 ) {
855 let MergeScratch {
856 d,
857 cands,
858 open,
859 out,
860 } = scratch;
861 self.d2_stream(t, list, d);
862 cands.clear();
863 cands.extend(
864 d.iter()
865 .zip(list)
866 .enumerate()
867 .map(|(i, (&d, &u))| (d, u, i < closed)),
868 );
869 cands.sort_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1)));
870 let a2 = alpha * alpha;
871 out.clear();
872 open.clear();
873 for &(dc, c, was_kept) in cands.iter() {
874 if out.len() == r {
875 break;
876 }
877 if !self.occluded(c, if was_kept { open } else { out }, dc, a2) {
878 out.push(c);
879 if !was_kept {
880 open.push(c);
881 }
882 }
883 }
884 }
885
886 #[allow(clippy::too_many_arguments)]
889 fn insert_pass(
890 &self,
891 graph: &mut Adjacency,
892 order: &[u32],
893 entry: u32,
894 l: usize,
895 alpha: f32,
896 r: usize,
897 ) {
898 let n = graph.len();
899 let mut closed = vec![0u32; n];
901 let max_batch = (n / 50).max(64);
902 let mut outs = Adjacency::new(max_batch.min(order.len()), r);
903 let (mut start, mut size) = (0, 1);
904 while start < order.len() {
905 let end = (start + size).min(order.len());
906 let batch = &order[start..end];
907 let g: &Adjacency = graph;
908 outs.par_set(batch.len(), |i, out| {
909 let p = batch[i];
910 INSERT_SCRATCH.with(|cell| {
911 let InsertScratch { cands, d } = &mut *cell.borrow_mut();
912 let t0 = prof::start();
913 BUILD_VISITED.with(|vis| {
914 self.greedy_into(g, entry, p, l, &mut vis.borrow_mut(), cands);
915 });
916 prof::add(&prof::GREEDY, t0);
917 let t0 = prof::start();
918 let adj = g.get(p as usize);
919 self.d2_stream(p, adj, d);
920 cands.extend(d.iter().copied().zip(adj.iter().copied()));
921 self.prune_into(p, cands, alpha, r, out);
922 prof::add(&prof::PRUNE, t0);
923 });
924 });
925 let trev = prof::start();
926 let mut rev: Vec<(u32, u32)> = batch
927 .iter()
928 .enumerate()
929 .flat_map(|(i, &p)| outs.get(i).iter().map(move |&u| (u, p)))
930 .collect();
931 for (i, &p) in batch.iter().enumerate() {
932 let out = outs.get(i);
933 closed[p as usize] = out.len() as u32;
934 graph.set(p as usize, out);
935 }
936 rev.par_sort_unstable();
937 let groups: Vec<_> = rev.chunk_by(|a, b| a.0 == b.0).collect();
939 let targets = || groups.iter().map(|grp| grp[0].0 as usize);
940 let closed_of = pick_mut(&mut closed, targets());
941 let jobs: Vec<_> = graph
942 .rows_mut(targets())
943 .zip(closed_of)
944 .zip(&groups)
945 .collect();
946 jobs.into_par_iter().for_each_init(
947 <(Vec<u32>, MergeScratch)>::default,
948 |(merged, scratch), (((row, len), closed), grp)| {
949 let t = grp[0].0;
950 merged.clear();
951 merged.extend_from_slice(&row[..*len as usize]);
952 for &(_, s) in grp.iter() {
953 if !merged.contains(&s) {
954 merged.push(s);
955 }
956 }
957 let list = if merged.len() > r {
958 self.prune_merged_into(t, merged, *closed as usize, alpha, r, scratch);
959 *closed = scratch.out.len() as u32;
960 &scratch.out
961 } else {
962 &*merged
963 };
964 row[..list.len()].copy_from_slice(list);
965 *len = list.len() as u32;
966 },
967 );
968 prof::add(&prof::REV, trev);
969 start = end;
970 size = (size * 2).min(max_batch);
971 }
972 }
973
974 fn fill(&self, graph: &mut Adjacency, r: usize) {
977 graph.par_append(|g, v, out| {
978 let have = g.get(v);
979 if have.len() >= r {
980 return;
981 }
982 FILL_SCRATCH.with(|cell| self.fill_one(g, v, r, &mut cell.borrow_mut(), out));
983 });
984 }
985
986 fn fill_one(&self, g: &Adjacency, v: usize, r: usize, s: &mut FillScratch, out: &mut Vec<u32>) {
988 let FillScratch {
989 ids,
990 skip,
991 bounds,
992 upper,
993 d,
994 near,
995 } = s;
996 let have = g.get(v);
997 ids.clear();
998 ids.extend(have.iter().flat_map(|&u| g.get(u as usize).iter().copied()));
999 ids.sort_unstable();
1000 ids.dedup();
1001 skip.clear();
1003 skip.extend_from_slice(have);
1004 skip.push(v as u32);
1005 skip.sort_unstable();
1006 let mut at = 0;
1007 ids.retain(|&u| {
1008 while at < skip.len() && skip[at] < u {
1009 at += 1;
1010 }
1011 skip.get(at) != Some(&u)
1012 });
1013 let k = r - have.len();
1014 if let Some(codes) = self.codes.filter(|_| ids.len() > k) {
1015 const AHEAD: usize = 8;
1018 ids.iter().take(AHEAD).for_each(|&u| codes.prefetch(u));
1019 bounds.clear();
1020 bounds.extend(ids.iter().enumerate().map(|(i, &u)| {
1021 if let Some(&w) = ids.get(i + AHEAD) {
1022 codes.prefetch(w);
1023 }
1024 codes.d2_bounds(v as u32, u)
1025 }));
1026 upper.clear();
1027 upper.extend(bounds.iter().map(|b| b.1));
1028 let cut = *upper.select_nth_unstable_by(k - 1, f32::total_cmp).1;
1029 let mut b = bounds.iter();
1030 ids.retain(|_| b.next().is_some_and(|b| b.0 <= cut));
1031 }
1032 self.d2_stream(v as u32, ids, d);
1033 near.clear();
1034 near.extend(d.iter().copied().zip(ids.iter().copied()));
1035 let order = |a: &(f32, u32), b: &(f32, u32)| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1));
1038 if near.len() > k {
1039 near.select_nth_unstable_by(k - 1, order);
1040 near.truncate(k);
1041 }
1042 near.sort_by(order);
1043 out.extend(near.iter().map(|x| x.1));
1044 }
1045}
1046
1047fn pick_mut<T>(values: &mut [T], ids: impl Iterator<Item = usize>) -> impl Iterator<Item = &mut T> {
1049 let mut rest = values;
1050 let mut next = 0;
1051 ids.map(move |v| {
1052 assert!(v >= next, "ids must be strictly increasing");
1053 let (x, tail) = std::mem::take(&mut rest)[v - next..]
1054 .split_first_mut()
1055 .expect("id in range");
1056 rest = tail;
1057 next = v + 1;
1058 x
1059 })
1060}
1061
1062fn edge_factors(pv: &[f32], pu: &[f32], code: &mut [u64]) -> (f32, f32, f32) {
1065 let (mut r2, mut l1, mut sv, mut vw) = (0f32, 0f32, 0f32, 0f32);
1066 for ((a64, b64), word) in pv
1068 .as_chunks::<64>()
1069 .0
1070 .iter()
1071 .zip(pu.as_chunks::<64>().0)
1072 .zip(code.iter_mut())
1073 {
1074 let mut bits = 0u64;
1075 for (i, (&a, &b)) in a64.iter().zip(b64).enumerate() {
1076 let w = b - a;
1077 r2 += w * w;
1078 l1 += w.abs();
1079 vw += a * w;
1080 let up = w > 0.0;
1081 sv += if up { a } else { -a };
1082 bits |= u64::from(up) << i;
1083 }
1084 *word = bits;
1085 }
1086 let m = if l1 > 0.0 { r2 / l1 } else { 0.0 };
1087 let pop = code.iter().map(|c| c.count_ones()).sum::<u32>() as f32;
1088 (vw - m * sv, m, pop)
1089}
1090
1091impl QGraph {
1092 pub fn build(data: &[f32], dim: usize, params: &BuildParams) -> Result<Self, WarmError> {
1095 Self::from_graph(&Graph::build(data, dim, params)?, params)
1096 }
1097
1098 pub fn from_graph(g: &Graph, params: &BuildParams) -> Result<Self, WarmError> {
1100 let (n, dim, r) = (g.n, g.dim, g.degree);
1101 if r != params.degree {
1102 return Err(invalid(format!(
1103 "graph built with degree {r}, parameters say {}",
1104 params.degree
1105 )));
1106 }
1107 if g.metric != Metric::Angular {
1108 return Err(invalid(
1109 "the edge index is angular only; use VGraph for Euclidean data".into(),
1110 ));
1111 }
1112 let rotation = Rotation::new(dim, params.code_bits, params.seed);
1113 let words = rotation.padded() / 64;
1114 let rows = g.rows();
1115 let graph = &g.adj;
1116 let codes_off = dim.div_ceil(2);
1117 let factors_off = codes_off + r * words;
1118 let ids_off = factors_off + 3 * r / 2;
1119 let stride = ids_off + r / 2;
1120 let padded = rotation.padded();
1122 let mut rotated = vec![0f32; n * padded];
1123 rotated
1124 .par_chunks_mut(padded)
1125 .enumerate()
1126 .for_each(|(v, out)| rotation.apply(rows.row(v as u32), out));
1127 let rot_row = |v: u32| &rotated[v as usize * padded..(v as usize + 1) * padded];
1128 let mut blocks = vec![0u64; n * stride];
1129 blocks
1130 .par_chunks_mut(stride)
1131 .enumerate()
1132 .for_each(|(v, b)| {
1133 let adj = graph.get(v);
1134 let (vec_w, rest) = b.split_at_mut(codes_off);
1135 let (codes, rest) = rest.split_at_mut(r * words);
1136 let (fac, ids) = rest.split_at_mut(3 * r / 2);
1137 as_f32_mut(vec_w)[..dim].copy_from_slice(rows.row(v as u32));
1138 let fac = as_f32_mut(fac);
1139 let ids = as_u32_mut(ids);
1140 ids.fill(NONE);
1141 let pv = rot_row(v as u32);
1142 for &u in adj.iter().take(4) {
1143 build_kernels::prefetch_row(rot_row(u));
1144 }
1145 for (j, &u) in adj.iter().enumerate() {
1146 if let Some(&next) = adj.get(j + 4) {
1147 build_kernels::prefetch_row(rot_row(next));
1148 }
1149 ids[j] = u;
1150 let (k, m, pop) =
1151 edge_factors(pv, rot_row(u), &mut codes[j * words..(j + 1) * words]);
1152 fac[j] = k;
1153 fac[r + j] = m;
1154 fac[2 * r + j] = pop;
1155 }
1156 });
1157 Ok(Self {
1158 n,
1159 dim,
1160 degree: r,
1161 words,
1162 rotation,
1163 stride,
1164 codes_off,
1165 factors_off,
1166 ids_off,
1167 blocks,
1168 entry: g.entry,
1169 upper_ids: g.upper_ids.clone(),
1170 layers: g.layers.clone(),
1171 })
1172 }
1173}
1174
1175pub struct Graph {
1179 n: usize,
1180 dim: usize,
1181 degree: usize,
1182 metric: Metric,
1183 vectors: Vec<f32>,
1185 norms: Vec<f32>,
1187 adj: Adjacency,
1188 entry: u32,
1189 upper_ids: Vec<u32>,
1190 layers: Vec<Vec<Vec<u32>>>,
1191 build: BuildParams,
1193}
1194
1195struct GraphStructure {
1196 adj: Adjacency,
1197 entry: u32,
1198 upper_ids: Vec<u32>,
1199 layers: Vec<Vec<Vec<u32>>>,
1200}
1201
1202#[derive(Clone, Copy)]
1203struct GraphView<'a> {
1204 n: usize,
1205 dim: usize,
1206 degree: usize,
1207 metric: Metric,
1208 vectors: &'a [f32],
1209 norms: &'a [f32],
1210 adj: &'a Adjacency,
1211 entry: u32,
1212 upper_ids: &'a [u32],
1213 layers: &'a [Vec<Vec<u32>>],
1214 backing: Option<&'a memmap2::Mmap>,
1215}
1216
1217impl GraphView<'_> {
1218 fn neighbors(&self, v: usize) -> &[u32] {
1219 self.adj.get(v)
1220 }
1221
1222 fn release_rows(&self) -> std::io::Result<()> {
1223 if let Some(map) = self.backing {
1224 #[cfg(target_os = "linux")]
1225 unsafe {
1228 map.unchecked_advise(memmap2::UncheckedAdvice::DontNeed)?
1229 };
1230 #[cfg(not(target_os = "linux"))]
1231 let _ = map;
1232 }
1233 Ok(())
1234 }
1235}
1236
1237#[cfg(not(target_arch = "wasm32"))]
1238pub(crate) struct PreparedRows {
1239 file: std::io::BufWriter<std::fs::File>,
1240 values: usize,
1241 expected: usize,
1242 dim: usize,
1243}
1244
1245#[cfg(not(target_arch = "wasm32"))]
1246impl PreparedRows {
1247 pub(crate) fn new(count: usize, dim: usize) -> std::io::Result<Self> {
1248 let invalid =
1249 || std::io::Error::new(std::io::ErrorKind::InvalidInput, "prepared row geometry");
1250 if dim == 0 || count >= NONE as usize {
1251 return Err(invalid());
1252 }
1253 let expected = count.checked_mul(dim).ok_or_else(invalid)?;
1254 expected
1255 .checked_mul(4)
1256 .filter(|&n| n <= isize::MAX as usize)
1257 .ok_or_else(invalid)?;
1258 Ok(Self {
1259 file: std::io::BufWriter::with_capacity(1 << 20, tempfile::tempfile()?),
1260 values: 0,
1261 expected,
1262 dim,
1263 })
1264 }
1265
1266 pub(crate) fn append(&mut self, rows: &[f32]) -> std::io::Result<()> {
1267 use std::io::Write;
1268 let end = self
1269 .values
1270 .checked_add(rows.len())
1271 .filter(|&n| n <= self.expected);
1272 if !rows.len().is_multiple_of(self.dim) || end.is_none() {
1273 return Err(std::io::Error::new(
1274 std::io::ErrorKind::InvalidInput,
1275 "prepared row count",
1276 ));
1277 }
1278 self.file.write_all(bytemuck::cast_slice(rows))?;
1279 self.values = end.expect("checked row count");
1280 Ok(())
1281 }
1282
1283 pub(crate) fn build(self, params: &BuildParams) -> Result<PreparedGraph, WarmError> {
1284 let dim = self.dim;
1285 self.build_with(params, |data, norms| {
1286 build_structure(data, norms, dim, params)
1287 })
1288 }
1289
1290 fn build_with(
1293 self,
1294 params: &BuildParams,
1295 structure: impl FnOnce(&[f32], &[f32]) -> Result<GraphStructure, WarmError>,
1296 ) -> Result<PreparedGraph, WarmError> {
1297 let io = |error: std::io::Error| WarmError::StorageFailure {
1298 context: format!("prepared graph rows: {error}"),
1299 };
1300 if self.values != self.expected || self.values == 0 {
1301 return Err(invalid("incomplete or empty prepared rows".into()));
1302 }
1303 let file = self.file.into_inner().map_err(|e| io(e.into_error()))?;
1304 let mut map = unsafe { memmap2::MmapMut::map_mut(&file) }.map_err(io)?;
1307 let data: &mut [f32] = bytemuck::cast_slice_mut(&mut map[..self.expected * 4]);
1308 validate_build(data, self.dim, params)?;
1309 let norms = prepare_rows(data, self.dim, params.metric);
1310 map.flush().map_err(io)?;
1311 let vectors = map.make_read_only().map_err(io)?;
1312 #[cfg(unix)]
1313 vectors.advise(memmap2::Advice::Random).map_err(io)?;
1314 let data: &[f32] = bytemuck::cast_slice(&vectors[..self.expected * 4]);
1315 let structure = structure(data, &norms)?;
1316 Ok(PreparedGraph {
1317 vectors,
1318 norms,
1319 dim: self.dim,
1320 metric: params.metric,
1321 degree: params.degree,
1322 structure,
1323 })
1324 }
1325}
1326
1327#[cfg(not(target_arch = "wasm32"))]
1328pub(crate) struct PreparedGraph {
1329 vectors: memmap2::Mmap,
1330 norms: Vec<f32>,
1331 dim: usize,
1332 metric: Metric,
1333 degree: usize,
1334 structure: GraphStructure,
1335}
1336
1337#[cfg(not(target_arch = "wasm32"))]
1338impl PreparedGraph {
1339 fn view(&self) -> GraphView<'_> {
1340 let vectors: &[f32] = bytemuck::cast_slice(&self.vectors[..]);
1341 GraphView {
1342 n: vectors.len() / self.dim,
1343 dim: self.dim,
1344 degree: self.degree,
1345 metric: self.metric,
1346 vectors,
1347 norms: &self.norms,
1348 adj: &self.structure.adj,
1349 entry: self.structure.entry,
1350 upper_ids: &self.structure.upper_ids,
1351 layers: &self.structure.layers,
1352 backing: Some(&self.vectors),
1353 }
1354 }
1355}
1356
1357fn prepare_rows(data: &mut [f32], dim: usize, metric: Metric) -> Vec<f32> {
1360 match metric {
1361 Metric::Angular => {
1362 data.par_chunks_mut(dim).for_each(normalize_row);
1363 Vec::new()
1364 }
1365 Metric::Euclidean => data.par_chunks(dim).map(|x| dot(x, x)).collect(),
1366 }
1367}
1368
1369fn normalize_row(x: &mut [f32]) {
1374 let squared = dot(x, x);
1375 if squared.is_finite() && squared >= f32::MIN_POSITIVE {
1376 let norm = squared.sqrt();
1377 x.iter_mut().for_each(|o| *o /= norm);
1378 return;
1379 }
1380 let largest = x.iter().fold(0.0f32, |m, v| m.max(v.abs()));
1381 if largest > 0.0 && largest.is_finite() {
1382 x.iter_mut().for_each(|o| *o /= largest);
1383 let norm = dot(x, x).sqrt();
1384 x.iter_mut().for_each(|o| *o /= norm);
1385 }
1386}
1387
1388fn norms_opt(norms: &[f32]) -> Option<&[f32]> {
1390 (!norms.is_empty()).then_some(norms)
1391}
1392
1393impl Graph {
1394 fn view(&self) -> GraphView<'_> {
1395 GraphView {
1396 n: self.n,
1397 dim: self.dim,
1398 degree: self.degree,
1399 metric: self.metric,
1400 vectors: &self.vectors,
1401 norms: &self.norms,
1402 adj: &self.adj,
1403 entry: self.entry,
1404 upper_ids: &self.upper_ids,
1405 layers: &self.layers,
1406 backing: None,
1407 }
1408 }
1409
1410 fn rows(&self) -> Rows<'_> {
1411 Rows {
1412 data: &self.vectors,
1413 dim: self.dim,
1414 norms: norms_opt(&self.norms),
1415 codes: None,
1416 quantized: false,
1417 }
1418 }
1419
1420 pub fn metric(&self) -> Metric {
1422 self.metric
1423 }
1424
1425 pub(crate) fn neighbors(&self, v: usize) -> &[u32] {
1427 self.adj.get(v)
1428 }
1429
1430 pub fn structure_hash(&self) -> u64 {
1433 let mut h = Fnv::new();
1434 h.word(self.entry);
1435 for list in self.adj.iter() {
1436 h.word(list.len() as u32);
1437 list.iter().for_each(|&u| h.word(u));
1438 }
1439 for layer in &self.layers {
1440 h.word(layer.len() as u32);
1441 for list in layer {
1442 h.word(list.len() as u32);
1443 list.iter().for_each(|&u| h.word(u));
1444 }
1445 }
1446 h.finish()
1447 }
1448
1449 pub fn build(data: &[f32], dim: usize, params: &BuildParams) -> Result<Self, WarmError> {
1453 Self::build_owned(data.to_vec(), dim, params)
1454 }
1455
1456 pub fn build_owned(
1460 mut data: Vec<f32>,
1461 dim: usize,
1462 params: &BuildParams,
1463 ) -> Result<Self, WarmError> {
1464 validate_build(&data, dim, params)?;
1465 let norms = prepare_rows(&mut data, dim, params.metric);
1466 let structure = build_structure(&data, &norms, dim, params)?;
1467 Ok(Self {
1468 n: data.len() / dim,
1469 dim,
1470 degree: params.degree,
1471 metric: params.metric,
1472 vectors: data,
1473 norms,
1474 adj: structure.adj,
1475 entry: structure.entry,
1476 upper_ids: structure.upper_ids,
1477 layers: structure.layers,
1478 build: params.clone(),
1479 })
1480 }
1481}
1482
1483fn validate_build(data: &[f32], dim: usize, params: &BuildParams) -> Result<(), WarmError> {
1484 if dim == 0 || !data.len().is_multiple_of(dim) {
1485 return Err(invalid(format!(
1486 "ragged data: {} values are not a multiple of dimension {dim}",
1487 data.len()
1488 )));
1489 }
1490 let n = data.len() / dim;
1491 if n == 0 {
1492 return Err(invalid("empty index: no vectors to build from".into()));
1493 }
1494 if n >= NONE as usize {
1495 return Err(invalid(format!(
1496 "index too large: {n} vectors, the limit is {}",
1497 NONE - 1
1498 )));
1499 }
1500 if let Some(at) = data.iter().position(|v| !v.is_finite()) {
1501 return Err(invalid(format!(
1502 "non-finite value in row {} (coordinate {})",
1503 at / dim,
1504 at % dim
1505 )));
1506 }
1507 if params.build_ef == 0 {
1508 return Err(invalid("build_ef must be positive".into()));
1509 }
1510 let r = params.degree;
1511 if r == 0 || !r.is_multiple_of(2) {
1512 return Err(invalid(format!(
1513 "degree must be positive and even, not {r}"
1514 )));
1515 }
1516 Ok(())
1517}
1518
1519fn build_structure(
1520 vectors: &[f32],
1521 norms: &[f32],
1522 dim: usize,
1523 params: &BuildParams,
1524) -> Result<GraphStructure, WarmError> {
1525 let tsetup = prof::start();
1526 let n = vectors.len() / dim;
1527 let r = params.degree;
1528 let metric = params.metric;
1529 if let Some(at) = norms.iter().position(|v| !v.is_finite()) {
1530 return Err(invalid(format!("row {at} overflows its squared norm")));
1531 }
1532 let codes = (params.quantized_search || dim >= BUILD_CODES_MIN_DIM)
1533 .then(|| BuildCodes::new(vectors, dim, norms_opt(norms)));
1534 let rows = Rows {
1535 data: vectors,
1536 dim,
1537 norms: norms_opt(norms),
1538 codes: codes.as_ref(),
1539 quantized: params.quantized_search,
1540 };
1541
1542 let entry = medoid(&rows, n, metric);
1543
1544 let mut rng = SplitMix(params.seed ^ 0xA5A5_A5A5);
1545 let mut order: Vec<u32> = (0..n as u32).collect();
1546 for i in (1..n).rev() {
1547 order.swap(i, (rng.next_u64() % (i as u64 + 1)) as usize);
1548 }
1549 let mut graph = Adjacency::new(n, r);
1550 prof::report("setup", tsetup);
1551 let first_ef = params.build_ef.min(FIRST_PASS_EF);
1552 for (l, alpha) in [(first_ef, 1.0), (params.build_ef, params.alpha)] {
1553 let t = prof::start();
1554 rows.insert_pass(&mut graph, &order, entry, l, alpha, r);
1555 prof::report("pass", t);
1556 }
1557 let t = prof::start();
1558 rows.fill(&mut graph, r);
1559 prof::report("fill", t);
1560 let tup = prof::start();
1561
1562 let layers = upper_layers(&rows, &order, params);
1563 prof::report("upper", tup);
1564 prof::graph_hash((&graph, &layers, entry));
1565 let upper_ids = order[..n / LAYER_RATIO].to_vec();
1566 Ok(GraphStructure {
1567 adj: graph,
1568 entry,
1569 upper_ids,
1570 layers,
1571 })
1572}
1573
1574fn medoid(rows: &Rows<'_>, n: usize, metric: Metric) -> u32 {
1577 medoid_by(rows, n, |i| i, metric)
1578}
1579
1580fn medoid_by(rows: &Rows<'_>, n: usize, row_at: impl Fn(u32) -> u32 + Sync, metric: Metric) -> u32 {
1581 let dim = rows.dim;
1582 let mut centroid = vec![0f32; dim];
1583 for i in 0..n as u32 {
1584 let x = rows.row(row_at(i));
1585 centroid.iter_mut().zip(x).for_each(|(c, v)| *c += v);
1586 }
1587 match metric {
1588 Metric::Angular => {
1589 (0..n as u32)
1590 .into_par_iter()
1591 .map(|i| (dot(rows.row(row_at(i)), ¢roid), i))
1592 .reduce(
1593 || (f32::NEG_INFINITY, NONE),
1594 |a, b| {
1595 if b.0 > a.0 || (b.0 == a.0 && b.1 < a.1) {
1596 b
1597 } else {
1598 a
1599 }
1600 },
1601 )
1602 .1
1603 }
1604 Metric::Euclidean => {
1605 centroid.iter_mut().for_each(|c| *c /= n as f32);
1608 (0..n as u32)
1609 .into_par_iter()
1610 .map(|i| {
1611 (
1612 rows.norms.expect("Euclidean rows carry norms")[row_at(i) as usize]
1613 - 2.0 * dot(rows.row(row_at(i)), ¢roid),
1614 i,
1615 )
1616 })
1617 .reduce(
1618 || (f32::INFINITY, NONE),
1619 |a, b| {
1620 if b.0 < a.0 || (b.0 == a.0 && b.1 < a.1) {
1621 b
1622 } else {
1623 a
1624 }
1625 },
1626 )
1627 .1
1628 }
1629 }
1630}
1631
1632fn upper_layers(rows: &Rows<'_>, order: &[u32], params: &BuildParams) -> Vec<Vec<Vec<u32>>> {
1636 let dim = rows.dim;
1637 let mut layers = Vec::new();
1638 let mut size = order.len() / LAYER_RATIO;
1639 while size >= MIN_LAYER {
1640 let sub: Vec<f32> = order[..size]
1641 .iter()
1642 .flat_map(|&g| rows.row(g).iter().copied())
1643 .collect();
1644 let sub_norms: Vec<f32> = match rows.norms {
1645 Some(n) => order[..size].iter().map(|&g| n[g as usize]).collect(),
1646 None => Vec::new(),
1647 };
1648 let sub_codes = (params.quantized_search || dim >= BUILD_CODES_MIN_DIM)
1649 .then(|| BuildCodes::new(&sub, dim, norms_opt(&sub_norms)));
1650 let sub_rows = Rows {
1651 data: &sub,
1652 dim,
1653 norms: norms_opt(&sub_norms),
1654 codes: sub_codes.as_ref(),
1655 quantized: params.quantized_search,
1656 };
1657 let local: Vec<u32> = (0..size as u32).collect();
1658 let mut g = Adjacency::new(size, UPPER_DEGREE);
1659 sub_rows.insert_pass(
1660 &mut g,
1661 &local,
1662 0,
1663 params.build_ef,
1664 params.alpha,
1665 UPPER_DEGREE,
1666 );
1667 layers.push(g.to_lists());
1668 size /= LAYER_RATIO;
1669 }
1670 layers
1671}
1672
1673struct Fnv(u64);
1675
1676impl Fnv {
1677 fn new() -> Self {
1678 Self(0xcbf2_9ce4_8422_2325)
1679 }
1680
1681 fn word(&mut self, w: u32) {
1682 self.0 = (self.0 ^ u64::from(w)).wrapping_mul(0x0000_0100_0000_01b3);
1683 }
1684
1685 fn finish(&self) -> u64 {
1686 self.0
1687 }
1688}
1689
1690fn upper_bytes(upper_ids: &[u32], layers: &[Vec<Vec<u32>>]) -> usize {
1692 let adj: usize = layers.iter().flatten().map(|a| a.len() * 4).sum();
1693 upper_ids.len() * 4 + adj
1694}
1695
1696impl QGraph {
1697 pub fn len(&self) -> usize {
1698 self.n
1699 }
1700
1701 pub fn is_empty(&self) -> bool {
1702 self.n == 0
1703 }
1704
1705 pub fn index_bytes(&self) -> usize {
1707 self.blocks.len() * 8 + upper_bytes(&self.upper_ids, &self.layers)
1708 }
1709
1710 fn vector(&self, v: u32) -> &[f32] {
1711 &as_f32(&self.block(v)[..self.codes_off])[..self.dim]
1712 }
1713
1714 fn descend(&self, pq: &[f32]) -> (u32, usize) {
1717 if self.layers.is_empty() {
1718 return (self.entry, 0);
1719 }
1720 let mut p = 0u32;
1721 let mut best = -dot(pq, self.vector(self.upper_ids[0]));
1722 let mut evals = 1;
1723 for layer in self.layers.iter().rev() {
1724 loop {
1725 let mut moved = false;
1726 for &u in &layer[p as usize] {
1727 let d = -dot(pq, self.vector(self.upper_ids[u as usize]));
1728 evals += 1;
1729 if d < best {
1730 best = d;
1731 p = u;
1732 moved = true;
1733 }
1734 }
1735 if !moved {
1736 break;
1737 }
1738 }
1739 }
1740 (self.upper_ids[p as usize], evals)
1741 }
1742
1743 pub fn searcher(&self) -> Searcher<'_> {
1744 Searcher {
1745 index: self,
1746 visited: Visited::new(self.n),
1747 pq: vec![0.0; self.rotation.padded()],
1748 code: QueryCode::new(self.words),
1749 raw: vec![0; self.degree],
1750 est: vec![0.0; self.degree],
1751 pool: Vec::new(),
1752 results: Vec::new(),
1753 }
1754 }
1755
1756 fn block(&self, v: u32) -> &[u64] {
1757 let v = v as usize;
1758 &self.blocks[v * self.stride..(v + 1) * self.stride]
1759 }
1760}
1761
1762pub struct Searcher<'a> {
1764 index: &'a QGraph,
1765 visited: Visited,
1766 pq: Vec<f32>,
1767 code: QueryCode,
1768 raw: Vec<u32>,
1769 est: Vec<f32>,
1770 pool: Vec<Candidate>,
1771 results: Vec<(f32, u32)>,
1772}
1773
1774impl Searcher<'_> {
1775 pub fn search(
1778 &mut self,
1779 query: &[f32],
1780 k: usize,
1781 params: SearchParams,
1782 out: &mut Vec<u32>,
1783 ) -> usize {
1784 let idx = self.index;
1785 let ef = params.ef.max(1);
1786 idx.rotation.apply(query, &mut self.pq);
1787 self.code.encode(&self.pq);
1788 let (lo2, d2q, sq) = (2.0 * self.code.lo, 2.0 * self.code.delta, self.code.sum);
1789
1790 self.visited.next();
1791 let (entry, mut exact) = idx.descend(query);
1792 let mut expanded = 0;
1793 self.visited.insert(entry);
1794 self.pool.clear();
1795 self.results.clear();
1796 self.pool.push(Candidate {
1797 dist: 0.0,
1798 id: entry,
1799 done: false,
1800 });
1801 let mut cur = 0;
1802 while cur < self.pool.len() {
1803 if params.max_exact > 0 && expanded >= params.max_exact {
1804 break;
1805 }
1806 let v = self.pool[cur].id;
1807 if let Some(next) = self.pool[cur + 1..].iter().find(|c| !c.done) {
1808 kernels::prefetch(idx.block(next.id));
1809 }
1810 let b = idx.block(v);
1811 let ipv = dot(query, &as_f32(&b[..idx.codes_off])[..idx.dim]);
1812 exact += 1;
1813 expanded += 1;
1814 push_top(&mut self.results, k, -ipv, v);
1815 let mut pos = self.pool.partition_point(|c| c.dist <= -ipv);
1820 if pos > cur {
1821 pos -= 1;
1822 self.pool[cur..=pos].rotate_left(1);
1823 } else {
1824 self.pool[pos..=cur].rotate_right(1);
1825 }
1826 self.pool[pos] = Candidate {
1827 dist: -ipv,
1828 id: v,
1829 done: true,
1830 };
1831
1832 kernels::raw_estimates(
1833 &b[idx.codes_off..idx.factors_off],
1834 &self.code.planes,
1835 idx.words,
1836 &mut self.raw,
1837 );
1838 let fac = as_f32(&b[idx.factors_off..idx.ids_off]);
1839 let (kf, rest) = fac.split_at(idx.degree);
1840 let (mf, pf) = rest.split_at(idx.degree);
1841 for ((((e, &kj), &mj), &pj), &rj) in
1842 self.est.iter_mut().zip(kf).zip(mf).zip(pf).zip(&self.raw)
1843 {
1844 *e = -(ipv + kj + mj * (lo2 * pj - sq + d2q * rj as f32));
1845 }
1846
1847 let mut best = usize::MAX;
1848 for (&u, &d) in as_u32(&b[idx.ids_off..]).iter().zip(&self.est) {
1849 if u == NONE || !self.visited.insert(u) {
1850 continue;
1851 }
1852 if self.pool.len() >= ef && d >= self.pool[ef - 1].dist {
1853 continue;
1854 }
1855 let pos = self.pool.partition_point(|c| c.dist <= d);
1856 self.pool.insert(
1857 pos,
1858 Candidate {
1859 dist: d,
1860 id: u,
1861 done: false,
1862 },
1863 );
1864 self.pool.truncate(ef);
1865 best = best.min(pos);
1866 }
1867 cur = best.min(cur);
1868 while cur < self.pool.len() && self.pool[cur].done {
1869 cur += 1;
1870 }
1871 }
1872 out.clear();
1873 out.extend(self.results.iter().map(|r| r.1));
1874 exact
1875 }
1876}
1877
1878fn push_top(top: &mut Vec<(f32, u32)>, k: usize, dist: f32, id: u32) {
1879 if top.len() >= k && top.last().is_none_or(|l| dist >= l.0) {
1880 return;
1881 }
1882 let pos = top.partition_point(|c| c.0 <= dist);
1883 top.insert(pos, (dist, id));
1884 top.truncate(k);
1885}
1886
1887#[cfg(test)]
1888mod tests {
1889 use super::*;
1890
1891 #[cfg(feature = "experimental-build-search")]
1892 #[test]
1893 fn build_pool_retirement_matches_ordered_merge_without_growth() {
1894 for limit in [
1895 1usize, 2, 3, 7, 8, 9, 31, 32, 33, 63, 64, 65, 199, 200, 201, 499, 500, 501,
1896 ] {
1897 for len in [0, limit - 1, limit] {
1898 let original: Vec<_> = (0..len)
1899 .map(|id| Candidate {
1900 dist: (id / 3) as f32,
1901 id: id as u32,
1902 done: id % 3 == 0,
1903 })
1904 .collect();
1905 let bound = original.last().map_or(0.0, |c| c.dist);
1906 for dist in [
1907 f32::NEG_INFINITY,
1908 -1.0,
1909 -0.0,
1910 0.0,
1911 0.25,
1912 1.0,
1913 bound - 0.25,
1914 bound,
1915 bound + 0.25,
1916 f32::INFINITY,
1917 f32::from_bits(0x7fc0_1234),
1918 ] {
1919 let candidate = Candidate {
1920 dist,
1921 id: limit as u32,
1922 done: false,
1923 };
1924 let expected_pos = if len == limit && dist >= bound {
1925 None
1926 } else {
1927 Some(original.iter().take_while(|c| c.dist <= dist).count())
1928 };
1929 let expected: Vec<_> = match expected_pos {
1930 Some(pos) => original[..pos]
1931 .iter()
1932 .copied()
1933 .chain(std::iter::once(candidate))
1934 .chain(original[pos..].iter().copied())
1935 .take(limit)
1936 .collect(),
1937 None => original.clone(),
1938 };
1939 let mut pool = Vec::with_capacity(limit);
1940 pool.extend_from_slice(&original);
1941 let capacity = pool.capacity();
1942 assert_eq!(
1943 insert_build_candidate(&mut pool, candidate, limit),
1944 expected_pos
1945 );
1946 assert_eq!(
1947 pool.capacity(),
1948 capacity,
1949 "limit={limit}, len={len}, dist={dist}"
1950 );
1951 let bits = |entries: &[Candidate]| {
1952 entries
1953 .iter()
1954 .map(|c| (c.dist.to_bits(), c.id, c.done))
1955 .collect::<Vec<_>>()
1956 };
1957 assert_eq!(
1958 bits(&pool),
1959 bits(&expected),
1960 "limit={limit}, len={len}, dist={dist}"
1961 );
1962 }
1963 }
1964 }
1965 }
1966
1967 #[test]
1970 fn angular_rows_normalize_at_extreme_magnitudes() {
1971 let mut rows = vec![
1972 1e20f32, 1e20, 0.0, 0.0, 1e-30, 1e-30, 1e-30, 1e-30, 3.0, 4.0, 0.0, 0.0,
1973 ];
1974 prepare_rows(&mut rows, 4, Metric::Angular);
1975 let half = std::f32::consts::FRAC_1_SQRT_2;
1976 for (got, want) in rows[..8]
1977 .iter()
1978 .zip([half, half, 0.0, 0.0, 0.5, 0.5, 0.5, 0.5])
1979 {
1980 assert!((got - want).abs() < 1e-6, "{rows:?}");
1981 }
1982 assert_eq!(&rows[8..], &[0.6, 0.8, 0.0, 0.0]);
1983 }
1984 use kernels::raw_estimates;
1985
1986 #[test]
1987 fn build_visited_reuses_words_across_resets_and_graph_sizes() {
1988 let mut visited = BuildVisited::default();
1989 for n in [4097, 129, 128, 127, 65, 64, 63, 1, 0, 1, 63, 64, 65, 4097] {
1990 for _ in 0..3 {
1991 visited.reset(n);
1992 let mut reference = std::collections::HashSet::new();
1993 for id in (0..n as u32).rev().chain(0..n as u32) {
1994 assert_eq!(visited.insert(id), reference.insert(id));
1995 }
1996 assert_eq!(visited.touched.len(), n.div_ceil(64));
1997 let mut unique = std::collections::HashSet::new();
1998 assert!(visited
1999 .touched
2000 .iter()
2001 .all(|&word| visited.bits[word as usize] != 0 && unique.insert(word)));
2002 visited.reset(n);
2003 assert!(visited.touched.is_empty());
2004 assert!(visited.bits.iter().all(|&word| word == 0));
2005 }
2006 }
2007 }
2008
2009 #[test]
2010 fn build_visited_records_each_sparse_word_once() {
2011 let mut visited = BuildVisited::default();
2012 visited.reset(8_421_866);
2013 for id in [0, 1, 63, 64, 65, 127, 128, 8_421_865] {
2014 assert!(visited.insert(id));
2015 assert!(!visited.insert(id));
2016 }
2017 assert_eq!(visited.touched, [0, 1, 2, 131_591]);
2018 visited.reset(0);
2019 visited.reset(8_421_866);
2020 for id in [0, 63, 64, 127, 128, 8_421_865] {
2021 assert!(visited.insert(id));
2022 }
2023 }
2024
2025 #[test]
2026 #[ignore = "synthetic visited-reset timing diagnostic; run in isolation with --release"]
2027 fn build_visited_reset_workload() {
2028 use std::hint::black_box;
2029 use std::time::{Duration, Instant};
2030
2031 let n = 8_421_866usize;
2032 for (visits, iterations) in [(64, 20_000), (1024, 10_000), (16_384, 1000), (n, 4)] {
2033 let ids: Vec<_> = (0..visits).map(|i| (i * 104_729 % n) as u32).collect();
2034 let words: std::collections::HashSet<_> = ids.iter().map(|id| id / 64).collect();
2035 let mut visited = BuildVisited::default();
2036 visited.reset(n);
2037 for &id in &ids {
2038 assert!(visited.insert(id));
2039 }
2040 let mut reset_time = Duration::ZERO;
2041 let start = Instant::now();
2042 for _ in 0..iterations {
2043 let reset_start = Instant::now();
2044 black_box(&mut visited).reset(black_box(n));
2045 reset_time += reset_start.elapsed();
2046 for &id in &ids {
2047 black_box(visited.insert(black_box(id)));
2048 }
2049 }
2050 println!(
2051 "{{\"vertices\":{n},\"visits\":{visits},\"touched_words\":{},\"full_words\":{},\"iterations\":{iterations},\"reset_ns\":{},\"total_ns\":{}}}",
2052 words.len(), n.div_ceil(64), reset_time.as_nanos(), start.elapsed().as_nanos()
2053 );
2054 }
2055 }
2056
2057 pub(super) fn gaussian(rng: &mut SplitMix) -> f32 {
2058 let u = rng.next_f32().max(1e-7);
2059 let v = rng.next_f32();
2060 (-2.0 * u.ln()).sqrt() * (std::f32::consts::TAU * v).cos()
2061 }
2062
2063 fn unit(mut x: Vec<f32>) -> Vec<f32> {
2064 let n = dot(&x, &x).sqrt();
2065 x.iter_mut().for_each(|v| *v /= n);
2066 x
2067 }
2068
2069 #[test]
2074 fn edge_estimator_is_unbiased_over_rotations() {
2075 let d = 100;
2076 let mut rng = SplitMix(42);
2077 for bits in [0, 0, 0, 512, 512, 512] {
2078 let q = unit((0..d).map(|_| gaussian(&mut rng)).collect());
2079 let v = unit((0..d).map(|_| gaussian(&mut rng)).collect());
2080 let u = unit(v.iter().map(|x| x + 0.15 * gaussian(&mut rng)).collect());
2081 let truth = dot(&q, &u);
2082 let (mut float_sum, mut quant_sum) = (0f64, 0f64);
2083 let seeds = 400;
2084 for seed in 0..seeds {
2085 let rot = Rotation::new(d, bits, seed);
2086 let p = rot.padded();
2087 let (mut pq, mut pv, mut pu) = (vec![0.0; p], vec![0.0; p], vec![0.0; p]);
2088 rot.apply(&q, &mut pq);
2089 rot.apply(&v, &mut pv);
2090 rot.apply(&u, &mut pu);
2091 let mut code = vec![0u64; p / 64];
2092 let (k, m, pop) = edge_factors(&pv, &pu, &mut code);
2093 let ipv = dot(&pq, &pv);
2094 let sq: f32 = pq
2095 .iter()
2096 .enumerate()
2097 .map(|(i, x)| {
2098 if code[i / 64] >> (i % 64) & 1 == 1 {
2099 *x
2100 } else {
2101 -*x
2102 }
2103 })
2104 .sum();
2105 float_sum += f64::from(ipv + k + m * sq);
2106 let mut qc = QueryCode::new(p / 64);
2107 qc.encode(&pq);
2108 let mut raw = [0u32];
2109 raw_estimates(&code, &qc.planes, p / 64, &mut raw);
2110 let est =
2111 ipv + k + m * (2.0 * qc.lo * pop - qc.sum + 2.0 * qc.delta * raw[0] as f32);
2112 quant_sum += f64::from(est);
2113 }
2114 let float_mean = float_sum / f64::from(seeds as u32);
2115 let quant_mean = quant_sum / f64::from(seeds as u32);
2116 assert!(
2117 (float_mean - f64::from(truth)).abs() < 0.02,
2118 "{float_mean} vs {truth}"
2119 );
2120 assert!(
2121 (quant_mean - f64::from(truth)).abs() < 0.03,
2122 "{quant_mean} vs {truth}"
2123 );
2124 }
2125 }
2126
2127 pub(super) fn clustered(n: usize, d: usize, seed: u64) -> Vec<f32> {
2128 let mut rng = SplitMix(seed);
2129 let centers: Vec<Vec<f32>> = (0..20)
2130 .map(|_| (0..d).map(|_| gaussian(&mut rng)).collect())
2131 .collect();
2132 (0..n)
2133 .flat_map(|i| {
2134 let c = ¢ers[i % centers.len()];
2135 c.iter()
2136 .map(|x| x + 0.4 * gaussian(&mut rng))
2137 .collect::<Vec<_>>()
2138 })
2139 .collect()
2140 }
2141
2142 pub(super) fn brute_force(data: &[f32], d: usize, q: &[f32], k: usize) -> Vec<u32> {
2143 let mut s: Vec<(f32, u32)> = data
2144 .chunks_exact(d)
2145 .enumerate()
2146 .map(|(i, x)| (-dot(x, q) / dot(x, x).sqrt(), i as u32))
2147 .collect();
2148 s.sort_by(|a, b| a.0.total_cmp(&b.0));
2149 s.into_iter().take(k).map(|x| x.1).collect()
2150 }
2151
2152 #[test]
2153 fn recall_on_clustered_data() {
2154 let d = 24;
2155 let data = clustered(4000, d, 1);
2156 let queries = clustered(100, d, 2);
2157 let index = QGraph::build(&data, d, &BuildParams::default()).expect("build");
2158 let mut s = index.searcher();
2159 let mut out = Vec::new();
2160 let mut hits = 0;
2161 for q in queries.chunks_exact(d) {
2162 let truth = brute_force(&data, d, q, 10);
2163 s.search(
2164 q,
2165 10,
2166 SearchParams {
2167 ef: 64,
2168 max_exact: 0,
2169 },
2170 &mut out,
2171 );
2172 hits += out.iter().filter(|x| truth.contains(x)).count();
2173 }
2174 let recall = hits as f64 / 1000.0;
2175 assert!(recall >= 0.9, "recall {recall}");
2177 }
2178
2179 #[test]
2182 fn prune_merged_matches_prune() {
2183 let d = 12;
2184 let mut data = clustered(600, d, 5);
2185 data.chunks_exact_mut(d).for_each(|x| {
2186 let n = dot(x, x).sqrt();
2187 x.iter_mut().for_each(|v| *v /= n);
2188 });
2189 let rows = Rows {
2190 data: &data,
2191 dim: d,
2192 norms: None,
2193 quantized: false,
2194 codes: None,
2195 };
2196 let mut rng = SplitMix(11);
2197 for trial in 0..200 {
2198 let t = (rng.next_u64() % 600) as u32;
2199 let (r, alpha) = (4 + trial % 13, [1.0, 1.2][trial % 2]);
2200 let all = |ids: &[u32]| -> Vec<(f32, u32)> {
2201 ids.iter().map(|&u| (rows.d2(t, u), u)).collect()
2202 };
2203 let pool: Vec<u32> = (0..60)
2204 .map(|_| (rng.next_u64() % 600) as u32)
2205 .filter(|&u| u != t)
2206 .collect();
2207 let mut list = rows.prune(t, all(&pool), alpha, r);
2208 let closed = list.len();
2209 for _ in 0..1 + trial % 9 {
2210 let u = (rng.next_u64() % 600) as u32;
2211 if u != t && !list.contains(&u) {
2212 list.push(u);
2213 }
2214 }
2215 assert_eq!(
2216 rows.prune_merged(t, &list, closed, alpha, r),
2217 rows.prune(t, all(&list), alpha, r),
2218 "trial {trial}"
2219 );
2220 }
2221 }
2222
2223 #[test]
2228 fn build_codes_leave_the_graph_unchanged() {
2229 for metric in [Metric::Angular, Metric::Euclidean] {
2230 for d in [100, 160] {
2231 let n = 1200;
2232 let mut data = clustered(n, d, 3);
2233 let mut rng = SplitMix(5);
2234 if metric == Metric::Euclidean {
2235 for (i, x) in data.chunks_exact_mut(d).enumerate() {
2236 let s = 0.1 * 100f32.powf((i * 7919 % n) as f32 / n as f32);
2237 x.iter_mut().for_each(|v| *v = s * (*v + 3.0));
2238 }
2239 }
2240 for i in (0..n).step_by(7) {
2241 let j = (rng.next_u64() % n as u64) as usize;
2242 for k in 0..d {
2243 data[i * d + k] = data[j * d + k] * (1.0 + 1e-6 * gaussian(&mut rng));
2244 }
2245 }
2246 data[5 * d..6 * d].fill(0.0);
2247 let norms = prepare_rows(&mut data, d, metric);
2248 let codes = BuildCodes::new(&data, d, norms_opt(&norms));
2249 let build = |codes: Option<&BuildCodes>| {
2250 let rows = Rows {
2251 data: &data,
2252 dim: d,
2253 norms: norms_opt(&norms),
2254 codes,
2255 quantized: false,
2256 };
2257 let order: Vec<u32> = (0..n as u32).rev().collect();
2258 let mut graph = Adjacency::new(n, 12);
2259 for alpha in [1.0, 1.2] {
2260 rows.insert_pass(&mut graph, &order, 0, 40, alpha, 12);
2261 }
2262 rows.fill(&mut graph, 12);
2263 graph
2264 };
2265 assert_eq!(build(Some(&codes)), build(None), "{metric:?}, d = {d}");
2266 }
2267 }
2268 }
2269
2270 #[test]
2271 fn build_and_search_are_deterministic() {
2272 let d = 16;
2273 let data = clustered(3000, d, 3);
2274 let params = BuildParams {
2275 seed: 9,
2276 ..BuildParams::default()
2277 };
2278 let a = QGraph::build(&data, d, ¶ms).expect("build");
2279 let b = QGraph::build(&data, d, ¶ms).expect("build");
2280 assert_eq!(a.entry, b.entry);
2281 assert!(a.blocks == b.blocks);
2282 let (mut sa, mut sb) = (a.searcher(), b.searcher());
2283 let (mut oa, mut ob) = (Vec::new(), Vec::new());
2284 let p = SearchParams {
2285 ef: 32,
2286 max_exact: 0,
2287 };
2288 for q in data.chunks_exact(d).take(50) {
2289 sa.search(q, 10, p, &mut oa);
2290 sb.search(q, 10, p, &mut ob);
2291 assert_eq!(oa, ob);
2292 }
2293 }
2294
2295 #[test]
2298 fn build_reproduces_the_pinned_graphs() {
2299 let cases = [
2300 (Metric::Angular, false, 16, 0x2b16_3721_02ba_f39b),
2301 (Metric::Euclidean, true, 128, 0xe014_c782_9bed_e9aa),
2302 ];
2303 for (metric, quantized_search, d, want) in cases {
2304 let data = clustered(2500, d, 11);
2305 let params = BuildParams {
2306 metric,
2307 quantized_search,
2308 degree: 16,
2309 build_ef: 48,
2310 seed: 7,
2311 ..BuildParams::default()
2312 };
2313 let g = Graph::build(&data, d, ¶ms).expect("build");
2314 assert_eq!(g.structure_hash(), want, "{metric:?}");
2315 }
2316 }
2317
2318 #[test]
2319 fn exact_budget_caps_work() {
2320 let d = 16;
2321 let data = clustered(2000, d, 4);
2322 let index = QGraph::build(&data, d, &BuildParams::default()).expect("build");
2323 let mut s = index.searcher();
2324 let mut out = Vec::new();
2325 let q = &data[..d];
2326 let capped = s.search(
2327 q,
2328 10,
2329 SearchParams {
2330 ef: 128,
2331 max_exact: 20,
2332 },
2333 &mut out,
2334 );
2335 assert_eq!(out.len(), 10);
2336 let free = s.search(
2337 q,
2338 10,
2339 SearchParams {
2340 ef: 128,
2341 max_exact: 0,
2342 },
2343 &mut out,
2344 );
2345 assert!(capped < free);
2346 assert_eq!(capped - index.descend(q).1, 20);
2347 }
2348}