citadel_vector/
ann.rs

1//! In-memory ANN index wrapping the vendored PRISM engine.
2
3use crate::prism::{Filter, Metric, PointStore, PrismConfig, PrismIndex};
4
5/// Request `k * OVER_FETCH` candidates to offset PRISM recall below 1.0.
6pub const OVER_FETCH: usize = 4;
7
8#[derive(Debug, thiserror::Error)]
9pub enum AnnError {
10    #[error("ANN build requires at least one row")]
11    EmptyInput,
12    #[error("ANN build vector dim mismatch: expected {expected}, got {got} for row_id {row_id}")]
13    DimMismatch {
14        expected: u16,
15        got: usize,
16        row_id: u64,
17    },
18    #[error(
19        "ANN build attribute arity mismatch: expected {expected}, got {got} for row_id {row_id}"
20    )]
21    AttrArityMismatch {
22        expected: usize,
23        got: usize,
24        row_id: u64,
25    },
26}
27
28// binary_rerank=0: the Hamming pre-filter kills recall on continuous vectors.
29// sigma_high low: stay in the fast HIGH search regime.
30fn prism_config(metric: Metric) -> PrismConfig {
31    PrismConfig {
32        metric,
33        binary_rerank: 0,
34        sigma_high: 0.001,
35        ..PrismConfig::default()
36    }
37}
38
39/// In-memory ANN index over a `(row_id, vector)` snapshot.
40pub struct AnnIndex {
41    prism: PrismIndex,
42    /// PRISM internal id -> external row_id.
43    id_map: Vec<u64>,
44    /// Highest row_id in the snapshot.
45    pub snapshot_max: u64,
46    pub metric: Metric,
47    pub dim: u16,
48}
49
50impl std::fmt::Debug for AnnIndex {
51    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
52        f.debug_struct("AnnIndex")
53            .field("snapshot_max", &self.snapshot_max)
54            .field("metric", &self.metric)
55            .field("dim", &self.dim)
56            .field("indexed_len", &self.id_map.len())
57            .finish()
58    }
59}
60
61impl AnnIndex {
62    /// Build an unfiltered index from `(row_id, vector)` pairs.
63    pub fn build(rows: Vec<(u64, Vec<f32>)>, metric: Metric, dim: u16) -> Result<Self, AnnError> {
64        let with_attrs = rows
65            .into_iter()
66            .map(|(id, v)| (id, v, Vec::new()))
67            .collect();
68        Self::build_with_attrs(with_attrs, 0, metric, dim)
69    }
70
71    /// Build a filtered index from `(row_id, vector, attr_codes)` triples. Each
72    /// attribute is a PRISM dimension; distinct tuples form the searchable cells.
73    pub fn build_with_attrs(
74        mut rows: Vec<(u64, Vec<f32>, Vec<u32>)>,
75        num_attrs: usize,
76        metric: Metric,
77        dim: u16,
78    ) -> Result<Self, AnnError> {
79        if rows.is_empty() {
80            return Err(AnnError::EmptyInput);
81        }
82        for (rid, v, a) in &rows {
83            if v.len() != dim as usize {
84                return Err(AnnError::DimMismatch {
85                    expected: dim,
86                    got: v.len(),
87                    row_id: *rid,
88                });
89            }
90            if a.len() != num_attrs {
91                return Err(AnnError::AttrArityMismatch {
92                    expected: num_attrs,
93                    got: a.len(),
94                    row_id: *rid,
95                });
96            }
97        }
98
99        rows.sort_unstable_by_key(|(id, _, _)| *id);
100        let snapshot_max = rows.last().map(|(id, _, _)| *id).unwrap_or(0);
101
102        let n = rows.len();
103        let mut flat: Vec<f32> = Vec::with_capacity(n * dim as usize);
104        let mut row_ids: Vec<u64> = Vec::with_capacity(n);
105        // PRISM needs >=1 attribute dim; an all-zero column = one cell.
106        let attr_dims = num_attrs.max(1);
107        let mut attr_cols: Vec<Vec<u32>> = vec![Vec::with_capacity(n); attr_dims];
108        for (rid, v, a) in &rows {
109            flat.extend_from_slice(v);
110            row_ids.push(*rid);
111            if num_attrs == 0 {
112                attr_cols[0].push(0);
113            } else {
114                for (j, &code) in a.iter().enumerate() {
115                    attr_cols[j].push(code);
116                }
117            }
118        }
119
120        let store = PointStore::from_parts(flat, dim as usize, attr_cols);
121        let prism = PrismIndex::build(store, prism_config(metric));
122
123        // PRISM reorders points by cell; remap to external row_ids.
124        let id_map: Vec<u64> = prism
125            .original_ids
126            .iter()
127            .map(|&old| row_ids[old as usize])
128            .collect();
129
130        Ok(Self {
131            prism,
132            id_map,
133            snapshot_max,
134            metric,
135            dim,
136        })
137    }
138
139    /// Top-k search returning `(row_id, distance)` ascending, at the default ef.
140    pub fn search(&self, query: &[f32], k: usize) -> Vec<(u64, f32)> {
141        let ef = (k * OVER_FETCH).max(self.prism.config.beam_width);
142        self.search_with_ef(query, k, ef)
143    }
144
145    /// Unfiltered search with an explicit beam width `ef`.
146    pub fn search_with_ef(&self, query: &[f32], k: usize, ef: usize) -> Vec<(u64, f32)> {
147        self.search_filtered(query, k, ef, &Filter::none())
148    }
149
150    /// Filtered search at the default ef.
151    pub fn search_filtered_default_ef(
152        &self,
153        query: &[f32],
154        k: usize,
155        filter: &Filter,
156    ) -> Vec<(u64, f32)> {
157        let ef = (k * OVER_FETCH).max(self.prism.config.beam_width);
158        self.search_filtered(query, k, ef, filter)
159    }
160
161    /// Filtered search; `filter` dims index the `build_with_attrs` attributes
162    /// (`Filter::none()` matches all).
163    pub fn search_filtered(
164        &self,
165        query: &[f32],
166        k: usize,
167        ef: usize,
168        filter: &Filter,
169    ) -> Vec<(u64, f32)> {
170        debug_assert_eq!(query.len(), self.dim as usize);
171        self.prism
172            .search(query, filter, k, ef)
173            .into_iter()
174            .map(|r| (self.id_map[r.id as usize], r.dist))
175            .collect()
176    }
177
178    /// Number of indexed rows.
179    pub fn indexed_len(&self) -> usize {
180        self.id_map.len()
181    }
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187
188    fn synth_rows(n: usize, dim: u16) -> Vec<(u64, Vec<f32>)> {
189        (0..n)
190            .map(|i| {
191                let row_id = (i as u64) + 1;
192                let v: Vec<f32> = (0..dim).map(|d| (i as f32 + d as f32) * 0.01).collect();
193                (row_id, v)
194            })
195            .collect()
196    }
197
198    #[test]
199    fn build_empty_input_errors() {
200        let err = AnnIndex::build(Vec::new(), Metric::L2, 4).unwrap_err();
201        assert!(matches!(err, AnnError::EmptyInput));
202    }
203
204    #[test]
205    fn build_dim_mismatch_errors() {
206        let rows = vec![(1u64, vec![1.0, 2.0])];
207        let err = AnnIndex::build(rows, Metric::L2, 4).unwrap_err();
208        assert!(matches!(
209            err,
210            AnnError::DimMismatch {
211                expected: 4,
212                got: 2,
213                row_id: 1
214            }
215        ));
216    }
217
218    #[test]
219    fn build_single_row_succeeds() {
220        let rows = vec![(7u64, vec![0.1, 0.2, 0.3, 0.4])];
221        let idx = AnnIndex::build(rows, Metric::L2, 4).unwrap();
222        assert_eq!(idx.indexed_len(), 1);
223        assert_eq!(idx.snapshot_max, 7);
224    }
225
226    #[test]
227    fn build_small_n_succeeds() {
228        let rows = synth_rows(5, 8);
229        let idx = AnnIndex::build(rows, Metric::L2, 8).unwrap();
230        assert_eq!(idx.indexed_len(), 5);
231    }
232
233    #[test]
234    fn build_large_n_succeeds() {
235        let rows = synth_rows(500, 16);
236        let idx = AnnIndex::build(rows, Metric::L2, 16).unwrap();
237        assert_eq!(idx.indexed_len(), 500);
238    }
239
240    #[test]
241    fn search_returns_row_ids_not_internal_ids() {
242        let n = 200;
243        let rows = synth_rows(n, 8);
244        let idx = AnnIndex::build(rows, Metric::L2, 8).unwrap();
245        let hits = idx.search(&[0.5; 8], 5);
246        assert!(!hits.is_empty());
247        for (rid, _d) in &hits {
248            assert!(*rid >= 1 && *rid <= n as u64);
249        }
250    }
251
252    #[test]
253    fn snapshot_max_tracks_highest_row_id() {
254        let rows = vec![
255            (5u64, vec![1.0, 0.0]),
256            (10u64, vec![0.0, 1.0]),
257            (3u64, vec![1.0, 1.0]),
258        ];
259        let idx = AnnIndex::build(rows, Metric::L2, 2).unwrap();
260        assert_eq!(idx.snapshot_max, 10);
261    }
262
263    #[test]
264    fn cosine_metric_propagates_to_prism() {
265        let rows = synth_rows(50, 16);
266        let idx = AnnIndex::build(rows, Metric::Cosine, 16).unwrap();
267        assert_eq!(idx.metric, Metric::Cosine);
268        assert_eq!(idx.prism.config.metric, Metric::Cosine);
269    }
270
271    #[test]
272    fn inner_metric_propagates_to_prism() {
273        let rows = synth_rows(50, 16);
274        let idx = AnnIndex::build(rows, Metric::InnerProduct, 16).unwrap();
275        assert_eq!(idx.metric, Metric::InnerProduct);
276        assert_eq!(idx.prism.config.metric, Metric::InnerProduct);
277    }
278
279    /// attr 0 = i % 2; row_id = i + 1.
280    fn attr_rows(n: u64, dim: u16) -> Vec<(u64, Vec<f32>, Vec<u32>)> {
281        (0..n)
282            .map(|i| {
283                let v: Vec<f32> = (0..dim).map(|d| (i as f32 + d as f32) * 0.01).collect();
284                (i + 1, v, vec![(i % 2) as u32])
285            })
286            .collect()
287    }
288
289    #[test]
290    fn build_with_attrs_filters_by_attribute() {
291        let idx = AnnIndex::build_with_attrs(attr_rows(100, 8), 1, Metric::L2, 8).unwrap();
292        let hits = idx.search_filtered(&[0.5; 8], 10, 200, &Filter::eq(0, 1));
293        assert!(!hits.is_empty());
294        assert!(hits.len() <= 10);
295        for (rid, _) in &hits {
296            // category 1 == odd i == even row_id.
297            assert_eq!(rid % 2, 0, "row {rid} is not category 1");
298        }
299    }
300
301    #[test]
302    fn build_with_attrs_unfiltered_spans_all_cells() {
303        let idx = AnnIndex::build_with_attrs(attr_rows(100, 8), 1, Metric::L2, 8).unwrap();
304        let hits = idx.search_with_ef(&[0.5; 8], 10, 200);
305        assert_eq!(hits.len(), 10);
306        for (rid, _) in &hits {
307            assert!(*rid >= 1 && *rid <= 100);
308        }
309    }
310
311    #[test]
312    fn build_with_attrs_two_dims_conjunctive_filter() {
313        let n = 180u64;
314        let dim = 8u16;
315        let rows: Vec<(u64, Vec<f32>, Vec<u32>)> = (0..n)
316            .map(|i| {
317                let v: Vec<f32> = (0..dim).map(|d| (i as f32 + d as f32) * 0.01).collect();
318                (i + 1, v, vec![(i % 2) as u32, (i % 3) as u32])
319            })
320            .collect();
321        let idx = AnnIndex::build_with_attrs(rows, 2, Metric::L2, dim).unwrap();
322        let filter = Filter::new(vec![(0, vec![1]), (1, vec![2])]);
323        let hits = idx.search_filtered(&[0.5; 8], 10, 200, &filter);
324        assert!(!hits.is_empty());
325        for (rid, _) in &hits {
326            let i = rid - 1;
327            assert_eq!(i % 2, 1, "row {rid} fails attr0 = 1");
328            assert_eq!(i % 3, 2, "row {rid} fails attr1 = 2");
329        }
330    }
331
332    #[test]
333    fn build_with_attrs_arity_mismatch_errors() {
334        let rows = vec![(1u64, vec![0.0; 4], vec![0u32])];
335        let err = AnnIndex::build_with_attrs(rows, 2, Metric::L2, 4).unwrap_err();
336        assert!(matches!(
337            err,
338            AnnError::AttrArityMismatch {
339                expected: 2,
340                got: 1,
341                row_id: 1
342            }
343        ));
344    }
345
346    #[test]
347    fn build_delegates_to_attrs_path() {
348        let idx = AnnIndex::build(synth_rows(50, 8), Metric::L2, 8).unwrap();
349        assert_eq!(idx.indexed_len(), 50);
350        let hits = idx.search(&[0.3; 8], 5);
351        assert!(!hits.is_empty());
352    }
353}
citadel_vector/ann.rs

citadel_vector/
ann.rs
