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