reddb_file/
vector_ivf_index.rs

1//! Persisted IVF (inverted-file) vector-index payload codec.
2//!
3//! The vector engine owns k-means training, probing, and assignment. This
4//! module owns only the durable byte layout of a serialized IVF index. The
5//! `"IVF1"` magic doubles as the format identifier; there is no separate
6//! version word.
7//!
8//! Byte layout (little-endian) — DO NOT change order/width; these bytes live in
9//! existing `.rdb` artifacts:
10//!
11//! ```text
12//! "IVF1"                 4 bytes magic
13//! n_lists                u32
14//! n_probes               u32
15//! dimension              u32
16//! max_iterations         u32
17//! convergence_threshold  f32
18//! trained                u8  (0 / 1)
19//! count                  u64
20//! next_id                u64
21//! list_count             u32
22//! repeated list_count times:
23//!   centroid_len         u32
24//!   centroid             f32 * centroid_len
25//!   id_count             u32
26//!   id                   u64 * id_count
27//!   vector_count         u32
28//!   repeated vector_count times:
29//!     vector_len         u32
30//!     value             f32 * vector_len
31//! ```
32
33/// Magic prefix for a serialized IVF index payload.
34pub const IVF_INDEX_MAGIC: [u8; 4] = *b"IVF1";
35/// Minimum length of a well-formed payload (header through `list_count`).
36pub const IVF_INDEX_HEADER_LEN: usize = 4  // magic
37    + 4  // n_lists
38    + 4  // n_probes
39    + 4  // dimension
40    + 4  // max_iterations
41    + 4  // convergence_threshold
42    + 1  // trained
43    + 8  // count
44    + 8  // next_id
45    + 4; // list_count
46
47/// A decoded IVF inverted list (one Voronoi cell).
48#[derive(Debug, Clone, PartialEq)]
49pub struct IvfListFrame {
50    pub centroid: Vec<f32>,
51    pub ids: Vec<u64>,
52    pub vectors: Vec<Vec<f32>>,
53}
54
55/// A decoded IVF index payload. Plain data only — the engine rebuilds derived
56/// state (the id→list map) from these fields.
57#[derive(Debug, Clone, PartialEq)]
58pub struct IvfIndexFrame {
59    pub n_lists: u32,
60    pub n_probes: u32,
61    pub dimension: u32,
62    pub max_iterations: u32,
63    pub convergence_threshold: f32,
64    pub trained: bool,
65    pub count: u64,
66    pub next_id: u64,
67    pub lists: Vec<IvfListFrame>,
68}
69
70/// Errors decoding an IVF index payload.
71#[derive(Debug, Clone, PartialEq, Eq)]
72pub enum IvfIndexFrameError {
73    TooShort,
74    InvalidMagic,
75    Truncated { offset: usize, reason: &'static str },
76}
77
78impl std::fmt::Display for IvfIndexFrameError {
79    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
80        match self {
81            Self::TooShort => write!(f, "data too short"),
82            Self::InvalidMagic => write!(f, "invalid IVF magic"),
83            Self::Truncated { offset, reason } => {
84                write!(f, "truncated IVF payload at offset {offset}: {reason}")
85            }
86        }
87    }
88}
89
90impl std::error::Error for IvfIndexFrameError {}
91
92/// Serialize an IVF index payload to bytes.
93pub fn encode_ivf_index_frame(frame: &IvfIndexFrame) -> Vec<u8> {
94    let mut bytes = Vec::new();
95    bytes.extend_from_slice(&IVF_INDEX_MAGIC);
96    bytes.extend_from_slice(&frame.n_lists.to_le_bytes());
97    bytes.extend_from_slice(&frame.n_probes.to_le_bytes());
98    bytes.extend_from_slice(&frame.dimension.to_le_bytes());
99    bytes.extend_from_slice(&frame.max_iterations.to_le_bytes());
100    bytes.extend_from_slice(&frame.convergence_threshold.to_le_bytes());
101    bytes.push(if frame.trained { 1 } else { 0 });
102    bytes.extend_from_slice(&frame.count.to_le_bytes());
103    bytes.extend_from_slice(&frame.next_id.to_le_bytes());
104    bytes.extend_from_slice(&(frame.lists.len() as u32).to_le_bytes());
105
106    for list in &frame.lists {
107        bytes.extend_from_slice(&(list.centroid.len() as u32).to_le_bytes());
108        for value in &list.centroid {
109            bytes.extend_from_slice(&value.to_le_bytes());
110        }
111
112        bytes.extend_from_slice(&(list.ids.len() as u32).to_le_bytes());
113        for id in &list.ids {
114            bytes.extend_from_slice(&id.to_le_bytes());
115        }
116
117        bytes.extend_from_slice(&(list.vectors.len() as u32).to_le_bytes());
118        for vector in &list.vectors {
119            bytes.extend_from_slice(&(vector.len() as u32).to_le_bytes());
120            for value in vector {
121                bytes.extend_from_slice(&value.to_le_bytes());
122            }
123        }
124    }
125
126    bytes
127}
128
129/// Deserialize an IVF index payload from bytes.
130pub fn decode_ivf_index_frame(bytes: &[u8]) -> Result<IvfIndexFrame, IvfIndexFrameError> {
131    if bytes.len() < 41 {
132        return Err(IvfIndexFrameError::TooShort);
133    }
134    if bytes[0..4] != IVF_INDEX_MAGIC {
135        return Err(IvfIndexFrameError::InvalidMagic);
136    }
137
138    let mut pos = 4usize;
139    let n_lists = read_u32(bytes, &mut pos, "n_lists")?;
140    let n_probes = read_u32(bytes, &mut pos, "n_probes")?;
141    let dimension = read_u32(bytes, &mut pos, "dimension")?;
142    let max_iterations = read_u32(bytes, &mut pos, "max_iterations")?;
143    let convergence_threshold = read_f32(bytes, &mut pos, "convergence_threshold")?;
144    let trained = read_u8(bytes, &mut pos, "trained")? == 1;
145    let count = read_u64(bytes, &mut pos, "count")?;
146    let next_id = read_u64(bytes, &mut pos, "next_id")?;
147    let list_count = read_u32(bytes, &mut pos, "list_count")?;
148
149    let mut lists = Vec::new();
150    for _ in 0..list_count {
151        let centroid_len = read_u32(bytes, &mut pos, "centroid_len")?;
152        let mut centroid = Vec::new();
153        for _ in 0..centroid_len {
154            centroid.push(read_f32(bytes, &mut pos, "centroid")?);
155        }
156
157        let id_count = read_u32(bytes, &mut pos, "id_count")?;
158        let mut ids = Vec::new();
159        for _ in 0..id_count {
160            ids.push(read_u64(bytes, &mut pos, "id")?);
161        }
162
163        let vector_count = read_u32(bytes, &mut pos, "vector_count")?;
164        let mut vectors = Vec::new();
165        for _ in 0..vector_count {
166            let vector_len = read_u32(bytes, &mut pos, "vector_len")?;
167            let mut vector = Vec::new();
168            for _ in 0..vector_len {
169                vector.push(read_f32(bytes, &mut pos, "vector value")?);
170            }
171            vectors.push(vector);
172        }
173
174        lists.push(IvfListFrame {
175            centroid,
176            ids,
177            vectors,
178        });
179    }
180
181    Ok(IvfIndexFrame {
182        n_lists,
183        n_probes,
184        dimension,
185        max_iterations,
186        convergence_threshold,
187        trained,
188        count,
189        next_id,
190        lists,
191    })
192}
193
194fn read_u8(bytes: &[u8], pos: &mut usize, reason: &'static str) -> Result<u8, IvfIndexFrameError> {
195    if *pos + 1 > bytes.len() {
196        return Err(IvfIndexFrameError::Truncated {
197            offset: *pos,
198            reason,
199        });
200    }
201    let value = bytes[*pos];
202    *pos += 1;
203    Ok(value)
204}
205
206fn read_u32(
207    bytes: &[u8],
208    pos: &mut usize,
209    reason: &'static str,
210) -> Result<u32, IvfIndexFrameError> {
211    if *pos + 4 > bytes.len() {
212        return Err(IvfIndexFrameError::Truncated {
213            offset: *pos,
214            reason,
215        });
216    }
217    let value = u32::from_le_bytes(
218        bytes[*pos..*pos + 4]
219            .try_into()
220            .expect("u32 length checked"),
221    );
222    *pos += 4;
223    Ok(value)
224}
225
226fn read_u64(
227    bytes: &[u8],
228    pos: &mut usize,
229    reason: &'static str,
230) -> Result<u64, IvfIndexFrameError> {
231    if *pos + 8 > bytes.len() {
232        return Err(IvfIndexFrameError::Truncated {
233            offset: *pos,
234            reason,
235        });
236    }
237    let value = u64::from_le_bytes(
238        bytes[*pos..*pos + 8]
239            .try_into()
240            .expect("u64 length checked"),
241    );
242    *pos += 8;
243    Ok(value)
244}
245
246fn read_f32(
247    bytes: &[u8],
248    pos: &mut usize,
249    reason: &'static str,
250) -> Result<f32, IvfIndexFrameError> {
251    if *pos + 4 > bytes.len() {
252        return Err(IvfIndexFrameError::Truncated {
253            offset: *pos,
254            reason,
255        });
256    }
257    let value = f32::from_le_bytes(
258        bytes[*pos..*pos + 4]
259            .try_into()
260            .expect("f32 length checked"),
261    );
262    *pos += 4;
263    Ok(value)
264}
265
266#[cfg(test)]
267mod tests {
268    use super::*;
269
270    fn sample_frame() -> IvfIndexFrame {
271        IvfIndexFrame {
272            n_lists: 4,
273            n_probes: 2,
274            dimension: 3,
275            max_iterations: 50,
276            convergence_threshold: 1e-4,
277            trained: true,
278            count: 3,
279            next_id: 9,
280            lists: vec![
281                IvfListFrame {
282                    centroid: vec![0.5, 0.5, 0.5],
283                    ids: vec![1, 2],
284                    vectors: vec![vec![0.4, 0.4, 0.4], vec![0.6, 0.6, 0.6]],
285                },
286                IvfListFrame {
287                    centroid: vec![9.0, 9.0, 9.0],
288                    ids: vec![8],
289                    vectors: vec![vec![9.1, 9.0, 8.9]],
290                },
291            ],
292        }
293    }
294
295    #[test]
296    fn ivf_index_frame_round_trips() {
297        let frame = sample_frame();
298        let encoded = encode_ivf_index_frame(&frame);
299        let decoded = decode_ivf_index_frame(&encoded).unwrap();
300        assert_eq!(decoded, frame);
301        assert_eq!(encode_ivf_index_frame(&decoded), encoded);
302    }
303
304    #[test]
305    fn ivf_index_frame_pins_byte_layout() {
306        let frame = sample_frame();
307        let encoded = encode_ivf_index_frame(&frame);
308        assert_eq!(&encoded[0..4], b"IVF1");
309        assert_eq!(&encoded[4..8], &4u32.to_le_bytes()); // n_lists
310                                                         // `trained` byte: after magic + 4 u32 + 1 f32.
311        let trained_off = 4 + 4 * 4 + 4;
312        assert_eq!(encoded[trained_off], 1);
313    }
314
315    #[test]
316    fn ivf_index_frame_rejects_bad_input() {
317        assert_eq!(
318            decode_ivf_index_frame(&[0u8; 8]),
319            Err(IvfIndexFrameError::TooShort)
320        );
321        let mut bad_magic = encode_ivf_index_frame(&sample_frame());
322        bad_magic[0] = b'X';
323        assert_eq!(
324            decode_ivf_index_frame(&bad_magic),
325            Err(IvfIndexFrameError::InvalidMagic)
326        );
327        let encoded = encode_ivf_index_frame(&sample_frame());
328        assert!(matches!(
329            decode_ivf_index_frame(&encoded[..encoded.len() - 1]),
330            Err(IvfIndexFrameError::Truncated { .. })
331        ));
332    }
333
334    #[test]
335    fn ivf_index_frame_does_not_preallocate_untrusted_counts() {
336        let mut frame = sample_frame();
337        frame.lists.clear();
338        let mut encoded = encode_ivf_index_frame(&frame);
339        let list_count_off = IVF_INDEX_HEADER_LEN - 4;
340        encoded[list_count_off..list_count_off + 4].copy_from_slice(&u32::MAX.to_le_bytes());
341
342        assert!(matches!(
343            decode_ivf_index_frame(&encoded),
344            Err(IvfIndexFrameError::Truncated { .. })
345        ));
346    }
347}
reddb_file/vector_ivf_index.rs

reddb_file/
vector_ivf_index.rs
