nodedb_vector/ivf/
checkpoint.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Checkpoint encoding for `IvfPqIndex`: centroids, the PQ codec, every
4//! cell's ids, codes and FP32 vectors, the tombstones, and the training stamp.
5//! Codes are stored, never recomputed on load.
6
7use std::collections::HashMap;
8
9use nodedb_mem::ScopedMemory;
10use roaring::RoaringBitmap;
11
12use crate::error::VectorError;
13use crate::quantize::pq::PqCodec;
14
15use super::index::{IvfCell, IvfPqIndex};
16use super::params::IvfPqParams;
17
18#[derive(zerompk::ToMessagePack, zerompk::FromMessagePack)]
19struct IvfSnapshot {
20    dim: usize,
21    params: IvfPqParams,
22    centroids: Vec<Vec<f32>>,
23    pq_bytes: Option<Vec<u8>>,
24    cells: Vec<IvfCell>,
25    deleted: Vec<u32>,
26    next_id: u32,
27    trained_on: usize,
28    trained_at_ms: u64,
29}
30
31fn corrupt(detail: String) -> VectorError {
32    VectorError::CheckpointDeserializationError { detail }
33}
34
35impl IvfPqIndex {
36    /// Encode the whole index as MessagePack.
37    pub fn to_bytes(&self) -> Result<Vec<u8>, VectorError> {
38        let snapshot = IvfSnapshot {
39            dim: self.dim,
40            params: self.params.clone(),
41            centroids: self.centroids.clone(),
42            pq_bytes: self.pq.as_ref().map(PqCodec::to_bytes).transpose()?,
43            cells: self.cells.clone(),
44            deleted: self.deleted.iter().collect(),
45            next_id: self.next_id,
46            trained_on: self.trained_on,
47            trained_at_ms: self.trained_at_ms,
48        };
49        zerompk::to_msgpack_vec(&snapshot).map_err(|e| VectorError::CheckpointSerializationError {
50            detail: format!("IVF-PQ index encode: {e}"),
51        })
52    }
53
54    /// Decode an index written by [`Self::to_bytes`], charging the PQ codec
55    /// to `memory`.
56    ///
57    /// Fails with [`VectorError::CheckpointDeserializationError`] when the
58    /// bytes do not decode or the decoded cells disagree with the dimension,
59    /// the PQ code width, or the centroid count.
60    pub fn from_bytes(bytes: &[u8], memory: ScopedMemory) -> Result<Self, VectorError> {
61        let snap: IvfSnapshot = zerompk::from_msgpack(bytes)
62            .map_err(|e| corrupt(format!("IVF-PQ index decode: {e}")))?;
63        let pq = snap
64            .pq_bytes
65            .as_deref()
66            .map(|b| PqCodec::from_bytes(b, memory))
67            .transpose()
68            .map_err(|e| corrupt(format!("IVF-PQ codec decode: {e}")))?;
69        let m = pq.as_ref().map_or(0, |pq| pq.m);
70        if snap.cells.len() != snap.centroids.len() {
71            return Err(corrupt(format!(
72                "IVF-PQ index has {} cells for {} centroids",
73                snap.cells.len(),
74                snap.centroids.len()
75            )));
76        }
77        let mut slots = HashMap::new();
78        for (cell_idx, cell) in snap.cells.iter().enumerate() {
79            let n = cell.ids.len();
80            if cell.codes.len() != n * m || cell.vectors.len() != n * snap.dim {
81                return Err(corrupt(format!(
82                    "IVF-PQ cell {cell_idx} holds {n} ids, {} code bytes and {} vector \
83                     components; expected {} and {}",
84                    cell.codes.len(),
85                    cell.vectors.len(),
86                    n * m,
87                    n * snap.dim
88                )));
89            }
90            for (pos, &id) in cell.ids.iter().enumerate() {
91                if slots.insert(id, (cell_idx as u32, pos as u32)).is_some() {
92                    return Err(corrupt(format!("IVF-PQ index holds vector id {id} twice")));
93                }
94            }
95        }
96        let deleted: RoaringBitmap = snap.deleted.into_iter().collect();
97        if let Some(id) = deleted.iter().find(|id| !slots.contains_key(id)) {
98            return Err(corrupt(format!(
99                "IVF-PQ index tombstones vector id {id} it does not hold"
100            )));
101        }
102        Ok(Self {
103            dim: snap.dim,
104            params: snap.params,
105            centroids: snap.centroids,
106            pq,
107            cells: snap.cells,
108            slots,
109            deleted,
110            next_id: snap.next_id,
111            trained_on: snap.trained_on,
112            trained_at_ms: snap.trained_at_ms,
113        })
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120    use crate::distance::DistanceMetric;
121    use crate::test_support::test_memory;
122
123    #[test]
124    fn round_trip_keeps_entries_tombstones_and_training() {
125        let vecs: Vec<Vec<f32>> = (0..32)
126            .map(|i| (0..8).map(|d| ((i * 8 + d) as f32) * 0.01).collect())
127            .collect();
128        let refs: Vec<&[f32]> = vecs.iter().map(|v| v.as_slice()).collect();
129        let mut idx = IvfPqIndex::new(
130            8,
131            IvfPqParams {
132                n_cells: 4,
133                pq_m: 4,
134                pq_k: 8,
135                nprobe: 4,
136                metric: DistanceMetric::L2,
137            },
138        );
139        idx.train(&refs, test_memory()).unwrap();
140        for (i, v) in vecs.iter().enumerate() {
141            idx.insert_with_id(10 + i as u32, v.clone()).unwrap();
142        }
143        idx.delete(12);
144        idx.set_trained_at_ms(1_700_000_000_000);
145
146        let restored = IvfPqIndex::from_bytes(&idx.to_bytes().unwrap(), test_memory()).unwrap();
147        assert_eq!(restored.len(), 32);
148        assert!(restored.is_deleted(12));
149        assert_eq!(restored.trained_on(), 32);
150        assert_eq!(restored.trained_at_ms(), 1_700_000_000_000);
151        assert_eq!(restored.get_vector(20), Some(vecs[10].as_slice()));
152        assert_eq!(
153            restored.search(&vecs[5], 3).unwrap()[0].id,
154            idx.search(&vecs[5], 3).unwrap()[0].id
155        );
156    }
157
158    #[test]
159    fn garbage_is_a_typed_error() {
160        assert!(matches!(
161            IvfPqIndex::from_bytes(b"not an index", test_memory()),
162            Err(VectorError::CheckpointDeserializationError { .. })
163        ));
164    }
165}
nodedb_vector/ivf/checkpoint.rs

nodedb_vector/ivf/
checkpoint.rs
