reddb_file/
ivf_index_codec.rs

1//! On-disk codec for the `IVF1` vector-index payload.
2//!
3//! `reddb-file` owns the byte layout of the persisted IVF artifact: the `IVF1`
4//! magic, the config block, training state, and per-list centroid/vector data.
5//! The clustering algorithm stays in the server engine; only the format-faithful
6//! encode/decode lives here (ADR 0046).
7//!
8//! Layout (all integers little-endian, no version field):
9//! ```text
10//! magic   : b"IVF1"            (4 bytes)
11//! config  : n_lists u32, n_probes u32, dimension u32, max_iterations u32,
12//!           convergence_threshold f32
13//! state   : trained u8, count u64, next_id u64
14//! lists   : list_count u32, then per list:
15//!             centroid_len u32, centroid f32[centroid_len],
16//!             id_count u32, ids u64[id_count],
17//!             vector_count u32, then per vector: len u32, f32[len]
18//! ```
19
20/// Magic prefixing a persisted IVF index payload.
21pub const IVF_INDEX_MAGIC: &[u8; 4] = b"IVF1";
22
23/// Plain, engine-agnostic view of one persisted IVF inverted list.
24#[derive(Debug, Clone, PartialEq)]
25pub struct IvfListLayout {
26    /// Cluster centroid.
27    pub centroid: Vec<f32>,
28    /// Member vector ids.
29    pub ids: Vec<u64>,
30    /// Member vectors (parallel to `ids`).
31    pub vectors: Vec<Vec<f32>>,
32}
33
34/// Plain, engine-agnostic view of a persisted IVF index.
35#[derive(Debug, Clone, PartialEq)]
36pub struct IvfIndexLayout {
37    /// Number of inverted lists / clusters.
38    pub n_lists: usize,
39    /// Number of lists probed at query time.
40    pub n_probes: usize,
41    /// Vector dimension.
42    pub dimension: usize,
43    /// Max k-means iterations.
44    pub max_iterations: usize,
45    /// k-means convergence threshold.
46    pub convergence_threshold: f32,
47    /// Whether the index has been trained.
48    pub trained: bool,
49    /// Total stored vector count.
50    pub count: usize,
51    /// Next auto-generated id.
52    pub next_id: u64,
53    /// Inverted lists in persistence order.
54    pub lists: Vec<IvfListLayout>,
55}
56
57/// Errors raised while decoding a persisted IVF payload.
58#[derive(Debug, Clone, PartialEq, Eq)]
59pub enum IvfCodecError {
60    /// Fewer than the minimum header bytes were present.
61    TooShort,
62    /// The leading magic was not `IVF1`.
63    InvalidMagic,
64    /// A field ran past the end of the buffer.
65    Truncated,
66}
67
68impl std::fmt::Display for IvfCodecError {
69    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
70        match self {
71            IvfCodecError::TooShort => write!(f, "data too short"),
72            IvfCodecError::InvalidMagic => write!(f, "invalid IVF magic"),
73            IvfCodecError::Truncated => write!(f, "truncated IVF payload"),
74        }
75    }
76}
77
78impl std::error::Error for IvfCodecError {}
79
80/// Encode a persisted IVF index, byte-faithful to the legacy server layout.
81pub fn encode_ivf_index(layout: &IvfIndexLayout) -> Vec<u8> {
82    let mut bytes = Vec::new();
83    bytes.extend_from_slice(IVF_INDEX_MAGIC);
84    bytes.extend_from_slice(&(layout.n_lists as u32).to_le_bytes());
85    bytes.extend_from_slice(&(layout.n_probes as u32).to_le_bytes());
86    bytes.extend_from_slice(&(layout.dimension as u32).to_le_bytes());
87    bytes.extend_from_slice(&(layout.max_iterations as u32).to_le_bytes());
88    bytes.extend_from_slice(&layout.convergence_threshold.to_le_bytes());
89    bytes.push(if layout.trained { 1 } else { 0 });
90    bytes.extend_from_slice(&(layout.count as u64).to_le_bytes());
91    bytes.extend_from_slice(&layout.next_id.to_le_bytes());
92    bytes.extend_from_slice(&(layout.lists.len() as u32).to_le_bytes());
93
94    for list in &layout.lists {
95        bytes.extend_from_slice(&(list.centroid.len() as u32).to_le_bytes());
96        for value in &list.centroid {
97            bytes.extend_from_slice(&value.to_le_bytes());
98        }
99
100        bytes.extend_from_slice(&(list.ids.len() as u32).to_le_bytes());
101        for id in &list.ids {
102            bytes.extend_from_slice(&id.to_le_bytes());
103        }
104
105        bytes.extend_from_slice(&(list.vectors.len() as u32).to_le_bytes());
106        for vector in &list.vectors {
107            bytes.extend_from_slice(&(vector.len() as u32).to_le_bytes());
108            for value in vector {
109                bytes.extend_from_slice(&value.to_le_bytes());
110            }
111        }
112    }
113
114    bytes
115}
116
117fn read_u32(buf: &[u8], pos: &mut usize) -> Result<u32, IvfCodecError> {
118    if *pos + 4 > buf.len() {
119        return Err(IvfCodecError::Truncated);
120    }
121    let value = u32::from_le_bytes([buf[*pos], buf[*pos + 1], buf[*pos + 2], buf[*pos + 3]]);
122    *pos += 4;
123    Ok(value)
124}
125
126fn read_u64(buf: &[u8], pos: &mut usize) -> Result<u64, IvfCodecError> {
127    if *pos + 8 > buf.len() {
128        return Err(IvfCodecError::Truncated);
129    }
130    let value = u64::from_le_bytes([
131        buf[*pos],
132        buf[*pos + 1],
133        buf[*pos + 2],
134        buf[*pos + 3],
135        buf[*pos + 4],
136        buf[*pos + 5],
137        buf[*pos + 6],
138        buf[*pos + 7],
139    ]);
140    *pos += 8;
141    Ok(value)
142}
143
144fn read_f32(buf: &[u8], pos: &mut usize) -> Result<f32, IvfCodecError> {
145    if *pos + 4 > buf.len() {
146        return Err(IvfCodecError::Truncated);
147    }
148    let value = f32::from_le_bytes([buf[*pos], buf[*pos + 1], buf[*pos + 2], buf[*pos + 3]]);
149    *pos += 4;
150    Ok(value)
151}
152
153/// Decode a persisted IVF index produced by [`encode_ivf_index`] or the legacy
154/// server `to_bytes`.
155pub fn decode_ivf_index(bytes: &[u8]) -> Result<IvfIndexLayout, IvfCodecError> {
156    if bytes.len() < 41 {
157        return Err(IvfCodecError::TooShort);
158    }
159    if &bytes[0..4] != IVF_INDEX_MAGIC {
160        return Err(IvfCodecError::InvalidMagic);
161    }
162
163    let mut pos = 4usize;
164    let n_lists = read_u32(bytes, &mut pos)? as usize;
165    let n_probes = read_u32(bytes, &mut pos)? as usize;
166    let dimension = read_u32(bytes, &mut pos)? as usize;
167    let max_iterations = read_u32(bytes, &mut pos)? as usize;
168    let convergence_threshold = read_f32(bytes, &mut pos)?;
169
170    if pos >= bytes.len() {
171        return Err(IvfCodecError::Truncated);
172    }
173    let trained = bytes[pos] == 1;
174    pos += 1;
175    let count = read_u64(bytes, &mut pos)? as usize;
176    let next_id = read_u64(bytes, &mut pos)?;
177    let list_count = read_u32(bytes, &mut pos)? as usize;
178
179    let mut lists = Vec::with_capacity(list_count);
180    for _ in 0..list_count {
181        let centroid_len = read_u32(bytes, &mut pos)? as usize;
182        let mut centroid = Vec::with_capacity(centroid_len);
183        for _ in 0..centroid_len {
184            centroid.push(read_f32(bytes, &mut pos)?);
185        }
186
187        let id_count = read_u32(bytes, &mut pos)? as usize;
188        let mut ids = Vec::with_capacity(id_count);
189        for _ in 0..id_count {
190            ids.push(read_u64(bytes, &mut pos)?);
191        }
192
193        let vector_count = read_u32(bytes, &mut pos)? as usize;
194        let mut vectors = Vec::with_capacity(vector_count);
195        for _ in 0..vector_count {
196            let vector_len = read_u32(bytes, &mut pos)? as usize;
197            let mut vector = Vec::with_capacity(vector_len);
198            for _ in 0..vector_len {
199                vector.push(read_f32(bytes, &mut pos)?);
200            }
201            vectors.push(vector);
202        }
203
204        lists.push(IvfListLayout {
205            centroid,
206            ids,
207            vectors,
208        });
209    }
210
211    Ok(IvfIndexLayout {
212        n_lists,
213        n_probes,
214        dimension,
215        max_iterations,
216        convergence_threshold,
217        trained,
218        count,
219        next_id,
220        lists,
221    })
222}
223
224#[cfg(test)]
225mod tests {
226    use super::*;
227
228    fn sample() -> IvfIndexLayout {
229        IvfIndexLayout {
230            n_lists: 4,
231            n_probes: 2,
232            dimension: 3,
233            max_iterations: 50,
234            convergence_threshold: 1e-4,
235            trained: true,
236            count: 3,
237            next_id: 3,
238            lists: vec![
239                IvfListLayout {
240                    centroid: vec![0.0, 1.0, 2.0],
241                    ids: vec![0, 2],
242                    vectors: vec![vec![0.0, 1.0, 2.0], vec![0.1, 1.1, 2.1]],
243                },
244                IvfListLayout {
245                    centroid: vec![9.0, 9.0, 9.0],
246                    ids: vec![1],
247                    vectors: vec![vec![9.0, 9.0, 9.0]],
248                },
249            ],
250        }
251    }
252
253    #[test]
254    fn round_trip_preserves_layout() {
255        let layout = sample();
256        let bytes = encode_ivf_index(&layout);
257        let decoded = decode_ivf_index(&bytes).expect("decode");
258        assert_eq!(decoded, layout);
259    }
260
261    #[test]
262    fn fixture_bytes_are_byte_identical() {
263        let layout = sample();
264        let bytes = encode_ivf_index(&layout);
265        assert_eq!(&bytes[0..4], b"IVF1", "magic must lead the payload");
266        // n_lists u32 directly after the magic (no version field).
267        assert_eq!(&bytes[4..8], &4u32.to_le_bytes());
268        // trained byte sits at offset 4 + 4*4 + 4 = 24.
269        assert_eq!(bytes[24], 1);
270    }
271
272    #[test]
273    fn rejects_short_and_bad_magic() {
274        assert_eq!(decode_ivf_index(&[0u8; 10]), Err(IvfCodecError::TooShort));
275        let mut bytes = encode_ivf_index(&sample());
276        bytes[0] = b'X';
277        assert_eq!(decode_ivf_index(&bytes), Err(IvfCodecError::InvalidMagic));
278    }
279}
reddb_file/ivf_index_codec.rs

reddb_file/
ivf_index_codec.rs
