nodedb_vector/ivf/
search.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! IVF-PQ search: probe the nearest cells, rank their entries by PQ
4//! distance, and rerank the best of them by exact FP32 distance.
5
6use roaring::RoaringBitmap;
7
8use crate::distance::{DistanceMetric, distance};
9use crate::error::{VectorError, check_dim};
10use crate::hnsw::SearchResult;
11
12use super::index::{IvfPqIndex, residual};
13
14/// PQ-ranked candidates reranked exactly, per result asked for.
15const RERANK_PER_RESULT: usize = 3;
16/// Fewest PQ-ranked candidates reranked exactly.
17const MIN_RERANK: usize = 20;
18
19/// One PQ-ranked entry: id, PQ distance, cell, position in the cell.
20struct Candidate {
21    id: u32,
22    pq_distance: f32,
23    cell: usize,
24    pos: usize,
25}
26
27impl IvfPqIndex {
28    /// The `top_k` live entries nearest to `query` under the index metric.
29    ///
30    /// A query without the index dimension fails with
31    /// [`VectorError::DimensionMismatch`]. A distance table over the memory
32    /// budget fails the search: skipping its cell would drop results.
33    pub fn search(&self, query: &[f32], top_k: usize) -> Result<Vec<SearchResult>, VectorError> {
34        self.search_with(query, top_k, self.params.metric, None)
35    }
36
37    /// The `top_k` live entries nearest to `query` under `metric`, keeping
38    /// only ids in `filter` when one is given. Fails as [`Self::search`].
39    pub fn search_with(
40        &self,
41        query: &[f32],
42        top_k: usize,
43        metric: DistanceMetric,
44        filter: Option<&RoaringBitmap>,
45    ) -> Result<Vec<SearchResult>, VectorError> {
46        check_dim(self.dim, query.len())?;
47        let Some(pq) = &self.pq else {
48            return Ok(Vec::new());
49        };
50        if top_k == 0 || self.slots.is_empty() {
51            return Ok(Vec::new());
52        }
53
54        let mut probe: Vec<(usize, f32)> = self
55            .centroids
56            .iter()
57            .enumerate()
58            .map(|(i, c)| (i, distance(query, c, self.params.metric)))
59            .collect();
60        probe.sort_by(|a, b| a.1.total_cmp(&b.1));
61        probe.truncate(self.params.nprobe.max(1));
62
63        let m = pq.m;
64        let mut candidates: Vec<Candidate> = Vec::new();
65        for &(cell_idx, _) in &probe {
66            let cell = &self.cells[cell_idx];
67            if cell.ids.is_empty() {
68                continue;
69            }
70            let table = pq.build_distance_table(&residual(query, &self.centroids[cell_idx]))?;
71            for (pos, &id) in cell.ids.iter().enumerate() {
72                if self.deleted.contains(id) || filter.is_some_and(|f| !f.contains(id)) {
73                    continue;
74                }
75                let code = &cell.codes[pos * m..(pos + 1) * m];
76                candidates.push(Candidate {
77                    id,
78                    pq_distance: pq.asymmetric_distance(&table, code),
79                    cell: cell_idx,
80                    pos,
81                });
82            }
83        }
84
85        let pool = top_k.saturating_mul(RERANK_PER_RESULT).max(MIN_RERANK);
86        if candidates.len() > pool {
87            candidates.select_nth_unstable_by(pool, |a, b| a.pq_distance.total_cmp(&b.pq_distance));
88            candidates.truncate(pool);
89        }
90
91        let dim = self.dim;
92        let mut results: Vec<SearchResult> = candidates
93            .into_iter()
94            .map(|c| {
95                let start = c.pos * dim;
96                let vector = &self.cells[c.cell].vectors[start..start + dim];
97                SearchResult {
98                    id: c.id,
99                    distance: distance(query, vector, metric),
100                }
101            })
102            .collect();
103        results.sort_by(|a, b| a.distance.total_cmp(&b.distance));
104        results.truncate(top_k);
105        Ok(results)
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use crate::ivf::IvfPqParams;
113    use crate::test_support::test_memory;
114
115    fn make_vectors(n: usize, dim: usize) -> Vec<Vec<f32>> {
116        (0..n)
117            .map(|i| (0..dim).map(|d| ((i * dim + d) as f32) * 0.01).collect())
118            .collect()
119    }
120
121    fn index_with(vecs: &[Vec<f32>], params: IvfPqParams) -> IvfPqIndex {
122        let refs: Vec<&[f32]> = vecs.iter().map(|v| v.as_slice()).collect();
123        let mut idx = IvfPqIndex::new(vecs[0].len(), params);
124        idx.train(&refs, test_memory()).unwrap();
125        idx.add_batch(&refs).unwrap();
126        idx
127    }
128
129    #[test]
130    fn train_and_search_finds_the_exact_match_first() {
131        let vecs = make_vectors(1000, 16);
132        let idx = index_with(
133            &vecs,
134            IvfPqParams {
135                n_cells: 32,
136                pq_m: 4,
137                pq_k: 32,
138                nprobe: 8,
139                metric: DistanceMetric::L2,
140            },
141        );
142        assert_eq!(idx.len(), 1000);
143        let results = idx.search(&vecs[500], 5).unwrap();
144        assert_eq!(results.len(), 5);
145        assert_eq!(results[0].id, 500, "the exact rerank puts the match first");
146        assert_eq!(results[0].distance, 0.0);
147    }
148
149    #[test]
150    fn probing_every_cell_returns_every_live_entry_once() {
151        let vecs = make_vectors(64, 8);
152        let mut idx = index_with(
153            &vecs,
154            IvfPqParams {
155                n_cells: 4,
156                pq_m: 4,
157                pq_k: 8,
158                nprobe: 4,
159                metric: DistanceMetric::L2,
160            },
161        );
162        idx.delete(7);
163        let mut ids: Vec<u32> = idx
164            .search(&vecs[0], 100)
165            .unwrap()
166            .iter()
167            .map(|r| r.id)
168            .collect();
169        ids.sort_unstable();
170        let expected: Vec<u32> = (0..64).filter(|id| *id != 7).collect();
171        assert_eq!(ids, expected);
172    }
173
174    #[test]
175    fn a_filter_keeps_only_its_ids() {
176        let vecs = make_vectors(64, 8);
177        let idx = index_with(
178            &vecs,
179            IvfPqParams {
180                n_cells: 4,
181                pq_m: 4,
182                pq_k: 8,
183                nprobe: 4,
184                metric: DistanceMetric::L2,
185            },
186        );
187        let filter: RoaringBitmap = [3u32, 40, 41].into_iter().collect();
188        let mut ids: Vec<u32> = idx
189            .search_with(&vecs[0], 10, DistanceMetric::L2, Some(&filter))
190            .unwrap()
191            .iter()
192            .map(|r| r.id)
193            .collect();
194        ids.sort_unstable();
195        assert_eq!(ids, vec![3, 40, 41]);
196    }
197
198    #[test]
199    fn an_untrained_index_finds_nothing() {
200        let idx = IvfPqIndex::new(8, IvfPqParams::default());
201        assert!(idx.search(&[0.0; 8], 5).unwrap().is_empty());
202    }
203}
nodedb_vector/ivf/search.rs

nodedb_vector/ivf/
search.rs
