1use std::collections::BinaryHeap;
18use std::sync::atomic::{AtomicU64, Ordering};
19use std::sync::OnceLock;
20use std::time::Instant;
21
22use rayon::prelude::*;
23
24use super::adjacency::Adjacency;
25use super::fastscan::{pack32, scan32, BLOCK};
26use super::kernels::{
27 dot_i4_query, dot_i4_query_x4, dot_i8_query, dot_i8_query_x4, interleave_for_i4, prefetch,
28 prefetch_bytes, raw_one, QueryCode, Rotation, SplitMix,
29};
30use super::rerank::{dot_f32_i4, dot_f32_i8};
31use super::vnni::query_sum;
32use super::{invalid, prepare_rows, upper_bytes, BuildParams, Graph, GraphView, Metric, NONE};
33use crate::core::error::WarmError;
34
35mod file;
36mod gated;
37pub use file::Residency;
38use file::Slab;
39
40#[derive(Clone, Copy, Debug, PartialEq, Eq)]
42pub enum VMetric {
43 Angular,
44 Euclidean,
45 InnerProduct,
46}
47
48impl From<Metric> for VMetric {
49 fn from(metric: Metric) -> Self {
50 match metric {
51 Metric::Angular => Self::Angular,
52 Metric::Euclidean => Self::Euclidean,
53 }
54 }
55}
56
57impl VMetric {
58 fn code(self) -> u64 {
59 match self {
60 Self::Angular => 0,
61 Self::Euclidean => 1,
62 Self::InnerProduct => 2,
63 }
64 }
65}
66
67#[derive(Default, serde::Serialize)]
68pub struct QueryProfile {
69 pub query_encoding_ns: u64,
70 pub entry_ns: u64,
71 pub adjacency_visited_prefetch_ns: u64,
72 pub scoring_pool_ns: u64,
73 pub rerank_sort_ns: u64,
74 pub total_ns: u64,
75 pub expansions: usize,
76 pub neighbors_read: usize,
77 pub entry_evaluations: usize,
78 pub traversal_evaluations: usize,
79 pub rerank_evaluations: usize,
80 pub beam_retained: usize,
81 pub retained: usize,
82}
83
84#[derive(Clone, Copy, Debug)]
86pub struct VSearchParams {
87 pub ef: usize,
89 pub reservoir: usize,
93 pub rerank: usize,
96 pub patience: usize,
99 pub screen_sigmas: f32,
103 pub screen: bool,
105 pub fuse: bool,
109 pub neighbour_cap: usize,
112}
113
114impl Default for VSearchParams {
115 fn default() -> Self {
116 Self {
117 ef: 64,
118 reservoir: 0,
119 rerank: 0,
120 patience: 0,
121 screen_sigmas: 1.0,
122 screen: true,
123 fuse: false,
124 neighbour_cap: 0,
125 }
126 }
127}
128
129#[derive(Clone, Copy, Debug, PartialEq)]
152pub struct SearchKnobs {
153 pub prefetch: u8,
155 pub batch_est: bool,
157 pub frontier: f32,
159 pub block_screen: bool,
161 pub rerank_z: f32,
166}
167
168impl SearchKnobs {
169 pub const OFF: Self = Self {
171 prefetch: 0,
172 batch_est: false,
173 frontier: 1.0,
174 block_screen: false,
175 rerank_z: 0.0,
176 };
177
178 pub fn from_env() -> Self {
181 static KNOBS: OnceLock<SearchKnobs> = OnceLock::new();
182 *KNOBS.get_or_init(|| Self {
183 prefetch: env_knob::<u8>("WARM_QGRAPH_PREFETCH").unwrap_or(0).min(3),
184 batch_est: env_knob::<u8>("WARM_QGRAPH_BATCH_EST").unwrap_or(0) != 0,
185 frontier: env_knob::<f32>("WARM_QGRAPH_FRONTIER")
186 .filter(|f| f.is_finite())
187 .map_or(1.0, |f| f.clamp(0.0, 1.0)),
188 block_screen: env_knob::<u8>("WARM_QGRAPH_BLOCK_SCREEN").unwrap_or(0) != 0,
189 rerank_z: env_knob::<f32>("WARM_QGRAPH_RERANK_Z")
190 .filter(|z| z.is_finite() && *z > 0.0)
191 .unwrap_or(0.0),
192 })
193 }
194}
195
196impl Default for SearchKnobs {
197 fn default() -> Self {
198 Self::OFF
199 }
200}
201
202fn env_knob<T: std::str::FromStr>(name: &str) -> Option<T> {
205 let value = crate::compatibility::var(name).ok()?;
206 match value.trim().parse::<T>() {
207 Ok(v) => Some(v),
208 Err(_) => {
209 eprintln!("{name}={value:?} is not a valid value; using the default");
210 None
211 }
212 }
213}
214
215fn huge_pages() -> bool {
223 static HUGE: OnceLock<bool> = OnceLock::new();
224 *HUGE.get_or_init(|| env_knob::<u8>("WARM_QGRAPH_HUGEPAGES").unwrap_or(1) != 0)
225}
226
227static HUGE_ADVISED_BYTES: AtomicU64 = AtomicU64::new(0);
229static HUGE_ADVICE_FAILED: AtomicU64 = AtomicU64::new(0);
230#[cfg(test)]
233static FAIL_HUGE_ADVICE: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false);
234
235#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Serialize)]
237pub struct HugePageStatus {
238 pub enabled: bool,
240 pub advised_bytes: u64,
242 pub failed_calls: u64,
244 pub anon_huge_bytes: Option<u64>,
249}
250
251pub fn huge_page_status() -> HugePageStatus {
253 let anon_huge_bytes = std::fs::read_to_string("/proc/self/smaps_rollup")
254 .ok()
255 .and_then(|text| {
256 let line = text.lines().find(|l| l.starts_with("AnonHugePages:"))?;
257 let kib: u64 = line.split_whitespace().nth(1)?.parse().ok()?;
258 Some(kib * 1024)
259 });
260 HugePageStatus {
261 enabled: huge_pages(),
262 advised_bytes: HUGE_ADVISED_BYTES.load(Ordering::Relaxed),
263 failed_calls: HUGE_ADVICE_FAILED.load(Ordering::Relaxed),
264 anon_huge_bytes,
265 }
266}
267
268fn huge_bytes_behind(ptr: *const u8, len: usize) -> Option<u64> {
272 let text = std::fs::read_to_string("/proc/self/smaps").ok()?;
273 let (lo, hi) = (ptr as u64, ptr as u64 + len as u64);
274 let mut total = 0u64;
275 let mut inside = false;
276 for line in text.lines() {
277 let head = line.split_whitespace().next().unwrap_or("");
278 if let Some((a, b)) = head.split_once('-') {
279 if let (Ok(a), Ok(b)) = (u64::from_str_radix(a, 16), u64::from_str_radix(b, 16)) {
280 inside = a < hi && lo < b;
281 continue;
282 }
283 }
284 if inside {
285 if let Some(rest) = line.strip_prefix("AnonHugePages:") {
286 let kib: u64 = rest.split_whitespace().next()?.parse().ok()?;
287 total += kib * 1024;
288 }
289 }
290 }
291 Some(total)
292}
293
294fn advise_huge_pages(ptr: *const u8, len: usize) {
301 #[cfg(target_os = "linux")]
302 {
303 const HUGE: usize = 2 << 20;
304 const MADV_HUGEPAGE: i32 = 14;
305 extern "C" {
306 fn madvise(addr: *mut std::ffi::c_void, len: usize, advice: i32) -> i32;
307 }
308 let start = (ptr as usize).next_multiple_of(HUGE);
309 let end = (ptr as usize + len) & !(HUGE - 1);
310 if start < end {
311 #[cfg(test)]
312 if FAIL_HUGE_ADVICE.load(Ordering::Relaxed) {
313 HUGE_ADVICE_FAILED.fetch_add(1, Ordering::Relaxed);
314 return;
315 }
316 let rc = unsafe { madvise(start as *mut std::ffi::c_void, end - start, MADV_HUGEPAGE) };
319 if rc == 0 {
320 HUGE_ADVISED_BYTES.fetch_add((end - start) as u64, Ordering::Relaxed);
321 } else {
322 HUGE_ADVICE_FAILED.fetch_add(1, Ordering::Relaxed);
323 }
324 }
325 }
326 #[cfg(not(target_os = "linux"))]
327 let _ = (ptr, len);
328}
329
330fn huge_capacity<T>(ptr: *const T, capacity: usize) {
333 if huge_pages() {
334 advise_huge_pages(ptr.cast::<u8>(), capacity * std::mem::size_of::<T>());
335 }
336}
337
338#[repr(C, align(128))]
340#[derive(Clone, Copy)]
341struct Line([i8; 128]);
342
343struct Lines {
345 lines: Slab<Line>,
347 len: usize,
348}
349
350impl Lines {
351 fn zeroed(len: usize) -> Self {
352 let count = len.div_ceil(128);
353 let mut lines = Vec::with_capacity(count);
354 huge_capacity(lines.as_ptr(), count);
355 lines.resize(count, Line([0; 128]));
356 Self {
357 lines: lines.into(),
358 len,
359 }
360 }
361
362 #[inline(always)]
363 fn bytes(&self) -> &[i8] {
364 unsafe { std::slice::from_raw_parts(self.lines.as_ptr().cast::<i8>(), self.len) }
367 }
368
369 fn bytes_mut(&mut self) -> &mut [i8] {
370 let len = self.len;
371 let lines = self.lines.as_mut_slice();
372 unsafe { std::slice::from_raw_parts_mut(lines.as_mut_ptr().cast::<i8>(), len) }
374 }
375}
376
377#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Serialize)]
381pub struct VGraphMemoryBreakdown {
382 pub adjacency_ids_bytes: usize,
383 pub adjacency_offsets_bytes: usize,
384 pub traversal_codes_bytes: usize,
385 pub traversal_scales_bytes: usize,
386 pub residual_codes_bytes: usize,
387 pub residual_scales_bytes: usize,
388 pub mean_bytes: usize,
389 pub norms_bytes: usize,
390 pub original_ids_bytes: usize,
391 pub screen_bytes: usize,
392 pub neighbor_blocks_bytes: usize,
393 pub upper_layers_bytes: usize,
394}
395
396impl VGraphMemoryBreakdown {
397 pub fn total_bytes(&self) -> usize {
398 self.adjacency_ids_bytes
399 + self.adjacency_offsets_bytes
400 + self.traversal_codes_bytes
401 + self.traversal_scales_bytes
402 + self.residual_codes_bytes
403 + self.residual_scales_bytes
404 + self.mean_bytes
405 + self.norms_bytes
406 + self.original_ids_bytes
407 + self.screen_bytes
408 + self.neighbor_blocks_bytes
409 + self.upper_layers_bytes
410 }
411}
412
413pub struct VGraph {
414 n: usize,
415 dim: usize,
416 metric: VMetric,
417 padded: usize,
419 bits: u8,
421 cbytes: usize,
423 stride: usize,
426 degree: usize,
427 mean: Vec<f32>,
428 half_norms: Slab<f32>,
431 adj: Slab<u32>,
433 adj3: Slab<u8>,
435 adj_off: Vec<u32>,
437 adjv: Slab<u8>,
441 adjv_off: Slab<u32>,
442 codes: Lines,
443 scales: Slab<f32>,
444 res_bits: u8,
446 res_stride: usize,
448 residual: Slab<i8>,
449 res_scales: Slab<f32>,
450 entry: u32,
451 upper_ids: Vec<u32>,
452 layers: Vec<Vec<Vec<u32>>>,
453 ids: Slab<u32>,
455 screen: Slab<u64>,
459 screen_words: usize,
460 screen_rot: Option<Rotation>,
461 screen_sigma: f32,
464 block_bits: u8,
467 block_dim: usize,
469 block_m: usize,
471 blocks_per_vertex: usize,
473 blocks: Lines,
476 block_rot: Option<Rotation>,
477 block_kappa: f32,
480}
481
482const SCREEN_SEED: u64 = 0x5C12_EE00;
484const BLOCK_SEED: u64 = 0xB10C_5CA0;
486
487fn block_record_bytes(m: usize) -> usize {
489 16 * m + 8 * BLOCK
490}
491
492#[inline]
494fn block_level(nibbles: &[u8], bits: u8, i: usize) -> f32 {
495 match bits {
496 1 => {
497 if nibbles[i / 4] >> (i % 4) & 1 == 1 {
498 1.0
499 } else {
500 -1.0
501 }
502 }
503 2 => 2.0 * f32::from(nibbles[i / 2] >> (2 * (i % 2)) & 3) - 3.0,
504 _ => f32::from(nibbles[i]) - 7.0,
505 }
506}
507
508#[inline]
510fn le_f32(bytes: &[u8], i: usize) -> f32 {
511 f32::from_le_bytes([
512 bytes[4 * i],
513 bytes[4 * i + 1],
514 bytes[4 * i + 2],
515 bytes[4 * i + 3],
516 ])
517}
518
519fn encode_block_code(x: &[f32], bits: u8, nibbles: &mut [u8]) -> f32 {
529 nibbles.fill(0);
530 let m = x.iter().fold(0f32, |a, &v| a.max(v.abs()));
531 if m == 0.0 {
532 return 0.0;
533 }
534 match bits {
535 1 => {
536 let (mut norm2, mut abs) = (0f32, 0f32);
537 for (i, &v) in x.iter().enumerate() {
538 norm2 += v * v;
539 abs += v.abs();
540 if v >= 0.0 {
541 nibbles[i / 4] |= 1 << (i % 4);
542 }
543 }
544 norm2 / abs
545 }
546 2 => {
547 let err = |s: f32| -> f32 {
549 x.iter()
550 .map(|&v| {
551 let t = ((v / s + 3.0) / 2.0).round().clamp(0.0, 3.0);
552 let d = v - s * (2.0 * t - 3.0);
553 d * d
554 })
555 .sum()
556 };
557 let (mut best, mut best_err) = (m / 3.0, f32::INFINITY);
558 for k in 0..8 {
559 let s = m / 3.0 * (1.0 - 0.1 * k as f32);
560 let e = err(s);
561 if e < best_err {
562 (best, best_err) = (s, e);
563 }
564 }
565 let mut dot = 0f32;
566 for (i, &v) in x.iter().enumerate() {
567 let t = ((v / best + 3.0) / 2.0).round().clamp(0.0, 3.0) as u8;
568 nibbles[i / 2] |= t << (2 * (i % 2));
569 dot += v * (2.0 * f32::from(t) - 3.0);
570 }
571 unbiased_scale(x, dot, best)
572 }
573 _ => {
574 let err = |s: f32| -> f32 {
575 x.iter()
576 .map(|&v| {
577 let d = v - s * (v / s).round().clamp(-7.0, 7.0);
578 d * d
579 })
580 .sum()
581 };
582 let (mut best, mut best_err) = (m / 7.0, f32::INFINITY);
583 for k in 0..8 {
584 let s = m / 7.0 * (1.0 - 0.06 * k as f32);
585 let e = err(s);
586 if e < best_err {
587 (best, best_err) = (s, e);
588 }
589 }
590 let mut dot = 0f32;
591 for (i, &v) in x.iter().enumerate() {
592 let level = (v / best).round().clamp(-7.0, 7.0);
593 nibbles[i] = (level as i8 + 7) as u8;
594 dot += v * level;
595 }
596 unbiased_scale(x, dot, best)
597 }
598 }
599}
600
601fn unbiased_scale(x: &[f32], dot: f32, fallback: f32) -> f32 {
604 let norm2: f32 = x.iter().map(|v| v * v).sum();
605 if dot > 0.0 {
606 norm2 / dot
607 } else {
608 fallback
609 }
610}
611
612fn block_table(q: &[f32], bits: u8, m: usize, t: &mut [f32]) {
616 match bits {
617 1 => {
618 for j in 0..m {
619 let qs = &q[4 * j..4 * j + 4];
620 let total: f32 = qs.iter().sum();
621 for v in 0..16usize {
622 let on: f32 = (0..4).filter(|b| v >> b & 1 == 1).map(|b| qs[b]).sum();
623 t[16 * j + v] = 2.0 * on - total;
624 }
625 }
626 }
627 2 => {
628 for j in 0..m {
629 let (qa, qb) = (q[2 * j], q[2 * j + 1]);
630 for v in 0..16usize {
631 let (ta, tb) = ((v & 3) as f32, (v >> 2) as f32);
632 t[16 * j + v] = qa * (2.0 * ta - 3.0) + qb * (2.0 * tb - 3.0);
633 }
634 }
635 }
636 _ => {
637 for j in 0..m {
638 for v in 0..16usize {
639 t[16 * j + v] = q[j] * (v as f32 - 7.0);
640 }
641 }
642 }
643 }
644}
645
646pub struct BlockQuery {
650 table: Vec<f32>,
651 lut: Vec<u8>,
652 mins: Vec<f32>,
654 delta: f32,
655 min_sum: f32,
656 inv_sq: f32,
658}
659
660impl BlockQuery {
661 fn new(m: usize) -> Self {
662 Self {
663 table: vec![0.0; 16 * m],
664 lut: vec![0; 16 * m],
665 mins: vec![0.0; m],
666 delta: 1.0,
667 min_sum: 0.0,
668 inv_sq: 1.0,
669 }
670 }
671
672 fn quantize(&mut self, m: usize) {
675 debug_assert_eq!(self.lut.len(), 16 * m);
676 debug_assert_eq!(self.mins.len(), m);
677 let mut range = 0f32;
678 let mut min_sum = 0f32;
679 let groups = self.table.as_chunks::<16>().0;
680 for (g, lo) in groups.iter().zip(self.mins.iter_mut()) {
681 *lo = g.iter().fold(f32::INFINITY, |a, &v| a.min(v));
682 let hi = g.iter().fold(f32::NEG_INFINITY, |a, &v| a.max(v));
683 min_sum += *lo;
684 range = range.max(hi - *lo);
685 }
686 let delta = if range > 0.0 { range / 255.0 } else { 1.0 };
687 let luts = self.lut.as_chunks_mut::<16>().0;
688 for ((g, &lo), lut) in groups.iter().zip(&self.mins).zip(luts) {
689 for (o, &e) in lut.iter_mut().zip(g) {
690 *o = ((e - lo) / delta).round().clamp(0.0, 255.0) as u8;
691 }
692 }
693 self.delta = delta;
694 self.min_sum = min_sum;
695 }
696}
697
698#[inline]
700fn key_dist(k: u64) -> f32 {
701 let o = (k >> 32) as u32;
702 f32::from_bits(if o >> 31 == 1 { o & 0x7FFF_FFFF } else { !o })
703}
704
705pub(super) fn quantize(x: &[f32], code: &mut [i8]) -> f32 {
708 let m = x.iter().fold(0f32, |a, &v| a.max(v.abs()));
709 code.fill(0);
710 if m == 0.0 {
711 return 0.0;
712 }
713 let scale = m / 127.0;
714 for (c, &v) in code.iter_mut().zip(x) {
715 *c = (v / scale).round().clamp(-127.0, 127.0) as i8;
716 }
717 scale
718}
719
720pub(super) fn quantize4(x: &[f32], code: &mut [i8]) -> f32 {
723 let m = x.iter().fold(0f32, |a, &v| a.max(v.abs()));
724 code.fill(0);
725 if m == 0.0 {
726 return 0.0;
727 }
728 let scale = m / 7.0;
729 let q = |i: usize| {
730 x.get(i)
731 .map_or(0, |&v| (v / scale).round().clamp(-7.0, 7.0) as i8)
732 };
733 for (j, c) in code.iter_mut().enumerate() {
734 *c = ((q(2 * j) as u8 & 0x0F) | ((q(2 * j + 1) as u8) << 4)) as i8;
735 }
736 scale
737}
738
739pub(super) fn bfs_order(g: GraphView<'_>) -> Vec<u32> {
742 let mut seen = vec![false; g.n];
743 let mut order = Vec::with_capacity(g.n);
744 for s in std::iter::once(g.entry as usize).chain(0..g.n) {
745 if seen[s] {
746 continue;
747 }
748 seen[s] = true;
749 let mut head = order.len();
750 order.push(s as u32);
751 while head < order.len() {
752 let v = order[head] as usize;
753 head += 1;
754 for &u in g.neighbors(v) {
755 if !std::mem::replace(&mut seen[u as usize], true) {
756 order.push(u);
757 }
758 }
759 }
760 }
761 order
762}
763
764pub(super) fn code_at(c: &[i8], bits: u8, i: usize) -> i8 {
766 if bits == 8 {
767 c.get(i).copied().unwrap_or(0)
768 } else {
769 c.get(i / 2).map_or(0, |&b| {
770 if i.is_multiple_of(2) {
771 (((b as u8) << 4) as i8) >> 4
772 } else {
773 b >> 4
774 }
775 })
776 }
777}
778
779const NONE3: u32 = 0x00FF_FFFF;
781
782#[inline]
783fn put_varint(out: &mut Vec<u8>, mut v: u64) {
785 while v >= 0x80 {
786 out.push((v as u8) | 0x80);
787 v >>= 7;
788 }
789 out.push(v as u8);
790}
791
792#[inline]
794fn get_varint(bytes: &[u8]) -> (u64, usize) {
795 let mut v = 0u64;
796 for (i, &b) in bytes.iter().enumerate() {
797 v |= u64::from(b & 0x7f) << (7 * i);
798 if b & 0x80 == 0 {
799 return (v, i + 1);
800 }
801 }
802 (v, bytes.len())
803}
804
805fn zigzag(v: i64) -> u64 {
806 ((v << 1) ^ (v >> 63)) as u64
807}
808
809fn unzigzag(v: u64) -> i64 {
810 ((v >> 1) as i64) ^ -((v & 1) as i64)
811}
812
813fn prefetch_u8(bytes: &[u8]) {
814 prefetch_bytes(unsafe { std::slice::from_raw_parts(bytes.as_ptr().cast::<i8>(), bytes.len()) });
816}
817
818#[inline]
819fn prefetch_f32(x: &f32) {
820 prefetch_bytes(unsafe { std::slice::from_raw_parts((x as *const f32).cast::<i8>(), 4) });
822}
823
824#[inline]
825fn prefetch_ids(ids: &[u32]) {
826 prefetch_bytes(unsafe { std::slice::from_raw_parts(ids.as_ptr().cast::<i8>(), ids.len() * 4) });
828}
829
830#[inline]
832fn key(dist: f32, id: u32) -> u64 {
833 let b = dist.to_bits();
834 let o = if b >> 31 == 1 { !b } else { b | 0x8000_0000 };
835 (u64::from(o) << 32) | u64::from(id)
836}
837
838#[inline]
839fn key_id(k: u64) -> u32 {
840 k as u32
841}
842
843struct VertexLayout {
844 padded: usize,
845 cbytes: usize,
846 stride: usize,
847 res_stride: usize,
848}
849
850impl VertexLayout {
851 fn new(
852 n: usize,
853 dim: usize,
854 bits: u8,
855 residual_bits: u8,
856 align: bool,
857 id_bytes: u8,
858 ) -> Result<Self, WarmError> {
859 if !matches!(bits, 4 | 8) {
860 return Err(invalid(format!(
861 "vertex code bits must be 4 or 8, not {bits}"
862 )));
863 }
864 if !(id_bytes == 4 || id_bytes == 0 || (id_bytes == 3 && (n as u64) < u64::from(NONE3))) {
865 return Err(invalid(format!("id_bytes must be 4, 0 (delta varints), or 3 below 2^24 - 1 vertices (got {id_bytes} for {n} vertices)")));
866 }
867 if !matches!(residual_bits, 0 | 4 | 8) {
868 return Err(invalid(format!(
869 "residual code bits must be 0, 4 or 8, not {residual_bits}"
870 )));
871 }
872 let padded = dim.next_multiple_of(if bits == 4 { 32 } else { 16 });
873 let cbytes = if bits == 4 { padded / 2 } else { padded };
874 let stride = match (align, cbytes) {
875 (false, _) => cbytes,
876 (true, c) if c <= 64 => c.next_power_of_two(),
877 (true, c) => c.next_multiple_of(128),
878 };
879 let res_stride = match residual_bits {
880 8 => padded,
881 4 => padded / 2,
882 _ => 0,
883 };
884 Ok(Self {
885 padded,
886 cbytes,
887 stride,
888 res_stride,
889 })
890 }
891}
892
893impl VGraph {
894 pub fn from_graph(g: &Graph, residual: bool) -> Result<Self, WarmError> {
897 Self::from_graph_with(g, 8, if residual { 8 } else { 0 }, false, false, 4, false)
898 }
899
900 pub fn from_graph_with(
909 g: &Graph,
910 bits: u8,
911 residual_bits: u8,
912 align: bool,
913 reorder: bool,
914 id_bytes: u8,
915 screen: bool,
916 ) -> Result<Self, WarmError> {
917 Self::encode_graph(
918 g,
919 (g.dim, g.metric.into()),
920 (bits, residual_bits),
921 align,
922 reorder,
923 id_bytes,
924 screen,
925 )
926 }
927
928 pub fn from_lifted_ip_graph(
931 g: &Graph,
932 bits: u8,
933 residual_bits: u8,
934 align: bool,
935 reorder: bool,
936 id_bytes: u8,
937 ) -> Result<Self, WarmError> {
938 if g.metric != Metric::Euclidean || g.dim < 2 {
939 return Err(invalid(
940 "direct IP serving requires a lifted Euclidean graph".into(),
941 ));
942 }
943 Self::encode_graph(
944 g,
945 (g.dim - 1, VMetric::InnerProduct),
946 (bits, residual_bits),
947 align,
948 reorder,
949 id_bytes,
950 false,
951 )
952 }
953
954 fn encode_graph(
955 g: &Graph,
956 geometry: (usize, VMetric),
957 precision: (u8, u8),
958 align: bool,
959 reorder: bool,
960 id_bytes: u8,
961 screen: bool,
962 ) -> Result<Self, WarmError> {
963 let (dim, metric) = geometry;
964 let (bits, residual_bits) = precision;
965 let layout = VertexLayout::new(g.n, dim, bits, residual_bits, align, id_bytes)?;
966 let euclidean = metric == VMetric::Euclidean;
967 if metric != VMetric::Angular && screen {
968 return Err(invalid("the 1-bit screen is angular only".into()));
969 }
970 let n = g.n;
971 let VertexLayout {
972 padded,
973 cbytes,
974 stride,
975 res_stride,
976 } = layout;
977 let mut sum = vec![0f64; dim];
978 for x in g.vectors.chunks_exact(g.dim) {
979 sum.iter_mut().zip(x).for_each(|(s, &v)| *s += f64::from(v));
980 }
981 let mut mean: Vec<f32> = sum.iter().map(|s| (s / n as f64) as f32).collect();
982 mean.resize(padded, 0.0);
983
984 let ids = if reorder {
985 bfs_order(g.view())
986 } else {
987 Vec::new()
988 };
989 let old = |v: usize| if reorder { ids[v] as usize } else { v };
990 let mut new_id: Vec<u32> = Vec::new();
991 if reorder {
992 new_id = vec![0; n];
993 for (v, &o) in ids.iter().enumerate() {
994 new_id[o as usize] = v as u32;
995 }
996 }
997 let new = |o: u32| if reorder { new_id[o as usize] } else { o };
998 let half_norms: Vec<f32> = if euclidean {
999 let mut half_norms = Vec::with_capacity(n);
1000 huge_capacity(half_norms.as_ptr(), n);
1001 half_norms.extend((0..n).map(|v| 0.5 * g.norms[old(v)]));
1002 half_norms
1003 } else {
1004 Vec::new()
1005 };
1006
1007 let screen_rot = screen.then(|| Rotation::new(dim, 0, SCREEN_SEED));
1008 let screen_words = screen_rot.as_ref().map_or(0, |r| r.padded() / 64);
1009 let centred = |v: usize, out: &mut Vec<f32>| {
1012 let o = old(v);
1013 out.clear();
1014 out.extend(
1015 g.vectors[o * g.dim..o * g.dim + dim]
1016 .iter()
1017 .zip(&mean)
1018 .map(|(a, m)| a - m),
1019 );
1020 };
1021 let mut codes = Lines::zeroed(n * stride);
1023 let mut scales = Vec::with_capacity(n);
1024 huge_capacity(scales.as_ptr(), n);
1025 scales.resize(n, 0f32);
1026 codes
1027 .bytes_mut()
1028 .par_chunks_exact_mut(stride)
1029 .zip(scales.par_iter_mut())
1030 .enumerate()
1031 .for_each_init(Vec::<f32>::new, |x, (v, (dst, s))| {
1032 centred(v, x);
1033 *s = if bits == 8 {
1034 quantize(x, &mut dst[..cbytes])
1035 } else {
1036 quantize4(x, &mut dst[..cbytes])
1037 };
1038 });
1039 let mut residual = Vec::with_capacity(n * res_stride);
1040 huge_capacity(residual.as_ptr(), n * res_stride);
1041 residual.resize(n * res_stride, 0i8);
1042 let mut res_scales = vec![0f32; if residual_bits > 0 { n } else { 0 }];
1043 if residual_bits > 0 {
1044 let code_bytes = codes.bytes();
1045 residual
1046 .par_chunks_exact_mut(res_stride)
1047 .zip(res_scales.par_iter_mut())
1048 .enumerate()
1049 .for_each_init(Vec::<f32>::new, |x, (v, (r, sr))| {
1050 centred(v, x);
1051 let c = &code_bytes[v * stride..v * stride + cbytes];
1052 let s = scales[v];
1053 x.iter_mut()
1054 .enumerate()
1055 .for_each(|(i, a)| *a -= s * f32::from(code_at(c, bits, i)));
1056 *sr = if residual_bits == 8 {
1057 quantize(x, r)
1058 } else {
1059 quantize4(x, r)
1060 };
1061 });
1062 }
1063 let mut screen_rows = vec![0u64; if screen { n * (screen_words + 1) } else { 0 }];
1064 if let Some(rot) = &screen_rot {
1065 screen_rows
1066 .par_chunks_exact_mut(screen_words + 1)
1067 .enumerate()
1068 .for_each_init(
1069 || (Vec::<f32>::new(), vec![0f32; rot.padded()]),
1070 |(x, xr), (v, srow)| {
1071 centred(v, x);
1072 rot.apply(x, xr);
1073 let (mut norm2, mut abs) = (0f32, 0f32);
1074 let mut pop = 0u32;
1075 for (i, &a) in xr.iter().enumerate() {
1076 norm2 += a * a;
1077 abs += a.abs();
1078 if a >= 0.0 {
1079 srow[i / 64] |= 1u64 << (i % 64);
1080 pop += 1;
1081 }
1082 }
1083 let f = if abs > 0.0 { norm2 / abs } else { 0.0 };
1085 srow[screen_words] = u64::from(f.to_bits()) | (u64::from(pop) << 32);
1086 },
1087 );
1088 }
1089
1090 let r = g.degree;
1091 let list = |v: usize| g.neighbors(old(v)).iter().take(r).map(|&u| new(u));
1093 let (mut adj, mut adj3) = (Vec::new(), Vec::new());
1094 let (mut adjv, mut adjv_off) = (Vec::new(), Vec::new());
1095 if id_bytes == 3 {
1096 adj3 = Vec::<u8>::with_capacity(3 * n * r);
1097 huge_capacity(adj3.as_ptr(), adj3.capacity());
1098 for v in 0..n {
1099 let mut count = 0;
1100 for u in list(v) {
1101 adj3.extend_from_slice(&u.min(NONE3).to_le_bytes()[..3]);
1102 count += 1;
1103 }
1104 for _ in count..r {
1105 adj3.extend_from_slice(&NONE3.to_le_bytes()[..3]);
1106 }
1107 }
1108 } else if id_bytes == 4 {
1109 adj = vec![NONE; n * r];
1110 for (v, dst) in adj.chunks_exact_mut(r).enumerate() {
1111 for (d, u) in dst.iter_mut().zip(list(v)) {
1112 *d = u;
1113 }
1114 }
1115 } else {
1116 adjv_off.reserve(n + 1);
1117 let mut sorted: Vec<u32> = Vec::with_capacity(r);
1118 for v in 0..n {
1119 adjv_off.push(adjv.len() as u32);
1120 sorted.clear();
1121 sorted.extend(list(v).take_while(|&u| u != NONE));
1122 sorted.sort_unstable();
1123 let list = &sorted;
1124 let mut prev = v as u32;
1125 for (i, &u) in list.iter().enumerate() {
1126 let delta = if i == 0 {
1127 zigzag(i64::from(u) - i64::from(prev))
1128 } else {
1129 u64::from(u - prev - 1)
1130 };
1131 put_varint(&mut adjv, delta);
1132 prev = u;
1133 }
1134 }
1135 adjv_off.push(adjv.len() as u32);
1136 }
1137 let entry = new(g.entry);
1138 let upper_ids = g.upper_ids.iter().map(|&u| new(u)).collect();
1139 let mut out = Self {
1140 n,
1141 dim,
1142 metric,
1143 padded,
1144 bits,
1145 cbytes,
1146 stride,
1147 degree: r,
1148 mean,
1149 half_norms: half_norms.into(),
1150 adj: adj.into(),
1151 adj3: adj3.into(),
1152 adj_off: Vec::new(),
1153 adjv: adjv.into(),
1154 adjv_off: adjv_off.into(),
1155 codes,
1156 scales: scales.into(),
1157 res_bits: residual_bits,
1158 res_stride,
1159 residual: residual.into(),
1160 res_scales: res_scales.into(),
1161 entry,
1162 upper_ids,
1163 layers: g.layers.clone(),
1164 ids: ids.into(),
1165 screen: screen_rows.into(),
1166 screen_words,
1167 screen_rot,
1168 screen_sigma: 0.0,
1169 block_bits: 0,
1170 block_dim: 0,
1171 block_m: 0,
1172 blocks_per_vertex: 0,
1173 blocks: Lines::zeroed(0),
1174 block_rot: None,
1175 block_kappa: 0.0,
1176 };
1177 if screen {
1178 out.screen_sigma = out.measure_screen_sigma(g);
1179 }
1180 Ok(out)
1181 }
1182
1183 pub fn with_blocks(mut self, g: &Graph, bits: u8) -> Result<Self, WarmError> {
1189 if !matches!(bits, 1 | 2 | 4) {
1190 return Err(invalid(format!("block bits must be 1, 2 or 4, not {bits}")));
1191 }
1192 if self.metric != VMetric::Angular {
1193 return Err(invalid("neighbour blocks are angular only".into()));
1194 }
1195 let (n, dim, r) = (self.n, self.dim, self.degree);
1196 let rot = Rotation::new(dim, 0, BLOCK_SEED);
1197 let bd = rot.padded();
1198 let m = bd * usize::from(bits) / 4;
1199 let per_vertex = r.div_ceil(BLOCK);
1200 let record = block_record_bytes(m);
1201 let ids = &self.ids;
1202 let old = |v: usize| if ids.is_empty() { v } else { ids[v] as usize };
1203 let row = |v: usize| &g.vectors[old(v) * dim..(old(v) + 1) * dim];
1204 let mut blocks = Lines::zeroed(n * per_vertex * record);
1205 blocks
1206 .bytes_mut()
1207 .par_chunks_exact_mut(per_vertex * record)
1208 .enumerate()
1209 .for_each_init(
1210 || {
1211 let list = Vec::<u32>::with_capacity(r);
1212 (
1213 vec![0f32; bd],
1214 vec![0f32; dim],
1215 vec![0f32; bd],
1216 vec![0u8; BLOCK * m],
1217 list,
1218 )
1219 },
1220 |(xr, diff, rr, nibbles, list), (v, dst)| {
1221 let xv = row(v);
1230 rot.apply(xv, xr);
1231 let mut scales = [0f32; BLOCK];
1232 let mut offsets = [0f32; BLOCK];
1233 list.clear();
1234 self.for_each_neighbour(v, |u| list.push(u));
1235 for (b, rec) in dst.chunks_exact_mut(record).enumerate() {
1236 nibbles.fill(0);
1237 scales.fill(0.0);
1238 offsets.fill(0.0);
1239 for i in 0..BLOCK {
1240 let Some(&u) = list.get(b * BLOCK + i) else {
1241 break;
1242 };
1243 let xu = row(u as usize);
1244 for (d, (a, b)) in diff.iter_mut().zip(xu.iter().zip(xv)) {
1245 *d = a - b;
1246 }
1247 rot.apply(diff, rr);
1248 let nib = &mut nibbles[i * m..(i + 1) * m];
1249 let s = encode_block_code(rr, bits, nib);
1250 let dot_lv: f32 =
1251 (0..bd).map(|j| block_level(nib, bits, j) * xr[j]).sum();
1252 let dot_vr: f32 = xv.iter().zip(diff.iter()).map(|(a, b)| a * b).sum();
1253 scales[i] = s;
1254 offsets[i] = dot_vr - s * dot_lv;
1255 }
1256 let rec_u8: &mut [u8] = bytemuck::cast_slice_mut(rec);
1257 pack32(nibbles, m, &mut rec_u8[..16 * m]);
1258 for (i, (s, e)) in scales.iter().zip(&offsets).enumerate() {
1259 let at = 16 * m + 4 * i;
1260 rec_u8[at..at + 4].copy_from_slice(&s.to_le_bytes());
1261 let at = 16 * m + 4 * BLOCK + 4 * i;
1262 rec_u8[at..at + 4].copy_from_slice(&e.to_le_bytes());
1263 }
1264 }
1265 },
1266 );
1267 self.block_bits = bits;
1268 self.block_dim = bd;
1269 self.block_m = m;
1270 self.blocks_per_vertex = per_vertex;
1271 self.blocks = blocks;
1272 self.block_rot = Some(rot);
1273 self.block_kappa = self.measure_block_kappa(g);
1274 Ok(self)
1275 }
1276
1277 fn unit_row<'g>(&self, g: &'g Graph, v: usize) -> &'g [f32] {
1279 let o = if self.ids.is_empty() {
1280 v
1281 } else {
1282 self.ids[v] as usize
1283 };
1284 &g.vectors[o * self.dim..(o + 1) * self.dim]
1285 }
1286
1287 fn measure_block_kappa(&self, g: &Graph) -> f32 {
1291 let (nq, nv) = (32.min(self.n), 64.min(self.n));
1292 let mut rng = SplitMix(BLOCK_SEED ^ 1);
1293 let mut s = self.searcher();
1294 let (mut num, mut den) = (0f64, 0f64);
1295 let (mut diag_sum2, mut diag_n) = (0f64, 0f64);
1296 let mut scan = [0u32; BLOCK];
1297 for _ in 0..nq {
1298 let qi = (rng.next_u64() % self.n as u64) as usize;
1299 let q = self.unit_row(g, qi).to_vec();
1300 s.scratch.qf[..self.dim].copy_from_slice(&q);
1301 let sq = self.encode_query(
1302 &s.scratch.qf,
1303 &mut s.scratch.qc,
1304 &mut s.scratch.qtmp,
1305 &mut s.scratch.qrot,
1306 &mut s.scratch.qs,
1307 );
1308 if sq == 0.0 {
1309 continue;
1310 }
1311 self.encode_block_query(&s.scratch.qf, sq, &mut s.scratch.bqrot, &mut s.scratch.bq);
1312 let centred = |x: &[f32]| -> f32 {
1313 q.iter()
1314 .zip(x)
1315 .zip(&self.mean)
1316 .map(|((a, b), m)| a * (b - m))
1317 .sum()
1318 };
1319 let mut list = Vec::with_capacity(self.degree);
1320 for _ in 0..nv {
1321 let v = (rng.next_u64() % self.n as u64) as usize;
1322 let xv = self.unit_row(g, v);
1323 let ip_v = centred(xv);
1324 let d2: f32 = q.iter().zip(xv).map(|(a, b)| (a - b) * (a - b)).sum();
1325 list.clear();
1326 self.for_each_neighbour(v, |u| list.push(u));
1327 for b in 0..self.blocks_per_vertex {
1328 let (codes, scales, offsets) = self.block(v, b);
1329 scan32(codes, &s.scratch.bq.lut, self.block_m, &mut scan);
1330 for (i, &sum) in scan.iter().enumerate() {
1331 let Some(&u) = list.get(b * BLOCK + i) else {
1332 break;
1333 };
1334 let sc = le_f32(scales, i);
1335 let est = ip_v
1336 + sc * (s.scratch.bq.delta * sum as f32 + s.scratch.bq.min_sum)
1337 + le_f32(offsets, i);
1338 let exact = centred(self.unit_row(g, u as usize));
1339 let err = f64::from(est - exact);
1340 num += err * err;
1341 den += f64::from(sc * sc * d2);
1342 diag_sum2 += f64::from(exact - ip_v).powi(2);
1343 diag_n += 1.0;
1344 }
1345 }
1346 }
1347 }
1348 if crate::compatibility::var_os("WARM_BLOCK_DIAG").is_some() && diag_n > 0.0 {
1349 eprintln!(
1350 "block diag: bits {} kappa {:.4} rms_err {:.5} rms_residual_ip {:.5} pairs {}",
1351 self.block_bits,
1352 num / den.max(1e-30),
1353 (num / diag_n).sqrt(),
1354 (diag_sum2 / diag_n).sqrt(),
1355 diag_n
1356 );
1357 }
1358 if den > 0.0 {
1359 (num / den) as f32
1360 } else {
1361 0.0
1362 }
1363 }
1364
1365 pub fn has_blocks(&self) -> bool {
1367 self.block_bits > 0
1368 }
1369
1370 pub fn block_bytes(&self) -> usize {
1372 self.blocks.len
1373 }
1374
1375 #[inline]
1378 fn block(&self, v: usize, b: usize) -> (&[u8], &[u8], &[u8]) {
1379 let record = block_record_bytes(self.block_m);
1380 let at = (v * self.blocks_per_vertex + b) * record;
1381 let bytes = self.blocks.bytes();
1382 let rec =
1384 unsafe { std::slice::from_raw_parts(bytes.as_ptr().add(at).cast::<u8>(), record) };
1385 let (codes, rest) = rec.split_at(16 * self.block_m);
1386 let (scales, offsets) = rest.split_at(4 * BLOCK);
1387 (codes, scales, offsets)
1388 }
1389
1390 #[inline]
1391 fn prefetch_blocks(&self, v: usize) {
1392 let record = block_record_bytes(self.block_m) * self.blocks_per_vertex;
1393 let at = v * record;
1394 prefetch_bytes(&self.blocks.bytes()[at..at + record]);
1395 }
1396
1397 fn encode_block_query(&self, qf: &[f32], sq: f32, qrot: &mut [f32], bq: &mut BlockQuery) {
1399 let rot = self.block_rot.as_ref().expect("blocks present");
1400 rot.apply(&qf[..self.dim], qrot);
1401 block_table(qrot, self.block_bits, self.block_m, &mut bq.table);
1402 bq.quantize(self.block_m);
1403 bq.inv_sq = if sq > 0.0 { 1.0 / sq } else { 1.0 };
1404 }
1405
1406 fn measure_screen_sigma(&self, g: &Graph) -> f32 {
1409 let (nq, nv) = (128.min(self.n), 1024.min(self.n));
1410 let mut rng = SplitMix(SCREEN_SEED ^ 1);
1411 let mut qc = vec![0i8; self.padded];
1412 let mut qtmp = vec![0i8; self.padded];
1413 let mut qf = vec![0f32; self.padded];
1414 let mut qrot = vec![0f32; self.screen_words * 64];
1415 let mut qs = QueryCode::new(self.screen_words);
1416 let (mut sum, mut sum2, mut count) = (0f64, 0f64, 0usize);
1417 for _ in 0..nq {
1418 let q = (rng.next_u64() % self.n as u64) as usize;
1419 qf[..self.dim].copy_from_slice(&g.vectors[q * self.dim..(q + 1) * self.dim]);
1420 let sq = self.encode_query(&qf, &mut qc, &mut qtmp, &mut qrot, &mut qs);
1421 if sq == 0.0 {
1422 continue;
1423 }
1424 for _ in 0..nv {
1425 let v = (rng.next_u64() % self.n as u64) as u32;
1426 let d = f64::from(self.screen_est(&qs, v) - self.est(&qc, query_sum(&qc), v, 1.0));
1427 sum += d;
1428 sum2 += d * d;
1429 count += 1;
1430 }
1431 }
1432 if count == 0 {
1433 return 0.0;
1434 }
1435 let mean = sum / count as f64;
1436 ((sum2 / count as f64) - mean * mean).max(0.0).sqrt() as f32
1437 }
1438
1439 fn encode_query(
1442 &self,
1443 qf: &[f32],
1444 qc: &mut [i8],
1445 qtmp: &mut [i8],
1446 qrot: &mut [f32],
1447 qs: &mut QueryCode,
1448 ) -> f32 {
1449 let sq = if self.bits == 8 {
1450 quantize(qf, qc)
1451 } else {
1452 let s = quantize(qf, qtmp);
1453 interleave_for_i4(qtmp, qc);
1454 s
1455 };
1456 if let Some(rot) = &self.screen_rot {
1457 rot.apply(&qf[..self.dim], qrot);
1458 if sq > 0.0 {
1459 qrot.iter_mut().for_each(|v| *v /= sq);
1462 }
1463 qs.encode(qrot);
1464 }
1465 sq
1466 }
1467
1468 pub fn has_screen(&self) -> bool {
1470 !self.screen.is_empty()
1471 }
1472
1473 pub fn screen_bytes(&self) -> usize {
1475 self.screen.len() * 8
1476 }
1477
1478 #[inline]
1481 fn screen_est(&self, qs: &QueryCode, v: u32) -> f32 {
1482 let w = self.screen_words;
1483 let row = &self.screen[v as usize * (w + 1)..(v as usize + 1) * (w + 1)];
1484 let raw = raw_one(&row[..w], &qs.planes, w) as f32;
1485 let meta = row[w];
1486 let f = f32::from_bits(meta as u32);
1487 let pop = (meta >> 32) as f32;
1488 -(f * (2.0 * (qs.lo * pop + qs.delta * raw) - qs.sum))
1489 }
1490
1491 #[inline]
1492 fn prefetch_screen(&self, v: u32) {
1493 let w = self.screen_words;
1494 prefetch(&self.screen[v as usize * (w + 1)..(v as usize + 1) * (w + 1)]);
1495 }
1496
1497 pub fn len(&self) -> usize {
1498 self.n
1499 }
1500
1501 pub fn is_empty(&self) -> bool {
1502 self.n == 0
1503 }
1504
1505 pub fn code_rows_huge_bytes(&self) -> Option<u64> {
1508 let rows = self.codes.bytes();
1509 huge_bytes_behind(rows.as_ptr().cast(), rows.len())
1510 }
1511
1512 pub fn has_residual(&self) -> bool {
1513 self.res_bits > 0
1514 }
1515
1516 pub fn code_bits(&self) -> u8 {
1517 self.bits
1518 }
1519
1520 pub fn residual_bits(&self) -> u8 {
1521 self.res_bits
1522 }
1523
1524 pub fn index_bytes(&self) -> usize {
1527 self.memory_breakdown().total_bytes()
1528 }
1529
1530 pub fn memory_breakdown(&self) -> VGraphMemoryBreakdown {
1531 VGraphMemoryBreakdown {
1532 adjacency_ids_bytes: self.adj.len() * 4 + self.adj3.len() + self.adjv.len(),
1533 adjacency_offsets_bytes: (self.adj_off.len() + self.adjv_off.len()) * 4,
1534 traversal_codes_bytes: self.n * self.stride,
1535 traversal_scales_bytes: self.scales.len() * 4,
1536 residual_codes_bytes: self.residual.len(),
1537 residual_scales_bytes: self.res_scales.len() * 4,
1538 mean_bytes: self.mean.len() * 4,
1539 norms_bytes: self.half_norms.len() * 4,
1540 original_ids_bytes: self.ids.len() * 4,
1541 screen_bytes: self.screen_bytes(),
1542 neighbor_blocks_bytes: self.block_bytes(),
1543 upper_layers_bytes: upper_bytes(&self.upper_ids, &self.layers),
1544 }
1545 }
1546
1547 pub fn residual_bytes(&self) -> usize {
1549 self.residual.len() + self.res_scales.len() * 4
1550 }
1551
1552 #[inline]
1553 fn code(&self, v: u32) -> &[i8] {
1554 let at = v as usize * self.stride;
1555 &self.codes.bytes()[at..at + self.cbytes]
1556 }
1557
1558 pub(super) fn neighbours_into(&self, v: usize, out: &mut Vec<u32>) {
1560 self.for_each_neighbour(v, |u| out.push(u));
1561 }
1562
1563 pub(super) fn entry_vertex(&self) -> u32 {
1565 self.entry
1566 }
1567
1568 #[inline]
1570 fn for_each_neighbour(&self, v: usize, mut f: impl FnMut(u32)) {
1571 if !self.adjv_off.is_empty() {
1572 return self.for_each_delta_neighbour(v, f);
1573 }
1574 let (start, end) = self.adjacency_span(v);
1575 if self.adj3.is_empty() {
1576 for &u in &self.adj[start..end] {
1577 if u == NONE {
1578 break;
1579 }
1580 f(u);
1581 }
1582 } else {
1583 for b in self.adj3[3 * start..3 * end].as_chunks::<3>().0 {
1584 let u = u32::from_le_bytes([b[0], b[1], b[2], 0]);
1585 if u == NONE3 {
1586 break;
1587 }
1588 f(u);
1589 }
1590 }
1591 }
1592
1593 #[cfg(feature = "experimental-neighbor-admission")]
1594 #[inline(always)]
1595 fn admit_neighbour<const SCREEN: bool>(
1596 &self,
1597 u: u32,
1598 visited: &mut Marks,
1599 fresh: &mut Vec<u32>,
1600 ) {
1601 if visited.insert(u) {
1602 if SCREEN {
1603 self.prefetch_screen(u);
1604 } else {
1605 prefetch_bytes(self.code(u));
1606 self.prefetch_scale(u);
1607 }
1608 fresh.push(u);
1609 }
1610 }
1611
1612 #[cfg(feature = "experimental-neighbor-admission")]
1613 #[inline(always)]
1614 fn admit_neighbours<const PROFILE: bool, const SCREEN: bool>(
1615 &self,
1616 v: usize,
1617 cap: usize,
1618 visited: &mut Marks,
1619 fresh: &mut Vec<u32>,
1620 ) -> usize {
1621 let mut neighbors_read = 0;
1622 if cap == 0 {
1623 self.for_each_neighbour(v, |u| {
1624 if PROFILE {
1625 neighbors_read += 1;
1626 }
1627 self.admit_neighbour::<SCREEN>(u, visited, fresh);
1628 });
1629 } else {
1630 let mut seen = 0usize;
1631 self.for_each_neighbour(v, |u| {
1632 seen += 1;
1633 if seen <= cap {
1634 self.admit_neighbour::<SCREEN>(u, visited, fresh);
1635 }
1636 });
1637 }
1638 neighbors_read
1639 }
1640
1641 #[inline]
1642 fn adjacency_span(&self, v: usize) -> (usize, usize) {
1643 if self.adj_off.is_empty() {
1644 (v * self.degree, (v + 1) * self.degree)
1645 } else {
1646 (self.adj_off[v] as usize, self.adj_off[v + 1] as usize)
1647 }
1648 }
1649
1650 pub fn compact_adjacency(&mut self) -> Result<(), WarmError> {
1654 if !self.adj_off.is_empty() || !self.adjv_off.is_empty() {
1655 return Ok(());
1656 }
1657 let mut offsets = Vec::new();
1658 offsets
1659 .try_reserve_exact(self.n + 1)
1660 .map_err(|_| invalid("cannot allocate adjacency offsets".into()))?;
1661 offsets.push(0u32);
1662 for v in 0..self.n {
1663 let mut count = 0u32;
1664 self.for_each_neighbour(v, |_| count += 1);
1665 let end = offsets[v]
1666 .checked_add(count)
1667 .ok_or_else(|| invalid("compact adjacency exceeds u32 offsets".into()))?;
1668 offsets.push(end);
1669 }
1670 let edges = offsets[self.n] as usize;
1671 let width = if self.adj3.is_empty() { 4 } else { 3 };
1672 if (self.n * self.degree - edges).saturating_mul(width) <= offsets.len() * 4 {
1673 return Ok(());
1674 }
1675 if self.adj3.is_empty() {
1676 let mut compact = Vec::new();
1677 compact
1678 .try_reserve_exact(edges)
1679 .map_err(|_| invalid("cannot allocate compact adjacency".into()))?;
1680 for v in 0..self.n {
1681 self.for_each_neighbour(v, |u| compact.push(u));
1682 }
1683 self.adj = compact.into();
1684 } else {
1685 let bytes = edges
1686 .checked_mul(3)
1687 .ok_or_else(|| invalid("compact adjacency size overflows".into()))?;
1688 let mut compact = Vec::new();
1689 compact
1690 .try_reserve_exact(bytes)
1691 .map_err(|_| invalid("cannot allocate compact adjacency".into()))?;
1692 for v in 0..self.n {
1693 self.for_each_neighbour(v, |u| compact.extend_from_slice(&u.to_le_bytes()[..3]));
1694 }
1695 self.adj3 = compact.into();
1696 }
1697 self.adj_off = offsets;
1698 Ok(())
1699 }
1700
1701 #[cold]
1703 #[inline(never)]
1704 fn for_each_delta_neighbour(&self, v: usize, mut f: impl FnMut(u32)) {
1705 let span = &self.adjv[self.adjv_off[v] as usize..self.adjv_off[v + 1] as usize];
1706 let mut at = 0;
1707 let mut prev = v as u32;
1708 let mut first = true;
1709 while at < span.len() {
1710 let (delta, used) = get_varint(&span[at..]);
1711 at += used;
1712 let u = if first {
1713 (i64::from(prev) + unzigzag(delta)) as u32
1714 } else {
1715 prev + delta as u32 + 1
1716 };
1717 first = false;
1718 prev = u;
1719 f(u);
1720 }
1721 }
1722
1723 #[inline]
1724 fn prefetch_neighbours(&self, v: usize) {
1725 let (start, end) = self.adjacency_span(v);
1726 if !self.adjv_off.is_empty() {
1727 prefetch_u8(&self.adjv[self.adjv_off[v] as usize..self.adjv_off[v + 1] as usize]);
1728 } else if self.adj3.is_empty() {
1729 prefetch_ids(&self.adj[start..end]);
1730 } else {
1731 prefetch_u8(&self.adj3[3 * start..3 * end]);
1732 }
1733 }
1734
1735 #[inline]
1738 fn prefetch_scale(&self, v: u32) {
1739 prefetch_f32(&self.scales[v as usize]);
1740 if self.metric == VMetric::Euclidean {
1741 prefetch_f32(&self.half_norms[v as usize]);
1742 }
1743 }
1744
1745 #[inline]
1747 fn prefetch_rerank_rows(&self, v: u32) {
1748 prefetch_bytes(self.code(v));
1749 self.prefetch_scale(v);
1750 if self.has_residual() {
1751 let v = v as usize;
1752 prefetch_bytes(&self.residual[v * self.res_stride..(v + 1) * self.res_stride]);
1753 prefetch_f32(&self.res_scales[v]);
1754 }
1755 }
1756
1757 #[inline]
1758 fn original(&self, v: u32) -> u32 {
1759 if self.ids.is_empty() {
1760 v
1761 } else {
1762 self.ids[v as usize]
1763 }
1764 }
1765
1766 #[inline]
1771 fn est(&self, qc: &[i8], qsum: i32, v: u32, w: f32) -> f32 {
1772 let dot = if self.bits == 8 {
1773 dot_i8_query(qc, self.code(v), qsum)
1774 } else {
1775 dot_i4_query(qc, self.code(v), qsum)
1776 };
1777 self.order(v, self.scales[v as usize] * dot as f32, w)
1778 }
1779
1780 #[inline]
1783 fn order(&self, v: u32, ip: f32, w: f32) -> f32 {
1784 if self.metric == VMetric::Euclidean {
1785 self.half_norms[v as usize] * w - ip
1786 } else {
1787 -ip
1788 }
1789 }
1790
1791 #[inline]
1794 fn est4<const QUERY_X4: bool>(&self, qc: &[i8], qsum: i32, v: [u32; 4], w: f32) -> [f32; 4] {
1795 let dots = if self.bits == 8 {
1796 if QUERY_X4 {
1797 dot_i8_query_x4(qc, v.map(|u| self.code(u)), qsum)
1798 } else {
1799 v.map(|u| dot_i8_query(qc, self.code(u), qsum))
1800 }
1801 } else {
1802 dot_i4_query_x4(qc, v.map(|u| self.code(u)), qsum)
1803 };
1804 std::array::from_fn(|i| self.order(v[i], self.scales[v[i] as usize] * dots[i] as f32, w))
1805 }
1806
1807 fn fine(&self, qf: &[f32], v: u32) -> f32 {
1810 let primary = if self.bits == 8 {
1811 dot_f32_i8(qf, self.code(v))
1812 } else {
1813 dot_f32_i4(qf, self.code(v))
1814 };
1815 let mut ip = self.scales[v as usize] * primary;
1816 if self.has_residual() {
1817 let v = v as usize;
1818 let r = &self.residual[v * self.res_stride..(v + 1) * self.res_stride];
1819 let d = if self.res_bits == 8 {
1820 dot_f32_i8(qf, r)
1821 } else {
1822 dot_f32_i4(qf, r)
1823 };
1824 ip += self.res_scales[v] * d;
1825 }
1826 self.order(v, ip, 1.0)
1827 }
1828
1829 fn descend(&self, qc: &[i8], qsum: i32, w: f32, reservoir: &mut Reservoir) -> (u32, usize) {
1832 if self.layers.is_empty() {
1833 return (self.entry, 0);
1834 }
1835 let mut p = 0u32;
1836 let mut best = self.est(qc, qsum, self.upper_ids[0], w);
1837 reservoir.offer(key(best, self.upper_ids[0]));
1838 let mut evals = 1;
1839 for layer in self.layers.iter().rev() {
1840 loop {
1841 let list = &layer[p as usize];
1844 for &u in list {
1845 let id = self.upper_ids[u as usize];
1846 prefetch_bytes(self.code(id));
1847 self.prefetch_scale(id);
1848 }
1849 let mut moved = false;
1850 for &u in list {
1851 let id = self.upper_ids[u as usize];
1852 let d = self.est(qc, qsum, id, w);
1853 reservoir.offer(key(d, id));
1854 evals += 1;
1855 if d < best {
1856 best = d;
1857 p = u;
1858 moved = true;
1859 }
1860 }
1861 if !moved {
1862 break;
1863 }
1864 }
1865 }
1866 (self.upper_ids[p as usize], evals)
1867 }
1868
1869 pub fn searcher(&self) -> VSearcher<'_> {
1871 self.searcher_with(SearchKnobs::from_env())
1872 }
1873
1874 pub fn searcher_with(&self, knobs: SearchKnobs) -> VSearcher<'_> {
1876 VSearcher {
1877 index: self,
1878 scratch: SearchScratch {
1879 #[cfg(not(target_arch = "wasm32"))]
1880 dim: self.dim,
1881 knobs,
1882 visited: Marks::new(self.n, self.has_blocks()),
1883 qf: vec![0.0; self.padded],
1884 qc: vec![0; self.padded],
1885 qsum: 0,
1886 qtmp: vec![0; self.padded],
1887 qrot: vec![0.0; self.screen_words * 64],
1888 qs: QueryCode::new(self.screen_words.max(1)),
1889 bqrot: vec![0.0; self.block_dim],
1890 bq: BlockQuery::new(self.block_m),
1891 scan: [0; BLOCK],
1892 fresh: Vec::with_capacity(self.degree),
1893 survivors: Vec::with_capacity(self.degree),
1894 screened: 0,
1895 pool: Pool {
1896 keys: Vec::new(),
1897 expanded: Vec::new(),
1898 #[cfg(feature = "experimental-batch-pool")]
1899 merge_keys: Vec::new(),
1900 #[cfg(feature = "experimental-batch-pool")]
1901 merge_expanded: Vec::new(),
1902 cap: 1,
1903 cursor: 0,
1904 },
1905 expanded: BinaryHeap::new(),
1906 reservoir: Reservoir::default(),
1907 keys: Vec::new(),
1908 results: Vec::new(),
1909 },
1910 }
1911 }
1912
1913 #[cfg(not(target_arch = "wasm32"))]
1914 pub(super) fn update_encoding(&self) -> super::update::Encoding {
1915 super::update::Encoding {
1916 bits: self.bits,
1917 residual_bits: self.res_bits,
1918 align: self.stride != self.cbytes,
1919 reorder: !self.ids.is_empty(),
1920 id_bytes: if !self.adjv_off.is_empty() {
1921 0
1922 } else if !self.adj3.is_empty() {
1923 3
1924 } else {
1925 4
1926 },
1927 screen: self.has_screen(),
1928 block_bits: self.block_bits,
1929 }
1930 }
1931
1932 #[cfg(not(target_arch = "wasm32"))]
1933 pub(super) fn search_reusing(
1934 &self,
1935 cache: &mut Option<SearchScratch>,
1936 query: &[f32],
1937 k: usize,
1938 params: VSearchParams,
1939 out: &mut Vec<u32>,
1940 ) -> Result<usize, WarmError> {
1941 let compatible = cache.as_ref().is_some_and(|s| {
1942 s.dim == self.dim
1943 && s.visited.n == self.n
1944 && s.qf.len() == self.padded
1945 && s.qrot.len() == self.screen_words * 64
1946 && s.bqrot.len() == self.block_dim
1947 && s.bq.table.len() == 16 * self.block_m
1948 && s.visited.traces.len() == if self.has_blocks() { self.n } else { 0 }
1949 });
1950 if !compatible {
1951 *cache = Some(self.searcher().scratch);
1952 }
1953 cache
1954 .as_mut()
1955 .expect("initialized search scratch")
1956 .search_validated(self, query, k, params, out)
1957 }
1958}
1959
1960struct Marks {
1963 words: Vec<u64>,
1967 touched: Vec<u32>,
1968 n: usize,
1969 epoch: u8,
1970 traces: Vec<Trace>,
1972}
1973
1974#[derive(Clone, Copy, Default)]
1978struct Trace {
1979 epoch: u8,
1980 expanded: bool,
1981 w: f32,
1982 wsum: f32,
1983 key: u64,
1985}
1986
1987impl Marks {
1988 fn new(n: usize, traces: bool) -> Self {
1990 Self {
1991 words: vec![0; n.div_ceil(64)],
1992 touched: Vec::new(),
1993 n,
1994 epoch: 0,
1995 traces: vec![Trace::default(); if traces { n } else { 0 }],
1996 }
1997 }
1998
1999 fn next(&mut self) {
2000 for &w in &self.touched {
2001 self.words[w as usize] = 0;
2002 }
2003 self.touched.clear();
2004 self.epoch = self.epoch.wrapping_add(1);
2005 if self.epoch == 0 {
2006 self.traces.fill(Trace::default());
2007 self.epoch = 1;
2008 }
2009 }
2010
2011 #[inline]
2013 fn trace(&mut self, id: u32) -> (bool, &mut Trace) {
2014 let t = &mut self.traces[id as usize];
2015 let fresh = t.epoch != self.epoch;
2016 if fresh {
2017 *t = Trace {
2018 epoch: self.epoch,
2019 key: u64::MAX,
2020 ..Trace::default()
2021 };
2022 }
2023 (fresh, t)
2024 }
2025
2026 #[inline]
2028 fn seen(&self, id: u32) -> bool {
2029 self.words[(id >> 6) as usize] >> (id & 63) & 1 == 1
2030 }
2031
2032 #[inline]
2034 fn insert(&mut self, id: u32) -> bool {
2035 let w = &mut self.words[(id >> 6) as usize];
2036 let bit = 1u64 << (id & 63);
2037 if *w & bit != 0 {
2038 return false;
2039 }
2040 if *w == 0 {
2041 self.touched.push(id >> 6);
2042 }
2043 *w |= bit;
2044 true
2045 }
2046}
2047
2048struct Pool {
2056 keys: Vec<u64>,
2057 expanded: Vec<u64>,
2059 #[cfg(feature = "experimental-batch-pool")]
2060 merge_keys: Vec<u64>,
2061 #[cfg(feature = "experimental-batch-pool")]
2062 merge_expanded: Vec<u64>,
2063 cap: usize,
2064 cursor: usize,
2065}
2066
2067impl Pool {
2068 fn reset(&mut self, cap: usize) {
2069 self.keys.clear();
2070 self.expanded.clear();
2071 self.cap = cap.max(1);
2072 self.cursor = 0;
2073 }
2074
2075 fn is_full(&self) -> bool {
2076 self.keys.len() >= self.cap
2077 }
2078
2079 fn worst(&self) -> Option<u64> {
2081 self.is_full().then(|| self.keys[self.keys.len() - 1])
2082 }
2083
2084 #[cfg(test)]
2085 fn is_expanded(&self, i: usize) -> bool {
2086 self.expanded[i / 64] >> (i % 64) & 1 == 1
2087 }
2088
2089 #[inline]
2091 fn insert_flag(&mut self, pos: usize, len: usize) {
2092 let words = (len + 1).div_ceil(64);
2093 if self.expanded.len() < words {
2094 self.expanded.resize(words, 0);
2095 }
2096 let w = pos / 64;
2097 for i in (w + 1..words).rev() {
2098 self.expanded[i] = self.expanded[i] << 1 | self.expanded[i - 1] >> 63;
2099 }
2100 let low = (1u64 << (pos % 64)) - 1;
2101 let x = self.expanded[w];
2102 self.expanded[w] = x & low | (x & !low) << 1;
2103 }
2104
2105 #[inline]
2107 fn remove_flag(&mut self, pos: usize, len: usize) {
2108 let words = len.div_ceil(64);
2109 let w = pos / 64;
2110 let low = (1u64 << (pos % 64)) - 1;
2111 let carry = |e: &[u64], i: usize| if i + 1 < words { e[i + 1] << 63 } else { 0 };
2112 let x = self.expanded[w];
2113 self.expanded[w] = x & low | (x >> 1) & !low | carry(&self.expanded, w);
2114 for i in w + 1..words {
2115 self.expanded[i] = self.expanded[i] >> 1 | carry(&self.expanded, i);
2116 }
2117 }
2118
2119 #[inline]
2122 fn insert(&mut self, key: u64) -> Option<(usize, Option<u64>)> {
2123 let n = self.keys.len();
2124 if n >= self.cap && key >= self.keys[n - 1] {
2125 return None;
2126 }
2127 let pos = self.keys.partition_point(|&t| t < key);
2128 let mut evicted = None;
2129 let mut len = n;
2130 if n >= self.cap {
2131 evicted = self.keys.pop();
2132 len -= 1;
2133 self.expanded[len / 64] &= !(1u64 << (len % 64));
2134 }
2135 self.keys.insert(pos, key);
2136 self.insert_flag(pos, len);
2137 if pos < self.cursor {
2138 self.cursor = pos;
2139 }
2140 Some((pos, evicted))
2141 }
2142
2143 #[cfg(feature = "experimental-batch-pool")]
2146 fn insert_batch(&mut self, pending: &mut Vec<u64>, k: usize) -> bool {
2147 debug_assert!(pending.iter().enumerate().all(|(position, &candidate)| {
2148 let id = key_id(candidate);
2149 pending[..position].iter().all(|&other| key_id(other) != id)
2150 && self.keys.iter().all(|&other| key_id(other) != id)
2151 }));
2152 if let Some(worst) = self.worst() {
2153 pending.retain(|&key| key < worst);
2154 }
2155 if pending.len() < 8 {
2156 let mut improved = false;
2157 for &key in pending.iter() {
2158 if let Some((pos, _)) = self.insert(key) {
2159 improved |= pos < k;
2160 }
2161 }
2162 return improved;
2163 }
2164 pending.sort_unstable();
2165 let improved = k > 0 && (self.keys.len() < k || pending[0] <= self.keys[k - 1]);
2167 let len = self.cap.min(self.keys.len().saturating_add(pending.len()));
2168 self.merge_keys.clear();
2169 self.merge_keys.reserve(len);
2170 self.merge_expanded.clear();
2171 self.merge_expanded.resize(len.div_ceil(64), 0);
2172 let (mut old, mut new) = (0, 0);
2173 for pos in 0..len {
2174 if new < pending.len() && (old == self.keys.len() || pending[new] <= self.keys[old]) {
2175 self.merge_keys.push(pending[new]);
2176 self.cursor = self.cursor.min(pos);
2177 new += 1;
2178 } else {
2179 self.merge_keys.push(self.keys[old]);
2180 let flag = self.expanded[old / 64] >> (old % 64) & 1;
2181 self.merge_expanded[pos / 64] |= flag << (pos % 64);
2182 old += 1;
2183 }
2184 }
2185 std::mem::swap(&mut self.keys, &mut self.merge_keys);
2186 std::mem::swap(&mut self.expanded, &mut self.merge_expanded);
2187 improved
2188 }
2189
2190 fn remove(&mut self, key: u64) {
2192 let pos = self.keys.partition_point(|&t| t < key);
2193 if pos < self.keys.len() && self.keys[pos] == key {
2194 let len = self.keys.len();
2195 self.keys.remove(pos);
2196 self.remove_flag(pos, len);
2197 if pos < self.cursor {
2198 self.cursor -= 1;
2199 }
2200 }
2201 }
2202
2203 #[inline]
2205 fn unexpanded(&self, mut n: usize) -> Option<usize> {
2206 let len = self.keys.len();
2207 let mut i = self.cursor;
2208 while i < len {
2209 let free = !self.expanded[i / 64] >> (i % 64);
2210 if free == 0 {
2211 i = (i / 64 + 1) * 64;
2212 continue;
2213 }
2214 i += free.trailing_zeros() as usize;
2215 if i >= len {
2216 break;
2217 }
2218 if n == 0 {
2219 return Some(i);
2220 }
2221 n -= 1;
2222 i += 1;
2223 }
2224 None
2225 }
2226
2227 #[inline]
2229 fn next(&mut self) -> Option<u64> {
2230 match self.unexpanded(0) {
2231 Some(i) => {
2232 self.expanded[i / 64] |= 1u64 << (i % 64);
2233 self.cursor = i;
2234 Some(self.keys[i])
2235 }
2236 None => {
2237 self.cursor = self.keys.len();
2238 None
2239 }
2240 }
2241 }
2242
2243 #[inline]
2245 fn peek_next(&self) -> Option<u64> {
2246 self.peek_unexpanded(0)
2247 }
2248
2249 #[inline]
2251 fn peek_unexpanded(&self, n: usize) -> Option<u64> {
2252 self.unexpanded(n).map(|i| self.keys[i])
2253 }
2254}
2255
2256#[derive(Default)]
2257struct Reservoir {
2258 keys: Vec<u64>,
2259 cap: usize,
2260}
2261
2262impl Reservoir {
2263 fn reset(&mut self, cap: usize) {
2264 self.keys.clear();
2265 self.cap = cap;
2266 }
2267
2268 #[inline]
2269 fn offer(&mut self, key: u64) {
2270 if self.cap == 0 || (self.keys.len() == self.cap && key >= self.keys[self.keys.len() - 1]) {
2271 return;
2272 }
2273 let pos = match self.keys.binary_search(&key) {
2274 Ok(_) => return,
2275 Err(pos) => pos,
2276 };
2277 if self.keys.len() == self.cap {
2278 self.keys.pop();
2279 }
2280 self.keys.insert(pos, key);
2281 }
2282}
2283
2284pub struct VSearcher<'a> {
2286 index: &'a VGraph,
2287 scratch: SearchScratch,
2288}
2289
2290pub(super) struct SearchScratch {
2291 #[cfg(not(target_arch = "wasm32"))]
2292 dim: usize,
2293 visited: Marks,
2294 qf: Vec<f32>,
2295 qc: Vec<i8>,
2297 qsum: i32,
2298 qtmp: Vec<i8>,
2299 qrot: Vec<f32>,
2301 qs: QueryCode,
2302 bqrot: Vec<f32>,
2304 bq: BlockQuery,
2305 scan: [u32; BLOCK],
2306 fresh: Vec<u32>,
2307 survivors: Vec<u32>,
2309 screened: usize,
2311 pool: Pool,
2313 expanded: BinaryHeap<u64>,
2316 reservoir: Reservoir,
2317 keys: Vec<u64>,
2318 results: Vec<(f32, u32)>,
2319 knobs: SearchKnobs,
2320}
2321
2322impl VSearcher<'_> {
2323 pub fn last_screened(&self) -> usize {
2325 self.scratch.screened
2326 }
2327
2328 pub fn knobs(&self) -> SearchKnobs {
2330 self.scratch.knobs
2331 }
2332
2333 #[inline]
2337 pub fn search(
2338 &mut self,
2339 query: &[f32],
2340 k: usize,
2341 params: VSearchParams,
2342 out: &mut Vec<u32>,
2343 ) -> Result<usize, WarmError> {
2344 self.scratch.search(self.index, query, k, params, out)
2345 }
2346}
2347
2348impl SearchScratch {
2349 #[inline]
2354 fn score_fresh<const BATCH_POOL: bool, const QUERY_X4: bool>(
2355 &mut self,
2356 g: &VGraph,
2357 k: usize,
2358 patience: usize,
2359 stale: &mut usize,
2360 w: f32,
2361 ) -> usize {
2362 #[cfg(feature = "experimental-batch-pool")]
2363 if BATCH_POOL {
2364 return self.score_fresh_batch::<QUERY_X4>(g, k, patience, stale, w);
2365 }
2366 let qc = &self.qc[..];
2367 let n = self.fresh.len();
2368 let mut i = 0;
2369 if self.knobs.batch_est || (QUERY_X4 && g.bits == 8) {
2370 while i + 4 <= n {
2371 let ids = [
2372 self.fresh[i],
2373 self.fresh[i + 1],
2374 self.fresh[i + 2],
2375 self.fresh[i + 3],
2376 ];
2377 for (e, u) in g
2378 .est4::<QUERY_X4>(qc, self.qsum, ids, w)
2379 .into_iter()
2380 .zip(ids)
2381 {
2382 let kk = key(e, u);
2383 self.reservoir.offer(kk);
2384 if let Some((pos, _)) = self.pool.insert(kk) {
2385 if patience > 0 && pos < k {
2386 *stale = 0;
2387 }
2388 }
2389 }
2390 i += 4;
2391 }
2392 }
2393 for &u in &self.fresh[i..] {
2394 let kk = key(g.est(qc, self.qsum, u, w), u);
2395 self.reservoir.offer(kk);
2396 if let Some((pos, _)) = self.pool.insert(kk) {
2397 if patience > 0 && pos < k {
2398 *stale = 0;
2399 }
2400 }
2401 }
2402 n
2403 }
2404
2405 #[cfg(feature = "experimental-batch-pool")]
2406 fn score_fresh_batch<const QUERY_X4: bool>(
2407 &mut self,
2408 g: &VGraph,
2409 k: usize,
2410 patience: usize,
2411 stale: &mut usize,
2412 w: f32,
2413 ) -> usize {
2414 let qc = &self.qc[..];
2415 let n = self.fresh.len();
2416 let (reservoir, pending) = (&mut self.reservoir, &mut self.keys);
2417 pending.clear();
2418 let mut offer = |estimate, id| {
2419 let key = key(estimate, id);
2420 reservoir.offer(key);
2421 pending.push(key);
2422 };
2423 let mut i = 0;
2424 if self.knobs.batch_est || (QUERY_X4 && g.bits == 8) {
2425 while i + 4 <= n {
2426 let ids = [
2427 self.fresh[i],
2428 self.fresh[i + 1],
2429 self.fresh[i + 2],
2430 self.fresh[i + 3],
2431 ];
2432 for (estimate, id) in g
2433 .est4::<QUERY_X4>(qc, self.qsum, ids, w)
2434 .into_iter()
2435 .zip(ids)
2436 {
2437 offer(estimate, id);
2438 }
2439 i += 4;
2440 }
2441 }
2442 for &id in &self.fresh[i..] {
2443 offer(g.est(qc, self.qsum, id, w), id);
2444 }
2445 if self.pool.insert_batch(pending, k) && patience > 0 {
2446 *stale = 0;
2447 }
2448 n
2449 }
2450
2451 fn search(
2457 &mut self,
2458 g: &VGraph,
2459 query: &[f32],
2460 k: usize,
2461 params: VSearchParams,
2462 out: &mut Vec<u32>,
2463 ) -> Result<usize, WarmError> {
2464 Self::validate_query(g, query)?;
2465 self.search_validated(g, query, k, params, out)
2466 }
2467
2468 fn validate_query(g: &VGraph, query: &[f32]) -> Result<(), WarmError> {
2469 if query.len() != g.dim {
2470 return Err(invalid(format!(
2471 "query has {} dimensions, the index has {}",
2472 query.len(),
2473 g.dim
2474 )));
2475 }
2476 if let Some(at) = query.iter().position(|v| !v.is_finite()) {
2477 return Err(invalid(format!(
2478 "non-finite query value at coordinate {at}"
2479 )));
2480 }
2481 if g.metric == VMetric::Angular && query.iter().all(|&v| v == 0.0) {
2482 return Err(invalid("zero query: cosine similarity is undefined".into()));
2483 }
2484 Ok(())
2485 }
2486
2487 fn search_validated(
2489 &mut self,
2490 g: &VGraph,
2491 query: &[f32],
2492 k: usize,
2493 params: VSearchParams,
2494 out: &mut Vec<u32>,
2495 ) -> Result<usize, WarmError> {
2496 self.search_inner::<
2497 false,
2498 { cfg!(feature = "experimental-batch-pool") },
2499 { cfg!(feature = "experimental-query-x4") },
2500 >(
2501 g,
2502 query,
2503 k,
2504 params,
2505 out,
2506 &mut QueryProfile::default(),
2507 )
2508 }
2509
2510 fn search_inner<const PROFILE: bool, const BATCH_POOL: bool, const QUERY_X4: bool>(
2511 &mut self,
2512 g: &VGraph,
2513 query: &[f32],
2514 k: usize,
2515 params: VSearchParams,
2516 out: &mut Vec<u32>,
2517 profile: &mut QueryProfile,
2518 ) -> Result<usize, WarmError> {
2519 if params.reservoir > 0 && (g.has_blocks() || (g.has_screen() && params.screen)) {
2520 return Err(invalid(
2521 "independent retention requires unscreened row-code search".into(),
2522 ));
2523 }
2524 let started = PROFILE.then(Instant::now);
2525 let ef = params.ef.max(k).max(1);
2526 self.reservoir.reset(if params.reservoir == 0 {
2527 0
2528 } else {
2529 params.reservoir.max(k).min(g.n)
2530 });
2531 self.qf[..g.dim].copy_from_slice(query);
2534 let sq = g.encode_query(
2535 &self.qf,
2536 &mut self.qc,
2537 &mut self.qtmp,
2538 &mut self.qrot,
2539 &mut self.qs,
2540 );
2541 let w = if sq > 0.0 { 1.0 / sq } else { 1.0 };
2544 self.qsum = if cfg!(target_arch = "x86_64") {
2546 query_sum(&self.qc)
2547 } else {
2548 0
2549 };
2550 if let Some(started) = started {
2551 profile.query_encoding_ns = started.elapsed().as_nanos() as u64;
2552 }
2553 let qc = &self.qc[..];
2554 let blocks = g.has_blocks();
2555 if blocks {
2556 g.encode_block_query(&self.qf, sq, &mut self.bqrot, &mut self.bq);
2557 }
2558 let knobs = self.knobs;
2559 let block_screen = blocks && knobs.block_screen;
2560 let screen = g.has_screen() && params.screen && !blocks;
2561 let margin = params.screen_sigmas * g.screen_sigma;
2562 self.screened = 0;
2563
2564 self.visited.next();
2565 let entry_started = PROFILE.then(Instant::now);
2566 let (entry, mut evals) = g.descend(qc, self.qsum, w, &mut self.reservoir);
2567 if let Some(started) = entry_started {
2568 profile.entry_ns = started.elapsed().as_nanos() as u64;
2569 profile.entry_evaluations = evals;
2570 }
2571 self.visited.insert(entry);
2572 let fuse = blocks && params.fuse && !block_screen;
2573 if fuse {
2574 self.visited.trace(entry);
2575 }
2576 let (qn2, qm) = if blocks {
2579 let q = &self.qf[..g.dim];
2580 (
2581 q.iter().map(|a| a * a).sum::<f32>(),
2582 q.iter().zip(&g.mean).map(|(a, m)| a * m).sum::<f32>(),
2583 )
2584 } else {
2585 (0.0, 0.0)
2586 };
2587 self.pool.reset(ef);
2588 self.expanded.clear();
2589 let take = params.rerank.max(k).min(g.n);
2590 let k0 = key(g.est(qc, self.qsum, entry, w), entry);
2591 self.reservoir.offer(k0);
2592 self.pool.insert(k0);
2593 if fuse {
2594 self.visited.trace(entry).1.key = k0;
2595 }
2596 evals += 1;
2597 let patience = params.patience;
2602 let frontier = if knobs.frontier.is_finite() {
2603 knobs.frontier.clamp(0.0, 1.0)
2604 } else {
2605 1.0
2606 };
2607 let mut stale = 0usize;
2608 let target = k.min(g.n);
2609 let enough = |search: &Self| {
2610 let count = if params.reservoir > 0 {
2611 search.reservoir.keys.len()
2612 } else if blocks && !block_screen {
2613 search.expanded.len()
2614 } else {
2615 search.pool.keys.len()
2616 };
2617 count >= target
2618 };
2619 let mut seed_cursor = 0usize;
2620 loop {
2621 let c = match self.pool.next() {
2622 Some(c) => c,
2623 None => {
2624 if enough(self) {
2625 break;
2626 }
2627 let seed = (seed_cursor..g.n)
2630 .find(|&v| {
2631 if fuse {
2632 self.visited.traces[v].epoch != self.visited.epoch
2633 } else {
2634 !self.visited.seen(v as u32)
2635 }
2636 })
2637 .ok_or_else(|| {
2638 invalid(
2639 "search exhausted vertices before retaining enough candidates"
2640 .into(),
2641 )
2642 })?;
2643 seed_cursor = seed + 1;
2644 let seed = seed as u32;
2645 self.visited.insert(seed);
2646 let candidate = key(g.est(&self.qc, self.qsum, seed, w), seed);
2647 self.reservoir.offer(candidate);
2648 let (_, evicted) = self
2649 .pool
2650 .insert(candidate)
2651 .ok_or_else(|| invalid("search could not admit a recovery seed".into()))?;
2652 if fuse {
2653 self.visited.trace(seed).1.key = candidate;
2654 if let Some(evicted) = evicted {
2655 self.visited.trace(key_id(evicted)).1.key = u64::MAX;
2656 }
2657 }
2658 evals += 1;
2659 if PROFILE {
2660 profile.traversal_evaluations += 1;
2661 }
2662 stale = 0;
2663 continue;
2664 }
2665 };
2666 if patience > 0 && stale >= patience && enough(self) {
2667 break;
2668 }
2669 if frontier < 1.0 && k > 0 && self.pool.is_full() {
2670 let len = self.pool.keys.len();
2674 let kth = key_dist(self.pool.keys[k.min(len) - 1]);
2675 let worst = key_dist(self.pool.keys[len - 1]);
2676 if key_dist(c) > kth + frontier * (worst - kth) && enough(self) {
2677 break;
2678 }
2679 }
2680 if fuse {
2681 let (_, t) = self.visited.trace(key_id(c));
2683 if t.expanded {
2684 continue;
2685 }
2686 t.expanded = true;
2687 }
2688 stale += 1;
2689 if PROFILE {
2690 profile.expansions += 1;
2691 }
2692 let adjacency_started = PROFILE.then(Instant::now);
2693 if let Some(next) = self.pool.peek_next() {
2694 g.prefetch_neighbours(key_id(next) as usize);
2695 if blocks {
2696 g.prefetch_blocks(key_id(next) as usize);
2697 if !block_screen {
2698 prefetch_bytes(g.code(key_id(next)));
2699 }
2700 }
2701 }
2702 if knobs.prefetch >= 2 {
2703 if let Some(second) = self.pool.peek_unexpanded(1) {
2704 g.prefetch_neighbours(key_id(second) as usize);
2705 if blocks {
2706 g.prefetch_blocks(key_id(second) as usize);
2707 }
2708 }
2709 }
2710 if block_screen {
2711 let v = key_id(c) as usize;
2717 let ip_v = -key_dist(c) * sq;
2718 let d2_v = (qn2 + 1.0 - 2.0 * (ip_v + qm)).max(1e-6);
2719 let (delta, min_sum, inv_sq) = (self.bq.delta, self.bq.min_sum, self.bq.inv_sq);
2720 let worst = self.pool.worst().map(key_dist);
2721 let sigmas = params.screen_sigmas;
2722 self.fresh.clear();
2723 let list = &mut self.fresh;
2724 g.for_each_neighbour(v, |u| list.push(u));
2725 self.survivors.clear();
2726 for (b, ids) in self
2727 .fresh
2728 .chunks(BLOCK)
2729 .take(g.blocks_per_vertex)
2730 .enumerate()
2731 {
2732 let (codes, scales, offsets) = g.block(v, b);
2733 scan32(codes, &self.bq.lut, g.block_m, &mut self.scan);
2734 let (scales, offsets) = (scales.as_chunks::<4>().0, offsets.as_chunks::<4>().0);
2735 for (((&u, &raw), s), o) in ids.iter().zip(&self.scan).zip(scales).zip(offsets)
2736 {
2737 if !self.visited.insert(u) {
2738 continue;
2739 }
2740 let s = f32::from_le_bytes(*s);
2741 let ip = ip_v + s * (delta * raw as f32 + min_sum) + f32::from_le_bytes(*o);
2742 if let Some(w) = worst {
2743 let sigma = (g.block_kappa * s * s * d2_v).sqrt();
2744 if -(ip + sigmas * sigma) * inv_sq >= w {
2745 self.screened += 1;
2746 continue;
2747 }
2748 }
2749 prefetch_bytes(g.code(u));
2750 g.prefetch_scale(u);
2751 self.survivors.push(u);
2752 }
2753 }
2754 std::mem::swap(&mut self.fresh, &mut self.survivors);
2755 evals += self.score_fresh::<BATCH_POOL, QUERY_X4>(g, k, patience, &mut stale, w);
2756 continue;
2757 }
2758 if blocks {
2759 let v = key_id(c) as usize;
2765 let est_v = g.est(&self.qc, self.qsum, v as u32, w);
2766 let ek = key(est_v, v as u32);
2767 evals += 1;
2768 let ip_v = -est_v * sq;
2770 let d2_v = (qn2 + 1.0 - 2.0 * (ip_v + qm)).max(1e-6);
2771 if self.expanded.len() < take {
2772 self.expanded.push(ek);
2773 } else if let Some(mut worst) = self.expanded.peek_mut() {
2774 if ek < *worst {
2775 *worst = ek;
2776 }
2777 }
2778 let (delta, min_sum, inv_sq) = (self.bq.delta, self.bq.min_sum, self.bq.inv_sq);
2779 self.fresh.clear();
2780 let list = &mut self.fresh;
2781 g.for_each_neighbour(v, |u| list.push(u));
2782 for (b, ids) in self
2783 .fresh
2784 .chunks(BLOCK)
2785 .take(g.blocks_per_vertex)
2786 .enumerate()
2787 {
2788 let (codes, scales, offsets) = g.block(v, b);
2789 scan32(codes, &self.bq.lut, g.block_m, &mut self.scan);
2790 let (scales, offsets) = (scales.as_chunks::<4>().0, offsets.as_chunks::<4>().0);
2791 for (((&u, &raw), s), o) in ids.iter().zip(&self.scan).zip(scales).zip(offsets)
2792 {
2793 let s = f32::from_le_bytes(*s);
2794 let ip = ip_v + s * (delta * raw as f32 + min_sum) + f32::from_le_bytes(*o);
2795 if !fuse {
2796 if !self.visited.insert(u) {
2799 continue;
2800 }
2801 evals += 1;
2802 if let Some((pos, _)) = self.pool.insert(key(-ip * inv_sq, u)) {
2803 if patience > 0 && pos < k {
2804 stale = 0;
2805 }
2806 }
2807 continue;
2808 }
2809 let w = 1.0 / (g.block_kappa * s * s * d2_v + 1e-12);
2811 let (fresh, t) = self.visited.trace(u);
2812 let ip = if fresh {
2813 t.w = w;
2814 t.wsum = w * ip;
2815 ip
2816 } else {
2817 if t.expanded {
2818 continue;
2819 }
2820 t.w += w;
2821 t.wsum += w * ip;
2822 t.wsum / t.w
2823 };
2824 let kk = key(-ip * inv_sq, u);
2825 let old = t.key;
2826 evals += 1;
2827 if !fresh && kk >= old {
2828 continue;
2829 }
2830 if old != u64::MAX {
2831 self.pool.remove(old);
2832 }
2833 let Some((pos, evicted)) = self.pool.insert(kk) else {
2835 t.key = u64::MAX;
2836 continue;
2837 };
2838 t.key = kk;
2839 if let Some(e) = evicted {
2840 self.visited.trace(key_id(e)).1.key = u64::MAX;
2841 }
2842 if patience > 0 && pos < k {
2843 stale = 0;
2844 }
2845 }
2846 }
2847 continue;
2848 }
2849 self.fresh.clear();
2850 let (visited, fresh) = (&mut self.visited, &mut self.fresh);
2851 #[cfg(feature = "experimental-neighbor-admission")]
2852 let neighbors_read = if screen {
2853 g.admit_neighbours::<PROFILE, true>(
2854 key_id(c) as usize,
2855 params.neighbour_cap,
2856 visited,
2857 fresh,
2858 )
2859 } else {
2860 g.admit_neighbours::<PROFILE, false>(
2861 key_id(c) as usize,
2862 params.neighbour_cap,
2863 visited,
2864 fresh,
2865 )
2866 };
2867 #[cfg(not(feature = "experimental-neighbor-admission"))]
2868 let neighbors_read = {
2869 let mut neighbors_read = 0;
2870 let admit = |u: u32, visited: &mut Marks, fresh: &mut Vec<u32>| {
2871 if visited.insert(u) {
2872 if screen {
2873 g.prefetch_screen(u);
2874 } else {
2875 prefetch_bytes(g.code(u));
2876 g.prefetch_scale(u);
2877 }
2878 fresh.push(u);
2879 }
2880 };
2881 if params.neighbour_cap == 0 {
2882 g.for_each_neighbour(key_id(c) as usize, |u| {
2883 if PROFILE {
2884 neighbors_read += 1;
2885 }
2886 admit(u, visited, fresh);
2887 });
2888 } else {
2889 let mut seen = 0usize;
2890 g.for_each_neighbour(key_id(c) as usize, |u| {
2891 seen += 1;
2892 if seen <= params.neighbour_cap {
2893 admit(u, visited, fresh);
2894 }
2895 });
2896 }
2897 neighbors_read
2898 };
2899 if screen {
2900 let worst = self.pool.worst().map(key_dist);
2903 self.survivors.clear();
2904 for &u in &self.fresh {
2905 if let Some(w) = worst {
2906 if g.screen_est(&self.qs, u) - margin >= w {
2907 self.screened += 1;
2908 continue;
2909 }
2910 }
2911 prefetch_bytes(g.code(u));
2912 g.prefetch_scale(u);
2913 self.survivors.push(u);
2914 }
2915 std::mem::swap(&mut self.fresh, &mut self.survivors);
2916 }
2917 if let Some(started) = adjacency_started {
2918 profile.adjacency_visited_prefetch_ns += started.elapsed().as_nanos() as u64;
2919 profile.neighbors_read += neighbors_read;
2920 }
2921 let scoring_started = PROFILE.then(Instant::now);
2922 let count = self.score_fresh::<BATCH_POOL, QUERY_X4>(g, k, patience, &mut stale, w);
2923 evals += count;
2924 if let Some(started) = scoring_started {
2925 profile.scoring_pool_ns += started.elapsed().as_nanos() as u64;
2926 profile.traversal_evaluations += count;
2927 }
2928 if knobs.prefetch >= 3 {
2929 if let Some(next) = self.pool.peek_next() {
2932 let visited = &self.visited;
2933 g.for_each_neighbour(key_id(next) as usize, |u| {
2934 if !visited.seen(u) {
2935 if screen {
2936 g.prefetch_screen(u);
2937 } else {
2938 prefetch_bytes(g.code(u));
2939 g.prefetch_scale(u);
2940 }
2941 }
2942 });
2943 }
2944 }
2945 }
2946
2947 let rerank_started = PROFILE.then(Instant::now);
2948 if PROFILE {
2949 profile.beam_retained = self.pool.keys.len();
2950 profile.traversal_evaluations += 1;
2951 }
2952 self.keys.clear();
2953 if params.reservoir > 0 {
2954 self.keys.extend_from_slice(&self.reservoir.keys);
2955 } else if blocks && !block_screen {
2956 self.keys.extend(self.expanded.drain());
2957 self.keys.sort_unstable();
2958 } else {
2959 self.keys.extend_from_slice(&self.pool.keys);
2960 }
2961 if PROFILE {
2962 profile.retained = self.keys.len();
2963 }
2964 self.results.clear();
2965 let take = take.min(self.keys.len());
2966 if params.rerank > 0 && knobs.rerank_z > 0.0 && sq > 0.0 && k > 0 && take > k {
2967 evals += self.rerank_gated(g, k, take, sq, knobs.rerank_z);
2968 if PROFILE {
2969 profile.rerank_evaluations = self.results.len();
2970 }
2971 } else if params.rerank > 0 {
2972 if PROFILE {
2973 profile.rerank_evaluations = take;
2974 }
2975 for &c in &self.keys[..take] {
2976 g.prefetch_rerank_rows(key_id(c));
2977 }
2978 self.results.extend(self.keys[..take].iter().map(|&c| {
2980 let id = key_id(c);
2981 (g.fine(&self.qf, id), g.original(id))
2982 }));
2983 evals += take;
2984 self.results
2985 .sort_unstable_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1)));
2986 } else {
2987 self.results.extend(
2988 self.keys[..take]
2989 .iter()
2990 .map(|&c| (0.0, g.original(key_id(c)))),
2991 );
2992 }
2993 out.clear();
2994 out.extend(self.results.iter().take(k).map(|r| r.1));
2995 if let Some(started) = rerank_started {
2996 profile.rerank_sort_ns = started.elapsed().as_nanos() as u64;
2997 }
2998 if let Some(started) = started {
2999 profile.total_ns = started.elapsed().as_nanos() as u64;
3000 }
3001 Ok(evals)
3002 }
3003}
3004
3005#[cfg(feature = "experimental")]
3006impl VSearcher<'_> {
3007 pub fn profile_search(
3009 &mut self,
3010 query: &[f32],
3011 k: usize,
3012 params: VSearchParams,
3013 out: &mut Vec<u32>,
3014 ) -> Result<QueryProfile, WarmError> {
3015 let g = self.index;
3016 if g.has_screen()
3017 || g.has_blocks()
3018 || self.knobs() != SearchKnobs::OFF
3019 || params.neighbour_cap != 0
3020 {
3021 return Err(invalid("profiling requires the plain row-code path".into()));
3022 }
3023 SearchScratch::validate_query(g, query)?;
3024 let mut profile = QueryProfile::default();
3025 self.scratch
3026 .search_inner::<
3027 true,
3028 { cfg!(feature = "experimental-batch-pool") },
3029 { cfg!(feature = "experimental-query-x4") },
3030 >(
3031 g,
3032 query,
3033 k,
3034 params,
3035 out,
3036 &mut profile,
3037 )?;
3038 Ok(profile)
3039 }
3040}
3041
3042#[cfg(feature = "experimental")]
3043impl VGraph {
3044 pub fn encoding_diagnostic(&self, source: &Graph) -> Result<serde_json::Value, WarmError> {
3046 if source.n != self.n || source.dim < self.dim {
3047 return Err(invalid("encoding diagnostic source shape mismatch".into()));
3048 }
3049 let (mut coarse_error, mut fine_error) = (0f64, 0f64);
3050 for v in 0..self.n {
3051 let original = self.original(v as u32) as usize;
3052 for d in 0..self.dim {
3053 let expected =
3054 f64::from(source.vectors[original * source.dim + d]) - f64::from(self.mean[d]);
3055 let primary = f64::from(self.scales[v])
3056 * f64::from(code_at(self.code(v as u32), self.bits, d));
3057 let residual = if self.has_residual() {
3058 f64::from(self.res_scales[v])
3059 * f64::from(code_at(
3060 &self.residual[v * self.res_stride..(v + 1) * self.res_stride],
3061 self.res_bits,
3062 d,
3063 ))
3064 } else {
3065 0.0
3066 };
3067 coarse_error += (primary - expected).powi(2);
3068 fine_error += (primary + residual - expected).powi(2);
3069 }
3070 }
3071 let count = (self.n * self.dim) as f64;
3072 Ok(
3073 serde_json::json!({"metric": format!("{:?}", self.metric), "dimensions": self.dim,
3074 "padded": self.padded, "primary_stride": self.stride, "residual_stride": self.res_stride,
3075 "norm_bytes": self.half_norms.len() * 4, "index_bytes": self.index_bytes(),
3076 "primary_coordinate_rmse": (coarse_error / count).sqrt(),
3077 "refined_coordinate_rmse": (fine_error / count).sqrt()}),
3078 )
3079 }
3080
3081 pub fn diagnose(
3086 &mut self,
3087 source: &Graph,
3088 query: &[f32],
3089 k: usize,
3090 params: VSearchParams,
3091 oracle: Option<u32>,
3092 exact: bool,
3093 ) -> Result<super::diagnostics::TraversalDiagnostic, WarmError> {
3094 if source.n != self.n
3095 || source.dim != self.dim + usize::from(self.metric == VMetric::InnerProduct)
3096 || source.degree != self.degree
3097 || if self.metric == VMetric::InnerProduct {
3098 source.metric != Metric::Euclidean
3099 } else {
3100 self.metric != source.metric.into()
3101 }
3102 {
3103 return Err(invalid(
3104 "diagnostic source graph does not match index geometry and metric".into(),
3105 ));
3106 }
3107 if self.has_screen() || self.has_blocks() {
3108 return Err(invalid(
3109 "diagnostic traversal does not support screens or blocks".into(),
3110 ));
3111 }
3112 if query.len() != self.dim
3113 || query.iter().any(|x| !x.is_finite())
3114 || (self.metric == VMetric::Angular && query.iter().all(|&x| x == 0.0))
3115 || k == 0
3116 || k > self.n
3117 {
3118 return Err(invalid(
3119 "invalid diagnostic query dimension, coordinates or k".into(),
3120 ));
3121 }
3122 let stored_oracle = match oracle {
3123 Some(id) if id as usize >= self.n => {
3124 return Err(invalid("diagnostic oracle ID outside index".into()));
3125 }
3126 Some(id) if self.ids.is_empty() => Some(id),
3127 Some(id) => Some(
3128 self.ids
3129 .iter()
3130 .position(|&original| original == id)
3131 .expect("encoded IDs are a permutation") as u32,
3132 ),
3133 None => None,
3134 };
3135 let saved_entry = self.entry;
3136 let saved_layers = stored_oracle.map(|id| {
3137 self.entry = id;
3138 std::mem::take(&mut self.layers)
3139 });
3140 let result = if exact {
3141 self.diagnose_exact(source, query, k, params)
3142 } else {
3143 self.diagnose_coded(query, k, params)
3144 };
3145 self.entry = saved_entry;
3146 if let Some(layers) = saved_layers {
3147 self.layers = layers;
3148 }
3149 result
3150 }
3151
3152 fn diagnostic_entry(&self, mut score: impl FnMut(u32) -> f32, evaluated: &mut Vec<u32>) -> u32 {
3153 if self.layers.is_empty() {
3154 return self.entry;
3155 }
3156 let mut p = 0u32;
3157 let mut best = score(self.upper_ids[0]);
3158 evaluated.push(self.original(self.upper_ids[0]));
3159 for layer in self.layers.iter().rev() {
3160 loop {
3161 let mut moved = false;
3162 for &u in &layer[p as usize] {
3163 let id = self.upper_ids[u as usize];
3164 let candidate = score(id);
3165 evaluated.push(self.original(id));
3166 if candidate < best {
3167 best = candidate;
3168 p = u;
3169 moved = true;
3170 }
3171 }
3172 if !moved {
3173 break;
3174 }
3175 }
3176 }
3177 self.upper_ids[p as usize]
3178 }
3179
3180 fn diagnose_coded(
3181 &self,
3182 query: &[f32],
3183 k: usize,
3184 params: VSearchParams,
3185 ) -> Result<super::diagnostics::TraversalDiagnostic, WarmError> {
3186 let mut report = super::diagnostics::TraversalDiagnostic::default();
3187 let mut searcher = self.searcher_with(SearchKnobs::OFF);
3188 searcher.search(query, k, params, &mut report.results)?;
3189 let scratch = &mut searcher.scratch;
3190 let sq = self.encode_query(
3191 &scratch.qf,
3192 &mut scratch.qc,
3193 &mut scratch.qtmp,
3194 &mut scratch.qrot,
3195 &mut scratch.qs,
3196 );
3197 let w = if sq > 0.0 { 1.0 / sq } else { 1.0 };
3198 let qsum = query_sum(&scratch.qc);
3200 self.diagnostic_entry(
3201 |id| self.est(&scratch.qc, qsum, id, w),
3202 &mut report.evaluated,
3203 );
3204 report.evaluated.extend(
3205 (0..self.n as u32)
3206 .filter(|&id| scratch.visited.seen(id))
3207 .map(|id| self.original(id)),
3208 );
3209 report.evaluated.sort_unstable();
3210 report.evaluated.dedup();
3211 report
3212 .beam
3213 .extend(scratch.pool.keys.iter().map(|&c| self.original(key_id(c))));
3214 report.retained.extend(
3215 scratch
3216 .keys
3217 .iter()
3218 .map(|&candidate| self.original(key_id(candidate))),
3219 );
3220 if params.rerank > 0 {
3221 let take = params.rerank.max(k).min(self.n).min(report.retained.len());
3222 report
3223 .rerank_candidates
3224 .extend_from_slice(&report.retained[..take]);
3225 }
3226 Ok(report)
3227 }
3228
3229 fn diagnose_exact(
3230 &self,
3231 source: &Graph,
3232 query: &[f32],
3233 k: usize,
3234 params: VSearchParams,
3235 ) -> Result<super::diagnostics::TraversalDiagnostic, WarmError> {
3236 let mut report = super::diagnostics::TraversalDiagnostic::default();
3237 let score = |id: u32| {
3238 let original = self.original(id) as usize;
3239 let row = &source.vectors[original * source.dim..original * source.dim + self.dim];
3240 let ip: f64 = row
3241 .iter()
3242 .zip(query)
3243 .map(|(&x, &q)| f64::from(x) * f64::from(q))
3244 .sum();
3245 let norm = if self.metric == VMetric::Euclidean {
3246 row.iter().map(|&x| f64::from(x).powi(2)).sum::<f64>() * 0.5
3247 } else {
3248 0.0
3249 };
3250 (norm - ip) as f32
3251 };
3252 let mut reservoir = Reservoir::default();
3253 reservoir.reset(if params.reservoir == 0 {
3254 0
3255 } else {
3256 params.reservoir.max(k).min(self.n)
3257 });
3258 let entry = self.diagnostic_entry(
3259 |id| {
3260 let distance = score(id);
3261 reservoir.offer(key(distance, id));
3262 distance
3263 },
3264 &mut report.evaluated,
3265 );
3266 let mut pool = Pool {
3267 keys: Vec::new(),
3268 expanded: Vec::new(),
3269 #[cfg(feature = "experimental-batch-pool")]
3270 merge_keys: Vec::new(),
3271 #[cfg(feature = "experimental-batch-pool")]
3272 merge_expanded: Vec::new(),
3273 cap: params.ef.max(k).max(1),
3274 cursor: 0,
3275 };
3276 let mut seen = vec![false; self.n];
3277 seen[entry as usize] = true;
3278 report.evaluated.push(self.original(entry));
3279 let entry_key = key(score(entry), entry);
3280 reservoir.offer(entry_key);
3281 pool.insert(entry_key);
3282 let mut stale = 0usize;
3283 while let Some(candidate) = pool.next() {
3284 if params.patience > 0 && stale >= params.patience {
3285 break;
3286 }
3287 stale += 1;
3288 let mut count = 0usize;
3289 self.for_each_neighbour(key_id(candidate) as usize, |id| {
3290 count += 1;
3291 if (params.neighbour_cap > 0 && count > params.neighbour_cap) || seen[id as usize] {
3292 return;
3293 }
3294 seen[id as usize] = true;
3295 report.evaluated.push(self.original(id));
3296 let kk = key(score(id), id);
3297 reservoir.offer(kk);
3298 if let Some((pos, _)) = pool.insert(kk) {
3299 if params.patience > 0 && pos < k {
3300 stale = 0;
3301 }
3302 }
3303 });
3304 }
3305 report.evaluated.sort_unstable();
3306 report.evaluated.dedup();
3307 report
3308 .beam
3309 .extend(pool.keys.iter().map(|&c| self.original(key_id(c))));
3310 let keys = if params.reservoir > 0 {
3311 &reservoir.keys
3312 } else {
3313 &pool.keys
3314 };
3315 report
3316 .retained
3317 .extend(keys.iter().map(|&c| self.original(key_id(c))));
3318 let take = params.rerank.max(k).min(self.n).min(keys.len());
3319 if params.rerank > 0 {
3320 report
3321 .rerank_candidates
3322 .extend_from_slice(&report.retained[..take]);
3323 let mut qf = vec![0.0; self.padded];
3324 qf[..self.dim].copy_from_slice(query);
3325 let mut results: Vec<_> = keys[..take]
3326 .iter()
3327 .map(|&c| (self.fine(&qf, key_id(c)), self.original(key_id(c))))
3328 .collect();
3329 results.sort_unstable_by(|a, b| a.0.total_cmp(&b.0).then(a.1.cmp(&b.1)));
3330 report
3331 .results
3332 .extend(results.iter().take(k).map(|&(_, id)| id));
3333 } else {
3334 report
3335 .results
3336 .extend_from_slice(&report.retained[..k.min(take)]);
3337 }
3338 Ok(report)
3339 }
3340}
3341
3342#[cfg(test)]
3343mod serving_tests {
3344 use super::*;
3345
3346 #[test]
3347 fn connected_search_fingerprints_and_evaluation_counts() {
3348 let rows: Vec<f32> = (0..64)
3349 .flat_map(|v| (0..16).map(move |d| if v & (1 << (d % 6)) == 0 { -1.0 } else { 1.0 }))
3350 .collect();
3351 for metric in [VMetric::Angular, VMetric::Euclidean, VMetric::InnerProduct] {
3352 let mut graph = if metric == VMetric::InnerProduct {
3353 lifted(&rows, 16)
3354 } else {
3355 Graph::build(
3356 &rows,
3357 16,
3358 &BuildParams {
3359 metric: if metric == VMetric::Angular {
3360 Metric::Angular
3361 } else {
3362 Metric::Euclidean
3363 },
3364 degree: 8,
3365 build_ef: 16,
3366 ..Default::default()
3367 },
3368 )
3369 .unwrap()
3370 };
3371 graph.entry = 0;
3372 graph.layers.clear();
3373 graph.upper_ids.clear();
3374 for v in 0..graph.n {
3375 let neighbors: Vec<_> = [1, 3]
3376 .map(|offset| ((v + offset) % graph.n) as u32)
3377 .to_vec();
3378 graph.adj.set(v, &neighbors);
3379 }
3380 for bits in [4, 8] {
3381 let index = if metric == VMetric::InnerProduct {
3382 VGraph::from_lifted_ip_graph(&graph, bits, 8, true, true, 3).unwrap()
3383 } else {
3384 VGraph::from_graph_with(&graph, bits, 8, true, true, 3, false).unwrap()
3385 };
3386 for (knob_id, knobs) in [
3387 SearchKnobs::OFF,
3388 SearchKnobs {
3389 prefetch: 3,
3390 batch_est: true,
3391 frontier: 0.8,
3392 ..SearchKnobs::OFF
3393 },
3394 ]
3395 .into_iter()
3396 .enumerate()
3397 {
3398 let mut searcher = index.searcher_with(knobs);
3399 let mut out = Vec::new();
3400 let mut fingerprint = 0xcbf29ce484222325u64;
3401 let mut evaluations = 0usize;
3402 for ef in [8, 16] {
3403 for reservoir in [0, 16] {
3404 for patience in [0, 1] {
3405 for query in rows.as_chunks::<16>().0.iter().step_by(7) {
3406 let count = searcher
3407 .search(
3408 query,
3409 5,
3410 VSearchParams {
3411 ef,
3412 rerank: 16,
3413 reservoir,
3414 patience,
3415 ..Default::default()
3416 },
3417 &mut out,
3418 )
3419 .unwrap();
3420 assert_eq!(out.len(), 5);
3421 evaluations += count;
3422 for value in
3423 out.iter()
3424 .map(|&id| u64::from(id))
3425 .chain(searcher.scratch.results.iter().map(
3426 |&(distance, _)| u64::from(distance.to_bits()),
3427 ))
3428 .chain([count as u64])
3429 {
3430 fingerprint =
3431 (fingerprint ^ value).wrapping_mul(0x100000001b3);
3432 }
3433 }
3434 }
3435 }
3436 }
3437 println!("CONNECTED {metric:?} {bits} {knob_id} {fingerprint} {evaluations}");
3438 }
3439 }
3440 }
3441 }
3442
3443 fn disconnect(graph: &mut Graph, component_size: usize) {
3444 graph.entry = 0;
3445 graph.layers.clear();
3446 graph.upper_ids.clear();
3447 for v in 0..graph.n {
3448 let first = v / component_size * component_size;
3449 let neighbors: Vec<_> = (first..(first + component_size).min(graph.n))
3450 .filter(|&u| u != v)
3451 .map(|u| u as u32)
3452 .collect();
3453 graph.adj.set(v, &neighbors);
3454 }
3455 }
3456
3457 fn assert_complete(searcher: &VSearcher<'_>, out: &[u32], k: usize, rerank: usize) {
3458 assert_eq!(out.len(), k.min(searcher.index.n));
3459 assert_eq!(
3460 out.iter()
3461 .copied()
3462 .collect::<std::collections::BTreeSet<_>>()
3463 .len(),
3464 out.len()
3465 );
3466 assert!(out.iter().all(|&id| (id as usize) < searcher.index.n));
3467 assert!(searcher.scratch.keys.windows(2).all(|w| w[0] <= w[1]));
3468 if rerank > 0 {
3469 assert!(searcher.scratch.results.iter().all(|r| r.0.is_finite()));
3470 assert!(searcher
3471 .scratch
3472 .results
3473 .windows(2)
3474 .all(|w| { w[0].0.total_cmp(&w[1].0).then(w[0].1.cmp(&w[1].1)).is_le() }));
3475 }
3476 }
3477
3478 #[test]
3479 fn disconnected_graph_returns_the_requested_number_of_distinct_neighbors() {
3480 let rows: Vec<f32> = (0..67 * 2).map(|i| (i * 7 % 31 + 1) as f32).collect();
3481 for metric in [VMetric::Angular, VMetric::Euclidean, VMetric::InnerProduct] {
3482 let mut graph = if metric == VMetric::InnerProduct {
3483 lifted(&rows, 2)
3484 } else {
3485 Graph::build(
3486 &rows,
3487 2,
3488 &BuildParams {
3489 metric: if metric == VMetric::Angular {
3490 Metric::Angular
3491 } else {
3492 Metric::Euclidean
3493 },
3494 degree: 2,
3495 build_ef: 16,
3496 ..Default::default()
3497 },
3498 )
3499 .unwrap()
3500 };
3501 disconnect(&mut graph, 3);
3502 for bits in [4, 8] {
3503 for id_bytes in [0, 3, 4] {
3504 let index = if metric == VMetric::InnerProduct {
3505 VGraph::from_lifted_ip_graph(&graph, bits, 8, true, true, id_bytes).unwrap()
3506 } else {
3507 VGraph::from_graph_with(
3508 &graph,
3509 bits,
3510 8,
3511 true,
3512 true,
3513 id_bytes,
3514 metric == VMetric::Angular,
3515 )
3516 .unwrap()
3517 };
3518 for (ef, reservoir, rerank, patience, neighbour_cap) in [
3519 (0, 0, 0, 0, 0),
3520 (8, 0, 67, 1, 1),
3521 (8, 1, 67, 1, 0),
3522 (8, 68, 67, 1, 1),
3523 ] {
3524 let knobs = SearchKnobs {
3525 prefetch: 3,
3526 batch_est: true,
3527 frontier: 0.0,
3528 rerank_z: 1.0,
3529 ..SearchKnobs::OFF
3530 };
3531 let mut searcher = index.searcher_with(knobs);
3532 let params = VSearchParams {
3533 ef,
3534 reservoir,
3535 rerank,
3536 patience,
3537 neighbour_cap,
3538 screen: reservoir == 0,
3539 screen_sigmas: 0.0,
3540 ..Default::default()
3541 };
3542 let mut out = Vec::new();
3543 for k in [0, 1, 2, 3, 4, 63, 64, 65, 66, 67, 68, usize::MAX] {
3544 let evaluations =
3545 searcher.search(&[1.0, 2.0], k, params, &mut out).unwrap();
3546 assert_complete(&searcher, &out, k, rerank);
3547 let first = out.clone();
3548 assert_eq!(
3549 searcher.search(&[1.0, 2.0], k, params, &mut out).unwrap(),
3550 evaluations
3551 );
3552 assert_eq!(out, first);
3553 }
3554 }
3555 }
3556 }
3557 }
3558 }
3559
3560 #[test]
3561 fn recovery_expands_the_seed_and_counts_its_score() {
3562 let rows = [
3563 1.0, 0.0, 2.0, 0.0, 3.0, 0.0, -10.0, 0.0, 5.0, 0.0, 6.0, 0.0, 100.0, 0.0,
3564 ];
3565 let mut graph = lifted(&rows, 2);
3566 disconnect(&mut graph, 1);
3567 graph.adj.set(0, &[1]);
3568 graph.adj.set(1, &[2]);
3569 graph.adj.set(3, &[5]);
3570 graph.adj.set(5, &[4]);
3571 let index = VGraph::from_lifted_ip_graph(&graph, 8, 8, false, false, 4).unwrap();
3572 let mut searcher = index.searcher_with(SearchKnobs::OFF);
3573 for reservoir in [0, 4] {
3574 let params = VSearchParams {
3575 ef: 4,
3576 rerank: 4,
3577 reservoir,
3578 patience: 1,
3579 ..Default::default()
3580 };
3581 let mut out = Vec::new();
3582 let evaluations = searcher.search(&[1.0, 0.0], 4, params, &mut out).unwrap();
3583 assert_eq!(out, [5, 4, 2, 1]);
3584 assert_eq!(evaluations, 10);
3585 assert!(!searcher.scratch.visited.seen(6));
3586 #[cfg(feature = "experimental")]
3587 {
3588 let profile = searcher
3589 .profile_search(&[1.0, 0.0], 4, params, &mut out)
3590 .unwrap();
3591 assert_eq!(profile.traversal_evaluations, 6);
3592 assert_eq!(
3593 profile.entry_evaluations
3594 + profile.traversal_evaluations
3595 + profile.rerank_evaluations,
3596 evaluations
3597 );
3598 }
3599 }
3600 searcher
3601 .search(
3602 &[0.0, 0.0],
3603 7,
3604 VSearchParams {
3605 rerank: 7,
3606 ..Default::default()
3607 },
3608 &mut Vec::new(),
3609 )
3610 .unwrap();
3611 assert_eq!(
3612 searcher
3613 .scratch
3614 .results
3615 .iter()
3616 .map(|r| r.1)
3617 .collect::<Vec<_>>(),
3618 [0, 1, 2, 3, 4, 5, 6]
3619 );
3620 }
3621
3622 #[test]
3623 fn disconnected_blocks_retain_enough_expansions_with_stopping() {
3624 let rows = crate::core::qgraph::tests::clustered(67, 16, 77);
3625 let mut graph = Graph::build(
3626 &rows,
3627 16,
3628 &BuildParams {
3629 degree: 32,
3630 build_ef: 64,
3631 ..Default::default()
3632 },
3633 )
3634 .unwrap();
3635 for component_size in [3, 32, 33] {
3636 disconnect(&mut graph, component_size);
3637 for bits in [1, 2, 4] {
3638 for id_bytes in [0, 3, 4] {
3639 let index = VGraph::from_graph_with(&graph, 4, 8, true, true, id_bytes, false)
3640 .unwrap()
3641 .with_blocks(&graph, bits)
3642 .unwrap();
3643 for (fuse, block_screen) in [(false, false), (true, false), (false, true)] {
3644 let mut searcher = index.searcher_with(SearchKnobs {
3645 frontier: 0.0,
3646 block_screen,
3647 ..SearchKnobs::OFF
3648 });
3649 for rerank in [0, 67] {
3650 let params = VSearchParams {
3651 ef: 1,
3652 rerank,
3653 patience: 1,
3654 fuse,
3655 screen_sigmas: 0.0,
3656 ..Default::default()
3657 };
3658 let mut out = Vec::new();
3659 for k in [
3660 0,
3661 component_size - 1,
3662 component_size,
3663 component_size + 1,
3664 66,
3665 67,
3666 68,
3667 ] {
3668 searcher.search(&rows[..16], k, params, &mut out).unwrap();
3669 assert_complete(&searcher, &out, k, rerank);
3670 }
3671 }
3672 }
3673 }
3674 }
3675 }
3676 }
3677
3678 #[test]
3679 fn recovery_reuses_marks_and_fused_traces_across_epoch_rollover() {
3680 let rows = crate::core::qgraph::tests::clustered(67, 16, 79);
3681 let mut graph = Graph::build(
3682 &rows,
3683 16,
3684 &BuildParams {
3685 degree: 2,
3686 build_ef: 16,
3687 ..Default::default()
3688 },
3689 )
3690 .unwrap();
3691 disconnect(&mut graph, 3);
3692 let index = VGraph::from_graph(&graph, true)
3693 .unwrap()
3694 .with_blocks(&graph, 2)
3695 .unwrap();
3696 let mut searcher = index.searcher_with(SearchKnobs::OFF);
3697 let mut out = Vec::new();
3698 for iteration in 0..260 {
3699 for fuse in [false, true] {
3700 let k = if iteration % 2 == 0 { 4 } else { 65 };
3701 let query = &rows[(iteration % 67) * 16..][..16];
3702 let params = VSearchParams {
3703 ef: 1,
3704 rerank: 67,
3705 patience: 1,
3706 fuse,
3707 ..Default::default()
3708 };
3709 let mut fresh = index.searcher_with(SearchKnobs::OFF);
3710 let mut expected = Vec::new();
3711 let evaluations = fresh.search(query, k, params, &mut expected).unwrap();
3712 assert_eq!(
3713 searcher.search(query, k, params, &mut out).unwrap(),
3714 evaluations
3715 );
3716 assert_eq!(out, expected);
3717 assert_complete(&searcher, &out, k, params.rerank);
3718 }
3719 }
3720 }
3721
3722 #[test]
3723 fn singleton_search_bounds_k_without_recovery() {
3724 let graph = lifted(&[2.0, 0.0], 2);
3725 let index = VGraph::from_lifted_ip_graph(&graph, 8, 8, false, false, 4).unwrap();
3726 let mut searcher = index.searcher_with(SearchKnobs::OFF);
3727 for k in [0, 1, 2, usize::MAX] {
3728 let mut out = Vec::new();
3729 searcher
3730 .search(
3731 &[1.0, 0.0],
3732 k,
3733 VSearchParams {
3734 ef: 0,
3735 ..Default::default()
3736 },
3737 &mut out,
3738 )
3739 .unwrap();
3740 assert_complete(&searcher, &out, k, 0);
3741 }
3742 }
3743
3744 fn lifted(rows: &[f32], dim: usize) -> Graph {
3745 let norms: Vec<f64> = rows
3746 .chunks_exact(dim)
3747 .map(|row| row.iter().map(|&x| f64::from(x).powi(2)).sum())
3748 .collect();
3749 let max = norms.iter().copied().fold(0.0, f64::max);
3750 let data: Vec<f32> = rows
3751 .chunks_exact(dim)
3752 .zip(norms)
3753 .flat_map(|(row, norm)| {
3754 row.iter()
3755 .copied()
3756 .chain([(max - norm).max(0.0).sqrt() as f32])
3757 })
3758 .collect();
3759 Graph::build_owned(
3760 data,
3761 dim + 1,
3762 &BuildParams {
3763 metric: Metric::Euclidean,
3764 degree: 2,
3765 build_ef: 16,
3766 ..Default::default()
3767 },
3768 )
3769 .unwrap()
3770 }
3771
3772 #[test]
3773 fn reservoir_recovers_a_beam_rejection_without_changing_expansions() {
3774 let mut graph = lifted(&[0.0, 0.0, 0.8, 0.0, 1.0, 381.0], 2);
3775 graph.entry = 0;
3776 graph.layers.clear();
3777 graph.upper_ids.clear();
3778 graph.adj.set(0, &[1, 2]);
3779 graph.adj.set(1, &[]);
3780 graph.adj.set(2, &[]);
3781 let index = VGraph::from_lifted_ip_graph(&graph, 8, 8, false, false, 4).unwrap();
3782 let mut searcher = index.searcher_with(SearchKnobs::OFF);
3783 let mut out = Vec::new();
3784 let params = VSearchParams {
3785 ef: 1,
3786 rerank: 3,
3787 ..Default::default()
3788 };
3789 searcher.search(&[1.0, 0.0], 1, params, &mut out).unwrap();
3790 assert_eq!(out, [1]);
3791 let beam = searcher.scratch.pool.keys.clone();
3792 let marks = searcher.scratch.visited.words.clone();
3793 for capacity in [1, 2, 3, usize::MAX] {
3794 let params = VSearchParams {
3795 reservoir: capacity,
3796 ..params
3797 };
3798 searcher.search(&[1.0, 0.0], 1, params, &mut out).unwrap();
3799 assert_eq!(searcher.scratch.pool.keys, beam);
3800 assert_eq!(searcher.scratch.visited.words, marks);
3801 assert_eq!(out, if capacity == 1 { vec![1] } else { vec![2] });
3802 assert_eq!(searcher.scratch.results.len(), capacity.min(3));
3803 #[cfg(feature = "experimental")]
3804 {
3805 let profile = searcher
3806 .profile_search(&[1.0, 0.0], 1, params, &mut out)
3807 .unwrap();
3808 assert_eq!(profile.expansions, 2);
3809 assert_eq!(profile.rerank_evaluations, capacity.min(3));
3810 assert_eq!(profile.beam_retained, 1);
3811 assert_eq!(profile.retained, capacity.min(3));
3812 }
3813 }
3814 }
3815
3816 #[test]
3817 fn retention_preserves_traversal_with_stopping_and_batch_scoring() {
3818 let rows = crate::core::qgraph::tests::clustered(256, 32, 77);
3819 for metric in [Metric::Angular, Metric::Euclidean] {
3820 let graph = Graph::build(
3821 &rows,
3822 32,
3823 &BuildParams {
3824 metric,
3825 degree: 8,
3826 build_ef: 32,
3827 ..Default::default()
3828 },
3829 )
3830 .unwrap();
3831 let index = VGraph::from_graph_with(&graph, 4, 8, true, true, 3, false).unwrap();
3832 for patience in [0, 1, 4] {
3833 for knobs in [
3834 SearchKnobs::OFF,
3835 SearchKnobs {
3836 prefetch: 3,
3837 batch_est: true,
3838 frontier: 0.8,
3839 block_screen: false,
3840 rerank_z: 0.0,
3841 },
3842 ] {
3843 let mut a = index.searcher_with(knobs);
3844 let mut b = index.searcher_with(knobs);
3845 for q in rows.as_chunks::<32>().0.iter().take(8) {
3846 let mut out = Vec::new();
3847 let p = VSearchParams {
3848 ef: 12,
3849 rerank: 40,
3850 patience,
3851 ..Default::default()
3852 };
3853 let ea = a.search(q, 5, p, &mut out).unwrap();
3854 let eb = b
3855 .search(q, 5, VSearchParams { reservoir: 40, ..p }, &mut out)
3856 .unwrap();
3857 assert_eq!(a.scratch.pool.keys, b.scratch.pool.keys);
3858 assert_eq!(a.scratch.pool.expanded, b.scratch.pool.expanded);
3859 assert_eq!(a.scratch.visited.words, b.scratch.visited.words);
3860 assert_eq!(ea - a.scratch.results.len(), eb - b.scratch.results.len());
3861 }
3862 }
3863 }
3864 }
3865 }
3866
3867 #[test]
3868 fn reservoir_includes_upper_layer_scores_once() {
3869 let mut graph = lifted(&[0.0, 0.0, 0.8, 0.0, 1.0, 381.0], 2);
3870 graph.upper_ids = vec![0, 1, 2];
3871 graph.layers = vec![vec![vec![1, 2], vec![0, 2], vec![0, 1]]];
3872 for v in 0..3 {
3873 graph.adj.set(v, &[]);
3874 }
3875 let index = VGraph::from_lifted_ip_graph(&graph, 8, 8, false, false, 4).unwrap();
3876 let mut searcher = index.searcher_with(SearchKnobs::OFF);
3877 let mut out = Vec::new();
3878 for _ in 0..260 {
3879 searcher
3880 .search(
3881 &[1.0, 0.0],
3882 1,
3883 VSearchParams {
3884 ef: 1,
3885 rerank: 3,
3886 reservoir: 3,
3887 ..Default::default()
3888 },
3889 &mut out,
3890 )
3891 .unwrap();
3892 assert_eq!(out, [2]);
3893 assert_eq!(searcher.scratch.keys.len(), 3);
3894 }
3895 }
3896
3897 #[test]
3898 fn direct_ip_preserves_magnitudes_zero_queries_and_topology() {
3899 let rows = [1.0, 0.0, 0.0, 3.0, 0.0, 3.0, -2.0, 0.0];
3900 let mut graph = lifted(&rows, 2);
3901 for v in 0..graph.n {
3902 graph.adj.set(v, &[((v + 1) % graph.n) as u32]);
3903 }
3904 for bits in [4, 8] {
3905 for reorder in [false, true] {
3906 let index =
3907 VGraph::from_lifted_ip_graph(&graph, bits, 8, true, reorder, 3).unwrap();
3908 assert_eq!(index.metric(), VMetric::InnerProduct);
3909 assert!(index.half_norms.is_empty());
3910 for v in 0..index.n {
3911 let mut neighbors = Vec::new();
3912 index.for_each_neighbour(v, |u| neighbors.push(index.original(u)));
3913 assert_eq!(
3914 neighbors,
3915 graph.neighbors(index.original(v as u32) as usize)
3916 );
3917 }
3918 let mut searcher = index.searcher_with(SearchKnobs::OFF);
3919 let mut out = Vec::new();
3920 let params = VSearchParams {
3921 ef: 4,
3922 rerank: 4,
3923 reservoir: 4,
3924 ..Default::default()
3925 };
3926 for (query, expected) in [
3927 ([1.0, 1.0], vec![1, 2, 0, 3]),
3928 ([0.0, 0.0], vec![0, 1, 2, 3]),
3929 ([-1.0, 0.0], vec![3, 1, 2, 0]),
3930 ] {
3931 searcher.search(&query, 4, params, &mut out).unwrap();
3932 assert_eq!(out, expected);
3933 }
3934 assert!(searcher
3935 .search(&[f32::NAN, 0.0], 1, params, &mut out)
3936 .is_err());
3937 assert!(searcher.search(&[0.0], 1, params, &mut out).is_err());
3938 assert!(index.with_blocks(&graph, 2).is_err());
3939 }
3940 }
3941 }
3942
3943 #[test]
3944 fn direct_ip_removes_the_128_dimension_stride_cliff() {
3945 let mut rows = vec![0.0; 3 * 128];
3946 rows[0] = 1.0;
3947 rows[128] = 2.0;
3948 rows[256] = -1.0;
3949 let graph = lifted(&rows, 128);
3950 let old = VGraph::from_graph_with(&graph, 8, 8, true, true, 3, false).unwrap();
3951 let new = VGraph::from_lifted_ip_graph(&graph, 8, 8, true, true, 3).unwrap();
3952 assert_eq!(
3953 (old.dim, old.padded, old.stride, old.res_stride),
3954 (129, 144, 256, 144)
3955 );
3956 assert_eq!(
3957 (new.dim, new.padded, new.stride, new.res_stride),
3958 (128, 128, 128, 128)
3959 );
3960 assert_eq!(
3961 old.index_bytes() - new.index_bytes(),
3962 148 * graph.n + 16 * 4
3963 );
3964 assert_eq!(old.entry, new.entry);
3965 assert_eq!(old.upper_ids, new.upper_ids);
3966 assert_eq!(old.layers, new.layers);
3967 assert_eq!(&*old.ids, &*new.ids);
3968 assert_eq!(&*old.adj3, &*new.adj3);
3969 for dim in [127, 129] {
3970 let mut rows = vec![0.0; 3 * dim];
3971 rows[0] = 1.0;
3972 rows[dim] = 2.0;
3973 rows[2 * dim] = -1.0;
3974 let graph = lifted(&rows, dim);
3975 let old = VGraph::from_graph_with(&graph, 8, 8, true, true, 3, false).unwrap();
3976 let new = VGraph::from_lifted_ip_graph(&graph, 8, 8, true, true, 3).unwrap();
3977 assert_eq!((new.stride, new.res_stride), (old.stride, old.res_stride));
3978 assert_eq!(old.index_bytes() - new.index_bytes(), 4 * graph.n);
3979 }
3980 }
3981}
3982
3983const CACHE_MAGIC: &[u8; 8] = b"HMSGRF04";
3987
3988fn data_fingerprint(vectors: &[f32]) -> u64 {
3990 let mut h = super::Fnv::new();
3991 vectors.iter().for_each(|x| h.word(x.to_bits()));
3992 h.finish()
3993}
3994
3995fn put_u32s(w: &mut impl std::io::Write, xs: &[u32]) -> std::io::Result<()> {
3996 w.write_all(&(xs.len() as u64).to_le_bytes())?;
3997 for x in xs {
3998 w.write_all(&x.to_le_bytes())?;
3999 }
4000 Ok(())
4001}
4002
4003fn get_u64(r: &mut impl std::io::Read) -> std::io::Result<u64> {
4004 let mut b = [0u8; 8];
4005 r.read_exact(&mut b)?;
4006 Ok(u64::from_le_bytes(b))
4007}
4008
4009fn get_u32s(r: &mut impl std::io::Read, max: u64) -> std::io::Result<Vec<u32>> {
4010 let len = get_u64(r)?;
4011 if len > max {
4012 return Err(std::io::Error::new(
4013 std::io::ErrorKind::InvalidData,
4014 "graph cache: list longer than the graph",
4015 ));
4016 }
4017 let want = len.saturating_mul(4);
4020 let mut bytes = Vec::new();
4021 std::io::Read::read_to_end(&mut std::io::Read::take(&mut *r, want), &mut bytes)?;
4022 if bytes.len() as u64 != want {
4023 return Err(std::io::Error::new(
4024 std::io::ErrorKind::UnexpectedEof,
4025 "graph cache: list shorter than its lengths",
4026 ));
4027 }
4028 Ok(bytes
4029 .as_chunks::<4>()
4030 .0
4031 .iter()
4032 .map(|c| u32::from_le_bytes(*c))
4033 .collect())
4034}
4035
4036fn put_lists<'a>(
4037 w: &mut impl std::io::Write,
4038 lists: impl Iterator<Item = &'a [u32]> + Clone,
4039) -> std::io::Result<()> {
4040 let lens: Vec<u32> = lists.clone().map(|l| l.len() as u32).collect();
4041 put_u32s(w, &lens)?;
4042 put_u32s(w, &lists.flatten().copied().collect::<Vec<u32>>())
4043}
4044
4045fn get_lists(r: &mut impl std::io::Read, n: u64) -> std::io::Result<Vec<Vec<u32>>> {
4046 let lens = get_u32s(r, n)?;
4047 let flat = get_u32s(r, lens.iter().map(|&l| u64::from(l)).sum())?;
4048 let mut out = Vec::with_capacity(lens.len());
4049 let mut at = 0;
4050 for l in lens {
4051 out.push(flat[at..at + l as usize].to_vec());
4052 at += l as usize;
4053 }
4054 Ok(out)
4055}
4056
4057impl Graph {
4060 pub fn save(&self, path: &std::path::Path) -> std::io::Result<()> {
4064 if self.degree != self.build.degree {
4065 return Err(std::io::Error::new(
4066 std::io::ErrorKind::InvalidInput,
4067 "query-trained graph cannot use an untrained graph cache key; encode and save VGraph instead",
4068 ));
4069 }
4070 let mut w = std::io::BufWriter::new(std::fs::File::create(path)?);
4071 use std::io::Write;
4072 w.write_all(CACHE_MAGIC)?;
4073 let header = [
4074 self.n as u64,
4075 self.dim as u64,
4076 self.degree as u64,
4077 self.build.build_ef as u64,
4078 u64::from(self.build.alpha.to_bits()),
4079 self.build.seed,
4080 u64::from(self.entry),
4081 data_fingerprint(&self.vectors),
4082 u64::from(self.build.quantized_search),
4083 self.metric.code(),
4084 ];
4085 for x in header {
4086 w.write_all(&x.to_le_bytes())?;
4087 }
4088 put_lists(&mut w, self.adj.iter())?;
4089 put_u32s(&mut w, &self.upper_ids)?;
4090 w.write_all(&(self.layers.len() as u64).to_le_bytes())?;
4091 for l in &self.layers {
4092 put_lists(&mut w, l.iter().map(Vec::as_slice))?;
4093 }
4094 w.flush()
4095 }
4096
4097 pub fn load(
4103 data: &[f32],
4104 dim: usize,
4105 path: &std::path::Path,
4106 params: &BuildParams,
4107 ) -> std::io::Result<Self> {
4108 let bad = |m: String| std::io::Error::new(std::io::ErrorKind::InvalidData, m);
4109 let mut r = std::io::BufReader::new(std::fs::File::open(path)?);
4110 let mut magic = [0u8; 8];
4111 std::io::Read::read_exact(&mut r, &mut magic)?;
4112 if &magic != CACHE_MAGIC {
4113 return Err(bad(
4114 "graph cache: bad magic (an older format or not a cache file); delete it to rebuild"
4115 .into(),
4116 ));
4117 }
4118 let mut header = [0u64; 10];
4119 for x in &mut header {
4120 *x = get_u64(&mut r)?;
4121 }
4122 let [n, d, degree, build_ef, alpha_bits, seed, entry, fingerprint, quantized, metric] =
4123 header;
4124 if dim == 0 || !data.len().is_multiple_of(dim) {
4125 return Err(bad("graph cache: ragged data".into()));
4126 }
4127 let mut vectors = data.to_vec();
4128 let norms = prepare_rows(&mut vectors, dim, params.metric);
4129 let checks: [(&str, u64, u64); 9] = [
4130 ("metric", metric, params.metric.code()),
4131 ("n", n, (data.len() / dim) as u64),
4132 ("dim", d, dim as u64),
4133 ("degree", degree, params.degree as u64),
4134 ("build_ef", build_ef, params.build_ef as u64),
4135 ("alpha", alpha_bits, u64::from(params.alpha.to_bits())),
4136 ("seed", seed, params.seed),
4137 ("data fingerprint", fingerprint, data_fingerprint(&vectors)),
4138 (
4139 "quantized_search",
4140 quantized,
4141 u64::from(params.quantized_search),
4142 ),
4143 ];
4144 if let Some((what, file, want)) = checks.iter().find(|(_, a, b)| a != b) {
4145 let show = |x: u64| {
4146 if *what == "alpha" {
4147 f32::from_bits(x as u32).to_string()
4148 } else {
4149 format!("{x:#x}")
4150 }
4151 };
4152 return Err(bad(format!(
4153 "graph cache: {what} {} in the file, {} requested; delete the file to rebuild",
4154 show(*file),
4155 show(*want)
4156 )));
4157 }
4158 let n = n as usize;
4159 if entry >= n as u64 {
4160 return Err(bad("graph cache: entry out of range".into()));
4161 }
4162 let lens = get_u32s(&mut r, n as u64)?;
4163 let flat = get_u32s(&mut r, lens.iter().map(|&l| u64::from(l)).sum())?;
4164 let adj = Adjacency::from_flat(lens, &flat, degree as usize)
4165 .ok_or_else(|| bad("graph cache: list longer than the degree".into()))?;
4166 drop(flat);
4167 let upper_ids = get_u32s(&mut r, n as u64)?;
4168 let nl = get_u64(&mut r)?;
4169 if nl > 64 {
4170 return Err(bad("graph cache: too many layers".into()));
4171 }
4172 let layers = (0..nl)
4173 .map(|_| get_lists(&mut r, n as u64))
4174 .collect::<std::io::Result<Vec<_>>>()?;
4175 let ok = adj.len() == n
4176 && adj.iter().flatten().all(|&u| (u as usize) < n)
4177 && upper_ids.iter().all(|&u| (u as usize) < n)
4178 && layers.iter().all(|l| {
4179 l.len() <= upper_ids.len() && l.iter().flatten().all(|&u| (u as usize) < l.len())
4180 });
4181 if !ok {
4182 return Err(bad("graph cache: id out of range".into()));
4183 }
4184 Ok(Self {
4185 n,
4186 dim,
4187 degree: degree as usize,
4188 metric: params.metric,
4189 vectors,
4190 norms,
4191 adj,
4192 entry: entry as u32,
4193 upper_ids,
4194 layers,
4195 build: params.clone(),
4196 })
4197 }
4198}
4199
4200#[cfg(test)]
4201mod tests {
4202 use super::super::kernels::SplitMix;
4203 use super::super::tests::{brute_force, clustered, gaussian};
4204 use super::*;
4205
4206 #[test]
4207 #[cfg(feature = "experimental-query-x4")]
4208 fn query_x4_estimates_match_serial_at_padding_boundaries() {
4209 for dim in [
4210 15, 16, 17, 31, 32, 33, 63, 64, 65, 111, 112, 113, 127, 128, 129,
4211 ] {
4212 let data = clustered(40, dim, 173);
4213 let queries = clustered(2, dim, 174);
4214 for metric in [Metric::Angular, Metric::Euclidean] {
4215 let graph = Graph::build(
4216 &data,
4217 dim,
4218 &BuildParams {
4219 metric,
4220 degree: 8,
4221 build_ef: 16,
4222 quantized_search: true,
4223 ..Default::default()
4224 },
4225 )
4226 .unwrap();
4227 for bits in [4, 8] {
4228 for align in [false, true] {
4229 let index = VGraph::from_graph_with(&graph, bits, 8, align, true, 3, false)
4230 .unwrap();
4231 let mut searcher = index.searcher_with(SearchKnobs::OFF);
4232 let scratch = &mut searcher.scratch;
4233 for query in queries.chunks_exact(dim) {
4234 scratch.qf[..dim].copy_from_slice(query);
4235 index.encode_query(
4236 &scratch.qf,
4237 &mut scratch.qc,
4238 &mut scratch.qtmp,
4239 &mut scratch.qrot,
4240 &mut scratch.qs,
4241 );
4242 let qsum = query_sum(&scratch.qc);
4243 for ids in [[0, 1, 2, 3], [39, 16, 31, 0], [5, 5, 5, 5]] {
4244 for w in [0.0, 1.0, 3.25] {
4245 let expected =
4246 ids.map(|u| index.est(&scratch.qc, qsum, u, w).to_bits());
4247 let actual = index
4248 .est4::<true>(&scratch.qc, qsum, ids, w)
4249 .map(f32::to_bits);
4250 assert_eq!(actual, expected, "dim {dim}, bits {bits}");
4251 }
4252 }
4253 }
4254 }
4255 }
4256 }
4257 }
4258 }
4259
4260 #[test]
4261 #[cfg(feature = "experimental-query-x4")]
4262 fn query_x4_preserves_admission_state_at_batch_boundaries() {
4263 let dim = 17;
4264 let data = clustered(48, dim, 175);
4265 let query = clustered(1, dim, 176);
4266 for metric in [Metric::Angular, Metric::Euclidean] {
4267 let graph = Graph::build(
4268 &data,
4269 dim,
4270 &BuildParams {
4271 metric,
4272 degree: 8,
4273 build_ef: 16,
4274 quantized_search: true,
4275 ..Default::default()
4276 },
4277 )
4278 .unwrap();
4279 for bits in [4, 8] {
4280 let index = VGraph::from_graph_with(&graph, bits, 8, true, true, 3, false).unwrap();
4281 for count in (0..=9).chain(15..=17) {
4282 for cap in [3, 4, 5, 8] {
4283 for batch_est in [false, true] {
4284 for patience in [0, 4] {
4285 for reservoir in [0, 12] {
4286 let knobs = SearchKnobs {
4287 batch_est,
4288 ..SearchKnobs::OFF
4289 };
4290 let mut serial = index.searcher_with(knobs);
4291 let mut fused = index.searcher_with(knobs);
4292 for s in [&mut serial.scratch, &mut fused.scratch] {
4293 s.qf[..dim].copy_from_slice(&query);
4294 index.encode_query(
4295 &s.qf,
4296 &mut s.qc,
4297 &mut s.qtmp,
4298 &mut s.qrot,
4299 &mut s.qs,
4300 );
4301 s.qsum = query_sum(&s.qc);
4302 s.pool.reset(cap);
4303 s.reservoir.reset(reservoir);
4304 for id in 32..48 {
4305 let kk = key(index.est(&s.qc, s.qsum, id, 1.0), id);
4306 s.pool.insert(kk);
4307 s.reservoir.offer(kk);
4308 }
4309 s.pool.next();
4310 s.fresh.extend(0..count);
4311 }
4312 let (mut a, mut b) = (patience, patience);
4313 let expected = serial.scratch.score_fresh::<{
4314 cfg!(feature = "experimental-batch-pool")
4315 }, false>(
4316 &index,
4317 cap.min(3),
4318 patience,
4319 &mut a,
4320 1.0,
4321 );
4322 let actual = fused.scratch.score_fresh::<{
4323 cfg!(feature = "experimental-batch-pool")
4324 }, true>(
4325 &index,
4326 cap.min(3),
4327 patience,
4328 &mut b,
4329 1.0,
4330 );
4331 assert_eq!(actual, count as usize);
4332 assert_eq!(actual, expected);
4333 assert_eq!(a, b);
4334 assert_eq!(serial.scratch.pool.keys, fused.scratch.pool.keys);
4335 assert_eq!(
4336 serial.scratch.pool.expanded,
4337 fused.scratch.pool.expanded
4338 );
4339 assert_eq!(
4340 serial.scratch.pool.cursor,
4341 fused.scratch.pool.cursor
4342 );
4343 assert_eq!(
4344 serial.scratch.reservoir.keys,
4345 fused.scratch.reservoir.keys
4346 );
4347 }
4348 }
4349 }
4350 }
4351 }
4352 }
4353 }
4354 }
4355
4356 #[test]
4357 #[cfg(feature = "experimental-query-x4")]
4358 fn query_x4_preserves_search_results_work_and_reopening() {
4359 let directory = tempfile::tempdir().unwrap();
4360 for dim in [17, 33] {
4361 let mut data = clustered(80, dim, 177);
4362 let first = data[..dim].to_vec();
4363 data[7 * dim..8 * dim].copy_from_slice(&first);
4364 data[19 * dim..20 * dim].copy_from_slice(&first);
4365 let queries = clustered(2, dim, 178);
4366 for metric in [Metric::Angular, Metric::Euclidean] {
4367 let graph = Graph::build(
4368 &data,
4369 dim,
4370 &BuildParams {
4371 metric,
4372 degree: 16,
4373 build_ef: 32,
4374 quantized_search: true,
4375 ..Default::default()
4376 },
4377 )
4378 .unwrap();
4379 for bits in [4, 8] {
4380 for mode in 0..3 {
4381 if mode > 0 && metric != Metric::Angular {
4382 continue;
4383 }
4384 let mut index =
4385 VGraph::from_graph_with(&graph, bits, 8, true, true, 3, mode == 1)
4386 .unwrap();
4387 if mode == 2 {
4388 index = index.with_blocks(&graph, 4).unwrap();
4389 }
4390 let path = directory.path().join("query-x4.warm");
4391 index.save(&path).unwrap();
4392 let reopened = VGraph::open(&path, Residency::Mapped).unwrap();
4393 let knobs = SearchKnobs {
4394 block_screen: mode == 2,
4395 ..SearchKnobs::OFF
4396 };
4397 let mut serial = index.searcher_with(knobs);
4398 let mut fused = reopened.searcher_with(knobs);
4399 for (ef, neighbour_cap, patience, reservoir) in
4400 [(3, 0, 0, 0), (4, 3, 1, 0), (5, 4, 4, 0), (17, 5, 0, 24)]
4401 {
4402 if mode > 0 && reservoir > 0 {
4403 continue;
4404 }
4405 let params = VSearchParams {
4406 ef,
4407 rerank: 24,
4408 neighbour_cap,
4409 patience,
4410 reservoir,
4411 ..Default::default()
4412 };
4413 for k in [0, 1, 5] {
4414 for query in queries.chunks_exact(dim) {
4415 let (mut expected, mut actual) = (Vec::new(), Vec::new());
4416 let (mut a, mut b) =
4417 (QueryProfile::default(), QueryProfile::default());
4418 let ea = serial
4419 .scratch
4420 .search_inner::<
4421 true,
4422 { cfg!(feature = "experimental-batch-pool") },
4423 false,
4424 >(&index, query, k, params, &mut expected, &mut a)
4425 .unwrap();
4426 let eb = fused
4427 .scratch
4428 .search_inner::<
4429 true,
4430 { cfg!(feature = "experimental-batch-pool") },
4431 true,
4432 >(&reopened, query, k, params, &mut actual, &mut b)
4433 .unwrap();
4434 assert_eq!(expected, actual);
4435 assert_eq!(ea, eb);
4436 let work = |p: QueryProfile| {
4437 (
4438 p.expansions,
4439 p.neighbors_read,
4440 p.entry_evaluations,
4441 p.traversal_evaluations,
4442 p.rerank_evaluations,
4443 p.beam_retained,
4444 p.retained,
4445 )
4446 };
4447 assert_eq!(work(a), work(b));
4448 assert_eq!(serial.scratch.pool.keys, fused.scratch.pool.keys);
4449 assert_eq!(
4450 serial.scratch.pool.expanded,
4451 fused.scratch.pool.expanded
4452 );
4453 assert_eq!(
4454 serial.scratch.pool.cursor,
4455 fused.scratch.pool.cursor
4456 );
4457 assert_eq!(
4458 serial.scratch.visited.words,
4459 fused.scratch.visited.words
4460 );
4461 assert_eq!(
4462 serial.scratch.reservoir.keys,
4463 fused.scratch.reservoir.keys
4464 );
4465 assert_eq!(serial.last_screened(), fused.last_screened());
4466 let mut ordinary = Vec::new();
4467 assert_eq!(
4468 fused.search(query, k, params, &mut ordinary).unwrap(),
4469 eb
4470 );
4471 assert_eq!(ordinary, actual);
4472 }
4473 }
4474 }
4475 }
4476 }
4477 }
4478 }
4479 }
4480
4481 #[test]
4484 fn refused_huge_page_advice_changes_nothing_but_the_counter() {
4485 let (n, dim) = (40_000, 128);
4487 let data = clustered(n, dim, 3);
4488 let graph = Graph::build(
4489 &data,
4490 dim,
4491 &BuildParams {
4492 degree: 8,
4493 build_ef: 16,
4494 metric: Metric::Euclidean,
4495 ..BuildParams::default()
4496 },
4497 )
4498 .unwrap();
4499 let build = || VGraph::from_graph_with(&graph, 8, 8, true, true, 3, false).unwrap();
4500 FAIL_HUGE_ADVICE.store(true, Ordering::Relaxed);
4501 let failed_before = huge_page_status().failed_calls;
4502 let refused = build();
4503 FAIL_HUGE_ADVICE.store(false, Ordering::Relaxed);
4504 let advised = build();
4505 if cfg!(target_os = "linux") && huge_pages() {
4506 assert!(huge_page_status().failed_calls > failed_before);
4507 }
4508 assert_eq!(refused.codes.bytes(), advised.codes.bytes());
4509 assert_eq!(&refused.scales[..], &advised.scales[..]);
4510 assert_eq!(&refused.half_norms[..], &advised.half_norms[..]);
4511 let params = VSearchParams {
4512 ef: 40,
4513 rerank: 40,
4514 ..VSearchParams::default()
4515 };
4516 let (mut a, mut b) = (refused.searcher(), advised.searcher());
4517 let (mut ra, mut rb) = (Vec::new(), Vec::new());
4518 for q in data.chunks_exact(dim).step_by(997) {
4519 a.search(q, 10, params, &mut ra).unwrap();
4520 b.search(q, 10, params, &mut rb).unwrap();
4521 assert_eq!(ra, rb);
4522 }
4523 }
4524
4525 #[test]
4526 fn visited_bits_reset_between_searches() {
4527 let n = 64 * 5 + 3;
4529 let mut marks = Marks::new(n, false);
4530 let mut rng = SplitMix(5);
4531 for search in 0..300 {
4532 marks.next();
4533 let mut set = std::collections::HashSet::new();
4534 let fixed = [0, 63, 64, 127, n as u32 - 1];
4535 for i in 0..40 {
4536 let id = if i < fixed.len() && search % 2 == 0 {
4537 fixed[i]
4538 } else {
4539 (rng.next_u64() % n as u64) as u32
4540 };
4541 assert_eq!(marks.seen(id), set.contains(&id));
4542 assert_eq!(marks.insert(id), set.insert(id));
4543 assert!(marks.seen(id));
4544 }
4545 let ones: u32 = marks.words.iter().map(|w| w.count_ones()).sum();
4546 assert_eq!(ones as usize, set.len());
4547 }
4548 }
4549
4550 #[test]
4551 #[cfg(feature = "experimental-neighbor-admission")]
4552 fn neighbor_admission_matches_ordered_set_model_at_boundaries() {
4553 fn check<const PROFILE: bool, const SCREEN: bool>(
4554 index: &VGraph,
4555 vertex: usize,
4556 list: &[u32],
4557 cap: usize,
4558 ) {
4559 let mut visited = Marks::new(index.n, false);
4560 for initial in [&[][..], &[0, 63, 64, 127, 128, 130][..], list] {
4561 visited.next();
4562 let mut seen = std::collections::HashSet::new();
4563 for &u in initial {
4564 visited.insert(u);
4565 seen.insert(u);
4566 }
4567 let limit = if cap == 0 { list.len() } else { cap };
4568 let expected: Vec<_> = list
4569 .iter()
4570 .take(limit)
4571 .copied()
4572 .filter(|&u| seen.insert(u))
4573 .collect();
4574 let mut fresh = Vec::new();
4575 let read = index.admit_neighbours::<PROFILE, SCREEN>(
4576 vertex,
4577 cap,
4578 &mut visited,
4579 &mut fresh,
4580 );
4581 assert_eq!(fresh, expected, "cap={cap}, vertex={vertex}");
4582 assert_eq!(read, if PROFILE && cap == 0 { list.len() } else { 0 });
4584 for u in 0..index.n as u32 {
4585 assert_eq!(visited.seen(u), seen.contains(&u));
4586 }
4587 let mut words = std::collections::HashSet::new();
4588 let touched: Vec<_> = initial
4589 .iter()
4590 .chain(&expected)
4591 .map(|&u| u >> 6)
4592 .filter(|&word| words.insert(word))
4593 .collect();
4594 assert_eq!(visited.touched, touched);
4595 visited.next();
4596 assert!(visited.words.iter().all(|&word| word == 0));
4597 assert!(visited.touched.is_empty());
4598 }
4599 }
4600
4601 let dim = 17;
4602 let data = clustered(131, dim, 91);
4603 let lists = [
4604 Vec::new(),
4605 vec![130],
4606 vec![63, 64, 65, 127, 128, 129, 130, 63, 64, 0, 0],
4607 (0..63).rev().collect(),
4608 (0..64).rev().collect(),
4609 (0..65).rev().collect(),
4610 ];
4611 for metric in [Metric::Angular, Metric::Euclidean] {
4612 let mut graph = Graph::build(
4613 &data,
4614 dim,
4615 &BuildParams {
4616 metric,
4617 degree: 66,
4618 build_ef: 80,
4619 ..BuildParams::default()
4620 },
4621 )
4622 .unwrap();
4623 for bytes in [3, 4, 0] {
4624 let ordered: Vec<Vec<u32>> = lists
4625 .iter()
4626 .map(|list| {
4627 let mut list = list.clone();
4628 if bytes == 0 {
4629 list.sort_unstable();
4630 list.dedup();
4631 }
4632 list
4633 })
4634 .collect();
4635 for v in 0..graph.n {
4636 graph.adj.set(v, &ordered[v % ordered.len()]);
4637 }
4638 let mut index = VGraph::from_graph_with(
4639 &graph,
4640 8,
4641 8,
4642 true,
4643 false,
4644 bytes,
4645 metric == Metric::Angular,
4646 )
4647 .unwrap();
4648 for compact in [false, true] {
4649 if compact {
4650 index.compact_adjacency().unwrap();
4651 assert!(bytes == 0 || !index.adj_off.is_empty());
4652 }
4653 for (vertex, list) in ordered.iter().enumerate() {
4654 for cap in [0, 1, 2, 63, 64, 65, 66, usize::MAX] {
4655 check::<false, false>(&index, vertex, list, cap);
4656 check::<true, false>(&index, vertex, list, cap);
4657 if metric == Metric::Angular {
4658 check::<false, true>(&index, vertex, list, cap);
4659 check::<true, true>(&index, vertex, list, cap);
4660 }
4661 }
4662 }
4663 }
4664 }
4665 }
4666 }
4667
4668 struct ReferencePool {
4670 keys: Vec<u64>,
4671 expanded: Vec<bool>,
4672 cap: usize,
4673 cursor: usize,
4674 }
4675
4676 impl ReferencePool {
4677 fn insert(&mut self, key: u64) -> Option<(usize, Option<u64>)> {
4678 let n = self.keys.len();
4679 if n >= self.cap && key >= self.keys[n - 1] {
4680 return None;
4681 }
4682 let pos = self.keys.partition_point(|&t| t < key);
4683 let mut evicted = None;
4684 if n >= self.cap {
4685 evicted = self.keys.pop();
4686 self.expanded.pop();
4687 }
4688 self.keys.insert(pos, key);
4689 self.expanded.insert(pos, false);
4690 if pos < self.cursor {
4691 self.cursor = pos;
4692 }
4693 Some((pos, evicted))
4694 }
4695
4696 fn remove(&mut self, key: u64) {
4697 let pos = self.keys.partition_point(|&t| t < key);
4698 if pos < self.keys.len() && self.keys[pos] == key {
4699 self.keys.remove(pos);
4700 self.expanded.remove(pos);
4701 if pos < self.cursor {
4702 self.cursor -= 1;
4703 }
4704 }
4705 }
4706
4707 fn next(&mut self) -> Option<u64> {
4708 while self.cursor < self.keys.len() {
4709 if !self.expanded[self.cursor] {
4710 self.expanded[self.cursor] = true;
4711 return Some(self.keys[self.cursor]);
4712 }
4713 self.cursor += 1;
4714 }
4715 None
4716 }
4717
4718 fn peek_unexpanded(&self, n: usize) -> Option<u64> {
4719 (self.cursor..self.keys.len())
4720 .filter(|&i| !self.expanded[i])
4721 .nth(n)
4722 .map(|i| self.keys[i])
4723 }
4724 }
4725
4726 #[test]
4727 fn pool_bits_follow_the_bool_pool() {
4728 let mut rng = SplitMix(9);
4729 let mut pool = Pool {
4730 keys: Vec::new(),
4731 expanded: Vec::new(),
4732 #[cfg(feature = "experimental-batch-pool")]
4733 merge_keys: Vec::new(),
4734 #[cfg(feature = "experimental-batch-pool")]
4735 merge_expanded: Vec::new(),
4736 cap: 1,
4737 cursor: 0,
4738 };
4739 for cap in [1usize, 2, 10, 63, 64, 65, 127, 128, 129, 300, 10_000] {
4741 for _ in 0..20 {
4742 pool.reset(cap);
4743 let mut model = ReferencePool {
4744 keys: Vec::new(),
4745 expanded: Vec::new(),
4746 cap,
4747 cursor: 0,
4748 };
4749 for _ in 0..600 {
4750 match rng.next_u64() % 8 {
4751 0 | 1 => assert_eq!(pool.next(), model.next()),
4752 2 if !model.keys.is_empty() => {
4753 let k = model.keys[(rng.next_u64() % model.keys.len() as u64) as usize];
4754 pool.remove(k);
4755 model.remove(k);
4756 }
4757 3 => {
4758 let n = (rng.next_u64() % 3) as usize;
4759 assert_eq!(pool.peek_unexpanded(n), model.peek_unexpanded(n));
4760 }
4761 _ => {
4762 let k = rng.next_u64() % 5000;
4763 assert_eq!(pool.insert(k), model.insert(k));
4764 }
4765 }
4766 assert_eq!(pool.keys, model.keys);
4767 assert_eq!(pool.cursor, model.cursor);
4768 let flags: Vec<bool> =
4769 (0..pool.keys.len()).map(|i| pool.is_expanded(i)).collect();
4770 assert_eq!(flags, model.expanded);
4771 let set: u32 = pool.expanded.iter().map(|w| w.count_ones()).sum();
4772 assert_eq!(set as usize, flags.iter().filter(|&&f| f).count());
4773 }
4774 }
4775 }
4776 }
4777
4778 #[test]
4779 #[cfg(feature = "experimental-batch-pool")]
4780 fn batch_pool_matches_sequential_insertion() {
4781 fn check(pool: &Pool, model: &ReferencePool) {
4782 assert_eq!(pool.keys, model.keys);
4783 assert_eq!(pool.cursor, model.cursor);
4784 let flags: Vec<_> = (0..pool.keys.len()).map(|i| pool.is_expanded(i)).collect();
4785 assert_eq!(flags, model.expanded);
4786 assert_eq!(
4787 pool.expanded
4788 .iter()
4789 .map(|word| word.count_ones())
4790 .sum::<u32>() as usize,
4791 flags.iter().filter(|&&flag| flag).count()
4792 );
4793 for next in [0, 1, 2] {
4794 assert_eq!(pool.peek_unexpanded(next), model.peek_unexpanded(next));
4795 }
4796 }
4797
4798 let mut rng = SplitMix(539365412);
4799 for cap in [1usize, 2, 7, 8, 9, 63, 64, 65, 127, 128, 129] {
4800 for k in [0, 1, cap - 1, cap, cap + 1] {
4801 for count in [0, 1, 7, 8, 9, cap + 1] {
4802 for fill in [0, cap / 2, cap] {
4803 let mut pool = Pool {
4804 keys: Vec::new(),
4805 expanded: Vec::new(),
4806 #[cfg(feature = "experimental-batch-pool")]
4807 merge_keys: Vec::new(),
4808 #[cfg(feature = "experimental-batch-pool")]
4809 merge_expanded: Vec::new(),
4810 cap,
4811 cursor: 0,
4812 };
4813 let mut model = ReferencePool {
4814 keys: Vec::new(),
4815 expanded: Vec::new(),
4816 cap,
4817 cursor: 0,
4818 };
4819 let mut id = 0;
4820 for _ in 0..fill {
4821 let candidate = key(1000.0 + (rng.next_u64() % 16) as f32, id);
4822 id += 1;
4823 assert_eq!(pool.insert(candidate), model.insert(candidate));
4824 }
4825 for _ in 0..=fill {
4826 assert_eq!(pool.next(), model.next());
4827 }
4828 for round in 0..8 {
4829 let size = if round == 0 {
4830 count
4831 } else {
4832 (rng.next_u64() % 17) as usize
4833 };
4834 let mut pending = Vec::new();
4835 for _ in 0..size {
4836 let distance = if round == 0 {
4837 (rng.next_u64() % 16) as f32
4838 } else {
4839 (rng.next_u64() % 32) as f32 - 8.0
4840 };
4841 pending.push(key(distance, id));
4842 id += 1;
4843 }
4844 let mut improved = false;
4845 for &candidate in &pending {
4846 if let Some((position, _)) = model.insert(candidate) {
4847 improved |= position < k;
4848 }
4849 }
4850 assert_eq!(pool.insert_batch(&mut pending, k), improved);
4851 check(&pool, &model);
4852 for _ in 0..3 {
4853 assert_eq!(pool.next(), model.next());
4854 }
4855 if round % 2 == 1 && !model.keys.is_empty() {
4856 let removed =
4857 model.keys[(rng.next_u64() % model.keys.len() as u64) as usize];
4858 pool.remove(removed);
4859 model.remove(removed);
4860 }
4861 check(&pool, &model);
4862 }
4863 }
4864 }
4865 }
4866 }
4867 }
4868
4869 #[test]
4870 #[cfg(feature = "experimental-batch-pool")]
4871 fn batch_pool_equal_distances_preserve_distinct_ids_and_flags() {
4872 for cap in [1usize, 2, 7, 8, 9, 63, 64, 65, 127, 128, 129] {
4873 for fill in [cap - 1, cap] {
4874 for first_id in [10, 1000] {
4875 let mut pool = Pool {
4876 keys: Vec::new(),
4877 expanded: Vec::new(),
4878 #[cfg(feature = "experimental-batch-pool")]
4879 merge_keys: Vec::new(),
4880 #[cfg(feature = "experimental-batch-pool")]
4881 merge_expanded: Vec::new(),
4882 cap,
4883 cursor: 0,
4884 };
4885 let mut model = ReferencePool {
4886 keys: Vec::new(),
4887 expanded: Vec::new(),
4888 cap,
4889 cursor: 0,
4890 };
4891 for id in first_id..first_id + fill as u32 {
4892 assert_eq!(pool.insert(key(1.0, id)), model.insert(key(1.0, id)));
4893 }
4894 for _ in 0..=fill {
4895 assert_eq!(pool.next(), model.next());
4896 }
4897 let mut pending: Vec<_> = [0, 200, 201, 202, 203, 204, 205, 206, 207]
4898 .map(|id| key(1.0, id))
4899 .to_vec();
4900 let mut improved = false;
4901 for &candidate in &pending {
4902 if let Some((position, _)) = model.insert(candidate) {
4903 improved |= position < 1;
4904 }
4905 }
4906 assert_eq!(pool.insert_batch(&mut pending, 1), improved);
4907 assert_eq!(pool.keys, model.keys);
4908 assert_eq!(pool.cursor, model.cursor);
4909 for (position, &expanded) in model.expanded.iter().enumerate() {
4910 assert_eq!(pool.is_expanded(position), expanded);
4911 }
4912 while let Some(expected) = model.next() {
4913 assert_eq!(pool.next(), Some(expected));
4914 }
4915 assert_eq!(pool.next(), None);
4916 }
4917 }
4918 }
4919 }
4920
4921 #[test]
4922 #[cfg(feature = "experimental-batch-pool")]
4923 fn batch_pool_preserves_search_results_work_and_reopening() {
4924 let dim = 16;
4925 let data = clustered(160, dim, 53);
4926 let queries = clustered(3, dim, 54);
4927 let directory = tempfile::tempdir().unwrap();
4928 for metric in [Metric::Angular, Metric::Euclidean] {
4929 let graph = Graph::build(
4930 &data,
4931 dim,
4932 &BuildParams {
4933 metric,
4934 degree: 32,
4935 build_ef: 64,
4936 quantized_search: true,
4937 ..Default::default()
4938 },
4939 )
4940 .unwrap();
4941 for bits in [4, 8] {
4942 let index = VGraph::from_graph_with(&graph, bits, 8, true, true, 3, false).unwrap();
4943 let path = directory.path().join("batch-pool.warm");
4944 index.save(&path).unwrap();
4945 let reopened = VGraph::open(&path, Residency::Mapped).unwrap();
4946 for batch_est in [false, true] {
4947 let knobs = SearchKnobs {
4948 batch_est,
4949 ..SearchKnobs::OFF
4950 };
4951 let mut plain = index.searcher_with(knobs);
4952 let mut batched = reopened.searcher_with(knobs);
4953 for ef in [7, 8, 9, 64, 129] {
4954 for k in [0, 1, 7, 8, 9, 64] {
4955 for patience in [0, 1, 4] {
4956 for reservoir in [0, 80] {
4957 let params = VSearchParams {
4958 ef,
4959 rerank: 80,
4960 patience,
4961 reservoir,
4962 ..Default::default()
4963 };
4964 for query in queries.chunks_exact(dim) {
4965 let (mut expected, mut actual) = (Vec::new(), Vec::new());
4966 let (mut a, mut b) =
4967 (QueryProfile::default(), QueryProfile::default());
4968 let ea = plain
4969 .scratch
4970 .search_inner::<true, false, { cfg!(feature = "experimental-query-x4") }>(
4971 &index,
4972 query,
4973 k,
4974 params,
4975 &mut expected,
4976 &mut a,
4977 )
4978 .unwrap();
4979 let eb = batched
4980 .scratch
4981 .search_inner::<true, true, { cfg!(feature = "experimental-query-x4") }>(
4982 &reopened,
4983 query,
4984 k,
4985 params,
4986 &mut actual,
4987 &mut b,
4988 )
4989 .unwrap();
4990 assert_eq!(expected, actual);
4991 assert_eq!(ea, eb);
4992 let work = |p: QueryProfile| {
4993 (
4994 p.expansions,
4995 p.neighbors_read,
4996 p.entry_evaluations,
4997 p.traversal_evaluations,
4998 p.rerank_evaluations,
4999 p.beam_retained,
5000 p.retained,
5001 )
5002 };
5003 assert_eq!(work(a), work(b));
5004 assert_eq!(
5005 plain.scratch.pool.keys,
5006 batched.scratch.pool.keys
5007 );
5008 assert_eq!(
5009 plain.scratch.pool.expanded,
5010 batched.scratch.pool.expanded
5011 );
5012 assert_eq!(
5013 plain.scratch.pool.cursor,
5014 batched.scratch.pool.cursor
5015 );
5016 assert_eq!(
5017 plain.scratch.visited.words,
5018 batched.scratch.visited.words
5019 );
5020 assert_eq!(
5021 plain.scratch.reservoir.keys,
5022 batched.scratch.reservoir.keys
5023 );
5024 let mut unprofiled = Vec::new();
5025 assert_eq!(
5026 plain
5027 .search(query, k, params, &mut unprofiled)
5028 .unwrap(),
5029 ea
5030 );
5031 assert_eq!(unprofiled, expected);
5032 assert_eq!(
5033 batched
5034 .search(query, k, params, &mut unprofiled)
5035 .unwrap(),
5036 eb
5037 );
5038 assert_eq!(unprofiled, actual);
5039 }
5040 }
5041 }
5042 }
5043 }
5044 }
5045 }
5046 }
5047 }
5048
5049 #[test]
5052 #[ignore]
5053 fn microbench_pool() {
5054 use std::hint::black_box;
5055 use std::time::Instant;
5056 for cap in [10usize, 20, 40, 80, 160, 320] {
5057 let stream: Vec<u64> = {
5058 let mut rng = SplitMix(cap as u64);
5059 (0..200_000).map(|_| rng.next_u64() % 100_000).collect()
5060 };
5061 let mut best = [f64::INFINITY; 2];
5062 for _ in 0..7 {
5063 for (which, b) in best.iter_mut().enumerate() {
5064 let t = Instant::now();
5065 let mut pool = Pool {
5066 keys: Vec::new(),
5067 expanded: Vec::new(),
5068 #[cfg(feature = "experimental-batch-pool")]
5069 merge_keys: Vec::new(),
5070 #[cfg(feature = "experimental-batch-pool")]
5071 merge_expanded: Vec::new(),
5072 cap: 1,
5073 cursor: 0,
5074 };
5075 let mut model = ReferencePool {
5076 keys: Vec::new(),
5077 expanded: Vec::new(),
5078 cap,
5079 cursor: 0,
5080 };
5081 pool.reset(cap);
5082 let mut acc = 0u64;
5083 for (i, &k) in stream.iter().enumerate() {
5084 if which == 0 {
5085 if i % 21 == 0 {
5086 if pool.next().is_none() {
5087 pool.reset(cap);
5088 }
5089 } else if let Some((p, _)) = pool.insert(black_box(k)) {
5090 acc += p as u64;
5091 }
5092 } else if i % 21 == 0 {
5093 if model.next().is_none() {
5094 model = ReferencePool {
5095 keys: Vec::new(),
5096 expanded: Vec::new(),
5097 cap,
5098 cursor: 0,
5099 };
5100 }
5101 } else if let Some((p, _)) = model.insert(black_box(k)) {
5102 acc += p as u64;
5103 }
5104 }
5105 black_box(acc);
5106 *b = b.min(t.elapsed().as_nanos() as f64 / stream.len() as f64);
5107 }
5108 }
5109 println!(
5110 "{{\"cap\": {cap}, \"bits_ns_per_op\": {:.2}, \"bools_ns_per_op\": {:.2}}}",
5111 best[0], best[1]
5112 );
5113 }
5114 }
5115
5116 #[test]
5117 fn compact_lists_preserve_order_search_and_persisted_bytes() {
5118 let dim = 16;
5119 let data = clustered(64, dim, 7);
5120 let queries = clustered(8, dim, 8);
5121 let mut graph = Graph::build(
5122 &data,
5123 dim,
5124 &BuildParams {
5125 degree: 16,
5126 build_ef: 32,
5127 ..BuildParams::default()
5128 },
5129 )
5130 .unwrap();
5131 for v in 0..graph.n {
5132 let count = [0, 1, 15, 16][v % 4];
5133 let ids: Vec<u32> = (1..=count).map(|j| ((v + j) % graph.n) as u32).collect();
5134 graph.adj.set(v, &ids);
5135 }
5136 let directory = tempfile::tempdir().unwrap();
5137 for bytes in [3, 4] {
5138 for bits in [4, 8] {
5139 let mut index =
5140 VGraph::from_graph_with(&graph, bits, 8, true, true, bytes, false).unwrap();
5141 let lists: Vec<Vec<u32>> = (0..index.n)
5142 .map(|v| {
5143 let mut ids = Vec::new();
5144 index.for_each_neighbour(v, |u| ids.push(u));
5145 ids
5146 })
5147 .collect();
5148 let before = directory.path().join("before");
5149 let after = directory.path().join("after");
5150 index.save(&before).unwrap();
5151 let old_bytes = index.index_bytes();
5152 let old_memory = index.memory_breakdown();
5153 let params = VSearchParams {
5154 ef: 32,
5155 rerank: 16,
5156 ..VSearchParams::default()
5157 };
5158 let expected: Vec<Vec<u32>> = {
5159 let mut searcher = index.searcher_with(SearchKnobs::OFF);
5160 queries
5161 .chunks_exact(dim)
5162 .map(|q| {
5163 let mut out = Vec::new();
5164 searcher.search(q, 10, params, &mut out).unwrap();
5165 out
5166 })
5167 .collect()
5168 };
5169 index.compact_adjacency().unwrap();
5170 assert!(index.index_bytes() < old_bytes);
5171 let memory = index.memory_breakdown();
5172 assert_eq!(
5173 old_memory.adjacency_ids_bytes,
5174 graph.n * graph.degree * bytes as usize
5175 );
5176 assert_eq!(old_memory.adjacency_offsets_bytes, 0);
5177 assert_eq!(
5178 memory.adjacency_ids_bytes,
5179 lists.iter().map(Vec::len).sum::<usize>() * bytes as usize
5180 );
5181 assert_eq!(memory.adjacency_offsets_bytes, (graph.n + 1) * 4);
5182 assert_eq!(
5183 old_bytes - index.index_bytes(),
5184 old_memory.adjacency_ids_bytes
5185 - memory.adjacency_ids_bytes
5186 - memory.adjacency_offsets_bytes
5187 );
5188 for (v, want) in lists.iter().enumerate() {
5189 let mut actual = Vec::new();
5190 index.for_each_neighbour(v, |u| actual.push(u));
5191 assert_eq!(&actual, want);
5192 }
5193 let compact_bytes = index.index_bytes();
5194 index.compact_adjacency().unwrap();
5195 assert_eq!(index.index_bytes(), compact_bytes);
5196 index.save(&after).unwrap();
5197 assert_eq!(
5198 std::fs::read(&before).unwrap(),
5199 std::fs::read(&after).unwrap()
5200 );
5201 for residency in [Residency::Mapped, Residency::Prefaulted, Residency::Heap] {
5202 let reopened = VGraph::open(&after, residency).unwrap();
5203 let (mut a, mut b) = (
5204 index.searcher_with(SearchKnobs::OFF),
5205 reopened.searcher_with(SearchKnobs::OFF),
5206 );
5207 for i in 0..260 {
5208 let qi = i % expected.len();
5209 let q = &queries[qi * dim..(qi + 1) * dim];
5210 let (mut x, mut y) = (Vec::new(), Vec::new());
5211 let evals = a.search(q, 10, params, &mut x).unwrap();
5212 b.search(q, 10, params, &mut y).unwrap();
5213 assert_eq!(x, expected[qi]);
5214 assert_eq!(x, y);
5215 #[cfg(feature = "experimental")]
5216 {
5217 let profile = a.profile_search(q, 10, params, &mut y).unwrap();
5218 assert_eq!(x, y);
5219 assert_eq!(
5220 evals,
5221 profile.entry_evaluations
5222 + profile.traversal_evaluations
5223 + profile.rerank_evaluations
5224 );
5225 assert!(profile.neighbors_read >= profile.traversal_evaluations - 1);
5226 }
5227 #[cfg(not(feature = "experimental"))]
5228 let _ = evals;
5229 }
5230 }
5231 }
5232 }
5233 }
5234
5235 #[test]
5238 fn graph_cache_lengths_are_bounded_by_the_file() {
5239 let d = 16;
5240 let n = 64usize;
5241 let data = clustered(n, d, 4);
5242 let params = BuildParams {
5243 build_ef: 16,
5244 degree: 8,
5245 seed: 3,
5246 ..BuildParams::default()
5247 };
5248 let graph = Graph::build(&data, d, ¶ms).expect("build");
5249 let dir = tempfile::tempdir().unwrap();
5250 let path = dir.path().join("graph.cache");
5251 graph.save(&path).unwrap();
5252 let mut bytes = std::fs::read(&path).unwrap();
5253 let lens = (8..bytes.len() - 8)
5256 .step_by(8)
5257 .filter(|&at| u64::from_le_bytes(bytes[at..at + 8].try_into().unwrap()) == n as u64)
5258 .nth(1)
5259 .expect("lengths list");
5260 for i in 0..n {
5261 let at = lens + 8 + 4 * i;
5262 bytes[at..at + 4].copy_from_slice(&u32::MAX.to_le_bytes());
5263 }
5264 let flat = lens + 8 + 4 * n;
5265 bytes[flat..flat + 8].copy_from_slice(&(n as u64 * u64::from(u32::MAX)).to_le_bytes());
5266 std::fs::write(&path, &bytes).unwrap();
5267 let error = Graph::load(&data, d, &path, ¶ms)
5268 .err()
5269 .unwrap()
5270 .to_string();
5271 assert!(error.contains("shorter than its lengths"), "{error}");
5272 }
5273
5274 #[test]
5277 fn graph_cache_round_trips_and_rejects_mismatches() {
5278 let d = 16;
5279 let data = clustered(600, d, 4);
5280 let params = BuildParams {
5281 build_ef: 32,
5282 degree: 8,
5283 alpha: 1.1,
5284 seed: 9,
5285 ..BuildParams::default()
5286 };
5287 let graph = Graph::build(&data, d, ¶ms).expect("build");
5288 let path = std::env::temp_dir().join(format!("warm-qgraph-cache-{}", std::process::id()));
5289 graph.save(&path).unwrap();
5290 let back = Graph::load(&data, d, &path, ¶ms).unwrap();
5291 assert_eq!(back.structure_hash(), graph.structure_hash());
5292 assert!(back.adj == graph.adj && back.layers == graph.layers);
5293 assert_eq!(
5294 (back.entry, &back.upper_ids),
5295 (graph.entry, &graph.upper_ids)
5296 );
5297 assert_eq!(back.vectors, graph.vectors);
5298 let rejects = |data: &[f32], p: &BuildParams, what: &str| {
5299 let err = Graph::load(data, d, &path, p)
5300 .err()
5301 .expect("a mismatched cache was accepted")
5302 .to_string();
5303 assert!(err.contains(what), "{what}: {err}");
5304 };
5305 for (what, p) in [
5306 (
5307 "degree",
5308 BuildParams {
5309 degree: 10,
5310 ..params.clone()
5311 },
5312 ),
5313 (
5314 "build_ef",
5315 BuildParams {
5316 build_ef: 33,
5317 ..params.clone()
5318 },
5319 ),
5320 (
5321 "alpha",
5322 BuildParams {
5323 alpha: 1.0,
5324 ..params.clone()
5325 },
5326 ),
5327 (
5328 "seed",
5329 BuildParams {
5330 seed: 10,
5331 ..params.clone()
5332 },
5333 ),
5334 (
5335 "metric",
5336 BuildParams {
5337 metric: Metric::Euclidean,
5338 ..params.clone()
5339 },
5340 ),
5341 ] {
5342 rejects(&data, &p, what);
5343 }
5344 let mut other = data.clone();
5345 other[d * 7 + 3] += 0.5;
5346 rejects(&other, ¶ms, "data fingerprint");
5347 rejects(&data[..d * 599], ¶ms, "n");
5348 std::fs::write(&path, b"HMSGRF01").unwrap();
5349 rejects(&data, ¶ms, "bad magic");
5350 std::fs::remove_file(&path).unwrap();
5351 }
5352
5353 #[test]
5358 fn euclidean_index_finds_euclidean_neighbours() {
5359 let (d, n, k) = (32, 4000, 10);
5360 let spread = |rows: &mut [f32], salt: usize| {
5361 let count = rows.len() / d;
5362 for (i, x) in rows.chunks_exact_mut(d).enumerate() {
5363 let s = 0.1 * 100f32.powf(((i * 7919 + salt) % count) as f32 / count as f32);
5364 x.iter_mut().for_each(|v| *v *= s);
5365 }
5366 };
5367 let mut data = clustered(n + 100, d, 7);
5369 spread(&mut data, 0);
5370 let queries = data.split_off(n * d);
5371 let l2 = |a: &[f32], b: &[f32]| -> f64 {
5372 a.iter()
5373 .zip(b)
5374 .map(|(x, y)| (f64::from(*x) - f64::from(*y)).powi(2))
5375 .sum()
5376 };
5377 let recall = |metric: Metric| -> f64 {
5378 let params = BuildParams {
5379 metric,
5380 degree: 32,
5381 build_ef: 100,
5382 ..BuildParams::default()
5383 };
5384 let g = Graph::build(&data, d, ¶ms).expect("build");
5385 let index = VGraph::from_graph_with(&g, 4, 8, true, true, 3, false).expect("encode");
5386 let mut s = index.searcher_with(SearchKnobs::OFF);
5387 let mut out = Vec::new();
5388 let mut hits = 0;
5389 for q in queries.chunks_exact(d) {
5390 let mut all: Vec<(f64, u32)> = data
5391 .chunks_exact(d)
5392 .enumerate()
5393 .map(|(i, x)| (l2(q, x), i as u32))
5394 .collect();
5395 all.sort_by(|a, b| a.0.total_cmp(&b.0));
5396 let truth: Vec<u32> = all[..k].iter().map(|t| t.1).collect();
5397 let params = VSearchParams {
5398 ef: 64,
5399 rerank: 32,
5400 ..VSearchParams::default()
5401 };
5402 s.search(q, k, params, &mut out).expect("search");
5403 hits += out.iter().filter(|x| truth.contains(x)).count();
5404 }
5405 hits as f64 / (k * 100) as f64
5406 };
5407 let (euclidean, angular) = (recall(Metric::Euclidean), recall(Metric::Angular));
5408 assert!(euclidean >= 0.95, "Euclidean recall {euclidean}");
5409 assert!(
5410 angular < euclidean - 0.3,
5411 "angular recall {angular} against Euclidean {euclidean}"
5412 );
5413 }
5414
5415 #[test]
5418 fn euclidean_index_limits() {
5419 let d = 16;
5420 let mut data = clustered(500, d, 2);
5421 for (i, x) in data.chunks_exact_mut(d).enumerate() {
5422 x.iter_mut().for_each(|v| *v *= 1.0 + i as f32 / 50.0);
5423 }
5424 let params = BuildParams {
5425 metric: Metric::Euclidean,
5426 degree: 16,
5427 build_ef: 64,
5428 ..BuildParams::default()
5429 };
5430 let g = Graph::build(&data, d, ¶ms).expect("build");
5431 assert!(VGraph::from_graph_with(&g, 8, 0, false, false, 4, true).is_err());
5432 assert!(super::super::QGraph::from_graph(&g, ¶ms).is_err());
5433 let blocks = VGraph::from_graph(&g, false)
5434 .expect("encode")
5435 .with_blocks(&g, 4);
5436 assert!(blocks.is_err());
5437 let index = VGraph::from_graph(&g, true).expect("encode");
5438 let mut out = Vec::new();
5439 let params = VSearchParams {
5440 ef: 500,
5441 rerank: 500,
5442 ..VSearchParams::default()
5443 };
5444 index
5445 .searcher_with(SearchKnobs::OFF)
5446 .search(&[0.0; 16], 1, params, &mut out)
5447 .expect("zero query");
5448 let smallest = data
5449 .chunks_exact(d)
5450 .enumerate()
5451 .min_by(|a, b| {
5452 let n2 = |x: &[f32]| x.iter().map(|v| v * v).sum::<f32>();
5453 n2(a.1).total_cmp(&n2(b.1))
5454 })
5455 .map(|m| m.0 as u32);
5456 assert_eq!(out.first().copied(), smallest);
5457 }
5458
5459 #[test]
5462 fn quantization_error_is_bounded() {
5463 let mut rng = SplitMix(12);
5464 let x: Vec<f32> = (0..100).map(|_| gaussian(&mut rng)).collect();
5465 let mut c = vec![0i8; 112];
5466 let s = quantize(&x, &mut c);
5467 let mut r: Vec<f32> = x
5468 .iter()
5469 .zip(&c)
5470 .map(|(a, &q)| a - s * f32::from(q))
5471 .collect();
5472 assert!(r.iter().all(|e| e.abs() <= s / 2.0 + 1e-6));
5473 assert!(c[100..].iter().all(|&q| q == 0));
5474 let mut rc = vec![0i8; 112];
5475 let sr = quantize(&r, &mut rc);
5476 r.iter_mut()
5477 .zip(&rc)
5478 .for_each(|(a, &q)| *a -= sr * f32::from(q));
5479 assert!(r.iter().all(|e| e.abs() <= s / 254.0 + 1e-6));
5480 assert_eq!(quantize(&[0.0; 4], &mut c[..16]), 0.0);
5481 }
5482
5483 fn recall(index: &VGraph, data: &[f32], queries: &[f32], d: usize, p: VSearchParams) -> f64 {
5484 let mut s = index.searcher();
5485 let mut out = Vec::new();
5486 let mut hits = 0;
5487 for q in queries.chunks_exact(d) {
5488 let truth = brute_force(data, d, q, 10);
5489 s.search(q, 10, p, &mut out).expect("search");
5490 assert_eq!(out.len(), 10);
5491 hits += out.iter().filter(|x| truth.contains(x)).count();
5492 }
5493 hits as f64 / (queries.len() / d * 10) as f64
5494 }
5495
5496 #[test]
5497 fn recall_on_clustered_data() {
5498 let d = 24;
5499 let data = clustered(4000, d, 1);
5500 let queries = clustered(100, d, 2);
5501 let graph = Graph::build(&data, d, &BuildParams::default()).expect("build");
5502 let plain = VGraph::from_graph(&graph, false).expect("encode");
5503 let fine = VGraph::from_graph(&graph, true).expect("encode");
5504 assert!(plain.index_bytes() < fine.index_bytes());
5505 let pool = VSearchParams {
5506 ef: 64,
5507 ..VSearchParams::default()
5508 };
5509 let rr = VSearchParams {
5510 ef: 64,
5511 rerank: 32,
5512 ..VSearchParams::default()
5513 };
5514 let r0 = recall(&plain, &data, &queries, d, pool);
5515 let r1 = recall(&fine, &data, &queries, d, rr);
5516 assert!(r0 >= 0.9, "recall {r0}");
5518 assert!(r1 >= r0, "rerank {r1} below pool order {r0}");
5519 let fine4 = VGraph::from_graph_with(&graph, 8, 4, false, false, 4, false).expect("encode");
5520 assert!(plain.index_bytes() < fine4.index_bytes());
5521 assert!(fine4.index_bytes() < fine.index_bytes());
5522 let r2 = recall(&fine4, &data, &queries, d, rr);
5523 assert!(r2 >= r0 - 0.02, "4-bit rerank {r2} below pool order {r0}");
5524 let lvq48 = VGraph::from_graph_with(&graph, 4, 8, true, false, 4, false).expect("encode");
5525 assert!(lvq48.index_bytes() < fine.index_bytes());
5526 let wide = VSearchParams {
5528 ef: 96,
5529 rerank: 64,
5530 ..VSearchParams::default()
5531 };
5532 let r3 = recall(&lvq48, &data, &queries, d, wide);
5533 assert!(
5534 r3 >= r0 - 0.05,
5535 "4-bit traversal with 8-bit residual {r3} vs {r0}"
5536 );
5537 }
5538
5539 #[test]
5543 fn block_codes_keep_recall() {
5544 let d = 24;
5545 let data = clustered(4000, d, 1);
5546 let queries = clustered(100, d, 2);
5547 let graph = Graph::build(&data, d, &BuildParams::default()).expect("build");
5548 let plain = VGraph::from_graph(&graph, true).expect("encode");
5549 let rr = VSearchParams {
5550 ef: 96,
5551 rerank: 64,
5552 ..VSearchParams::default()
5553 };
5554 let r0 = recall(&plain, &data, &queries, d, rr);
5555 for bits in [4u8, 2, 1] {
5556 for (reorder, id_bytes) in [(false, 4u8), (true, 3)] {
5557 let b = VGraph::from_graph_with(&graph, 8, 8, true, reorder, id_bytes, false)
5558 .expect("encode")
5559 .with_blocks(&graph, bits)
5560 .expect("blocks");
5561 assert!(b.has_blocks());
5562 assert!(b.index_bytes() > plain.index_bytes());
5563 let r = recall(&b, &data, &queries, d, rr);
5564 let slack = match bits {
5568 4 => 0.03,
5569 2 => 0.06,
5570 _ => 0.30,
5571 };
5572 assert!(
5573 r >= r0 - slack,
5574 "{bits}-bit blocks (reorder {reorder}, ids {id_bytes}) recall {r} vs {r0}"
5575 );
5576 }
5577 }
5578 let e = match VGraph::from_graph(&graph, true)
5579 .expect("encode")
5580 .with_blocks(&graph, 3)
5581 {
5582 Ok(_) => panic!("3 bits accepted"),
5583 Err(e) => e.to_string(),
5584 };
5585 assert!(e.contains("block bits"), "{e}");
5586 }
5587
5588 #[test]
5591 fn block_estimates_track_exact_inner_products() {
5592 let d = 64;
5593 let data = clustered(3000, d, 5);
5594 let graph = Graph::build(&data, d, &BuildParams::default()).expect("build");
5595 let unit = &graph.vectors;
5597 for (bits, bar) in [(4u8, 0.998), (2, 0.98), (1, 0.90)] {
5598 let g = VGraph::from_graph(&graph, false)
5599 .expect("encode")
5600 .with_blocks(&graph, bits)
5601 .expect("blocks");
5602 let mut s = g.searcher();
5603 let q: Vec<f32> = unit[7 * d..8 * d].to_vec();
5604 s.scratch.qf[..d].copy_from_slice(&q);
5605 let sq = g.encode_query(
5606 &s.scratch.qf,
5607 &mut s.scratch.qc,
5608 &mut s.scratch.qtmp,
5609 &mut s.scratch.qrot,
5610 &mut s.scratch.qs,
5611 );
5612 g.encode_block_query(&s.scratch.qf, sq, &mut s.scratch.bqrot, &mut s.scratch.bq);
5613 let centred =
5614 |u: usize| -> f32 { (0..d).map(|j| q[j] * (unit[u * d + j] - g.mean[j])).sum() };
5615 let (mut xs, mut ys) = (Vec::new(), Vec::new());
5616 for v in 0..64usize {
5617 let ip_v = centred(v);
5619 let mut list = Vec::new();
5620 g.for_each_neighbour(v, |u| list.push(u));
5621 for b in 0..g.blocks_per_vertex {
5622 let (codes, scales, offsets) = g.block(v, b);
5623 scan32(codes, &s.scratch.bq.lut, g.block_m, &mut s.scratch.scan);
5624 for i in 0..BLOCK {
5625 let Some(&u) = list.get(b * BLOCK + i) else {
5626 break;
5627 };
5628 let ip = ip_v
5629 + le_f32(scales, i)
5630 * (s.scratch.bq.delta * s.scratch.scan[i] as f32
5631 + s.scratch.bq.min_sum)
5632 + le_f32(offsets, i);
5633 xs.push(f64::from(ip));
5634 ys.push(f64::from(centred(u as usize)));
5635 }
5636 }
5637 }
5638 let n = xs.len() as f64;
5639 let (mx, my) = (xs.iter().sum::<f64>() / n, ys.iter().sum::<f64>() / n);
5640 let (mut sxy, mut sxx, mut syy) = (0.0, 0.0, 0.0);
5641 for (x, y) in xs.iter().zip(&ys) {
5642 sxy += (x - mx) * (y - my);
5643 sxx += (x - mx) * (x - mx);
5644 syy += (y - my) * (y - my);
5645 }
5646 let r = sxy / (sxx * syy).sqrt();
5647 eprintln!(
5648 "{bits}-bit block estimates: correlation {r:.4} over {} pairs",
5649 xs.len()
5650 );
5651 assert!(
5652 r > bar,
5653 "{bits}-bit block estimates: correlation {r} below {bar}"
5654 );
5655 }
5656 }
5657
5658 #[test]
5662 fn delta_coded_adjacency_round_trips() {
5663 let d = 24;
5664 let data = clustered(4000, d, 1);
5665 let queries = clustered(50, d, 2);
5666 let graph = Graph::build(&data, d, &BuildParams::default()).expect("build");
5667 let wide = VGraph::from_graph_with(&graph, 4, 8, true, true, 4, false).expect("encode");
5668 let delta = VGraph::from_graph_with(&graph, 4, 8, true, true, 0, false).expect("encode");
5669 assert!(
5670 delta.index_bytes() < wide.index_bytes(),
5671 "{} vs {}",
5672 delta.index_bytes(),
5673 wide.index_bytes()
5674 );
5675 for v in 0..wide.n {
5676 let (mut a, mut b) = (Vec::new(), Vec::new());
5677 wide.for_each_neighbour(v, |u| a.push(u));
5678 delta.for_each_neighbour(v, |u| b.push(u));
5679 a.sort_unstable();
5680 assert_eq!(a, b, "vertex {v}");
5681 }
5682 let p = VSearchParams {
5683 ef: 48,
5684 rerank: 16,
5685 ..VSearchParams::default()
5686 };
5687 let (mut s1, mut s2) = (wide.searcher(), delta.searcher());
5688 let (mut o1, mut o2) = (Vec::new(), Vec::new());
5689 for q in queries.chunks_exact(d) {
5690 s1.search(q, 10, p, &mut o1).expect("search");
5691 s2.search(q, 10, p, &mut o2).expect("search");
5692 assert_eq!(o1, o2);
5693 }
5694 let blocks = VGraph::from_graph_with(&graph, 8, 8, true, true, 0, false)
5695 .expect("encode")
5696 .with_blocks(&graph, 2)
5697 .expect("blocks");
5698 assert!(recall(&blocks, &data, &queries, d, p) > 0.8);
5699 assert_eq!(zigzag(-1), 1);
5700 assert_eq!(unzigzag(zigzag(-123_456)), -123_456);
5701 let mut buf = Vec::new();
5702 put_varint(&mut buf, 300);
5703 assert_eq!(get_varint(&buf), (300, 2));
5704 }
5705
5706 #[test]
5709 fn reordered_aligned_index_returns_the_same_ids() {
5710 let d = 24;
5711 let data = clustered(3000, d, 5);
5712 let graph = Graph::build(&data, d, &BuildParams::default()).expect("build");
5713 let a = VGraph::from_graph(&graph, true).expect("encode");
5714 let b = VGraph::from_graph_with(&graph, 8, 8, true, true, 4, false).expect("encode");
5715 assert_eq!(b.index_bytes(), a.index_bytes() + b.n * 4);
5717 let (mut sa, mut sb) = (a.searcher(), b.searcher());
5718 let (mut oa, mut ob) = (Vec::new(), Vec::new());
5719 let c = VGraph::from_graph_with(&graph, 8, 8, true, true, 3, false).expect("encode");
5720 assert_eq!(c.index_bytes(), b.index_bytes() - c.n * c.degree);
5721 let mut sc = c.searcher();
5722 let mut oc = Vec::new();
5723 let mut same = 0;
5724 for q in data.chunks_exact(d).take(200) {
5725 for p in [
5726 VSearchParams {
5727 ef: 32,
5728 ..VSearchParams::default()
5729 },
5730 VSearchParams {
5731 ef: 48,
5732 rerank: 16,
5733 ..VSearchParams::default()
5734 },
5735 ] {
5736 sa.search(q, 10, p, &mut oa).expect("search");
5737 sb.search(q, 10, p, &mut ob).expect("search");
5738 sc.search(q, 10, p, &mut oc).expect("search");
5739 assert_eq!(ob, oc, "3-byte ids changed the search");
5740 same += usize::from(oa == ob);
5741 }
5742 }
5743 assert!(same >= 396, "{same} of 400 searches agree");
5744 }
5745
5746 #[test]
5747 fn build_and_search_are_deterministic() {
5748 let d = 20;
5749 let data = clustered(3000, d, 3);
5750 let params = BuildParams {
5751 seed: 9,
5752 ..BuildParams::default()
5753 };
5754 let a = VGraph::from_graph(&Graph::build(&data, d, ¶ms).expect("build"), true)
5755 .expect("encode");
5756 let b = VGraph::from_graph(&Graph::build(&data, d, ¶ms).expect("build"), true)
5757 .expect("encode");
5758 assert!(
5759 a.codes.bytes() == b.codes.bytes()
5760 && a.residual[..] == b.residual[..]
5761 && a.adj[..] == b.adj[..]
5762 );
5763 assert!(a.scales[..] == b.scales[..] && a.res_scales[..] == b.res_scales[..]);
5764 assert_eq!((a.entry, &a.mean), (b.entry, &b.mean));
5765 let (mut sa, mut sb) = (a.searcher(), b.searcher());
5766 let (mut oa, mut ob) = (Vec::new(), Vec::new());
5767 let p = VSearchParams {
5768 ef: 32,
5769 rerank: 16,
5770 ..VSearchParams::default()
5771 };
5772 for q in data.chunks_exact(d).take(50) {
5773 sa.search(q, 10, p, &mut oa).expect("search");
5774 sb.search(q, 10, p, &mut ob).expect("search");
5775 assert_eq!(oa, ob);
5776 }
5777 }
5778
5779 #[test]
5781 fn invalid_inputs_are_errors_that_name_the_cause() {
5782 let d = 16;
5783 let data = clustered(500, d, 4);
5784 let err = |r: Result<Graph, WarmError>| match r {
5785 Err(e) => e.to_string(),
5786 Ok(_) => panic!("an invalid input was accepted"),
5787 };
5788 assert!(err(Graph::build(&[], d, &BuildParams::default())).contains("empty"));
5789 assert!(err(Graph::build(&data[..d + 1], d, &BuildParams::default())).contains("ragged"));
5790 assert!(err(Graph::build(&data, 0, &BuildParams::default())).contains("ragged"));
5791 let mut bad = data.clone();
5792 bad[3 * d + 2] = f32::NAN;
5793 let e = err(Graph::build(&bad, d, &BuildParams::default()));
5794 assert!(e.contains("non-finite") && e.contains("row 3"), "{e}");
5795 let odd = BuildParams {
5796 degree: 7,
5797 ..BuildParams::default()
5798 };
5799 assert!(err(Graph::build(&data, d, &odd)).contains("even"));
5800 let graph = Graph::build(&data, d, &BuildParams::default()).expect("build");
5801 let e = match VGraph::from_graph_with(&graph, 5, 8, false, false, 4, false) {
5802 Err(e) => e.to_string(),
5803 Ok(_) => panic!("bits 5 was accepted"),
5804 };
5805 assert!(e.contains("bits"), "{e}");
5806 let index = VGraph::from_graph(&graph, true).expect("encode");
5807 let mut s = index.searcher();
5808 let mut out = Vec::new();
5809 let p = VSearchParams::default();
5810 let e = s
5811 .search(&data[..d - 1], 10, p, &mut out)
5812 .expect_err("short query")
5813 .to_string();
5814 assert!(e.contains("dimension"), "{e}");
5815 let mut q = data[..d].to_vec();
5816 q[5] = f32::INFINITY;
5817 let e = s.search(&q, 10, p, &mut out).expect_err("inf").to_string();
5818 assert!(e.contains("non-finite") && e.contains('5'), "{e}");
5819 let e = s
5820 .search(&vec![0.0; d], 10, p, &mut out)
5821 .expect_err("zero")
5822 .to_string();
5823 assert!(e.contains("zero"), "{e}");
5824 let huge = VSearchParams {
5827 ef: usize::MAX,
5828 rerank: usize::MAX,
5829 ..VSearchParams::default()
5830 };
5831 let evals = s.search(&data[..d], 10, huge, &mut out).expect("search");
5832 assert_eq!(out.len(), 10);
5833 assert!(evals <= 2 * index.len() + 10);
5834 s.search(&data[..d], 5000, huge, &mut out).expect("search");
5835 assert!(out.len() <= index.len());
5836 let mut with_zero = data.clone();
5838 with_zero[..d].fill(0.0);
5839 let g0 = Graph::build(&with_zero, d, &BuildParams::default()).expect("build");
5840 let i0 = VGraph::from_graph(&g0, true).expect("encode");
5841 let mut s0 = i0.searcher();
5842 s0.search(&data[d..2 * d], 10, p, &mut out).expect("search");
5843 assert_eq!(out.len(), 10);
5844 assert!(out.iter().all(|&v| (v as usize) < i0.len()));
5845 }
5846
5847 #[test]
5850 fn screen_is_transparent_at_a_wide_margin() {
5851 let d = 24;
5852 let data = clustered(3000, d, 5);
5853 let queries = clustered(100, d, 2);
5854 let graph = Graph::build(&data, d, &BuildParams::default()).expect("build");
5855 let plain = VGraph::from_graph(&graph, true).expect("encode");
5856 let screened =
5857 VGraph::from_graph_with(&graph, 8, 8, false, false, 4, true).expect("encode");
5858 assert!(screened.has_screen() && screened.index_bytes() > plain.index_bytes());
5859 let (mut sa, mut sb) = (plain.searcher(), screened.searcher());
5860 let (mut oa, mut ob) = (Vec::new(), Vec::new());
5861 let wide = VSearchParams {
5862 ef: 64,
5863 rerank: 16,
5864 screen_sigmas: 1e6,
5865 ..VSearchParams::default()
5866 };
5867 for q in queries.chunks_exact(d) {
5868 let ea = sa.search(q, 10, wide, &mut oa).expect("search");
5869 let eb = sb.search(q, 10, wide, &mut ob).expect("search");
5870 assert_eq!((ea, &oa), (eb, &ob));
5871 assert_eq!(sb.last_screened(), 0);
5872 }
5873 let one = VSearchParams {
5874 ef: 64,
5875 rerank: 16,
5876 ..VSearchParams::default()
5877 };
5878 let r0 = recall(&plain, &data, &queries, d, one);
5879 let r1 = recall(&screened, &data, &queries, d, one);
5880 assert!(r1 >= r0 - 0.02, "screened {r1} vs {r0}");
5881 let mut skipped = 0;
5882 for q in queries.chunks_exact(d) {
5883 sb.search(q, 10, one, &mut ob).expect("search");
5884 skipped += sb.last_screened();
5885 }
5886 assert!(skipped > 0, "the screen never skipped a neighbour");
5887 }
5888}