nodedb_vector/sieve/
collection.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! SieveCollection — a keyed map of specialized HNSW subindices, one per
4//! stable predicate signature (e.g. "tenant_id=42", "lang=en").
5
6use std::collections::HashMap;
7
8use crate::error::VectorError;
9use crate::hnsw::graph::HnswIndex;
10use nodedb_types::hnsw::HnswParams;
11use nodedb_types::vector_distance::DistanceMetric;
12
13/// A stable predicate signature string, e.g. `"tenant_id=42"` or `"lang=en"`.
14pub type PredicateSignature = String;
15
16/// A collection of specialized HNSW subindices, one per stable predicate.
17///
18/// Each subindex is built over the subset of vectors that match its predicate.
19/// Smaller datasets mean a lower `sub_m` is sufficient, keeping memory overhead
20/// proportional to the per-predicate population rather than the global index.
21pub struct SieveCollection {
22    /// Map from predicate signature → specialized HNSW subindex.
23    subindices: HashMap<PredicateSignature, HnswIndex>,
24    /// `M` parameter used when creating subindices.
25    /// Smaller than the global M because each subindex covers fewer vectors.
26    sub_m: usize,
27}
28
29impl SieveCollection {
30    /// Create an empty collection.  `sub_m` is the HNSW `M` parameter used for
31    /// every subindex built by this collection.
32    pub fn new(sub_m: usize) -> Self {
33        Self {
34            subindices: HashMap::new(),
35            sub_m,
36        }
37    }
38
39    /// Build (or rebuild) a specialized subindex for `signature` from the
40    /// provided `(id, vector)` pairs.
41    ///
42    /// All vectors must have length `dim`.  The subindex uses `metric` for
43    /// distance computation.
44    ///
45    /// # Errors
46    ///
47    /// Returns `VectorError` if any insertion fails (e.g. dimension mismatch).
48    pub fn build_subindex(
49        &mut self,
50        signature: PredicateSignature,
51        vectors: &[(u32, Vec<f32>)],
52        dim: usize,
53        metric: DistanceMetric,
54    ) -> Result<(), VectorError> {
55        let params = HnswParams {
56            m: self.sub_m,
57            m0: self.sub_m * 2,
58            ef_construction: 200,
59            metric,
60            dtype: nodedb_types::vector_dtype::VectorStorageDtype::F32,
61        };
62        let mut index = HnswIndex::new(dim, params);
63        for (_, vec) in vectors {
64            index.insert(vec.clone())?;
65        }
66        self.subindices.insert(signature, index);
67        Ok(())
68    }
69
70    /// Returns `true` if a subindex exists for the given signature.
71    pub fn has(&self, signature: &PredicateSignature) -> bool {
72        self.subindices.contains_key(signature)
73    }
74
75    /// Returns a shared reference to the subindex for `signature`, or `None`.
76    pub fn get(&self, signature: &PredicateSignature) -> Option<&HnswIndex> {
77        self.subindices.get(signature)
78    }
79
80    /// Remove the subindex for `signature`, freeing its memory.
81    pub fn drop(&mut self, signature: &PredicateSignature) {
82        self.subindices.remove(signature);
83    }
84
85    /// All predicate signatures currently held in this collection.
86    pub fn signatures(&self) -> Vec<&PredicateSignature> {
87        self.subindices.keys().collect()
88    }
89}
90
91// Expose SearchResult at crate level through this module for callers that
92// import from `sieve::collection`.
93pub use crate::hnsw::graph::SearchResult as SubindexSearchResult;
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    fn sample_vectors(n: usize, dim: usize) -> Vec<(u32, Vec<f32>)> {
100        (0..n).map(|i| (i as u32, vec![i as f32; dim])).collect()
101    }
102
103    #[test]
104    fn build_subindex_has_and_get() {
105        let mut coll = SieveCollection::new(8);
106        let vecs = sample_vectors(5, 3);
107        coll.build_subindex("tenant_id=42".to_string(), &vecs, 3, DistanceMetric::L2)
108            .expect("build should succeed");
109
110        assert!(coll.has(&"tenant_id=42".to_string()));
111        let idx = coll.get(&"tenant_id=42".to_string());
112        assert!(idx.is_some());
113        assert_eq!(idx.unwrap().len(), 5);
114    }
115
116    #[test]
117    fn drop_removes_subindex() {
118        let mut coll = SieveCollection::new(8);
119        let vecs = sample_vectors(5, 3);
120        coll.build_subindex("lang=en".to_string(), &vecs, 3, DistanceMetric::Cosine)
121            .expect("build should succeed");
122        assert!(coll.has(&"lang=en".to_string()));
123
124        coll.drop(&"lang=en".to_string());
125        assert!(!coll.has(&"lang=en".to_string()));
126        assert!(coll.get(&"lang=en".to_string()).is_none());
127    }
128
129    #[test]
130    fn signatures_lists_all_keys() {
131        let mut coll = SieveCollection::new(8);
132        let vecs = sample_vectors(3, 2);
133        coll.build_subindex("a".to_string(), &vecs, 2, DistanceMetric::L2)
134            .expect("build a");
135        coll.build_subindex("b".to_string(), &vecs, 2, DistanceMetric::L2)
136            .expect("build b");
137
138        let mut sigs: Vec<String> = coll.signatures().into_iter().cloned().collect();
139        sigs.sort();
140        assert_eq!(sigs, vec!["a".to_string(), "b".to_string()]);
141    }
142
143    #[test]
144    fn search_on_subindex() {
145        let mut coll = SieveCollection::new(8);
146        let vecs = sample_vectors(5, 3);
147        coll.build_subindex("tenant_id=1".to_string(), &vecs, 3, DistanceMetric::L2)
148            .expect("build");
149
150        let idx = coll.get(&"tenant_id=1".to_string()).unwrap();
151        let results = idx.search(&[2.0, 2.0, 2.0], 2, 32);
152        assert!(!results.is_empty());
153    }
154}
nodedb_vector/sieve/collection.rs

nodedb_vector/sieve/
collection.rs
