1use crate::def::{Predicate, MAX_KEYMATCH_LIST};
2use crate::hnsw::HnswIndex;
3use core_storage::{list_tokens, Value, ValueKey};
4use std::collections::{BTreeMap, BTreeSet};
5
6pub const IVF_K_MIN: usize = 4;
13
14pub const IVF_K_MAX: usize = 1024;
16
17pub const IVF_ITERATIONS: usize = 12;
19
20pub const IVF_PROBE_DENOM: usize = 16;
23
24pub const IVF_DRIFT_REBUILD: u64 = 256;
28
29thread_local! {
30 static IVF_DRIFT_REBUILD_OVERRIDE: std::cell::Cell<Option<u64>> =
31 const { std::cell::Cell::new(None) };
32}
33
34pub(crate) fn ivf_drift_rebuild_threshold() -> u64 {
35 IVF_DRIFT_REBUILD_OVERRIDE.with(|c| c.get().unwrap_or(IVF_DRIFT_REBUILD))
36}
37
38pub fn with_ivf_drift_rebuild<R>(threshold: u64, f: impl FnOnce() -> R) -> R {
41 IVF_DRIFT_REBUILD_OVERRIDE.with(|c| {
42 let prev = c.replace(Some(threshold));
43 let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
44 c.set(prev);
45 match out {
46 Ok(v) => v,
47 Err(p) => std::panic::resume_unwind(p),
48 }
49 })
50}
51
52pub const EF_MAX: usize = 4_096;
57
58const BEAM_FLOOR_SLACK: f64 = 1e-5;
67
68thread_local! {
69 static EF_MAX_OVERRIDE: std::cell::Cell<Option<usize>> = const { std::cell::Cell::new(None) };
70}
71
72pub(crate) fn ef_max() -> usize {
73 EF_MAX_OVERRIDE.with(|c| c.get().unwrap_or(EF_MAX)).max(1)
74}
75
76pub fn with_ef_max<R>(cap: usize, f: impl FnOnce() -> R) -> R {
82 EF_MAX_OVERRIDE.with(|c| {
83 let prev = c.replace(Some(cap));
84 let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
85 c.set(prev);
86 match out {
87 Ok(v) => v,
88 Err(p) => std::panic::resume_unwind(p),
89 }
90 })
91}
92
93pub const HNSW_BUILD_BATCH: usize = 2_048;
102
103thread_local! {
104 static HNSW_BUILD_BATCH_OVERRIDE: std::cell::Cell<Option<usize>> =
105 const { std::cell::Cell::new(None) };
106}
107
108pub(crate) fn hnsw_build_batch() -> usize {
109 HNSW_BUILD_BATCH_OVERRIDE
110 .with(|c| c.get().unwrap_or(HNSW_BUILD_BATCH))
111 .max(1)
112}
113
114pub fn with_hnsw_build_batch<R>(batch: usize, f: impl FnOnce() -> R) -> R {
121 HNSW_BUILD_BATCH_OVERRIDE.with(|c| {
122 let prev = c.replace(Some(batch));
123 let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
124 c.set(prev);
125 match out {
126 Ok(v) => v,
127 Err(p) => std::panic::resume_unwind(p),
128 }
129 })
130}
131
132enum HnswLeg<'a> {
136 All,
138 Skip(&'a BTreeSet<u32>),
140 Defer,
142}
143
144pub fn hnsw_vector_present(spec: &CandidateSpec, get: &dyn Fn(&str) -> Option<Value>) -> bool {
148 match spec {
149 CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) => {
150 specs.iter().any(|s| hnsw_vector_present(s, get))
151 }
152 CandidateSpec::Hnsw { field, .. } => {
153 get(field).as_ref().and_then(as_numeric_list).is_some()
154 }
155 _ => false,
156 }
157}
158
159pub fn cluster_k(n: usize) -> usize {
161 if n == 0 {
162 return IVF_K_MIN;
163 }
164 let k = (n as f64).sqrt().ceil() as usize;
165 k.clamp(IVF_K_MIN, IVF_K_MAX)
166}
167
168pub fn probe_count(k: usize) -> usize {
170 k.div_ceil(IVF_PROBE_DENOM).max(1)
171}
172
173fn l2_normalize(xs: &[f64]) -> Option<Vec<f64>> {
175 let n = xs.iter().map(|x| x * x).sum::<f64>().sqrt();
176 if n == 0.0 {
177 return None;
178 }
179 Some(xs.iter().map(|x| x / n).collect())
180}
181
182fn l2_sq(a: &[f64], b: &[f64]) -> f64 {
185 if a.len() != b.len() {
186 return f64::MAX;
187 }
188 a.iter().zip(b.iter()).map(|(x, y)| (x - y) * (x - y)).sum()
189}
190
191pub fn nearest_centroid(centroids: &[Vec<f64>], xs: &[f64]) -> usize {
194 centroids
195 .iter()
196 .enumerate()
197 .min_by(|(_, a), (_, b)| {
198 l2_sq(xs, a)
199 .partial_cmp(&l2_sq(xs, b))
200 .unwrap_or(std::cmp::Ordering::Equal)
201 })
202 .map(|(i, _)| i)
203 .unwrap_or(0)
204}
205
206pub fn fnv1a_u64(data: &[u8]) -> u64 {
210 const FNV_OFFSET: u64 = 14_695_981_039_346_656_037;
211 const FNV_PRIME: u64 = 1_099_511_628_211;
212 let mut h = FNV_OFFSET;
213 for &b in data {
214 h ^= b as u64;
215 h = h.wrapping_mul(FNV_PRIME);
216 }
217 h
218}
219
220#[inline]
223fn lcg_next(state: u64) -> u64 {
224 state
225 .wrapping_mul(6_364_136_223_846_793_005)
226 .wrapping_add(1_442_695_040_888_963_407)
227}
228
229pub fn kmeans_fit(vecs: &[(u32, Vec<f64>)], k: usize, seed: u64) -> Vec<Vec<f64>> {
239 let vecs: Vec<(u32, Vec<f64>)> = vecs
240 .iter()
241 .filter_map(|(id, xs)| l2_normalize(xs).map(|n| (*id, n)))
242 .collect();
243 if vecs.is_empty() || k == 0 {
244 return vec![];
245 }
246 let n = vecs.len();
247 let k = k.min(n);
248 let dim = vecs[0].1.len();
249 if dim == 0 {
250 return vec![];
251 }
252
253 let mut state = seed;
255 let mut used = vec![false; n];
256 let mut init_idxs: Vec<usize> = Vec::with_capacity(k);
257 let mut attempts = 0usize;
258 while init_idxs.len() < k && attempts < n * 4 {
259 state = lcg_next(state);
260 let idx = (state >> 33) as usize % n;
261 if !used[idx] {
262 used[idx] = true;
263 init_idxs.push(idx);
264 }
265 attempts += 1;
266 }
267 if init_idxs.len() < k {
270 for (i, in_use) in used.iter().enumerate().take(n) {
271 if !in_use {
272 init_idxs.push(i);
273 if init_idxs.len() == k {
274 break;
275 }
276 }
277 }
278 }
279 let mut centroids: Vec<Vec<f64>> = init_idxs.iter().map(|&i| vecs[i].1.clone()).collect();
280 let mut assignments = vec![0usize; n];
281
282 for iter in 0..IVF_ITERATIONS {
284 for (j, (_, xs)) in vecs.iter().enumerate() {
286 assignments[j] = nearest_centroid(¢roids, xs);
287 }
288
289 let mut sums = vec![vec![0.0f64; dim]; k];
291 let mut counts = vec![0usize; k];
292 for (j, (_, xs)) in vecs.iter().enumerate() {
293 let c = assignments[j];
294 counts[c] += 1;
295 for d in 0..dim {
296 sums[c][d] += xs[d];
297 }
298 }
299
300 let mut new_centroids = vec![vec![0.0f64; dim]; k];
302 let mut empty: Vec<usize> = Vec::new();
303 for c in 0..k {
304 if counts[c] == 0 {
305 empty.push(c);
306 } else {
307 for d in 0..dim {
308 new_centroids[c][d] = sums[c][d] / counts[c] as f64;
309 }
310 }
311 }
312
313 for (ei, ec) in empty.into_iter().enumerate() {
316 let reseed =
317 seed ^ (iter as u64).wrapping_mul(0x9E37) ^ (ei as u64).wrapping_mul(0x1234_5679);
318 let mut rs = lcg_next(reseed);
319 rs = lcg_next(rs);
320 let pick = (rs >> 33) as usize % n;
321 new_centroids[ec] = vecs[pick].1.clone();
322 }
323
324 centroids = new_centroids;
325 }
326
327 centroids
328}
329
330#[cfg(test)]
331thread_local! {
332 static VECTOR_DIM_REJECT: std::cell::Cell<bool> = const { std::cell::Cell::new(true) };
333 static VECTOR_EARLY_EXIT: std::cell::Cell<bool> = const { std::cell::Cell::new(true) };
334}
335
336fn vector_dim_reject_enabled() -> bool {
337 #[cfg(test)]
338 {
339 VECTOR_DIM_REJECT.with(|c| c.get())
340 }
341 #[cfg(not(test))]
342 {
343 true
344 }
345}
346
347pub(crate) fn vector_early_exit_enabled() -> bool {
348 #[cfg(test)]
349 {
350 VECTOR_EARLY_EXIT.with(|c| c.get())
351 }
352 #[cfg(not(test))]
353 {
354 true
355 }
356}
357
358thread_local! {
359 static VECTOR_SCAN: std::cell::Cell<Option<bool>> = const { std::cell::Cell::new(None) };
362}
363
364pub fn vector_scan_forced() -> bool {
373 VECTOR_SCAN.with(|c| match c.get() {
374 Some(v) => v,
375 None => {
376 let v = std::env::var("MUSHROOMDB_VECTOR_SCAN")
377 .map(|s| s == "1" || s.eq_ignore_ascii_case("true"))
378 .unwrap_or(false);
379 c.set(Some(v));
380 v
381 }
382 })
383}
384
385pub fn with_vector_scan<R>(enabled: bool, f: impl FnOnce() -> R) -> R {
400 VECTOR_SCAN.with(|c| {
401 let prev = c.replace(Some(enabled));
402 let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
403 c.set(prev);
404 match out {
405 Ok(v) => v,
406 Err(p) => std::panic::resume_unwind(p),
407 }
408 })
409}
410
411#[cfg(test)]
413pub fn with_vector_dim_reject<R>(enabled: bool, f: impl FnOnce() -> R) -> R {
414 VECTOR_DIM_REJECT.with(|c| {
415 let prev = c.replace(enabled);
416 let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
417 c.set(prev);
418 match out {
419 Ok(v) => v,
420 Err(p) => std::panic::resume_unwind(p),
421 }
422 })
423}
424
425#[cfg(test)]
427pub fn with_vector_early_exit<R>(enabled: bool, f: impl FnOnce() -> R) -> R {
428 VECTOR_EARLY_EXIT.with(|c| {
429 let prev = c.replace(enabled);
430 let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
431 c.set(prev);
432 match out {
433 Ok(v) => v,
434 Err(p) => std::panic::resume_unwind(p),
435 }
436 })
437}
438
439#[derive(Debug, Default)]
440pub struct SideIndex {
441 by_key: BTreeMap<ValueKey, BTreeSet<u32>>,
442 vec_meta: BTreeMap<u32, (u32, f64)>,
447 vec_checkpoints: BTreeMap<u32, [f64; 8]>,
454 vec_anchor: BTreeMap<u32, f64>,
461
462 ivf_raw: BTreeMap<u32, Vec<f64>>,
467 ivf_centroids: Vec<Vec<f64>>,
469 ivf_clusters: BTreeMap<u32, usize>,
472 pub ivf_drift: u64,
476
477 hnsw: Option<HnswIndex>,
480 hnsw_tracked: BTreeSet<u32>,
483}
484
485#[derive(Debug, Default)]
486pub struct RuleIndex {
487 pub src_side: SideIndex,
488 pub dst_side: SideIndex,
489}
490
491#[derive(Debug)]
492pub enum CandidateSpec<'a> {
493 ByKey,
494 Scalar {
495 field: &'a str,
496 },
497 Tokens {
498 field: &'a str,
499 },
500 ScalarOrElements {
509 field: &'a str,
510 },
511 NumericBucket {
512 field: &'a str,
513 tolerance: f64,
514 },
515 GeoGrid {
516 field: &'a str,
517 km: f64,
518 },
519 ScanAll {
520 field: &'a str,
521 },
522 VectorClusters {
529 field: &'a str,
530 min: f64,
531 },
532 Hnsw {
539 field: &'a str,
540 k: usize,
543 floor: Option<f64>,
548 },
549 Union(Vec<CandidateSpec<'a>>),
556 Intersect(Vec<CandidateSpec<'a>>),
563}
564
565pub fn candidate_spec(p: &Predicate) -> CandidateSpec<'_> {
578 match p {
579 Predicate::KeyMatch { .. } => CandidateSpec::ByKey,
580 Predicate::FieldEqual { field } => CandidateSpec::Scalar { field },
581 Predicate::Overlap { field, .. } => CandidateSpec::Tokens { field },
582 Predicate::NumericWithin { field, tolerance } => CandidateSpec::NumericBucket {
583 field,
584 tolerance: *tolerance,
585 },
586 Predicate::GeoRadius { field, km } => CandidateSpec::GeoGrid { field, km: *km },
587 Predicate::VectorSimilar { field, .. } => CandidateSpec::ScanAll { field },
588 Predicate::All(parts) => {
589 debug_assert!(
590 !parts.is_empty(),
591 "candidate_spec requires a validated predicate"
592 );
593 CandidateSpec::Intersect(parts.iter().map(candidate_spec).collect())
594 }
595 Predicate::Any(parts) => {
596 debug_assert!(
597 !parts.is_empty(),
598 "candidate_spec requires a validated predicate"
599 );
600 CandidateSpec::Union(parts.iter().map(candidate_spec).collect())
601 }
602 }
603}
604
605pub fn candidate_spec_approx(p: &Predicate) -> CandidateSpec<'_> {
624 candidate_spec_approx_with_k(p, 64)
625}
626
627pub fn candidate_spec_approx_with_k(p: &Predicate, k: usize) -> CandidateSpec<'_> {
633 candidate_spec_approx_with_floor(p, k, false)
634}
635
636pub fn candidate_spec_approx_with_floor(
643 p: &Predicate,
644 k: usize,
645 floored: bool,
646) -> CandidateSpec<'_> {
647 match p {
648 Predicate::VectorSimilar { field, min } => CandidateSpec::Hnsw {
649 field,
650 k,
651 floor: floored.then_some(*min),
652 },
653 Predicate::All(parts) => {
654 debug_assert!(
655 !parts.is_empty(),
656 "candidate_spec_approx requires a validated predicate"
657 );
658 CandidateSpec::Intersect(
659 parts
660 .iter()
661 .map(|p| candidate_spec_approx_with_floor(p, k, floored))
662 .collect(),
663 )
664 }
665 other => candidate_spec(other),
666 }
667}
668
669pub fn spec_has_hnsw(spec: &CandidateSpec<'_>) -> bool {
673 match spec {
674 CandidateSpec::Hnsw { .. } => true,
675 CandidateSpec::Union(parts) | CandidateSpec::Intersect(parts) => {
676 parts.iter().any(spec_has_hnsw)
677 }
678 _ => false,
679 }
680}
681
682pub(crate) fn as_finite_f64(v: &Value) -> Option<f64> {
683 match v {
684 Value::Int(i) => Some(*i as f64),
685 Value::Float(f) if f.is_finite() => Some(*f),
686 _ => None,
687 }
688}
689
690fn as_latlon(v: &Value) -> Option<(f64, f64)> {
691 let Value::List(items) = v else {
692 return None;
693 };
694 if items.len() != 2 {
695 return None;
696 }
697 let lat = as_finite_f64(&items[0])?;
698 let lon = as_finite_f64(&items[1])?;
699 if (-90.0..=90.0).contains(&lat) && (-180.0..=180.0).contains(&lon) {
700 Some((lat, lon))
701 } else {
702 None
703 }
704}
705
706pub(crate) fn as_numeric_list(v: &Value) -> Option<Vec<f64>> {
707 let Value::List(items) = v else {
708 return None;
709 };
710 if items.is_empty() {
711 return None;
712 }
713 items.iter().map(as_finite_f64).collect()
714}
715
716fn vec_dim_norm(v: &Value) -> Option<(u32, f64)> {
717 let xs = as_numeric_list(v)?;
718 let mut n2 = 0.0;
719 for x in &xs {
720 n2 += *x * *x;
721 }
722 Some((xs.len() as u32, n2.sqrt()))
723}
724
725fn compute_ckpts(xs: &[f64]) -> [f64; 8] {
731 let dim = xs.len();
732 let mut ckpts = [0.0f64; 8];
733 if dim == 0 {
734 return ckpts;
735 }
736 let boundaries: [usize; 8] = std::array::from_fn(|i| i * dim / 8);
738 let mut suffix_sq = 0.0f64;
739 let mut ci = 7i32;
741 for j in (0..dim).rev() {
742 suffix_sq += xs[j] * xs[j];
743 while ci >= 0 && boundaries[ci as usize] == j {
745 ckpts[ci as usize] = suffix_sq.sqrt();
746 ci -= 1;
747 }
748 }
749 ckpts
750}
751
752fn floor_to_i64(x: f64) -> i64 {
753 let floored = x.floor();
754 if !floored.is_finite() {
755 return 0;
756 }
757 if floored >= i64::MAX as f64 {
758 i64::MAX
759 } else if floored <= i64::MIN as f64 {
760 i64::MIN
761 } else {
762 floored as i64
763 }
764}
765
766fn numeric_index_key(v: f64, tolerance: f64) -> Option<ValueKey> {
770 if !tolerance.is_finite() || tolerance < 0.0 {
771 return None;
772 }
773 if tolerance == 0.0 {
774 let v = if v == 0.0 { 0.0_f64 } else { v };
775 return Some(ValueKey::FloatBits(v.to_bits()));
776 }
777 Some(ValueKey::Int(floor_to_i64(v / tolerance)))
778}
779
780fn numeric_probe_keys(v: f64, tolerance: f64) -> BTreeSet<ValueKey> {
781 match numeric_index_key(v, tolerance) {
782 None => BTreeSet::new(),
783 Some(k @ ValueKey::FloatBits(_)) => BTreeSet::from([k]),
784 Some(ValueKey::Int(b)) => BTreeSet::from([
785 ValueKey::Int(b.saturating_sub(1)),
786 ValueKey::Int(b),
787 ValueKey::Int(b.saturating_add(1)),
788 ]),
789 Some(other) => BTreeSet::from([other]),
790 }
791}
792
793fn geo_cell(lat: f64, lon: f64, km: f64) -> Option<(i64, i64, f64, i64)> {
794 if !km.is_finite() || km <= 0.0 {
795 return None;
796 }
797 let cell_deg = (km / 111.0).max(1e-6);
798 let gx = floor_to_i64(lat / cell_deg);
799 let lon_cells = (360.0 / cell_deg).ceil() as i64;
802 let lon_cells = lon_cells.max(1);
803 let gy = floor_to_i64(lon / cell_deg).rem_euclid(lon_cells);
804 Some((gx, gy, cell_deg, lon_cells))
805}
806
807fn geo_index_key(lat: f64, lon: f64, km: f64) -> Option<ValueKey> {
808 let (gx, gy, _, _) = geo_cell(lat, lon, km)?;
809 Some(ValueKey::Str(format!("{gx}|{gy}")))
810}
811
812fn geo_probe_keys(lat: f64, lon: f64, km: f64) -> BTreeSet<ValueKey> {
813 let Some((gx, gy, cell_deg, lon_cells)) = geo_cell(lat, lon, km) else {
814 return BTreeSet::new();
815 };
816 let cos_lat = lat.to_radians().cos().max(0.05);
818 let n = ((km / (111.0 * cos_lat)) / cell_deg).ceil();
819 let n = if n.is_finite() {
820 floor_to_i64(n).max(0)
821 } else {
822 0
823 };
824 let mut out = BTreeSet::new();
825 for dx in -1..=1 {
826 for dy in -n..=n {
827 let cx = gx.saturating_add(dx);
828 let cy = gy.saturating_add(dy).rem_euclid(lon_cells);
829 out.insert(ValueKey::Str(format!("{cx}|{cy}")));
830 }
831 }
832 out
833}
834
835const SCAN_ALL_SENTINEL: ValueKey = ValueKey::Bool(true);
838
839fn ivf_cluster_key(cluster: usize) -> ValueKey {
842 ValueKey::Str(format!("\u{1}ivf:{cluster}"))
843}
844
845fn spec_is_scan_all_universe(spec: &CandidateSpec<'_>) -> bool {
848 match spec {
849 CandidateSpec::ScanAll { .. } => true,
850 CandidateSpec::Intersect(parts) => {
851 !parts.is_empty() && parts.iter().all(spec_is_scan_all_universe)
852 }
853 _ => false,
854 }
855}
856
857fn spec_is_bykey_external(spec: &CandidateSpec<'_>) -> bool {
860 match spec {
861 CandidateSpec::ByKey => true,
862 CandidateSpec::Intersect(parts) => {
863 !parts.is_empty() && parts.iter().all(spec_is_bykey_external)
864 }
865 _ => false,
866 }
867}
868
869impl SideIndex {
870 fn index_keys(spec: &CandidateSpec, get: &dyn Fn(&str) -> Option<Value>) -> BTreeSet<ValueKey> {
871 match spec {
872 CandidateSpec::ByKey => BTreeSet::new(),
873 CandidateSpec::Scalar { field } => get(field)
874 .as_ref()
875 .and_then(ValueKey::from_value)
876 .into_iter()
877 .collect(),
878 CandidateSpec::Tokens { field } => get(field)
879 .as_ref()
880 .and_then(list_tokens)
881 .unwrap_or_default(),
882 CandidateSpec::ScalarOrElements { field } => match get(field) {
883 Some(Value::List(items)) => items
884 .iter()
885 .take(MAX_KEYMATCH_LIST)
886 .filter(|v| matches!(v, Value::Str(_)))
887 .filter_map(ValueKey::from_value)
888 .collect(),
889 Some(v) => ValueKey::from_value(&v).into_iter().collect(),
890 None => BTreeSet::new(),
891 },
892 CandidateSpec::NumericBucket { field, tolerance } => get(field)
893 .as_ref()
894 .and_then(as_finite_f64)
895 .and_then(|v| numeric_index_key(v, *tolerance))
896 .into_iter()
897 .collect(),
898 CandidateSpec::GeoGrid { field, km } => get(field)
899 .as_ref()
900 .and_then(as_latlon)
901 .and_then(|(lat, lon)| geo_index_key(lat, lon, *km))
902 .into_iter()
903 .collect(),
904 CandidateSpec::ScanAll { field } => get(field)
905 .as_ref()
906 .and_then(as_numeric_list)
907 .map(|_| SCAN_ALL_SENTINEL)
908 .into_iter()
909 .collect(),
910 CandidateSpec::VectorClusters { .. } => BTreeSet::new(),
915 CandidateSpec::Hnsw { .. } => BTreeSet::new(),
917 CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) => {
921 let mut out = BTreeSet::new();
922 for s in specs {
923 out.extend(Self::index_keys(s, get));
924 }
925 out
926 }
927 }
928 }
929
930 fn probe_keys(spec: &CandidateSpec, get: &dyn Fn(&str) -> Option<Value>) -> BTreeSet<ValueKey> {
931 match spec {
932 CandidateSpec::ByKey
933 | CandidateSpec::Scalar { .. }
934 | CandidateSpec::Tokens { .. }
935 | CandidateSpec::ScalarOrElements { .. } => Self::index_keys(spec, get),
936 CandidateSpec::NumericBucket { field, tolerance } => get(field)
937 .as_ref()
938 .and_then(as_finite_f64)
939 .map(|v| numeric_probe_keys(v, *tolerance))
940 .unwrap_or_default(),
941 CandidateSpec::GeoGrid { field, km } => get(field)
942 .as_ref()
943 .and_then(as_latlon)
944 .map(|(lat, lon)| geo_probe_keys(lat, lon, *km))
945 .unwrap_or_default(),
946 CandidateSpec::ScanAll { field } => get(field)
947 .as_ref()
948 .and_then(as_numeric_list)
949 .map(|_| SCAN_ALL_SENTINEL)
950 .into_iter()
951 .collect(),
952 CandidateSpec::VectorClusters { .. } => BTreeSet::new(),
954 CandidateSpec::Hnsw { .. } => BTreeSet::new(),
956 CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) => {
959 let mut out = BTreeSet::new();
960 for s in specs {
961 out.extend(Self::probe_keys(s, get));
962 }
963 out
964 }
965 }
966 }
967
968 pub fn insert(&mut self, spec: &CandidateSpec, node: u32, get: &dyn Fn(&str) -> Option<Value>) {
969 self.insert_with(spec, node, &HnswLeg::All, get);
970 }
971
972 pub fn insert_skipping(
979 &mut self,
980 spec: &CandidateSpec,
981 node: u32,
982 already: &BTreeSet<u32>,
983 get: &dyn Fn(&str) -> Option<Value>,
984 ) {
985 self.insert_with(spec, node, &HnswLeg::Skip(already), get);
986 }
987
988 pub fn insert_deferring_hnsw(
997 &mut self,
998 spec: &CandidateSpec,
999 node: u32,
1000 get: &dyn Fn(&str) -> Option<Value>,
1001 ) {
1002 self.insert_with(spec, node, &HnswLeg::Defer, get);
1003 }
1004
1005 pub fn insert_hnsw_only(
1011 &mut self,
1012 spec: &CandidateSpec,
1013 node: u32,
1014 get: &dyn Fn(&str) -> Option<Value>,
1015 ) -> bool {
1016 match spec {
1017 CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) => {
1018 let mut any = false;
1019 for s in specs {
1020 any |= self.insert_hnsw_only(s, node, get);
1021 }
1022 any
1023 }
1024 CandidateSpec::Hnsw { field, .. } => {
1025 let Some(xs) = get(field).as_ref().and_then(as_numeric_list) else {
1026 return false;
1027 };
1028 self.record_vector_meta(node, &xs);
1029 self.hnsw_tracked.insert(node);
1030 if let Some(h) = &mut self.hnsw {
1031 h.insert(node, &xs);
1032 }
1033 true
1034 }
1035 _ => false,
1036 }
1037 }
1038
1039 fn insert_with(
1040 &mut self,
1041 spec: &CandidateSpec,
1042 node: u32,
1043 leg: &HnswLeg<'_>,
1044 get: &dyn Fn(&str) -> Option<Value>,
1045 ) {
1046 if let CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) = spec {
1049 for s in specs {
1050 self.insert_with(s, node, leg, get);
1051 }
1052 return;
1053 }
1054 if let CandidateSpec::Hnsw { field, .. } = spec {
1056 if let Some(xs) = get(field).as_ref().and_then(as_numeric_list) {
1057 self.record_vector_meta(node, &xs);
1061 self.hnsw_tracked.insert(node);
1062 match leg {
1063 HnswLeg::Skip(already) if already.contains(&node) => return,
1066 HnswLeg::Defer => return,
1067 _ => {}
1068 }
1069 if let Some(h) = &mut self.hnsw {
1070 h.insert(node, &xs);
1071 }
1072 }
1073 return;
1074 }
1075 if let CandidateSpec::VectorClusters { field, .. } = spec {
1077 if let Some(xs) = get(field).as_ref().and_then(as_numeric_list) {
1078 self.ivf_raw.insert(node, xs.clone());
1079 if !self.ivf_centroids.is_empty() {
1080 if let Some(unit) = l2_normalize(&xs) {
1082 let c = nearest_centroid(&self.ivf_centroids, &unit);
1083 self.ivf_clusters.insert(node, c);
1084 self.by_key
1085 .entry(ivf_cluster_key(c))
1086 .or_default()
1087 .insert(node);
1088 }
1089 self.ivf_drift = self.ivf_drift.saturating_add(1);
1090 }
1091 }
1092 return;
1093 }
1094
1095 for k in Self::index_keys(spec, get) {
1096 self.by_key.entry(k).or_default().insert(node);
1097 }
1098 if let CandidateSpec::ScanAll { field } = spec {
1099 if let Some(xs) = get(field).as_ref().and_then(as_numeric_list) {
1100 self.record_vector_meta(node, &xs);
1101 }
1102 }
1103 }
1104
1105 fn record_vector_meta(&mut self, node: u32, xs: &[f64]) {
1113 let mut n2 = 0.0f64;
1114 for x in xs {
1115 n2 += x * x;
1116 }
1117 self.vec_meta.insert(node, (xs.len() as u32, n2.sqrt()));
1118 self.vec_checkpoints.insert(node, compute_ckpts(xs));
1119 self.vec_anchor.insert(node, xs[0]);
1121 }
1122
1123 fn forget_vector_meta(&mut self, node: u32) {
1125 self.vec_meta.remove(&node);
1126 self.vec_checkpoints.remove(&node);
1127 self.vec_anchor.remove(&node);
1128 }
1129
1130 pub fn remove(&mut self, spec: &CandidateSpec, node: u32, get: &dyn Fn(&str) -> Option<Value>) {
1131 if let CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) = spec {
1133 for s in specs {
1134 self.remove(s, node, get);
1135 }
1136 return;
1137 }
1138 if let CandidateSpec::Hnsw { field, .. } = spec {
1143 if get(field).as_ref().and_then(as_numeric_list).is_some() {
1144 self.forget_vector_meta(node);
1145 self.hnsw_tracked.remove(&node);
1146 if let Some(h) = &mut self.hnsw {
1147 h.remove(node);
1148 }
1149 self.ivf_drift = self.ivf_drift.saturating_add(1);
1150 }
1151 return;
1152 }
1153 if let CandidateSpec::VectorClusters { .. } = spec {
1158 if self.ivf_raw.remove(&node).is_some() {
1159 self.ivf_drift = self.ivf_drift.saturating_add(1);
1160 if let Some(c) = self.ivf_clusters.remove(&node) {
1161 let key = ivf_cluster_key(c);
1162 if let Some(s) = self.by_key.get_mut(&key) {
1163 s.remove(&node);
1164 if s.is_empty() {
1165 self.by_key.remove(&key);
1166 }
1167 }
1168 }
1169 }
1170 return;
1171 }
1172
1173 for k in Self::index_keys(spec, get) {
1174 if let Some(set) = self.by_key.get_mut(&k) {
1175 set.remove(&node);
1176 if set.is_empty() {
1177 self.by_key.remove(&k);
1178 }
1179 }
1180 }
1181 if let CandidateSpec::ScanAll { field } = spec {
1182 if get(field).as_ref().and_then(as_numeric_list).is_some() {
1183 self.forget_vector_meta(node);
1184 }
1185 }
1186 }
1187
1188 pub fn vec_dim(&self, node: u32) -> Option<u32> {
1190 self.vec_meta.get(&node).map(|(d, _)| *d)
1191 }
1192
1193 pub fn vec_meta(&self, node: u32) -> Option<(u32, f64)> {
1195 self.vec_meta.get(&node).copied()
1196 }
1197
1198 pub fn vec_ckpts(&self, node: u32) -> Option<&[f64; 8]> {
1200 self.vec_checkpoints.get(&node)
1201 }
1202
1203 pub(crate) fn fresh_ckpts_for<'a>(
1230 &'a self,
1231 node: u32,
1232 live: &[f64],
1233 ) -> Option<(f64, &'a [f64; 8])> {
1234 let &(dim, norm) = self.vec_meta.get(&node)?;
1235 if dim != live.len() as u32 {
1236 return None;
1237 }
1238 let live_norm = {
1241 let mut n2 = 0.0f64;
1242 for x in live {
1243 n2 += x * x;
1244 }
1245 n2.sqrt()
1246 };
1247 if norm != live_norm {
1248 return None; }
1250 let live_anchor = live[0];
1254 let &cached_anchor = self.vec_anchor.get(&node)?;
1255 if live_anchor != cached_anchor {
1256 return None;
1257 }
1258 let ckpts = self.vec_checkpoints.get(&node)?;
1259 Some((norm, ckpts))
1260 }
1261
1262 pub fn candidates(
1263 &self,
1264 spec: &CandidateSpec,
1265 get: &dyn Fn(&str) -> Option<Value>,
1266 ) -> BTreeSet<u32> {
1267 if let CandidateSpec::Hnsw { field, k, floor } = spec {
1269 return self.hnsw_candidates(field, *k, *floor, get);
1270 }
1271 if let CandidateSpec::VectorClusters { field, .. } = spec {
1273 return self.ivf_candidates(field, get);
1274 }
1275 if let CandidateSpec::Union(specs) = spec {
1277 return specs.iter().flat_map(|s| self.candidates(s, get)).collect();
1278 }
1279 if let CandidateSpec::Intersect(specs) = spec {
1280 return self.intersect_candidates(specs, get);
1281 }
1282
1283 let mut out = BTreeSet::new();
1284 for k in Self::probe_keys(spec, get) {
1285 if let Some(set) = self.by_key.get(&k) {
1286 out.extend(set.iter().copied());
1287 }
1288 }
1289 if vector_dim_reject_enabled() {
1291 if let CandidateSpec::ScanAll { field } = spec {
1292 if let Some((dim, _)) = get(field).as_ref().and_then(vec_dim_norm) {
1293 out.retain(|id| self.vec_meta.get(id).is_none_or(|(d, _)| *d == dim));
1294 }
1295 }
1296 }
1297 out
1298 }
1299
1300 fn intersect_candidates(
1304 &self,
1305 specs: &[CandidateSpec<'_>],
1306 get: &dyn Fn(&str) -> Option<Value>,
1307 ) -> BTreeSet<u32> {
1308 let mut restrictive = Vec::new();
1309 let mut scan_alls = Vec::new();
1310 for s in specs {
1311 if spec_is_scan_all_universe(s) {
1312 scan_alls.push(s);
1313 } else if spec_is_bykey_external(s) {
1314 continue;
1315 } else {
1316 restrictive.push(s);
1317 }
1318 }
1319 let to_intersect: &[&CandidateSpec<'_>] = if !restrictive.is_empty() {
1320 &restrictive
1321 } else if !scan_alls.is_empty() {
1322 &scan_alls
1323 } else {
1324 return BTreeSet::new();
1325 };
1326 let mut iter = to_intersect.iter();
1327 let Some(first) = iter.next() else {
1328 return BTreeSet::new();
1329 };
1330 let mut acc = self.candidates(first, get);
1331 if acc.is_empty() {
1332 return acc;
1333 }
1334 for s in iter {
1335 let other = self.candidates(s, get);
1336 if other.is_empty() {
1337 return BTreeSet::new();
1338 }
1339 acc = acc.intersection(&other).copied().collect();
1340 if acc.is_empty() {
1341 return acc;
1342 }
1343 }
1344 acc
1345 }
1346
1347 fn ivf_candidates(&self, field: &str, get: &dyn Fn(&str) -> Option<Value>) -> BTreeSet<u32> {
1350 let Some(xs) = get(field).as_ref().and_then(as_numeric_list) else {
1351 return BTreeSet::new();
1352 };
1353 if self.ivf_centroids.is_empty() {
1354 return self.ivf_raw.keys().copied().collect();
1358 }
1359 if self.ivf_raw.len() <= self.ivf_centroids.len() {
1364 return self.ivf_raw.keys().copied().collect();
1365 }
1366 let k = self.ivf_centroids.len();
1367 let p = probe_count(k);
1368
1369 let Some(xs) = l2_normalize(&xs) else {
1372 return BTreeSet::new();
1373 };
1374
1375 let mut dists: Vec<(usize, f64)> = self
1377 .ivf_centroids
1378 .iter()
1379 .enumerate()
1380 .map(|(i, c)| (i, l2_sq(&xs, c)))
1381 .collect();
1382 dists.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
1383
1384 let mut out = BTreeSet::new();
1385 for (ci, _) in dists.iter().take(p) {
1386 let key = ivf_cluster_key(*ci);
1387 if let Some(nodes) = self.by_key.get(&key) {
1388 out.extend(nodes.iter().copied());
1389 }
1390 }
1391 out
1392 }
1393
1394 pub fn fit_ivf_clusters(&mut self, rule_name: &str) {
1405 if self.ivf_raw.is_empty() {
1406 self.ivf_centroids.clear();
1407 self.ivf_clusters.clear();
1408 self.ivf_drift = 0;
1409 return;
1410 }
1411
1412 for c in self.ivf_clusters.values() {
1414 self.by_key.remove(&ivf_cluster_key(*c));
1415 }
1416 self.ivf_clusters.clear();
1417
1418 let vecs: Vec<(u32, Vec<f64>)> = self
1420 .ivf_raw
1421 .iter()
1422 .map(|(&id, xs)| (id, xs.clone()))
1423 .collect();
1424
1425 let n = vecs.len();
1426 let k = cluster_k(n);
1427 let seed = fnv1a_u64(rule_name.as_bytes());
1428
1429 self.ivf_centroids = kmeans_fit(&vecs, k, seed);
1430
1431 for (node, xs) in &vecs {
1433 let Some(unit) = l2_normalize(xs) else {
1434 continue;
1435 };
1436 let c = nearest_centroid(&self.ivf_centroids, &unit);
1437 self.ivf_clusters.insert(*node, c);
1438 self.by_key
1439 .entry(ivf_cluster_key(c))
1440 .or_default()
1441 .insert(*node);
1442 }
1443 self.ivf_drift = 0;
1444 }
1445
1446 pub fn ivf_k(&self) -> usize {
1448 self.ivf_centroids.len()
1449 }
1450
1451 pub fn ivf_cluster_of(&self, node: u32) -> Option<usize> {
1453 self.ivf_clusters.get(&node).copied()
1454 }
1455
1456 pub fn export_ivf_state(&self) -> (Vec<Vec<f64>>, BTreeMap<u32, usize>, u64) {
1461 (
1462 self.ivf_centroids.clone(),
1463 self.ivf_clusters.clone(),
1464 self.ivf_drift,
1465 )
1466 }
1467
1468 pub fn load_ivf_state(
1480 &mut self,
1481 centroids: Vec<Vec<f64>>,
1482 clusters: BTreeMap<u32, usize>,
1483 drift: u64,
1484 ) {
1485 for c in self.ivf_clusters.values() {
1488 self.by_key.remove(&ivf_cluster_key(*c));
1489 }
1490 self.ivf_clusters.clear();
1491
1492 self.ivf_centroids = centroids;
1493 self.ivf_drift = drift;
1494
1495 for (&node, &c) in &clusters {
1497 if !self.ivf_raw.contains_key(&node) {
1498 continue;
1500 }
1501 self.ivf_clusters.insert(node, c);
1502 self.by_key
1503 .entry(ivf_cluster_key(c))
1504 .or_default()
1505 .insert(node);
1506 }
1507 }
1508
1509 pub fn init_hnsw(&mut self, rule_name: &str) {
1518 let seed = fnv1a_u64(rule_name.as_bytes());
1519 self.hnsw = Some(HnswIndex::new(seed));
1520 }
1521
1522 fn hnsw_candidates(
1560 &self,
1561 field: &str,
1562 k: usize,
1563 floor: Option<f64>,
1564 get: &dyn Fn(&str) -> Option<Value>,
1565 ) -> BTreeSet<u32> {
1566 let Some(xs) = get(field).as_ref().and_then(as_numeric_list) else {
1567 return BTreeSet::new();
1568 };
1569 if let Some(h) = &self.hnsw {
1570 if h.can_answer(xs.len()) {
1577 let Some(min) = floor else {
1578 return h.search(&xs, k).into_iter().map(|(id, _)| id).collect();
1579 };
1580 let cap = ef_max();
1581 let mut ef = h.ef_for(k);
1582 while ef < h.len() {
1583 let hits = h.search_with_ef(&xs, ef, ef);
1587 let full = hits.len() == ef;
1589 if full && hits[hits.len() - 1].1 < min - BEAM_FLOOR_SLACK {
1590 return hits.into_iter().map(|(id, _)| id).collect();
1591 }
1592 if !full || ef >= cap {
1597 break;
1598 }
1599 ef = ef.saturating_mul(2);
1600 }
1601 }
1602 }
1603 self.hnsw_tracked.clone()
1605 }
1606
1607 pub fn export_hnsw_blob(&self, complete: bool) -> Vec<u8> {
1617 self.hnsw
1618 .as_ref()
1619 .and_then(|h| crate::hnsw::encode_hnsw_blob(h, complete))
1620 .unwrap_or_default()
1621 }
1622
1623 pub fn load_hnsw_blob(&mut self, blob: &[u8]) {
1630 if let Ok(h) = crate::hnsw::decode_hnsw_blob(blob) {
1631 self.adopt_hnsw(h);
1632 }
1633 }
1634
1635 pub fn init_or_adopt_hnsw(&mut self, rule_name: &str, blob: &[u8]) -> (BTreeSet<u32>, bool) {
1645 self.hnsw = None;
1646 if !blob.is_empty() {
1647 match crate::hnsw::decode_hnsw_blob(blob) {
1648 Ok(h) => self.adopt_hnsw(h),
1649 Err(e) => eprintln!(
1650 "[mushroomdb] rule {rule_name:?}: a persisted HNSW index failed to load \
1651 ({e}); rebuilding it from the node scan"
1652 ),
1653 }
1654 }
1655 match &self.hnsw {
1656 Some(h) => (h.node_ids(), true),
1657 None => {
1658 self.init_hnsw(rule_name);
1659 (BTreeSet::new(), false)
1660 }
1661 }
1662 }
1663
1664 pub fn adopt_hnsw(&mut self, mut h: HnswIndex) {
1670 h.mark_complete();
1677 self.hnsw_tracked = h.node_ids();
1678 self.hnsw = Some(h);
1679 }
1680
1681 pub fn has_hnsw(&self) -> bool {
1683 self.hnsw.as_ref().is_some_and(|h| !h.is_empty())
1684 }
1685
1686 pub fn hnsw_ref(&self) -> Option<&HnswIndex> {
1688 self.hnsw.as_ref()
1689 }
1690
1691 pub fn take_hnsw(&mut self) -> Option<HnswIndex> {
1697 self.hnsw.take()
1698 }
1699}
1700
1701#[cfg(test)]
1702mod tests {
1703 use super::*;
1704 use crate::def::Predicate;
1705 use core_storage::Value;
1706 use std::collections::{BTreeMap, HashMap};
1707
1708 fn getter(map: &HashMap<String, Value>) -> impl Fn(&str) -> Option<Value> + '_ {
1709 move |f: &str| map.get(f).cloned()
1710 }
1711
1712 #[test]
1713 fn kmeans_centroids_are_unit_norm() {
1714 let vecs = vec![(0, vec![3.0, 0.0, 0.0]), (1, vec![0.0, 4.0, 0.0])];
1715 let cents = kmeans_fit(&vecs, 2, 1);
1716 for c in cents {
1717 let n = c.iter().map(|x| x * x).sum::<f64>().sqrt();
1718 assert!((n - 1.0).abs() < 1e-9, "{n}");
1719 }
1720 }
1721
1722 #[test]
1729 fn scaled_vector_joins_same_ivf_cluster_as_unit() {
1730 let spec = CandidateSpec::VectorClusters {
1732 field: "emb",
1733 min: 0.5,
1734 };
1735 let mut idx = SideIndex::default();
1736 idx.load_ivf_state(
1737 vec![vec![1.0, 0.0, 0.0], vec![2.5, 0.1, 0.0]],
1738 BTreeMap::new(),
1739 0,
1740 );
1741 idx.insert(&spec, 1, &getter(&emb(&[1.0, 0.0, 0.0])));
1742 idx.insert(&spec, 2, &getter(&emb(&[3.0, 0.0, 0.0])));
1743 assert_eq!(
1744 idx.ivf_cluster_of(1),
1745 idx.ivf_cluster_of(2),
1746 "scale-equivalent vectors must share an IVF cluster; got {:?} vs {:?}",
1747 idx.ivf_cluster_of(1),
1748 idx.ivf_cluster_of(2)
1749 );
1750 assert_eq!(idx.ivf_cluster_of(1), Some(0));
1751 }
1752
1753 #[test]
1754 fn scalar_index_buckets_by_value() {
1755 let pred = Predicate::FieldEqual {
1756 field: "ind".into(),
1757 };
1758 let spec = candidate_spec(&pred);
1759 let mut idx = SideIndex::default();
1760 let a: HashMap<_, _> = [("ind".to_string(), Value::Str("arch".into()))].into();
1761 let b: HashMap<_, _> = [("ind".to_string(), Value::Str("law".into()))].into();
1762 idx.insert(&spec, 1, &getter(&a));
1763 idx.insert(&spec, 2, &getter(&b));
1764 idx.insert(&spec, 3, &getter(&a));
1765 let c = idx.candidates(&spec, &getter(&a));
1766 assert_eq!(c.into_iter().collect::<Vec<_>>(), vec![1, 3]);
1767 idx.remove(&spec, 3, &getter(&a));
1768 assert_eq!(idx.candidates(&spec, &getter(&a)).len(), 1);
1769 let empty: HashMap<String, Value> = HashMap::new();
1771 idx.insert(&spec, 9, &getter(&empty));
1772 assert!(idx.candidates(&spec, &getter(&empty)).is_empty());
1773 }
1774
1775 #[test]
1776 fn token_index_unions_buckets() {
1777 let mk =
1778 |items: &[&str]| Value::List(items.iter().map(|s| Value::Str((*s).into())).collect());
1779 let pred = Predicate::Overlap {
1780 field: "tags".into(),
1781 min: 0.5,
1782 };
1783 let spec = candidate_spec(&pred);
1784 let mut idx = SideIndex::default();
1785 let a: HashMap<_, _> = [("tags".to_string(), mk(&["x", "y"]))].into();
1786 let b: HashMap<_, _> = [("tags".to_string(), mk(&["y", "z"]))].into();
1787 let c: HashMap<_, _> = [("tags".to_string(), mk(&["q"]))].into();
1788 idx.insert(&spec, 1, &getter(&a));
1789 idx.insert(&spec, 2, &getter(&b));
1790 idx.insert(&spec, 3, &getter(&c));
1791 let probe: HashMap<_, _> = [("tags".to_string(), mk(&["y"]))].into();
1792 assert_eq!(
1793 idx.candidates(&spec, &getter(&probe))
1794 .into_iter()
1795 .collect::<Vec<_>>(),
1796 vec![1, 2]
1797 );
1798 idx.remove(&spec, 2, &getter(&b));
1799 assert_eq!(
1800 idx.candidates(&spec, &getter(&probe))
1801 .into_iter()
1802 .collect::<Vec<_>>(),
1803 vec![1]
1804 );
1805 }
1806
1807 #[test]
1808 fn all_intersects_parts_and_bykey_indexes_nothing() {
1809 let all = Predicate::All(vec![
1810 Predicate::FieldEqual {
1811 field: "ind".into(),
1812 },
1813 Predicate::Overlap {
1814 field: "tags".into(),
1815 min: 0.5,
1816 },
1817 ]);
1818 match candidate_spec(&all) {
1819 CandidateSpec::Intersect(v) => assert_eq!(v.len(), 2),
1820 other => panic!("{other:?}"),
1821 }
1822 let km = Predicate::KeyMatch { field: "fk".into() };
1823 assert!(matches!(candidate_spec(&km), CandidateSpec::ByKey));
1824 let mut idx = SideIndex::default();
1825 let a: HashMap<_, _> = [("fk".to_string(), Value::Str("c1".into()))].into();
1826 idx.insert(&candidate_spec(&km), 1, &getter(&a));
1827 assert!(idx.candidates(&candidate_spec(&km), &getter(&a)).is_empty());
1828 }
1829
1830 fn year(v: Value) -> HashMap<String, Value> {
1831 [("year".to_string(), v)].into()
1832 }
1833
1834 fn loc(lat: f64, lon: f64) -> HashMap<String, Value> {
1835 [(
1836 "loc".to_string(),
1837 Value::List(vec![Value::Float(lat), Value::Float(lon)]),
1838 )]
1839 .into()
1840 }
1841
1842 fn emb(vals: &[f64]) -> HashMap<String, Value> {
1843 [(
1844 "emb".to_string(),
1845 Value::List(vals.iter().copied().map(Value::Float).collect()),
1846 )]
1847 .into()
1848 }
1849
1850 fn bucket_int(spec: &CandidateSpec, map: &HashMap<String, Value>) -> Option<i64> {
1851 match SideIndex::index_keys(spec, &getter(map)).into_iter().next() {
1852 Some(ValueKey::Int(b)) => Some(b),
1853 _ => None,
1854 }
1855 }
1856
1857 #[test]
1858 fn numeric_bucket_adjacency_and_far_value() {
1859 let pred = Predicate::NumericWithin {
1860 field: "year".into(),
1861 tolerance: 2.0,
1862 };
1863 let spec = candidate_spec(&pred);
1864 assert!(matches!(
1865 spec,
1866 CandidateSpec::NumericBucket {
1867 field: "year",
1868 tolerance
1869 } if tolerance == 2.0
1870 ));
1871
1872 let v10 = year(Value::Float(10.0));
1873 let v119 = year(Value::Float(11.9));
1874 let v99 = year(Value::Float(9.9));
1875 let v141 = year(Value::Float(14.1));
1876
1877 let b10 = bucket_int(&spec, &v10).unwrap();
1878 let b119 = bucket_int(&spec, &v119).unwrap();
1879 let b99 = bucket_int(&spec, &v99).unwrap();
1880 assert!((b10 - b119).abs() <= 1);
1882 assert!((b10 - b99).abs() <= 1);
1883
1884 let mut idx = SideIndex::default();
1885 idx.insert(&spec, 1, &getter(&v10));
1886 idx.insert(&spec, 2, &getter(&v119));
1887 idx.insert(&spec, 3, &getter(&v141));
1888 idx.insert(&spec, 4, &getter(&v99));
1889 let hits = idx.candidates(&spec, &getter(&v10));
1890 assert_eq!(hits.into_iter().collect::<Vec<_>>(), vec![1, 2, 4]);
1891 }
1892
1893 #[test]
1894 fn numeric_tol_zero_int_float_collide() {
1895 let pred = Predicate::NumericWithin {
1896 field: "year".into(),
1897 tolerance: 0.0,
1898 };
1899 let spec = candidate_spec(&pred);
1900 let mut idx = SideIndex::default();
1901 idx.insert(&spec, 1, &getter(&year(Value::Int(2))));
1902 assert_eq!(
1903 idx.candidates(&spec, &getter(&year(Value::Float(2.0))))
1904 .into_iter()
1905 .collect::<Vec<_>>(),
1906 vec![1]
1907 );
1908 assert!(idx
1909 .candidates(&spec, &getter(&year(Value::Float(2.1))))
1910 .is_empty());
1911 }
1912
1913 #[test]
1914 fn numeric_tol_zero_signed_zero_collides() {
1915 let pred = Predicate::NumericWithin {
1916 field: "year".into(),
1917 tolerance: 0.0,
1918 };
1919 let spec = candidate_spec(&pred);
1920 let neg = year(Value::Float(-0.0));
1921 let pos = year(Value::Float(0.0));
1922 let mut idx = SideIndex::default();
1923 idx.insert(&spec, 1, &getter(&neg));
1924 assert_eq!(
1925 idx.candidates(&spec, &getter(&pos))
1926 .into_iter()
1927 .collect::<Vec<_>>(),
1928 vec![1]
1929 );
1930 let mut idx2 = SideIndex::default();
1931 idx2.insert(&spec, 2, &getter(&pos));
1932 assert_eq!(
1933 idx2.candidates(&spec, &getter(&neg))
1934 .into_iter()
1935 .collect::<Vec<_>>(),
1936 vec![2]
1937 );
1938 }
1939
1940 #[test]
1941 fn geo_grid_same_cell_cross_cell_and_far_city() {
1942 let pred = Predicate::GeoRadius {
1943 field: "loc".into(),
1944 km: 400.0,
1945 };
1946 let spec = candidate_spec(&pred);
1947 assert!(matches!(
1948 spec,
1949 CandidateSpec::GeoGrid {
1950 field: "loc",
1951 km
1952 } if km == 400.0
1953 ));
1954
1955 let paris = loc(48.8566, 2.3522);
1956 let london = loc(51.5074, -0.1278);
1957 let nearby = loc(48.9, 2.4); let ny = loc(40.7128, -74.0060);
1959
1960 let mut idx = SideIndex::default();
1961 idx.insert(&spec, 1, &getter(&paris));
1962 idx.insert(&spec, 2, &getter(&london));
1963 idx.insert(&spec, 3, &getter(&nearby));
1964 idx.insert(&spec, 4, &getter(&ny));
1965
1966 let from_paris = idx.candidates(&spec, &getter(&paris));
1967 assert!(from_paris.contains(&1), "same-cell self");
1968 assert!(from_paris.contains(&3), "same-cell neighbor");
1969 assert!(from_paris.contains(&2), "cross-cell Paris↔London ~343.5 km");
1970 assert!(!from_paris.contains(&4), "New York not in 400 km probe");
1971 }
1972
1973 #[test]
1974 fn geo_grid_high_latitude_probe_is_superset() {
1975 let pred = Predicate::GeoRadius {
1976 field: "loc".into(),
1977 km: 340.0,
1978 };
1979 let spec = candidate_spec(&pred);
1980 let reyk = loc(64.1466, -21.9426);
1981 let lat = 64.0_f64;
1982 let dlon = 300.0 / (111.0 * lat.to_radians().cos());
1983 let east = loc(lat, -21.9426 + dlon);
1984
1985 let mut idx = SideIndex::default();
1986 idx.insert(&spec, 1, &getter(&reyk));
1987 idx.insert(&spec, 2, &getter(&east));
1988 let hits = idx.candidates(&spec, &getter(&reyk));
1989 assert!(
1990 hits.contains(&2),
1991 "300 km east of Reykjavik must stay in the high-lat probe"
1992 );
1993 }
1994
1995 #[test]
1996 fn geo_grid_antimeridian_wrap_and_evaluate_agree() {
1997 let pred = Predicate::GeoRadius {
1998 field: "loc".into(),
1999 km: 400.0,
2000 };
2001 let spec = candidate_spec(&pred);
2002 let east = loc(70.0, 179.9);
2003 let west = loc(70.0, -179.9);
2004
2005 let mut idx = SideIndex::default();
2006 idx.insert(&spec, 1, &getter(&east));
2007 assert!(
2008 idx.candidates(&spec, &getter(&west)).contains(&1),
2009 "±180 pair at lat 70 must land in the wrapped probe"
2010 );
2011
2012 let sp = |f: &str| east.get(f).cloned();
2013 let dp = |f: &str| west.get(f).cloned();
2014 let score = crate::def::evaluate(
2015 &pred,
2016 &crate::def::NodeView {
2017 key: "e",
2018 props: &sp,
2019 },
2020 &crate::def::NodeView {
2021 key: "w",
2022 props: &dp,
2023 },
2024 );
2025 assert!(
2026 score.is_some(),
2027 "haversine must match across the antimeridian"
2028 );
2029
2030 let paris = loc(48.8566, 2.3522);
2032 let ny = loc(40.7128, -74.0060);
2033 let mut idx2 = SideIndex::default();
2034 idx2.insert(&spec, 4, &getter(&ny));
2035 assert!(
2036 !idx2.candidates(&spec, &getter(&paris)).contains(&4),
2037 "New York still not in the Paris probe after wrap"
2038 );
2039 }
2040
2041 #[test]
2042 fn scan_all_returns_vector_nodes_skips_malformed() {
2043 let pred = Predicate::VectorSimilar {
2044 field: "emb".into(),
2045 min: 0.5,
2046 };
2047 let spec = candidate_spec(&pred);
2048 assert!(matches!(spec, CandidateSpec::ScanAll { field: "emb" }));
2049
2050 let mut idx = SideIndex::default();
2051 idx.insert(&spec, 1, &getter(&emb(&[1.0, 0.0])));
2052 idx.insert(&spec, 2, &getter(&emb(&[0.0, 1.0])));
2053 idx.insert(&spec, 3, &getter(&emb(&[1.0, 2.0, 3.0])));
2054 let empty: HashMap<_, _> = [("emb".to_string(), Value::List(vec![]))].into();
2055 let text: HashMap<_, _> =
2056 [("emb".to_string(), Value::List(vec![Value::Str("x".into())]))].into();
2057 let missing: HashMap<String, Value> = HashMap::new();
2058 idx.insert(&spec, 4, &getter(&empty));
2059 idx.insert(&spec, 5, &getter(&text));
2060 idx.insert(&spec, 6, &getter(&missing));
2061
2062 let hits = idx.candidates(&spec, &getter(&emb(&[1.0, 0.0])));
2063 assert_eq!(
2064 hits.into_iter().collect::<Vec<_>>(),
2065 vec![1, 2],
2066 "dim-2 probe must drop the dim-3 member"
2067 );
2068 assert_eq!(
2069 idx.candidates(&spec, &getter(&emb(&[1.0, 2.0, 3.0])))
2070 .into_iter()
2071 .collect::<Vec<_>>(),
2072 vec![3]
2073 );
2074 with_vector_dim_reject(false, || {
2075 assert_eq!(
2076 idx.candidates(&spec, &getter(&emb(&[1.0, 0.0])))
2077 .into_iter()
2078 .collect::<Vec<_>>(),
2079 vec![1, 2, 3],
2080 "unfiltered ScanAll still returns every vector node"
2081 );
2082 });
2083 assert_eq!(idx.vec_dim(1), Some(2));
2084 assert_eq!(idx.vec_dim(3), Some(3));
2085 assert!(idx.vec_meta(1).is_some());
2086 assert!(idx.vec_dim(4).is_none());
2087 assert!(idx.candidates(&spec, &getter(&empty)).is_empty());
2088 assert!(idx.candidates(&spec, &getter(&text)).is_empty());
2089 assert!(idx.candidates(&spec, &getter(&missing)).is_empty());
2090 idx.remove(&spec, 1, &getter(&emb(&[1.0, 0.0])));
2091 assert!(idx.vec_dim(1).is_none());
2092 }
2093
2094 #[test]
2095 fn legacy_specs_probe_keys_equal_index_keys() {
2096 let a: HashMap<_, _> = [
2097 ("ind".to_string(), Value::Str("arch".into())),
2098 (
2099 "tags".to_string(),
2100 Value::List(vec![Value::Str("x".into()), Value::Str("y".into())]),
2101 ),
2102 ("fk".to_string(), Value::Str("c1".into())),
2103 ]
2104 .into();
2105 let get = getter(&a);
2106 for pred in [
2107 Predicate::KeyMatch { field: "fk".into() },
2108 Predicate::FieldEqual {
2109 field: "ind".into(),
2110 },
2111 Predicate::Overlap {
2112 field: "tags".into(),
2113 min: 0.5,
2114 },
2115 ] {
2116 let spec = candidate_spec(&pred);
2117 assert_eq!(
2118 SideIndex::index_keys(&spec, &get),
2119 SideIndex::probe_keys(&spec, &get)
2120 );
2121 }
2122 }
2123
2124 #[test]
2125 fn all_vector_then_field_equal_does_not_scan_all() {
2126 let p = Predicate::All(vec![
2127 Predicate::VectorSimilar {
2128 field: "e".into(),
2129 min: 0.8,
2130 },
2131 Predicate::FieldEqual {
2132 field: "industry".into(),
2133 },
2134 ]);
2135 match candidate_spec(&p) {
2136 CandidateSpec::Intersect(v) => assert_eq!(v.len(), 2),
2137 other => panic!("{other:?}"),
2138 }
2139
2140 let spec = candidate_spec(&p);
2141 let mut idx = SideIndex::default();
2142 let mk = |industry: &str, e: &[f64]| {
2143 [
2144 ("industry".to_string(), Value::Str(industry.into())),
2145 (
2146 "e".to_string(),
2147 Value::List(e.iter().copied().map(Value::Float).collect()),
2148 ),
2149 ]
2150 .into()
2151 };
2152 let same: HashMap<_, _> = mk("tech", &[1.0, 0.0]);
2153 let other_ind: HashMap<_, _> = mk("law", &[1.0, 0.0]);
2154 let no_vec: HashMap<_, _> = [("industry".to_string(), Value::Str("tech".into()))].into();
2155 idx.insert(&spec, 1, &getter(&same));
2156 idx.insert(&spec, 2, &getter(&other_ind));
2157 idx.insert(&spec, 3, &getter(&no_vec));
2158
2159 let hits = idx.candidates(&spec, &getter(&same));
2160 assert!(hits.contains(&1), "matching industry must stay a candidate");
2161 assert!(
2162 !hits.contains(&2),
2163 "different industry must not be scanned in via VectorSimilar"
2164 );
2165 assert!(
2166 hits.contains(&3),
2167 "ScanAll is universe: extra Scalar-only candidates are allowed"
2168 );
2169
2170 let empty_ind: HashMap<_, _> = mk("finance", &[1.0, 0.0]);
2171 assert!(
2172 idx.candidates(&spec, &getter(&empty_ind)).is_empty(),
2173 "empty Scalar child → empty intersect"
2174 );
2175 }
2176
2177 #[test]
2178 fn all_approx_vector_then_field_equal_is_intersect() {
2179 let p = Predicate::All(vec![
2180 Predicate::VectorSimilar {
2181 field: "e".into(),
2182 min: 0.8,
2183 },
2184 Predicate::FieldEqual {
2185 field: "industry".into(),
2186 },
2187 ]);
2188 match candidate_spec_approx(&p) {
2189 CandidateSpec::Intersect(v) => assert_eq!(v.len(), 2),
2190 other => panic!("{other:?}"),
2191 }
2192
2193 let spec = candidate_spec_approx(&p);
2194 let mut idx = SideIndex::default();
2195 idx.init_hnsw("test-rule");
2197 let mk = |industry: &str, e: &[f64]| {
2198 [
2199 ("industry".to_string(), Value::Str(industry.into())),
2200 (
2201 "e".to_string(),
2202 Value::List(e.iter().copied().map(Value::Float).collect()),
2203 ),
2204 ]
2205 .into()
2206 };
2207 let same: HashMap<_, _> = mk("tech", &[1.0, 0.0]);
2208 let other_ind: HashMap<_, _> = mk("law", &[1.0, 0.0]);
2209 idx.insert(&spec, 1, &getter(&same));
2210 idx.insert(&spec, 2, &getter(&other_ind));
2211 let hits = idx.candidates(&spec, &getter(&same));
2212 assert!(hits.contains(&1), "matching industry must stay a candidate");
2213 assert!(
2214 !hits.contains(&2),
2215 "FieldEqual must be probed on the approximate All path"
2216 );
2217 }
2218
2219 #[test]
2220 fn all_of_scan_all_stays_scan_all() {
2221 let p = Predicate::All(vec![
2222 Predicate::VectorSimilar {
2223 field: "emb".into(),
2224 min: 0.5,
2225 },
2226 Predicate::VectorSimilar {
2227 field: "emb".into(),
2228 min: 0.9,
2229 },
2230 ]);
2231 match candidate_spec(&p) {
2232 CandidateSpec::Intersect(v) => assert_eq!(v.len(), 2),
2233 other => panic!("{other:?}"),
2234 }
2235 let spec = candidate_spec(&p);
2236 let mut idx = SideIndex::default();
2237 idx.insert(&spec, 1, &getter(&emb(&[1.0, 0.0])));
2238 idx.insert(&spec, 2, &getter(&emb(&[0.0, 1.0])));
2239 let hits = idx.candidates(&spec, &getter(&emb(&[1.0, 0.0])));
2240 assert_eq!(hits.into_iter().collect::<Vec<_>>(), vec![1, 2]);
2241 }
2242
2243 #[test]
2244 fn any_stays_union() {
2245 let p = Predicate::Any(vec![
2246 Predicate::FieldEqual {
2247 field: "industry".into(),
2248 },
2249 Predicate::Overlap {
2250 field: "tags".into(),
2251 min: 0.5,
2252 },
2253 ]);
2254 match candidate_spec(&p) {
2255 CandidateSpec::Union(v) => assert_eq!(v.len(), 2),
2256 other => panic!("{other:?}"),
2257 }
2258 }
2259
2260 #[test]
2263 fn checkpoint_populated_and_consistent_with_norm() {
2264 let pred = Predicate::VectorSimilar {
2265 field: "emb".into(),
2266 min: 0.8,
2267 };
2268 let spec = candidate_spec(&pred);
2269 let xs = [3.0f64, 4.0]; let mut idx = SideIndex::default();
2271 idx.insert(&spec, 1, &getter(&emb(&xs)));
2272
2273 let ckpts = idx
2274 .vec_ckpts(1)
2275 .expect("checkpoints must exist after insert");
2276 let (_, norm) = idx.vec_meta(1).unwrap();
2277 assert!(
2278 (ckpts[0] - norm).abs() < 1e-12,
2279 "ckpts[0] must equal the full L2 norm; got {} vs {}",
2280 ckpts[0],
2281 norm
2282 );
2283 assert!(
2284 (norm - 5.0).abs() < 1e-12,
2285 "norm of [3,4] must be 5.0, got {norm}"
2286 );
2287
2288 idx.remove(&spec, 1, &getter(&emb(&xs)));
2290 assert!(
2291 idx.vec_ckpts(1).is_none(),
2292 "checkpoints must be removed after remove()"
2293 );
2294 }
2295
2296 #[test]
2299 fn fresh_ckpts_for_freshness_gate() {
2300 let pred = Predicate::VectorSimilar {
2301 field: "emb".into(),
2302 min: 0.8,
2303 };
2304 let spec = candidate_spec(&pred);
2305 let xs = [1.0f64, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0];
2306 let mut idx = SideIndex::default();
2307 idx.insert(&spec, 7, &getter(&emb(&xs)));
2308
2309 let result = idx.fresh_ckpts_for(7, &xs);
2311 assert!(
2312 result.is_some(),
2313 "fresh_ckpts_for must succeed with matching live vector"
2314 );
2315 let (norm, ckpts) = result.unwrap();
2316 assert!((norm - 1.0).abs() < 1e-12);
2317 assert!((ckpts[0] - 1.0).abs() < 1e-12);
2318
2319 let wrong = [2.0f64, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]; assert!(
2322 idx.fresh_ckpts_for(7, &wrong).is_none(),
2323 "freshness gate must reject mismatched norm"
2324 );
2325
2326 let short = [1.0f64, 0.0];
2328 assert!(
2329 idx.fresh_ckpts_for(7, &short).is_none(),
2330 "freshness gate must reject mismatched dim"
2331 );
2332
2333 assert!(idx.fresh_ckpts_for(99, &xs).is_none());
2335 }
2336
2337 #[test]
2340 fn checkpoint_suffix_norms_non_increasing() {
2341 let pred = Predicate::VectorSimilar {
2342 field: "emb".into(),
2343 min: 0.5,
2344 };
2345 let spec = candidate_spec(&pred);
2346 let xs: Vec<f64> = (1..=16).map(|i| i as f64).collect();
2347 let mut idx = SideIndex::default();
2348 idx.insert(&spec, 42, &getter(&emb(&xs)));
2349
2350 let ckpts = *idx.vec_ckpts(42).unwrap();
2351 for c in 0..7 {
2352 assert!(
2353 ckpts[c] >= ckpts[c + 1] - 1e-12,
2354 "suffix norm must be non-increasing: ckpts[{c}]={} < ckpts[{}]={}",
2355 ckpts[c],
2356 c + 1,
2357 ckpts[c + 1]
2358 );
2359 }
2360 let expected_last = (15.0f64 * 15.0 + 16.0 * 16.0).sqrt();
2362 assert!(
2363 (ckpts[7] - expected_last).abs() < 1e-9,
2364 "ckpts[7] should be norm of last segment; got {} vs {}",
2365 ckpts[7],
2366 expected_last
2367 );
2368 }
2369 fn hnsw_side() -> (SideIndex, CandidateSpec<'static>) {
2375 let spec = CandidateSpec::Hnsw {
2376 field: "emb",
2377 k: 8,
2378 floor: None,
2379 };
2380 let mut side = SideIndex::default();
2381 side.init_hnsw("sim");
2382 for (id, xs) in [
2383 (1u32, vec![1.0, 0.0]),
2384 (2, vec![0.0, 1.0]),
2385 (3, vec![0.7, 0.7]),
2386 ] {
2387 side.insert(&spec, id, &getter(&emb(&xs)));
2388 }
2389 (side, spec)
2390 }
2391
2392 #[test]
2395 fn init_or_adopt_hnsw_adopts_a_usable_blob() {
2396 let (side, spec) = hnsw_side();
2397 let blob = side.export_hnsw_blob(true);
2398
2399 let mut fresh = SideIndex::default();
2400 let (ids, adopted) = fresh.init_or_adopt_hnsw("sim", &blob);
2401 assert!(adopted, "a usable blob must be adopted, not rebuilt");
2402 assert_eq!(ids, BTreeSet::from([1, 2, 3]));
2403 assert!(fresh.has_hnsw());
2404 assert_eq!(
2405 fresh.candidates(&spec, &getter(&emb(&[1.0, 0.0]))),
2406 side.candidates(&spec, &getter(&emb(&[1.0, 0.0]))),
2407 "the adopted graph must answer as the original did"
2408 );
2409 }
2410
2411 #[test]
2415 fn an_unknown_version_leaves_the_graph_empty() {
2416 let (side, spec) = hnsw_side();
2417 let mut blob = side.export_hnsw_blob(true);
2418 blob[4] = 99; let mut fresh = SideIndex::default();
2421 let (ids, adopted) = fresh.init_or_adopt_hnsw("sim", &blob);
2422 assert!(!adopted, "an unreadable blob must not count as adopted");
2423 assert!(ids.is_empty(), "nothing may be skipped by the scan");
2424 assert!(!fresh.has_hnsw(), "the graph must be empty");
2425
2426 for (id, xs) in [
2428 (1u32, vec![1.0, 0.0]),
2429 (2, vec![0.0, 1.0]),
2430 (3, vec![0.7, 0.7]),
2431 ] {
2432 fresh.insert_skipping(&spec, id, &ids, &getter(&emb(&xs)));
2433 }
2434 assert!(fresh.has_hnsw());
2435 assert_eq!(
2436 fresh.candidates(&spec, &getter(&emb(&[1.0, 0.0]))),
2437 BTreeSet::from([1, 2, 3])
2438 );
2439 }
2440
2441 #[test]
2444 fn an_unreadable_blob_leaves_the_graph_empty() {
2445 let (side, spec) = hnsw_side();
2446 let mut blob = side.export_hnsw_blob(true);
2447 blob.truncate(blob.len() / 2);
2448
2449 let mut fresh = SideIndex::default();
2450 let (ids, adopted) = fresh.init_or_adopt_hnsw("sim", &blob);
2451 assert!(!adopted, "an unreadable blob must not count as adopted");
2452 assert!(ids.is_empty(), "nothing may be skipped by the scan");
2453 assert!(!fresh.has_hnsw(), "the graph must be empty");
2454
2455 for (id, xs) in [
2457 (1u32, vec![1.0, 0.0]),
2458 (2, vec![0.0, 1.0]),
2459 (3, vec![0.7, 0.7]),
2460 ] {
2461 fresh.insert_skipping(&spec, id, &ids, &getter(&emb(&xs)));
2462 }
2463 assert!(fresh.has_hnsw());
2464 assert_eq!(
2465 fresh.candidates(&spec, &getter(&emb(&[1.0, 0.0]))),
2466 BTreeSet::from([1, 2, 3])
2467 );
2468 }
2469
2470 #[test]
2473 fn insert_skipping_tracks_but_does_not_reinsert() {
2474 let (side, spec) = hnsw_side();
2475 let blob = side.export_hnsw_blob(true);
2476
2477 let mut fresh = SideIndex::default();
2478 let (already, _) = fresh.init_or_adopt_hnsw("sim", &blob);
2479 let before = fresh.hnsw_ref().map(|h| h.len());
2480
2481 fresh.insert_skipping(&spec, 3, &already, &getter(&emb(&[0.7, 0.7])));
2483 assert_eq!(
2484 fresh.hnsw_ref().map(|h| h.len()),
2485 before,
2486 "an adopted id must not be re-inserted"
2487 );
2488 fresh.insert_skipping(&spec, 4, &already, &getter(&emb(&[-1.0, 0.0])));
2489 assert_eq!(
2490 fresh.hnsw_ref().map(|h| h.len()),
2491 before.map(|n| n + 1),
2492 "a post-snapshot id must be inserted"
2493 );
2494 assert_eq!(
2495 fresh.candidates(&spec, &getter(&emb(&[1.0, 0.0]))),
2496 BTreeSet::from([1, 2, 3, 4])
2497 );
2498 }
2499}