nodedb_vector/collection/
ivf_mode.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! The IVF-PQ mode of a `VectorCollection`.
4//!
5//! An `IvfPq` collection buffers its vectors in the growing segment until it
6//! holds the training threshold, `max(ivf_cells, pq_k)`. The growing segment
7//! is searched exactly, checkpointed, and rebuilt by WAL replay, so the
8//! buffer is durable and searchable with no extra machinery. It never seals.
9//!
10//! [`VectorCollection::train_ivf`] trains the IVF centroids and PQ codebooks
11//! on every live vector, moves each one into the trained index under its
12//! global id, and empties the segments. The move happens in one call on the
13//! owning core, so no search sees a vector twice or misses one. Later inserts
14//! land in the trained index.
15
16use nodedb_mem::ScopedMemory;
17
18use crate::error::VectorError;
19use crate::index_config::{IndexConfig, IndexType};
20use crate::ivf::IvfPqIndex;
21
22use super::lifecycle::VectorCollection;
23use super::lifecycle_insert_ops::sealed_vector;
24
25impl VectorCollection {
26    /// Whether the collection is configured as an IVF-PQ index.
27    pub fn is_ivf(&self) -> bool {
28        self.index_config.index_type == IndexType::IvfPq
29    }
30
31    /// The full index configuration.
32    pub fn index_config(&self) -> &IndexConfig {
33        &self.index_config
34    }
35
36    /// Replace the index configuration. The HNSW parameters the collection
37    /// holds stay: its segments were built with them.
38    pub fn set_index_config(&mut self, config: IndexConfig) {
39        self.index_config = IndexConfig {
40            hnsw: self.params.clone(),
41            ..config
42        };
43    }
44
45    /// Live vectors an `IvfPq` collection needs before it trains.
46    pub fn ivf_training_threshold(&self) -> usize {
47        self.index_config.to_ivf_params().training_threshold()
48    }
49
50    /// Whether an `IvfPq` collection is untrained and holds the training
51    /// threshold of live vectors.
52    pub fn needs_ivf_training(&self) -> bool {
53        self.is_ivf() && self.ivf.is_none() && self.live_count() >= self.ivf_training_threshold()
54    }
55
56    /// The trained IVF-PQ index, if any.
57    pub fn ivf_index(&self) -> Option<&IvfPqIndex> {
58        self.ivf.as_ref()
59    }
60
61    /// Train the IVF-PQ index on every live vector and move them all into it.
62    /// `trained_at_ms` stamps the training, in Unix milliseconds.
63    ///
64    /// Fails with [`VectorError::InvalidInput`] when the collection is not
65    /// `IvfPq`, is already trained, or holds fewer live vectors than the
66    /// training threshold. A training error propagates. Any failure leaves
67    /// the collection unchanged.
68    pub fn train_ivf(
69        &mut self,
70        memory: ScopedMemory,
71        trained_at_ms: u64,
72    ) -> Result<(), VectorError> {
73        if !self.is_ivf() || self.ivf.is_some() {
74            return Err(VectorError::InvalidInput {
75                detail: "IVF-PQ training needs an untrained ivf_pq collection".into(),
76            });
77        }
78        let live = self.gather_live_vectors();
79        let threshold = self.ivf_training_threshold();
80        if live.len() < threshold {
81            return Err(VectorError::InvalidInput {
82                detail: format!(
83                    "IVF-PQ training needs {threshold} live vectors; the collection holds {}",
84                    live.len()
85                ),
86            });
87        }
88        let mut ivf = IvfPqIndex::new(self.dim, self.index_config.to_ivf_params());
89        {
90            let refs: Vec<&[f32]> = live.iter().map(|(_, v)| v.as_slice()).collect();
91            ivf.train(&refs, memory)?;
92        }
93        for (id, vector) in live {
94            ivf.insert_with_id(id, vector)?;
95        }
96        ivf.set_trained_at_ms(trained_at_ms);
97        self.clear_segments();
98        self.ivf = Some(ivf);
99        Ok(())
100    }
101
102    /// Every live FP32 vector outside the IVF index, keyed by global id.
103    fn gather_live_vectors(&self) -> Vec<(u32, Vec<f32>)> {
104        let mut out = Vec::with_capacity(self.live_count());
105        for seg in &self.sealed {
106            for local in 0..seg.index.len() as u32 {
107                if !seg.index.is_deleted(local)
108                    && let Some(v) = sealed_vector(seg, local)
109                {
110                    out.push((seg.base_id + local, v));
111                }
112            }
113        }
114        for seg in &self.building {
115            for local in 0..seg.flat.len() as u32 {
116                if let Some(v) = seg.flat.get_vector(local) {
117                    out.push((seg.base_id + local, v.to_vec()));
118                }
119            }
120        }
121        for local in 0..self.growing.len() as u32 {
122            if let Some(v) = self.growing.get_vector(local) {
123                out.push((self.growing_base_id + local, v.to_vec()));
124            }
125        }
126        out
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use nodedb_types::Surrogate;
133
134    use super::*;
135    use crate::test_support::test_memory;
136
137    const DIM: usize = 8;
138
139    /// An L2 `IvfPq` collection with 4 cells, all probed, and 256 PQ
140    /// centroids: the threshold is 256 vectors.
141    fn ivf_collection() -> VectorCollection {
142        let config = IndexConfig {
143            hnsw: crate::hnsw::HnswParams {
144                metric: crate::distance::DistanceMetric::L2,
145                ..crate::hnsw::HnswParams::default()
146            },
147            index_type: IndexType::IvfPq,
148            pq_m: 4,
149            ivf_cells: 4,
150            ivf_nprobe: 4,
151            ..IndexConfig::default()
152        };
153        VectorCollection::with_index_config(DIM, config)
154    }
155
156    /// Distinct per `i`: the first component is `i + 1`.
157    fn vector(i: usize) -> Vec<f32> {
158        [1, 7, 11, 13, 17, 19, 23, 29]
159            .iter()
160            .map(|&m| (if m == 1 { i + 1 } else { i % m + 1 }) as f32)
161            .collect()
162    }
163
164    fn all_ids(coll: &VectorCollection, query: &[f32]) -> Vec<u32> {
165        let mut ids: Vec<u32> = coll
166            .search(query, 10_000, 64)
167            .unwrap()
168            .iter()
169            .map(|r| r.id)
170            .collect();
171        ids.sort_unstable();
172        ids
173    }
174
175    #[test]
176    fn below_the_threshold_vectors_wait_in_the_exact_buffer() {
177        let mut coll = ivf_collection();
178        assert_eq!(coll.ivf_training_threshold(), 256);
179        for i in 0..10 {
180            coll.insert_with_surrogate(vector(i), Surrogate::new(i as u32 + 1))
181                .unwrap();
182        }
183        assert!(!coll.needs_ivf_training());
184        assert!(!coll.needs_seal());
185        let hit = &coll.search(&vector(3), 1, 64).unwrap()[0];
186        assert_eq!(hit.id, 3);
187        assert_eq!(hit.distance, 0.0, "the buffer is searched exactly");
188    }
189
190    #[test]
191    fn training_moves_every_vector_once_and_later_inserts_follow() {
192        let mut coll = ivf_collection();
193        for i in 0..256 {
194            coll.insert_with_surrogate(vector(i), Surrogate::new(i as u32 + 1))
195                .unwrap();
196        }
197        coll.delete(5);
198        coll.insert(vector(256)).unwrap();
199        assert!(coll.needs_ivf_training());
200        let before = all_ids(&coll, &vector(0));
201
202        coll.train_ivf(test_memory(), 42).unwrap();
203
204        assert!(!coll.needs_ivf_training());
205        assert!(coll.growing_is_empty());
206        let ivf = coll.ivf_index().unwrap();
207        assert_eq!(ivf.trained_on(), 256);
208        assert_eq!(ivf.trained_at_ms(), 42);
209        assert_eq!(
210            all_ids(&coll, &vector(0)),
211            before,
212            "no vector lost or doubled"
213        );
214
215        let id = coll
216            .insert_with_surrogate(vector(300), Surrogate::new(9_999))
217            .unwrap();
218        assert_eq!(id, 257);
219        assert!(
220            coll.growing_is_empty(),
221            "a trained collection inserts into IVF"
222        );
223        assert_eq!(coll.search(&vector(300), 1, 64).unwrap()[0].id, 257);
224        assert_eq!(
225            coll.vector_for_surrogate(Surrogate::new(9_999)),
226            Some(vector(300))
227        );
228        assert!(coll.delete_by_surrogate(Surrogate::new(9_999)));
229        assert!(coll.search(&vector(300), 1, 64).unwrap()[0].id != 257);
230    }
231
232    #[test]
233    fn training_below_the_threshold_is_refused() {
234        let mut coll = ivf_collection();
235        coll.insert(vector(0)).unwrap();
236        assert!(matches!(
237            coll.train_ivf(test_memory(), 0),
238            Err(VectorError::InvalidInput { .. })
239        ));
240        assert!(coll.ivf_index().is_none());
241        assert_eq!(coll.live_count(), 1);
242    }
243
244    #[test]
245    fn a_trained_collection_survives_a_checkpoint() {
246        let mut coll = ivf_collection();
247        for i in 0..260 {
248            coll.insert(vector(i)).unwrap();
249        }
250        coll.train_ivf(test_memory(), 7).unwrap();
251        coll.insert(vector(400)).unwrap();
252        coll.delete(9);
253
254        let bytes = coll.checkpoint_to_bytes(None).unwrap();
255        let restored = VectorCollection::from_checkpoint(&bytes, None, test_memory()).unwrap();
256
257        assert!(restored.is_ivf());
258        assert_eq!(restored.ivf_index().unwrap().trained_at_ms(), 7);
259        assert_eq!(all_ids(&restored, &vector(0)), all_ids(&coll, &vector(0)));
260        assert_eq!(restored.search(&vector(400), 1, 64).unwrap()[0].id, 260);
261    }
262
263    #[test]
264    fn an_untrained_buffer_survives_a_checkpoint_and_trains_after() {
265        let mut coll = ivf_collection();
266        for i in 0..100 {
267            coll.insert(vector(i)).unwrap();
268        }
269        let bytes = coll.checkpoint_to_bytes(None).unwrap();
270        let mut restored = VectorCollection::from_checkpoint(&bytes, None, test_memory()).unwrap();
271        assert!(restored.is_ivf());
272        assert!(restored.ivf_index().is_none());
273        for i in 100..256 {
274            restored.insert(vector(i)).unwrap();
275        }
276        assert!(restored.needs_ivf_training());
277        restored.train_ivf(test_memory(), 1).unwrap();
278        assert_eq!(
279            all_ids(&restored, &vector(0)),
280            (0..256).collect::<Vec<u32>>()
281        );
282    }
283}
nodedb_vector/collection/ivf_mode.rs

nodedb_vector/collection/
ivf_mode.rs
