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 fn cluster_k(n: usize) -> usize {
54 if n == 0 {
55 return IVF_K_MIN;
56 }
57 let k = (n as f64).sqrt().ceil() as usize;
58 k.clamp(IVF_K_MIN, IVF_K_MAX)
59}
60
61pub fn probe_count(k: usize) -> usize {
63 k.div_ceil(IVF_PROBE_DENOM).max(1)
64}
65
66fn l2_normalize(xs: &[f64]) -> Option<Vec<f64>> {
68 let n = xs.iter().map(|x| x * x).sum::<f64>().sqrt();
69 if n == 0.0 {
70 return None;
71 }
72 Some(xs.iter().map(|x| x / n).collect())
73}
74
75fn l2_sq(a: &[f64], b: &[f64]) -> f64 {
78 if a.len() != b.len() {
79 return f64::MAX;
80 }
81 a.iter().zip(b.iter()).map(|(x, y)| (x - y) * (x - y)).sum()
82}
83
84pub fn nearest_centroid(centroids: &[Vec<f64>], xs: &[f64]) -> usize {
87 centroids
88 .iter()
89 .enumerate()
90 .min_by(|(_, a), (_, b)| {
91 l2_sq(xs, a)
92 .partial_cmp(&l2_sq(xs, b))
93 .unwrap_or(std::cmp::Ordering::Equal)
94 })
95 .map(|(i, _)| i)
96 .unwrap_or(0)
97}
98
99pub fn fnv1a_u64(data: &[u8]) -> u64 {
103 const FNV_OFFSET: u64 = 14_695_981_039_346_656_037;
104 const FNV_PRIME: u64 = 1_099_511_628_211;
105 let mut h = FNV_OFFSET;
106 for &b in data {
107 h ^= b as u64;
108 h = h.wrapping_mul(FNV_PRIME);
109 }
110 h
111}
112
113#[inline]
116fn lcg_next(state: u64) -> u64 {
117 state
118 .wrapping_mul(6_364_136_223_846_793_005)
119 .wrapping_add(1_442_695_040_888_963_407)
120}
121
122pub fn kmeans_fit(vecs: &[(u32, Vec<f64>)], k: usize, seed: u64) -> Vec<Vec<f64>> {
132 let vecs: Vec<(u32, Vec<f64>)> = vecs
133 .iter()
134 .filter_map(|(id, xs)| l2_normalize(xs).map(|n| (*id, n)))
135 .collect();
136 if vecs.is_empty() || k == 0 {
137 return vec![];
138 }
139 let n = vecs.len();
140 let k = k.min(n);
141 let dim = vecs[0].1.len();
142 if dim == 0 {
143 return vec![];
144 }
145
146 let mut state = seed;
148 let mut used = vec![false; n];
149 let mut init_idxs: Vec<usize> = Vec::with_capacity(k);
150 let mut attempts = 0usize;
151 while init_idxs.len() < k && attempts < n * 4 {
152 state = lcg_next(state);
153 let idx = (state >> 33) as usize % n;
154 if !used[idx] {
155 used[idx] = true;
156 init_idxs.push(idx);
157 }
158 attempts += 1;
159 }
160 if init_idxs.len() < k {
163 for (i, in_use) in used.iter().enumerate().take(n) {
164 if !in_use {
165 init_idxs.push(i);
166 if init_idxs.len() == k {
167 break;
168 }
169 }
170 }
171 }
172 let mut centroids: Vec<Vec<f64>> = init_idxs.iter().map(|&i| vecs[i].1.clone()).collect();
173 let mut assignments = vec![0usize; n];
174
175 for iter in 0..IVF_ITERATIONS {
177 for (j, (_, xs)) in vecs.iter().enumerate() {
179 assignments[j] = nearest_centroid(¢roids, xs);
180 }
181
182 let mut sums = vec![vec![0.0f64; dim]; k];
184 let mut counts = vec![0usize; k];
185 for (j, (_, xs)) in vecs.iter().enumerate() {
186 let c = assignments[j];
187 counts[c] += 1;
188 for d in 0..dim {
189 sums[c][d] += xs[d];
190 }
191 }
192
193 let mut new_centroids = vec![vec![0.0f64; dim]; k];
195 let mut empty: Vec<usize> = Vec::new();
196 for c in 0..k {
197 if counts[c] == 0 {
198 empty.push(c);
199 } else {
200 for d in 0..dim {
201 new_centroids[c][d] = sums[c][d] / counts[c] as f64;
202 }
203 }
204 }
205
206 for (ei, ec) in empty.into_iter().enumerate() {
209 let reseed =
210 seed ^ (iter as u64).wrapping_mul(0x9E37) ^ (ei as u64).wrapping_mul(0x1234_5679);
211 let mut rs = lcg_next(reseed);
212 rs = lcg_next(rs);
213 let pick = (rs >> 33) as usize % n;
214 new_centroids[ec] = vecs[pick].1.clone();
215 }
216
217 centroids = new_centroids;
218 }
219
220 centroids
221}
222
223#[cfg(test)]
224thread_local! {
225 static VECTOR_DIM_REJECT: std::cell::Cell<bool> = const { std::cell::Cell::new(true) };
226 static VECTOR_EARLY_EXIT: std::cell::Cell<bool> = const { std::cell::Cell::new(true) };
227}
228
229fn vector_dim_reject_enabled() -> bool {
230 #[cfg(test)]
231 {
232 VECTOR_DIM_REJECT.with(|c| c.get())
233 }
234 #[cfg(not(test))]
235 {
236 true
237 }
238}
239
240pub(crate) fn vector_early_exit_enabled() -> bool {
241 #[cfg(test)]
242 {
243 VECTOR_EARLY_EXIT.with(|c| c.get())
244 }
245 #[cfg(not(test))]
246 {
247 true
248 }
249}
250
251#[cfg(test)]
253pub fn with_vector_dim_reject<R>(enabled: bool, f: impl FnOnce() -> R) -> R {
254 VECTOR_DIM_REJECT.with(|c| {
255 let prev = c.replace(enabled);
256 let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
257 c.set(prev);
258 match out {
259 Ok(v) => v,
260 Err(p) => std::panic::resume_unwind(p),
261 }
262 })
263}
264
265#[cfg(test)]
267pub fn with_vector_early_exit<R>(enabled: bool, f: impl FnOnce() -> R) -> R {
268 VECTOR_EARLY_EXIT.with(|c| {
269 let prev = c.replace(enabled);
270 let out = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
271 c.set(prev);
272 match out {
273 Ok(v) => v,
274 Err(p) => std::panic::resume_unwind(p),
275 }
276 })
277}
278
279#[derive(Debug, Default)]
280pub struct SideIndex {
281 by_key: BTreeMap<ValueKey, BTreeSet<u32>>,
282 vec_meta: BTreeMap<u32, (u32, f64)>,
287 vec_checkpoints: BTreeMap<u32, [f64; 8]>,
294 vec_anchor: BTreeMap<u32, f64>,
301
302 ivf_raw: BTreeMap<u32, Vec<f64>>,
307 ivf_centroids: Vec<Vec<f64>>,
309 ivf_clusters: BTreeMap<u32, usize>,
312 pub ivf_drift: u64,
316
317 hnsw: Option<HnswIndex>,
320 hnsw_tracked: BTreeSet<u32>,
323}
324
325#[derive(Debug, Default)]
326pub struct RuleIndex {
327 pub src_side: SideIndex,
328 pub dst_side: SideIndex,
329}
330
331#[derive(Debug)]
332pub enum CandidateSpec<'a> {
333 ByKey,
334 Scalar {
335 field: &'a str,
336 },
337 Tokens {
338 field: &'a str,
339 },
340 ScalarOrElements {
349 field: &'a str,
350 },
351 NumericBucket {
352 field: &'a str,
353 tolerance: f64,
354 },
355 GeoGrid {
356 field: &'a str,
357 km: f64,
358 },
359 ScanAll {
360 field: &'a str,
361 },
362 VectorClusters {
369 field: &'a str,
370 min: f64,
371 },
372 Hnsw {
379 field: &'a str,
380 k: usize,
383 },
384 Union(Vec<CandidateSpec<'a>>),
391 Intersect(Vec<CandidateSpec<'a>>),
398}
399
400pub fn candidate_spec(p: &Predicate) -> CandidateSpec<'_> {
413 match p {
414 Predicate::KeyMatch { .. } => CandidateSpec::ByKey,
415 Predicate::FieldEqual { field } => CandidateSpec::Scalar { field },
416 Predicate::Overlap { field, .. } => CandidateSpec::Tokens { field },
417 Predicate::NumericWithin { field, tolerance } => CandidateSpec::NumericBucket {
418 field,
419 tolerance: *tolerance,
420 },
421 Predicate::GeoRadius { field, km } => CandidateSpec::GeoGrid { field, km: *km },
422 Predicate::VectorSimilar { field, .. } => CandidateSpec::ScanAll { field },
423 Predicate::All(parts) => {
424 debug_assert!(
425 !parts.is_empty(),
426 "candidate_spec requires a validated predicate"
427 );
428 CandidateSpec::Intersect(parts.iter().map(candidate_spec).collect())
429 }
430 Predicate::Any(parts) => {
431 debug_assert!(
432 !parts.is_empty(),
433 "candidate_spec requires a validated predicate"
434 );
435 CandidateSpec::Union(parts.iter().map(candidate_spec).collect())
436 }
437 }
438}
439
440pub fn candidate_spec_approx(p: &Predicate) -> CandidateSpec<'_> {
459 candidate_spec_approx_with_k(p, 64)
460}
461
462pub fn candidate_spec_approx_with_k(p: &Predicate, k: usize) -> CandidateSpec<'_> {
464 match p {
465 Predicate::VectorSimilar { field, .. } => CandidateSpec::Hnsw { field, k },
466 Predicate::All(parts) => {
467 debug_assert!(
468 !parts.is_empty(),
469 "candidate_spec_approx requires a validated predicate"
470 );
471 CandidateSpec::Intersect(
472 parts
473 .iter()
474 .map(|p| candidate_spec_approx_with_k(p, k))
475 .collect(),
476 )
477 }
478 other => candidate_spec(other),
479 }
480}
481
482pub(crate) fn as_finite_f64(v: &Value) -> Option<f64> {
483 match v {
484 Value::Int(i) => Some(*i as f64),
485 Value::Float(f) if f.is_finite() => Some(*f),
486 _ => None,
487 }
488}
489
490fn as_latlon(v: &Value) -> Option<(f64, f64)> {
491 let Value::List(items) = v else {
492 return None;
493 };
494 if items.len() != 2 {
495 return None;
496 }
497 let lat = as_finite_f64(&items[0])?;
498 let lon = as_finite_f64(&items[1])?;
499 if (-90.0..=90.0).contains(&lat) && (-180.0..=180.0).contains(&lon) {
500 Some((lat, lon))
501 } else {
502 None
503 }
504}
505
506pub(crate) fn as_numeric_list(v: &Value) -> Option<Vec<f64>> {
507 let Value::List(items) = v else {
508 return None;
509 };
510 if items.is_empty() {
511 return None;
512 }
513 items.iter().map(as_finite_f64).collect()
514}
515
516fn vec_dim_norm(v: &Value) -> Option<(u32, f64)> {
517 let xs = as_numeric_list(v)?;
518 let mut n2 = 0.0;
519 for x in &xs {
520 n2 += *x * *x;
521 }
522 Some((xs.len() as u32, n2.sqrt()))
523}
524
525fn compute_ckpts(xs: &[f64]) -> [f64; 8] {
531 let dim = xs.len();
532 let mut ckpts = [0.0f64; 8];
533 if dim == 0 {
534 return ckpts;
535 }
536 let boundaries: [usize; 8] = std::array::from_fn(|i| i * dim / 8);
538 let mut suffix_sq = 0.0f64;
539 let mut ci = 7i32;
541 for j in (0..dim).rev() {
542 suffix_sq += xs[j] * xs[j];
543 while ci >= 0 && boundaries[ci as usize] == j {
545 ckpts[ci as usize] = suffix_sq.sqrt();
546 ci -= 1;
547 }
548 }
549 ckpts
550}
551
552fn floor_to_i64(x: f64) -> i64 {
553 let floored = x.floor();
554 if !floored.is_finite() {
555 return 0;
556 }
557 if floored >= i64::MAX as f64 {
558 i64::MAX
559 } else if floored <= i64::MIN as f64 {
560 i64::MIN
561 } else {
562 floored as i64
563 }
564}
565
566fn numeric_index_key(v: f64, tolerance: f64) -> Option<ValueKey> {
570 if !tolerance.is_finite() || tolerance < 0.0 {
571 return None;
572 }
573 if tolerance == 0.0 {
574 let v = if v == 0.0 { 0.0_f64 } else { v };
575 return Some(ValueKey::FloatBits(v.to_bits()));
576 }
577 Some(ValueKey::Int(floor_to_i64(v / tolerance)))
578}
579
580fn numeric_probe_keys(v: f64, tolerance: f64) -> BTreeSet<ValueKey> {
581 match numeric_index_key(v, tolerance) {
582 None => BTreeSet::new(),
583 Some(k @ ValueKey::FloatBits(_)) => BTreeSet::from([k]),
584 Some(ValueKey::Int(b)) => BTreeSet::from([
585 ValueKey::Int(b.saturating_sub(1)),
586 ValueKey::Int(b),
587 ValueKey::Int(b.saturating_add(1)),
588 ]),
589 Some(other) => BTreeSet::from([other]),
590 }
591}
592
593fn geo_cell(lat: f64, lon: f64, km: f64) -> Option<(i64, i64, f64, i64)> {
594 if !km.is_finite() || km <= 0.0 {
595 return None;
596 }
597 let cell_deg = (km / 111.0).max(1e-6);
598 let gx = floor_to_i64(lat / cell_deg);
599 let lon_cells = (360.0 / cell_deg).ceil() as i64;
602 let lon_cells = lon_cells.max(1);
603 let gy = floor_to_i64(lon / cell_deg).rem_euclid(lon_cells);
604 Some((gx, gy, cell_deg, lon_cells))
605}
606
607fn geo_index_key(lat: f64, lon: f64, km: f64) -> Option<ValueKey> {
608 let (gx, gy, _, _) = geo_cell(lat, lon, km)?;
609 Some(ValueKey::Str(format!("{gx}|{gy}")))
610}
611
612fn geo_probe_keys(lat: f64, lon: f64, km: f64) -> BTreeSet<ValueKey> {
613 let Some((gx, gy, cell_deg, lon_cells)) = geo_cell(lat, lon, km) else {
614 return BTreeSet::new();
615 };
616 let cos_lat = lat.to_radians().cos().max(0.05);
618 let n = ((km / (111.0 * cos_lat)) / cell_deg).ceil();
619 let n = if n.is_finite() {
620 floor_to_i64(n).max(0)
621 } else {
622 0
623 };
624 let mut out = BTreeSet::new();
625 for dx in -1..=1 {
626 for dy in -n..=n {
627 let cx = gx.saturating_add(dx);
628 let cy = gy.saturating_add(dy).rem_euclid(lon_cells);
629 out.insert(ValueKey::Str(format!("{cx}|{cy}")));
630 }
631 }
632 out
633}
634
635const SCAN_ALL_SENTINEL: ValueKey = ValueKey::Bool(true);
638
639fn ivf_cluster_key(cluster: usize) -> ValueKey {
642 ValueKey::Str(format!("\u{1}ivf:{cluster}"))
643}
644
645fn spec_is_scan_all_universe(spec: &CandidateSpec<'_>) -> bool {
648 match spec {
649 CandidateSpec::ScanAll { .. } => true,
650 CandidateSpec::Intersect(parts) => {
651 !parts.is_empty() && parts.iter().all(spec_is_scan_all_universe)
652 }
653 _ => false,
654 }
655}
656
657fn spec_is_bykey_external(spec: &CandidateSpec<'_>) -> bool {
660 match spec {
661 CandidateSpec::ByKey => true,
662 CandidateSpec::Intersect(parts) => {
663 !parts.is_empty() && parts.iter().all(spec_is_bykey_external)
664 }
665 _ => false,
666 }
667}
668
669impl SideIndex {
670 fn index_keys(spec: &CandidateSpec, get: &dyn Fn(&str) -> Option<Value>) -> BTreeSet<ValueKey> {
671 match spec {
672 CandidateSpec::ByKey => BTreeSet::new(),
673 CandidateSpec::Scalar { field } => get(field)
674 .as_ref()
675 .and_then(ValueKey::from_value)
676 .into_iter()
677 .collect(),
678 CandidateSpec::Tokens { field } => get(field)
679 .as_ref()
680 .and_then(list_tokens)
681 .unwrap_or_default(),
682 CandidateSpec::ScalarOrElements { field } => match get(field) {
683 Some(Value::List(items)) => items
684 .iter()
685 .take(MAX_KEYMATCH_LIST)
686 .filter(|v| matches!(v, Value::Str(_)))
687 .filter_map(ValueKey::from_value)
688 .collect(),
689 Some(v) => ValueKey::from_value(&v).into_iter().collect(),
690 None => BTreeSet::new(),
691 },
692 CandidateSpec::NumericBucket { field, tolerance } => get(field)
693 .as_ref()
694 .and_then(as_finite_f64)
695 .and_then(|v| numeric_index_key(v, *tolerance))
696 .into_iter()
697 .collect(),
698 CandidateSpec::GeoGrid { field, km } => get(field)
699 .as_ref()
700 .and_then(as_latlon)
701 .and_then(|(lat, lon)| geo_index_key(lat, lon, *km))
702 .into_iter()
703 .collect(),
704 CandidateSpec::ScanAll { field } => get(field)
705 .as_ref()
706 .and_then(as_numeric_list)
707 .map(|_| SCAN_ALL_SENTINEL)
708 .into_iter()
709 .collect(),
710 CandidateSpec::VectorClusters { .. } => BTreeSet::new(),
715 CandidateSpec::Hnsw { .. } => BTreeSet::new(),
717 CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) => {
721 let mut out = BTreeSet::new();
722 for s in specs {
723 out.extend(Self::index_keys(s, get));
724 }
725 out
726 }
727 }
728 }
729
730 fn probe_keys(spec: &CandidateSpec, get: &dyn Fn(&str) -> Option<Value>) -> BTreeSet<ValueKey> {
731 match spec {
732 CandidateSpec::ByKey
733 | CandidateSpec::Scalar { .. }
734 | CandidateSpec::Tokens { .. }
735 | CandidateSpec::ScalarOrElements { .. } => Self::index_keys(spec, get),
736 CandidateSpec::NumericBucket { field, tolerance } => get(field)
737 .as_ref()
738 .and_then(as_finite_f64)
739 .map(|v| numeric_probe_keys(v, *tolerance))
740 .unwrap_or_default(),
741 CandidateSpec::GeoGrid { field, km } => get(field)
742 .as_ref()
743 .and_then(as_latlon)
744 .map(|(lat, lon)| geo_probe_keys(lat, lon, *km))
745 .unwrap_or_default(),
746 CandidateSpec::ScanAll { field } => get(field)
747 .as_ref()
748 .and_then(as_numeric_list)
749 .map(|_| SCAN_ALL_SENTINEL)
750 .into_iter()
751 .collect(),
752 CandidateSpec::VectorClusters { .. } => BTreeSet::new(),
754 CandidateSpec::Hnsw { .. } => BTreeSet::new(),
756 CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) => {
759 let mut out = BTreeSet::new();
760 for s in specs {
761 out.extend(Self::probe_keys(s, get));
762 }
763 out
764 }
765 }
766 }
767
768 pub fn insert(&mut self, spec: &CandidateSpec, node: u32, get: &dyn Fn(&str) -> Option<Value>) {
769 self.insert_skipping(spec, node, &BTreeSet::new(), get);
770 }
771
772 pub fn insert_skipping(
779 &mut self,
780 spec: &CandidateSpec,
781 node: u32,
782 already: &BTreeSet<u32>,
783 get: &dyn Fn(&str) -> Option<Value>,
784 ) {
785 if let CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) = spec {
788 for s in specs {
789 self.insert_skipping(s, node, already, get);
790 }
791 return;
792 }
793 if let CandidateSpec::Hnsw { field, .. } = spec {
795 if let Some(xs) = get(field).as_ref().and_then(as_numeric_list) {
796 self.hnsw_tracked.insert(node);
797 if already.contains(&node) {
798 return; }
800 if let Some(h) = &mut self.hnsw {
801 h.insert(node, &xs);
802 }
803 }
804 return;
805 }
806 if let CandidateSpec::VectorClusters { field, .. } = spec {
808 if let Some(xs) = get(field).as_ref().and_then(as_numeric_list) {
809 self.ivf_raw.insert(node, xs.clone());
810 if !self.ivf_centroids.is_empty() {
811 if let Some(unit) = l2_normalize(&xs) {
813 let c = nearest_centroid(&self.ivf_centroids, &unit);
814 self.ivf_clusters.insert(node, c);
815 self.by_key
816 .entry(ivf_cluster_key(c))
817 .or_default()
818 .insert(node);
819 }
820 self.ivf_drift = self.ivf_drift.saturating_add(1);
821 }
822 }
823 return;
824 }
825
826 for k in Self::index_keys(spec, get) {
827 self.by_key.entry(k).or_default().insert(node);
828 }
829 if let CandidateSpec::ScanAll { field } = spec {
830 if let Some(xs) = get(field).as_ref().and_then(as_numeric_list) {
831 let mut n2 = 0.0f64;
832 for x in &xs {
833 n2 += x * x;
834 }
835 let norm = n2.sqrt();
836 self.vec_meta.insert(node, (xs.len() as u32, norm));
837 self.vec_checkpoints.insert(node, compute_ckpts(&xs));
838 self.vec_anchor.insert(node, xs[0]);
840 }
841 }
842 }
843
844 pub fn remove(&mut self, spec: &CandidateSpec, node: u32, get: &dyn Fn(&str) -> Option<Value>) {
845 if let CandidateSpec::Union(specs) | CandidateSpec::Intersect(specs) = spec {
847 for s in specs {
848 self.remove(s, node, get);
849 }
850 return;
851 }
852 if let CandidateSpec::Hnsw { field, .. } = spec {
857 if get(field).as_ref().and_then(as_numeric_list).is_some() {
858 self.hnsw_tracked.remove(&node);
859 if let Some(h) = &mut self.hnsw {
860 h.remove(node);
861 }
862 self.ivf_drift = self.ivf_drift.saturating_add(1);
863 }
864 return;
865 }
866 if let CandidateSpec::VectorClusters { .. } = spec {
871 if self.ivf_raw.remove(&node).is_some() {
872 self.ivf_drift = self.ivf_drift.saturating_add(1);
873 if let Some(c) = self.ivf_clusters.remove(&node) {
874 let key = ivf_cluster_key(c);
875 if let Some(s) = self.by_key.get_mut(&key) {
876 s.remove(&node);
877 if s.is_empty() {
878 self.by_key.remove(&key);
879 }
880 }
881 }
882 }
883 return;
884 }
885
886 for k in Self::index_keys(spec, get) {
887 if let Some(set) = self.by_key.get_mut(&k) {
888 set.remove(&node);
889 if set.is_empty() {
890 self.by_key.remove(&k);
891 }
892 }
893 }
894 if let CandidateSpec::ScanAll { field } = spec {
895 if get(field).as_ref().and_then(as_numeric_list).is_some() {
896 self.vec_meta.remove(&node);
897 self.vec_checkpoints.remove(&node);
898 self.vec_anchor.remove(&node);
899 }
900 }
901 }
902
903 pub fn vec_dim(&self, node: u32) -> Option<u32> {
905 self.vec_meta.get(&node).map(|(d, _)| *d)
906 }
907
908 pub fn vec_meta(&self, node: u32) -> Option<(u32, f64)> {
910 self.vec_meta.get(&node).copied()
911 }
912
913 pub fn vec_ckpts(&self, node: u32) -> Option<&[f64; 8]> {
915 self.vec_checkpoints.get(&node)
916 }
917
918 pub(crate) fn fresh_ckpts_for<'a>(
945 &'a self,
946 node: u32,
947 live: &[f64],
948 ) -> Option<(f64, &'a [f64; 8])> {
949 let &(dim, norm) = self.vec_meta.get(&node)?;
950 if dim != live.len() as u32 {
951 return None;
952 }
953 let live_norm = {
956 let mut n2 = 0.0f64;
957 for x in live {
958 n2 += x * x;
959 }
960 n2.sqrt()
961 };
962 if norm != live_norm {
963 return None; }
965 let live_anchor = live[0];
969 let &cached_anchor = self.vec_anchor.get(&node)?;
970 if live_anchor != cached_anchor {
971 return None;
972 }
973 let ckpts = self.vec_checkpoints.get(&node)?;
974 Some((norm, ckpts))
975 }
976
977 pub fn candidates(
978 &self,
979 spec: &CandidateSpec,
980 get: &dyn Fn(&str) -> Option<Value>,
981 ) -> BTreeSet<u32> {
982 if let CandidateSpec::Hnsw { field, k } = spec {
984 return self.hnsw_candidates(field, *k, get);
985 }
986 if let CandidateSpec::VectorClusters { field, .. } = spec {
988 return self.ivf_candidates(field, get);
989 }
990 if let CandidateSpec::Union(specs) = spec {
992 return specs.iter().flat_map(|s| self.candidates(s, get)).collect();
993 }
994 if let CandidateSpec::Intersect(specs) = spec {
995 return self.intersect_candidates(specs, get);
996 }
997
998 let mut out = BTreeSet::new();
999 for k in Self::probe_keys(spec, get) {
1000 if let Some(set) = self.by_key.get(&k) {
1001 out.extend(set.iter().copied());
1002 }
1003 }
1004 if vector_dim_reject_enabled() {
1006 if let CandidateSpec::ScanAll { field } = spec {
1007 if let Some((dim, _)) = get(field).as_ref().and_then(vec_dim_norm) {
1008 out.retain(|id| self.vec_meta.get(id).is_none_or(|(d, _)| *d == dim));
1009 }
1010 }
1011 }
1012 out
1013 }
1014
1015 fn intersect_candidates(
1019 &self,
1020 specs: &[CandidateSpec<'_>],
1021 get: &dyn Fn(&str) -> Option<Value>,
1022 ) -> BTreeSet<u32> {
1023 let mut restrictive = Vec::new();
1024 let mut scan_alls = Vec::new();
1025 for s in specs {
1026 if spec_is_scan_all_universe(s) {
1027 scan_alls.push(s);
1028 } else if spec_is_bykey_external(s) {
1029 continue;
1030 } else {
1031 restrictive.push(s);
1032 }
1033 }
1034 let to_intersect: &[&CandidateSpec<'_>] = if !restrictive.is_empty() {
1035 &restrictive
1036 } else if !scan_alls.is_empty() {
1037 &scan_alls
1038 } else {
1039 return BTreeSet::new();
1040 };
1041 let mut iter = to_intersect.iter();
1042 let Some(first) = iter.next() else {
1043 return BTreeSet::new();
1044 };
1045 let mut acc = self.candidates(first, get);
1046 if acc.is_empty() {
1047 return acc;
1048 }
1049 for s in iter {
1050 let other = self.candidates(s, get);
1051 if other.is_empty() {
1052 return BTreeSet::new();
1053 }
1054 acc = acc.intersection(&other).copied().collect();
1055 if acc.is_empty() {
1056 return acc;
1057 }
1058 }
1059 acc
1060 }
1061
1062 fn ivf_candidates(&self, field: &str, get: &dyn Fn(&str) -> Option<Value>) -> BTreeSet<u32> {
1065 let Some(xs) = get(field).as_ref().and_then(as_numeric_list) else {
1066 return BTreeSet::new();
1067 };
1068 if self.ivf_centroids.is_empty() {
1069 return self.ivf_raw.keys().copied().collect();
1073 }
1074 if self.ivf_raw.len() <= self.ivf_centroids.len() {
1079 return self.ivf_raw.keys().copied().collect();
1080 }
1081 let k = self.ivf_centroids.len();
1082 let p = probe_count(k);
1083
1084 let Some(xs) = l2_normalize(&xs) else {
1087 return BTreeSet::new();
1088 };
1089
1090 let mut dists: Vec<(usize, f64)> = self
1092 .ivf_centroids
1093 .iter()
1094 .enumerate()
1095 .map(|(i, c)| (i, l2_sq(&xs, c)))
1096 .collect();
1097 dists.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
1098
1099 let mut out = BTreeSet::new();
1100 for (ci, _) in dists.iter().take(p) {
1101 let key = ivf_cluster_key(*ci);
1102 if let Some(nodes) = self.by_key.get(&key) {
1103 out.extend(nodes.iter().copied());
1104 }
1105 }
1106 out
1107 }
1108
1109 pub fn fit_ivf_clusters(&mut self, rule_name: &str) {
1120 if self.ivf_raw.is_empty() {
1121 self.ivf_centroids.clear();
1122 self.ivf_clusters.clear();
1123 self.ivf_drift = 0;
1124 return;
1125 }
1126
1127 for c in self.ivf_clusters.values() {
1129 self.by_key.remove(&ivf_cluster_key(*c));
1130 }
1131 self.ivf_clusters.clear();
1132
1133 let vecs: Vec<(u32, Vec<f64>)> = self
1135 .ivf_raw
1136 .iter()
1137 .map(|(&id, xs)| (id, xs.clone()))
1138 .collect();
1139
1140 let n = vecs.len();
1141 let k = cluster_k(n);
1142 let seed = fnv1a_u64(rule_name.as_bytes());
1143
1144 self.ivf_centroids = kmeans_fit(&vecs, k, seed);
1145
1146 for (node, xs) in &vecs {
1148 let Some(unit) = l2_normalize(xs) else {
1149 continue;
1150 };
1151 let c = nearest_centroid(&self.ivf_centroids, &unit);
1152 self.ivf_clusters.insert(*node, c);
1153 self.by_key
1154 .entry(ivf_cluster_key(c))
1155 .or_default()
1156 .insert(*node);
1157 }
1158 self.ivf_drift = 0;
1159 }
1160
1161 pub fn ivf_k(&self) -> usize {
1163 self.ivf_centroids.len()
1164 }
1165
1166 pub fn ivf_cluster_of(&self, node: u32) -> Option<usize> {
1168 self.ivf_clusters.get(&node).copied()
1169 }
1170
1171 pub fn export_ivf_state(&self) -> (Vec<Vec<f64>>, BTreeMap<u32, usize>, u64) {
1176 (
1177 self.ivf_centroids.clone(),
1178 self.ivf_clusters.clone(),
1179 self.ivf_drift,
1180 )
1181 }
1182
1183 pub fn load_ivf_state(
1195 &mut self,
1196 centroids: Vec<Vec<f64>>,
1197 clusters: BTreeMap<u32, usize>,
1198 drift: u64,
1199 ) {
1200 for c in self.ivf_clusters.values() {
1203 self.by_key.remove(&ivf_cluster_key(*c));
1204 }
1205 self.ivf_clusters.clear();
1206
1207 self.ivf_centroids = centroids;
1208 self.ivf_drift = drift;
1209
1210 for (&node, &c) in &clusters {
1212 if !self.ivf_raw.contains_key(&node) {
1213 continue;
1215 }
1216 self.ivf_clusters.insert(node, c);
1217 self.by_key
1218 .entry(ivf_cluster_key(c))
1219 .or_default()
1220 .insert(node);
1221 }
1222 }
1223
1224 pub fn init_hnsw(&mut self, rule_name: &str) {
1233 let seed = fnv1a_u64(rule_name.as_bytes());
1234 self.hnsw = Some(HnswIndex::new(seed));
1235 }
1236
1237 fn hnsw_candidates(
1242 &self,
1243 field: &str,
1244 k: usize,
1245 get: &dyn Fn(&str) -> Option<Value>,
1246 ) -> BTreeSet<u32> {
1247 let Some(xs) = get(field).as_ref().and_then(as_numeric_list) else {
1248 return BTreeSet::new();
1249 };
1250 if let Some(h) = &self.hnsw {
1251 if !h.is_empty() {
1252 return h.search(&xs, k).into_iter().map(|(id, _)| id).collect();
1253 }
1254 }
1255 self.hnsw_tracked.clone()
1257 }
1258
1259 pub fn export_hnsw_blob(&self) -> Vec<u8> {
1263 self.hnsw
1264 .as_ref()
1265 .and_then(|h| bincode::serialize(h).ok())
1266 .unwrap_or_default()
1267 }
1268
1269 pub fn load_hnsw_blob(&mut self, blob: &[u8]) {
1275 if blob.is_empty() {
1276 return;
1277 }
1278 if let Ok(h) = bincode::deserialize::<HnswIndex>(blob) {
1279 self.adopt_hnsw(h);
1280 }
1281 }
1282
1283 pub fn init_or_adopt_hnsw(&mut self, rule_name: &str, blob: &[u8]) -> (BTreeSet<u32>, bool) {
1293 self.hnsw = None;
1294 if !blob.is_empty() {
1295 match bincode::deserialize::<HnswIndex>(blob) {
1296 Ok(h) => self.adopt_hnsw(h),
1297 Err(e) => eprintln!(
1298 "[mushroomdb] rule {rule_name:?}: a persisted HNSW index failed to load \
1299 ({e}); rebuilding it from the node scan"
1300 ),
1301 }
1302 }
1303 match &self.hnsw {
1304 Some(h) => (h.node_ids(), true),
1305 None => {
1306 self.init_hnsw(rule_name);
1307 (BTreeSet::new(), false)
1308 }
1309 }
1310 }
1311
1312 pub fn adopt_hnsw(&mut self, h: HnswIndex) {
1318 self.hnsw_tracked = h.node_ids();
1319 self.hnsw = Some(h);
1320 }
1321
1322 pub fn has_hnsw(&self) -> bool {
1324 self.hnsw.as_ref().is_some_and(|h| !h.is_empty())
1325 }
1326
1327 pub fn hnsw_ref(&self) -> Option<&HnswIndex> {
1329 self.hnsw.as_ref()
1330 }
1331}
1332
1333#[cfg(test)]
1334mod tests {
1335 use super::*;
1336 use crate::def::Predicate;
1337 use core_storage::Value;
1338 use std::collections::{BTreeMap, HashMap};
1339
1340 fn getter(map: &HashMap<String, Value>) -> impl Fn(&str) -> Option<Value> + '_ {
1341 move |f: &str| map.get(f).cloned()
1342 }
1343
1344 #[test]
1345 fn kmeans_centroids_are_unit_norm() {
1346 let vecs = vec![(0, vec![3.0, 0.0, 0.0]), (1, vec![0.0, 4.0, 0.0])];
1347 let cents = kmeans_fit(&vecs, 2, 1);
1348 for c in cents {
1349 let n = c.iter().map(|x| x * x).sum::<f64>().sqrt();
1350 assert!((n - 1.0).abs() < 1e-9, "{n}");
1351 }
1352 }
1353
1354 #[test]
1361 fn scaled_vector_joins_same_ivf_cluster_as_unit() {
1362 let spec = CandidateSpec::VectorClusters {
1364 field: "emb",
1365 min: 0.5,
1366 };
1367 let mut idx = SideIndex::default();
1368 idx.load_ivf_state(
1369 vec![vec![1.0, 0.0, 0.0], vec![2.5, 0.1, 0.0]],
1370 BTreeMap::new(),
1371 0,
1372 );
1373 idx.insert(&spec, 1, &getter(&emb(&[1.0, 0.0, 0.0])));
1374 idx.insert(&spec, 2, &getter(&emb(&[3.0, 0.0, 0.0])));
1375 assert_eq!(
1376 idx.ivf_cluster_of(1),
1377 idx.ivf_cluster_of(2),
1378 "scale-equivalent vectors must share an IVF cluster; got {:?} vs {:?}",
1379 idx.ivf_cluster_of(1),
1380 idx.ivf_cluster_of(2)
1381 );
1382 assert_eq!(idx.ivf_cluster_of(1), Some(0));
1383 }
1384
1385 #[test]
1386 fn scalar_index_buckets_by_value() {
1387 let pred = Predicate::FieldEqual {
1388 field: "ind".into(),
1389 };
1390 let spec = candidate_spec(&pred);
1391 let mut idx = SideIndex::default();
1392 let a: HashMap<_, _> = [("ind".to_string(), Value::Str("arch".into()))].into();
1393 let b: HashMap<_, _> = [("ind".to_string(), Value::Str("law".into()))].into();
1394 idx.insert(&spec, 1, &getter(&a));
1395 idx.insert(&spec, 2, &getter(&b));
1396 idx.insert(&spec, 3, &getter(&a));
1397 let c = idx.candidates(&spec, &getter(&a));
1398 assert_eq!(c.into_iter().collect::<Vec<_>>(), vec![1, 3]);
1399 idx.remove(&spec, 3, &getter(&a));
1400 assert_eq!(idx.candidates(&spec, &getter(&a)).len(), 1);
1401 let empty: HashMap<String, Value> = HashMap::new();
1403 idx.insert(&spec, 9, &getter(&empty));
1404 assert!(idx.candidates(&spec, &getter(&empty)).is_empty());
1405 }
1406
1407 #[test]
1408 fn token_index_unions_buckets() {
1409 let mk =
1410 |items: &[&str]| Value::List(items.iter().map(|s| Value::Str((*s).into())).collect());
1411 let pred = Predicate::Overlap {
1412 field: "tags".into(),
1413 min: 0.5,
1414 };
1415 let spec = candidate_spec(&pred);
1416 let mut idx = SideIndex::default();
1417 let a: HashMap<_, _> = [("tags".to_string(), mk(&["x", "y"]))].into();
1418 let b: HashMap<_, _> = [("tags".to_string(), mk(&["y", "z"]))].into();
1419 let c: HashMap<_, _> = [("tags".to_string(), mk(&["q"]))].into();
1420 idx.insert(&spec, 1, &getter(&a));
1421 idx.insert(&spec, 2, &getter(&b));
1422 idx.insert(&spec, 3, &getter(&c));
1423 let probe: HashMap<_, _> = [("tags".to_string(), mk(&["y"]))].into();
1424 assert_eq!(
1425 idx.candidates(&spec, &getter(&probe))
1426 .into_iter()
1427 .collect::<Vec<_>>(),
1428 vec![1, 2]
1429 );
1430 idx.remove(&spec, 2, &getter(&b));
1431 assert_eq!(
1432 idx.candidates(&spec, &getter(&probe))
1433 .into_iter()
1434 .collect::<Vec<_>>(),
1435 vec![1]
1436 );
1437 }
1438
1439 #[test]
1440 fn all_intersects_parts_and_bykey_indexes_nothing() {
1441 let all = Predicate::All(vec![
1442 Predicate::FieldEqual {
1443 field: "ind".into(),
1444 },
1445 Predicate::Overlap {
1446 field: "tags".into(),
1447 min: 0.5,
1448 },
1449 ]);
1450 match candidate_spec(&all) {
1451 CandidateSpec::Intersect(v) => assert_eq!(v.len(), 2),
1452 other => panic!("{other:?}"),
1453 }
1454 let km = Predicate::KeyMatch { field: "fk".into() };
1455 assert!(matches!(candidate_spec(&km), CandidateSpec::ByKey));
1456 let mut idx = SideIndex::default();
1457 let a: HashMap<_, _> = [("fk".to_string(), Value::Str("c1".into()))].into();
1458 idx.insert(&candidate_spec(&km), 1, &getter(&a));
1459 assert!(idx.candidates(&candidate_spec(&km), &getter(&a)).is_empty());
1460 }
1461
1462 fn year(v: Value) -> HashMap<String, Value> {
1463 [("year".to_string(), v)].into()
1464 }
1465
1466 fn loc(lat: f64, lon: f64) -> HashMap<String, Value> {
1467 [(
1468 "loc".to_string(),
1469 Value::List(vec![Value::Float(lat), Value::Float(lon)]),
1470 )]
1471 .into()
1472 }
1473
1474 fn emb(vals: &[f64]) -> HashMap<String, Value> {
1475 [(
1476 "emb".to_string(),
1477 Value::List(vals.iter().copied().map(Value::Float).collect()),
1478 )]
1479 .into()
1480 }
1481
1482 fn bucket_int(spec: &CandidateSpec, map: &HashMap<String, Value>) -> Option<i64> {
1483 match SideIndex::index_keys(spec, &getter(map)).into_iter().next() {
1484 Some(ValueKey::Int(b)) => Some(b),
1485 _ => None,
1486 }
1487 }
1488
1489 #[test]
1490 fn numeric_bucket_adjacency_and_far_value() {
1491 let pred = Predicate::NumericWithin {
1492 field: "year".into(),
1493 tolerance: 2.0,
1494 };
1495 let spec = candidate_spec(&pred);
1496 assert!(matches!(
1497 spec,
1498 CandidateSpec::NumericBucket {
1499 field: "year",
1500 tolerance
1501 } if tolerance == 2.0
1502 ));
1503
1504 let v10 = year(Value::Float(10.0));
1505 let v119 = year(Value::Float(11.9));
1506 let v99 = year(Value::Float(9.9));
1507 let v141 = year(Value::Float(14.1));
1508
1509 let b10 = bucket_int(&spec, &v10).unwrap();
1510 let b119 = bucket_int(&spec, &v119).unwrap();
1511 let b99 = bucket_int(&spec, &v99).unwrap();
1512 assert!((b10 - b119).abs() <= 1);
1514 assert!((b10 - b99).abs() <= 1);
1515
1516 let mut idx = SideIndex::default();
1517 idx.insert(&spec, 1, &getter(&v10));
1518 idx.insert(&spec, 2, &getter(&v119));
1519 idx.insert(&spec, 3, &getter(&v141));
1520 idx.insert(&spec, 4, &getter(&v99));
1521 let hits = idx.candidates(&spec, &getter(&v10));
1522 assert_eq!(hits.into_iter().collect::<Vec<_>>(), vec![1, 2, 4]);
1523 }
1524
1525 #[test]
1526 fn numeric_tol_zero_int_float_collide() {
1527 let pred = Predicate::NumericWithin {
1528 field: "year".into(),
1529 tolerance: 0.0,
1530 };
1531 let spec = candidate_spec(&pred);
1532 let mut idx = SideIndex::default();
1533 idx.insert(&spec, 1, &getter(&year(Value::Int(2))));
1534 assert_eq!(
1535 idx.candidates(&spec, &getter(&year(Value::Float(2.0))))
1536 .into_iter()
1537 .collect::<Vec<_>>(),
1538 vec![1]
1539 );
1540 assert!(idx
1541 .candidates(&spec, &getter(&year(Value::Float(2.1))))
1542 .is_empty());
1543 }
1544
1545 #[test]
1546 fn numeric_tol_zero_signed_zero_collides() {
1547 let pred = Predicate::NumericWithin {
1548 field: "year".into(),
1549 tolerance: 0.0,
1550 };
1551 let spec = candidate_spec(&pred);
1552 let neg = year(Value::Float(-0.0));
1553 let pos = year(Value::Float(0.0));
1554 let mut idx = SideIndex::default();
1555 idx.insert(&spec, 1, &getter(&neg));
1556 assert_eq!(
1557 idx.candidates(&spec, &getter(&pos))
1558 .into_iter()
1559 .collect::<Vec<_>>(),
1560 vec![1]
1561 );
1562 let mut idx2 = SideIndex::default();
1563 idx2.insert(&spec, 2, &getter(&pos));
1564 assert_eq!(
1565 idx2.candidates(&spec, &getter(&neg))
1566 .into_iter()
1567 .collect::<Vec<_>>(),
1568 vec![2]
1569 );
1570 }
1571
1572 #[test]
1573 fn geo_grid_same_cell_cross_cell_and_far_city() {
1574 let pred = Predicate::GeoRadius {
1575 field: "loc".into(),
1576 km: 400.0,
1577 };
1578 let spec = candidate_spec(&pred);
1579 assert!(matches!(
1580 spec,
1581 CandidateSpec::GeoGrid {
1582 field: "loc",
1583 km
1584 } if km == 400.0
1585 ));
1586
1587 let paris = loc(48.8566, 2.3522);
1588 let london = loc(51.5074, -0.1278);
1589 let nearby = loc(48.9, 2.4); let ny = loc(40.7128, -74.0060);
1591
1592 let mut idx = SideIndex::default();
1593 idx.insert(&spec, 1, &getter(&paris));
1594 idx.insert(&spec, 2, &getter(&london));
1595 idx.insert(&spec, 3, &getter(&nearby));
1596 idx.insert(&spec, 4, &getter(&ny));
1597
1598 let from_paris = idx.candidates(&spec, &getter(&paris));
1599 assert!(from_paris.contains(&1), "same-cell self");
1600 assert!(from_paris.contains(&3), "same-cell neighbor");
1601 assert!(from_paris.contains(&2), "cross-cell Paris↔London ~343.5 km");
1602 assert!(!from_paris.contains(&4), "New York not in 400 km probe");
1603 }
1604
1605 #[test]
1606 fn geo_grid_high_latitude_probe_is_superset() {
1607 let pred = Predicate::GeoRadius {
1608 field: "loc".into(),
1609 km: 340.0,
1610 };
1611 let spec = candidate_spec(&pred);
1612 let reyk = loc(64.1466, -21.9426);
1613 let lat = 64.0_f64;
1614 let dlon = 300.0 / (111.0 * lat.to_radians().cos());
1615 let east = loc(lat, -21.9426 + dlon);
1616
1617 let mut idx = SideIndex::default();
1618 idx.insert(&spec, 1, &getter(&reyk));
1619 idx.insert(&spec, 2, &getter(&east));
1620 let hits = idx.candidates(&spec, &getter(&reyk));
1621 assert!(
1622 hits.contains(&2),
1623 "300 km east of Reykjavik must stay in the high-lat probe"
1624 );
1625 }
1626
1627 #[test]
1628 fn geo_grid_antimeridian_wrap_and_evaluate_agree() {
1629 let pred = Predicate::GeoRadius {
1630 field: "loc".into(),
1631 km: 400.0,
1632 };
1633 let spec = candidate_spec(&pred);
1634 let east = loc(70.0, 179.9);
1635 let west = loc(70.0, -179.9);
1636
1637 let mut idx = SideIndex::default();
1638 idx.insert(&spec, 1, &getter(&east));
1639 assert!(
1640 idx.candidates(&spec, &getter(&west)).contains(&1),
1641 "±180 pair at lat 70 must land in the wrapped probe"
1642 );
1643
1644 let sp = |f: &str| east.get(f).cloned();
1645 let dp = |f: &str| west.get(f).cloned();
1646 let score = crate::def::evaluate(
1647 &pred,
1648 &crate::def::NodeView {
1649 key: "e",
1650 props: &sp,
1651 },
1652 &crate::def::NodeView {
1653 key: "w",
1654 props: &dp,
1655 },
1656 );
1657 assert!(
1658 score.is_some(),
1659 "haversine must match across the antimeridian"
1660 );
1661
1662 let paris = loc(48.8566, 2.3522);
1664 let ny = loc(40.7128, -74.0060);
1665 let mut idx2 = SideIndex::default();
1666 idx2.insert(&spec, 4, &getter(&ny));
1667 assert!(
1668 !idx2.candidates(&spec, &getter(&paris)).contains(&4),
1669 "New York still not in the Paris probe after wrap"
1670 );
1671 }
1672
1673 #[test]
1674 fn scan_all_returns_vector_nodes_skips_malformed() {
1675 let pred = Predicate::VectorSimilar {
1676 field: "emb".into(),
1677 min: 0.5,
1678 };
1679 let spec = candidate_spec(&pred);
1680 assert!(matches!(spec, CandidateSpec::ScanAll { field: "emb" }));
1681
1682 let mut idx = SideIndex::default();
1683 idx.insert(&spec, 1, &getter(&emb(&[1.0, 0.0])));
1684 idx.insert(&spec, 2, &getter(&emb(&[0.0, 1.0])));
1685 idx.insert(&spec, 3, &getter(&emb(&[1.0, 2.0, 3.0])));
1686 let empty: HashMap<_, _> = [("emb".to_string(), Value::List(vec![]))].into();
1687 let text: HashMap<_, _> =
1688 [("emb".to_string(), Value::List(vec![Value::Str("x".into())]))].into();
1689 let missing: HashMap<String, Value> = HashMap::new();
1690 idx.insert(&spec, 4, &getter(&empty));
1691 idx.insert(&spec, 5, &getter(&text));
1692 idx.insert(&spec, 6, &getter(&missing));
1693
1694 let hits = idx.candidates(&spec, &getter(&emb(&[1.0, 0.0])));
1695 assert_eq!(
1696 hits.into_iter().collect::<Vec<_>>(),
1697 vec![1, 2],
1698 "dim-2 probe must drop the dim-3 member"
1699 );
1700 assert_eq!(
1701 idx.candidates(&spec, &getter(&emb(&[1.0, 2.0, 3.0])))
1702 .into_iter()
1703 .collect::<Vec<_>>(),
1704 vec![3]
1705 );
1706 with_vector_dim_reject(false, || {
1707 assert_eq!(
1708 idx.candidates(&spec, &getter(&emb(&[1.0, 0.0])))
1709 .into_iter()
1710 .collect::<Vec<_>>(),
1711 vec![1, 2, 3],
1712 "unfiltered ScanAll still returns every vector node"
1713 );
1714 });
1715 assert_eq!(idx.vec_dim(1), Some(2));
1716 assert_eq!(idx.vec_dim(3), Some(3));
1717 assert!(idx.vec_meta(1).is_some());
1718 assert!(idx.vec_dim(4).is_none());
1719 assert!(idx.candidates(&spec, &getter(&empty)).is_empty());
1720 assert!(idx.candidates(&spec, &getter(&text)).is_empty());
1721 assert!(idx.candidates(&spec, &getter(&missing)).is_empty());
1722 idx.remove(&spec, 1, &getter(&emb(&[1.0, 0.0])));
1723 assert!(idx.vec_dim(1).is_none());
1724 }
1725
1726 #[test]
1727 fn legacy_specs_probe_keys_equal_index_keys() {
1728 let a: HashMap<_, _> = [
1729 ("ind".to_string(), Value::Str("arch".into())),
1730 (
1731 "tags".to_string(),
1732 Value::List(vec![Value::Str("x".into()), Value::Str("y".into())]),
1733 ),
1734 ("fk".to_string(), Value::Str("c1".into())),
1735 ]
1736 .into();
1737 let get = getter(&a);
1738 for pred in [
1739 Predicate::KeyMatch { field: "fk".into() },
1740 Predicate::FieldEqual {
1741 field: "ind".into(),
1742 },
1743 Predicate::Overlap {
1744 field: "tags".into(),
1745 min: 0.5,
1746 },
1747 ] {
1748 let spec = candidate_spec(&pred);
1749 assert_eq!(
1750 SideIndex::index_keys(&spec, &get),
1751 SideIndex::probe_keys(&spec, &get)
1752 );
1753 }
1754 }
1755
1756 #[test]
1757 fn all_vector_then_field_equal_does_not_scan_all() {
1758 let p = Predicate::All(vec![
1759 Predicate::VectorSimilar {
1760 field: "e".into(),
1761 min: 0.8,
1762 },
1763 Predicate::FieldEqual {
1764 field: "industry".into(),
1765 },
1766 ]);
1767 match candidate_spec(&p) {
1768 CandidateSpec::Intersect(v) => assert_eq!(v.len(), 2),
1769 other => panic!("{other:?}"),
1770 }
1771
1772 let spec = candidate_spec(&p);
1773 let mut idx = SideIndex::default();
1774 let mk = |industry: &str, e: &[f64]| {
1775 [
1776 ("industry".to_string(), Value::Str(industry.into())),
1777 (
1778 "e".to_string(),
1779 Value::List(e.iter().copied().map(Value::Float).collect()),
1780 ),
1781 ]
1782 .into()
1783 };
1784 let same: HashMap<_, _> = mk("tech", &[1.0, 0.0]);
1785 let other_ind: HashMap<_, _> = mk("law", &[1.0, 0.0]);
1786 let no_vec: HashMap<_, _> = [("industry".to_string(), Value::Str("tech".into()))].into();
1787 idx.insert(&spec, 1, &getter(&same));
1788 idx.insert(&spec, 2, &getter(&other_ind));
1789 idx.insert(&spec, 3, &getter(&no_vec));
1790
1791 let hits = idx.candidates(&spec, &getter(&same));
1792 assert!(hits.contains(&1), "matching industry must stay a candidate");
1793 assert!(
1794 !hits.contains(&2),
1795 "different industry must not be scanned in via VectorSimilar"
1796 );
1797 assert!(
1798 hits.contains(&3),
1799 "ScanAll is universe: extra Scalar-only candidates are allowed"
1800 );
1801
1802 let empty_ind: HashMap<_, _> = mk("finance", &[1.0, 0.0]);
1803 assert!(
1804 idx.candidates(&spec, &getter(&empty_ind)).is_empty(),
1805 "empty Scalar child → empty intersect"
1806 );
1807 }
1808
1809 #[test]
1810 fn all_approx_vector_then_field_equal_is_intersect() {
1811 let p = Predicate::All(vec![
1812 Predicate::VectorSimilar {
1813 field: "e".into(),
1814 min: 0.8,
1815 },
1816 Predicate::FieldEqual {
1817 field: "industry".into(),
1818 },
1819 ]);
1820 match candidate_spec_approx(&p) {
1821 CandidateSpec::Intersect(v) => assert_eq!(v.len(), 2),
1822 other => panic!("{other:?}"),
1823 }
1824
1825 let spec = candidate_spec_approx(&p);
1826 let mut idx = SideIndex::default();
1827 idx.init_hnsw("test-rule");
1829 let mk = |industry: &str, e: &[f64]| {
1830 [
1831 ("industry".to_string(), Value::Str(industry.into())),
1832 (
1833 "e".to_string(),
1834 Value::List(e.iter().copied().map(Value::Float).collect()),
1835 ),
1836 ]
1837 .into()
1838 };
1839 let same: HashMap<_, _> = mk("tech", &[1.0, 0.0]);
1840 let other_ind: HashMap<_, _> = mk("law", &[1.0, 0.0]);
1841 idx.insert(&spec, 1, &getter(&same));
1842 idx.insert(&spec, 2, &getter(&other_ind));
1843 let hits = idx.candidates(&spec, &getter(&same));
1844 assert!(hits.contains(&1), "matching industry must stay a candidate");
1845 assert!(
1846 !hits.contains(&2),
1847 "FieldEqual must be probed on the approximate All path"
1848 );
1849 }
1850
1851 #[test]
1852 fn all_of_scan_all_stays_scan_all() {
1853 let p = Predicate::All(vec![
1854 Predicate::VectorSimilar {
1855 field: "emb".into(),
1856 min: 0.5,
1857 },
1858 Predicate::VectorSimilar {
1859 field: "emb".into(),
1860 min: 0.9,
1861 },
1862 ]);
1863 match candidate_spec(&p) {
1864 CandidateSpec::Intersect(v) => assert_eq!(v.len(), 2),
1865 other => panic!("{other:?}"),
1866 }
1867 let spec = candidate_spec(&p);
1868 let mut idx = SideIndex::default();
1869 idx.insert(&spec, 1, &getter(&emb(&[1.0, 0.0])));
1870 idx.insert(&spec, 2, &getter(&emb(&[0.0, 1.0])));
1871 let hits = idx.candidates(&spec, &getter(&emb(&[1.0, 0.0])));
1872 assert_eq!(hits.into_iter().collect::<Vec<_>>(), vec![1, 2]);
1873 }
1874
1875 #[test]
1876 fn any_stays_union() {
1877 let p = Predicate::Any(vec![
1878 Predicate::FieldEqual {
1879 field: "industry".into(),
1880 },
1881 Predicate::Overlap {
1882 field: "tags".into(),
1883 min: 0.5,
1884 },
1885 ]);
1886 match candidate_spec(&p) {
1887 CandidateSpec::Union(v) => assert_eq!(v.len(), 2),
1888 other => panic!("{other:?}"),
1889 }
1890 }
1891
1892 #[test]
1895 fn checkpoint_populated_and_consistent_with_norm() {
1896 let pred = Predicate::VectorSimilar {
1897 field: "emb".into(),
1898 min: 0.8,
1899 };
1900 let spec = candidate_spec(&pred);
1901 let xs = [3.0f64, 4.0]; let mut idx = SideIndex::default();
1903 idx.insert(&spec, 1, &getter(&emb(&xs)));
1904
1905 let ckpts = idx
1906 .vec_ckpts(1)
1907 .expect("checkpoints must exist after insert");
1908 let (_, norm) = idx.vec_meta(1).unwrap();
1909 assert!(
1910 (ckpts[0] - norm).abs() < 1e-12,
1911 "ckpts[0] must equal the full L2 norm; got {} vs {}",
1912 ckpts[0],
1913 norm
1914 );
1915 assert!(
1916 (norm - 5.0).abs() < 1e-12,
1917 "norm of [3,4] must be 5.0, got {norm}"
1918 );
1919
1920 idx.remove(&spec, 1, &getter(&emb(&xs)));
1922 assert!(
1923 idx.vec_ckpts(1).is_none(),
1924 "checkpoints must be removed after remove()"
1925 );
1926 }
1927
1928 #[test]
1931 fn fresh_ckpts_for_freshness_gate() {
1932 let pred = Predicate::VectorSimilar {
1933 field: "emb".into(),
1934 min: 0.8,
1935 };
1936 let spec = candidate_spec(&pred);
1937 let xs = [1.0f64, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0];
1938 let mut idx = SideIndex::default();
1939 idx.insert(&spec, 7, &getter(&emb(&xs)));
1940
1941 let result = idx.fresh_ckpts_for(7, &xs);
1943 assert!(
1944 result.is_some(),
1945 "fresh_ckpts_for must succeed with matching live vector"
1946 );
1947 let (norm, ckpts) = result.unwrap();
1948 assert!((norm - 1.0).abs() < 1e-12);
1949 assert!((ckpts[0] - 1.0).abs() < 1e-12);
1950
1951 let wrong = [2.0f64, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]; assert!(
1954 idx.fresh_ckpts_for(7, &wrong).is_none(),
1955 "freshness gate must reject mismatched norm"
1956 );
1957
1958 let short = [1.0f64, 0.0];
1960 assert!(
1961 idx.fresh_ckpts_for(7, &short).is_none(),
1962 "freshness gate must reject mismatched dim"
1963 );
1964
1965 assert!(idx.fresh_ckpts_for(99, &xs).is_none());
1967 }
1968
1969 #[test]
1972 fn checkpoint_suffix_norms_non_increasing() {
1973 let pred = Predicate::VectorSimilar {
1974 field: "emb".into(),
1975 min: 0.5,
1976 };
1977 let spec = candidate_spec(&pred);
1978 let xs: Vec<f64> = (1..=16).map(|i| i as f64).collect();
1979 let mut idx = SideIndex::default();
1980 idx.insert(&spec, 42, &getter(&emb(&xs)));
1981
1982 let ckpts = *idx.vec_ckpts(42).unwrap();
1983 for c in 0..7 {
1984 assert!(
1985 ckpts[c] >= ckpts[c + 1] - 1e-12,
1986 "suffix norm must be non-increasing: ckpts[{c}]={} < ckpts[{}]={}",
1987 ckpts[c],
1988 c + 1,
1989 ckpts[c + 1]
1990 );
1991 }
1992 let expected_last = (15.0f64 * 15.0 + 16.0 * 16.0).sqrt();
1994 assert!(
1995 (ckpts[7] - expected_last).abs() < 1e-9,
1996 "ckpts[7] should be norm of last segment; got {} vs {}",
1997 ckpts[7],
1998 expected_last
1999 );
2000 }
2001 fn hnsw_side() -> (SideIndex, CandidateSpec<'static>) {
2007 let spec = CandidateSpec::Hnsw { field: "emb", k: 8 };
2008 let mut side = SideIndex::default();
2009 side.init_hnsw("sim");
2010 for (id, xs) in [
2011 (1u32, vec![1.0, 0.0]),
2012 (2, vec![0.0, 1.0]),
2013 (3, vec![0.7, 0.7]),
2014 ] {
2015 side.insert(&spec, id, &getter(&emb(&xs)));
2016 }
2017 (side, spec)
2018 }
2019
2020 #[test]
2023 fn init_or_adopt_hnsw_adopts_a_usable_blob() {
2024 let (side, spec) = hnsw_side();
2025 let blob = side.export_hnsw_blob();
2026
2027 let mut fresh = SideIndex::default();
2028 let (ids, adopted) = fresh.init_or_adopt_hnsw("sim", &blob);
2029 assert!(adopted, "a usable blob must be adopted, not rebuilt");
2030 assert_eq!(ids, BTreeSet::from([1, 2, 3]));
2031 assert!(fresh.has_hnsw());
2032 assert_eq!(
2033 fresh.candidates(&spec, &getter(&emb(&[1.0, 0.0]))),
2034 side.candidates(&spec, &getter(&emb(&[1.0, 0.0]))),
2035 "the adopted graph must answer as the original did"
2036 );
2037 }
2038
2039 #[test]
2043 fn an_unreadable_blob_leaves_the_graph_empty() {
2044 let (side, spec) = hnsw_side();
2045 let mut blob = side.export_hnsw_blob();
2046 blob.truncate(blob.len() / 2);
2047
2048 let mut fresh = SideIndex::default();
2049 let (ids, adopted) = fresh.init_or_adopt_hnsw("sim", &blob);
2050 assert!(!adopted, "an unreadable blob must not count as adopted");
2051 assert!(ids.is_empty(), "nothing may be skipped by the scan");
2052 assert!(!fresh.has_hnsw(), "the graph must be empty");
2053
2054 for (id, xs) in [
2056 (1u32, vec![1.0, 0.0]),
2057 (2, vec![0.0, 1.0]),
2058 (3, vec![0.7, 0.7]),
2059 ] {
2060 fresh.insert_skipping(&spec, id, &ids, &getter(&emb(&xs)));
2061 }
2062 assert!(fresh.has_hnsw());
2063 assert_eq!(
2064 fresh.candidates(&spec, &getter(&emb(&[1.0, 0.0]))),
2065 BTreeSet::from([1, 2, 3])
2066 );
2067 }
2068
2069 #[test]
2072 fn insert_skipping_tracks_but_does_not_reinsert() {
2073 let (side, spec) = hnsw_side();
2074 let blob = side.export_hnsw_blob();
2075
2076 let mut fresh = SideIndex::default();
2077 let (already, _) = fresh.init_or_adopt_hnsw("sim", &blob);
2078 let before = fresh.hnsw_ref().map(|h| h.len());
2079
2080 fresh.insert_skipping(&spec, 3, &already, &getter(&emb(&[0.7, 0.7])));
2082 assert_eq!(
2083 fresh.hnsw_ref().map(|h| h.len()),
2084 before,
2085 "an adopted id must not be re-inserted"
2086 );
2087 fresh.insert_skipping(&spec, 4, &already, &getter(&emb(&[-1.0, 0.0])));
2088 assert_eq!(
2089 fresh.hnsw_ref().map(|h| h.len()),
2090 before.map(|n| n + 1),
2091 "a post-snapshot id must be inserted"
2092 );
2093 assert_eq!(
2094 fresh.candidates(&spec, &getter(&emb(&[1.0, 0.0]))),
2095 BTreeSet::from([1, 2, 3, 4])
2096 );
2097 }
2098}