nodedb_vector/rerank/
sidecar.rs

1// SPDX-License-Identifier: Apache-2.0
2
3use std::collections::HashMap;
4use std::sync::Arc;
5
6use super::codec::{CodecName, PreparedQuery, RerankCodec};
7use super::types::RerankError;
8
9const SIDECAR_MAGIC: [u8; 4] = *b"NDCC";
10const SIDECAR_VERSION: u8 = 1;
11
12/// Per-collection encoded-vector storage keyed by surrogate id, paired with
13/// the trained codec. Encoded vectors live alongside (not inside) the HNSW
14/// index — HNSW keeps full-precision vectors for graph traversal; the sidecar
15/// is consulted only during base-layer rerank.
16pub struct CodecSidecar {
17    codec: Arc<dyn RerankCodec>,
18    encoded: HashMap<u32, Vec<u8>>,
19}
20
21impl CodecSidecar {
22    pub fn new(codec: Arc<dyn RerankCodec>) -> Self {
23        Self {
24            codec,
25            encoded: HashMap::new(),
26        }
27    }
28
29    pub fn codec_name(&self) -> CodecName {
30        self.codec.name()
31    }
32
33    /// Encode a vector and insert it under `id`. Overwrites any existing entry.
34    pub fn encode_and_insert(&mut self, id: u32, vector: &[f32]) -> Result<(), RerankError> {
35        let bytes = self.codec.encode(vector)?;
36        self.encoded.insert(id, bytes);
37        Ok(())
38    }
39
40    pub fn remove(&mut self, id: u32) {
41        self.encoded.remove(&id);
42    }
43
44    pub fn get(&self, id: u32) -> Option<&[u8]> {
45        self.encoded.get(&id).map(|v| v.as_slice())
46    }
47
48    pub fn len(&self) -> usize {
49        self.encoded.len()
50    }
51
52    pub fn is_empty(&self) -> bool {
53        self.encoded.is_empty()
54    }
55
56    /// Serialize the sidecar (codec state + all encoded vectors) to bytes.
57    ///
58    /// Format: `[NDCC (4 bytes)][version: u8 = 1][msgpack payload]`
59    pub fn to_bytes(&self) -> Result<Vec<u8>, RerankError> {
60        let codec_bytes = self.codec.to_bytes()?;
61        let codec_name_byte = codec_name_to_u8(self.codec.name());
62
63        #[derive(
64            serde::Serialize, serde::Deserialize, zerompk::ToMessagePack, zerompk::FromMessagePack,
65        )]
66        struct Payload {
67            codec_name: u8,
68            codec_bytes: Vec<u8>,
69            encoded: Vec<(u32, Vec<u8>)>,
70        }
71
72        let payload = Payload {
73            codec_name: codec_name_byte,
74            codec_bytes,
75            encoded: self.encoded.iter().map(|(k, v)| (*k, v.clone())).collect(),
76        };
77
78        let body = zerompk::to_msgpack_vec(&payload)
79            .map_err(|e| RerankError::BadInput(format!("sidecar serialize: {e}")))?;
80        let mut buf = Vec::with_capacity(5 + body.len());
81        buf.extend_from_slice(&SIDECAR_MAGIC);
82        buf.push(SIDECAR_VERSION);
83        buf.extend_from_slice(&body);
84        Ok(buf)
85    }
86
87    /// Deserialize a sidecar from bytes produced by [`Self::to_bytes`].
88    pub fn from_bytes(bytes: &[u8], memory: nodedb_mem::ScopedMemory) -> Result<Self, RerankError> {
89        if bytes.len() < 5 {
90            return Err(RerankError::BadInput(
91                "sidecar from_bytes: too short".into(),
92            ));
93        }
94        if bytes[..4] != SIDECAR_MAGIC {
95            return Err(RerankError::BadInput(
96                "sidecar from_bytes: bad magic".into(),
97            ));
98        }
99        let version = bytes[4];
100        if version != SIDECAR_VERSION {
101            return Err(RerankError::BadInput(format!(
102                "sidecar from_bytes: unknown version {version}"
103            )));
104        }
105
106        #[derive(
107            serde::Serialize, serde::Deserialize, zerompk::ToMessagePack, zerompk::FromMessagePack,
108        )]
109        struct Payload {
110            codec_name: u8,
111            codec_bytes: Vec<u8>,
112            encoded: Vec<(u32, Vec<u8>)>,
113        }
114
115        let payload: Payload = zerompk::from_msgpack(&bytes[5..])
116            .map_err(|e| RerankError::BadInput(format!("sidecar deserialize: {e}")))?;
117
118        let codec_name = codec_name_from_u8(payload.codec_name).ok_or_else(|| {
119            RerankError::BadInput(format!(
120                "sidecar from_bytes: unknown codec_name byte {}",
121                payload.codec_name
122            ))
123        })?;
124        let codec =
125            super::codec::rerank_codec_from_bytes(codec_name, &payload.codec_bytes, memory)?;
126        let encoded = payload.encoded.into_iter().collect();
127        Ok(CodecSidecar { codec, encoded })
128    }
129
130    pub fn prepare_query(&self, query: &[f32]) -> Result<PreparedQuery, RerankError> {
131        self.codec.prepare_query(query)
132    }
133
134    /// Compute distance from a prepared query to the encoded vector at `id`.
135    /// Returns `Ok(None)` when the id isn't in the sidecar (lost / not yet
136    /// encoded); returns the distance otherwise.
137    pub fn distance_prepared(
138        &self,
139        prepared: &PreparedQuery,
140        id: u32,
141    ) -> Result<Option<f32>, RerankError> {
142        match self.encoded.get(&id) {
143            None => Ok(None),
144            Some(bytes) => self.codec.distance_prepared(prepared, bytes).map(Some),
145        }
146    }
147}
148
149fn codec_name_to_u8(name: CodecName) -> u8 {
150    match name {
151        CodecName::Sq8 => 0,
152        CodecName::Pq => 1,
153        CodecName::Binary => 2,
154        CodecName::RaBitQ => 3,
155        CodecName::Bbq => 4,
156    }
157}
158
159fn codec_name_from_u8(b: u8) -> Option<CodecName> {
160    match b {
161        0 => Some(CodecName::Sq8),
162        1 => Some(CodecName::Pq),
163        2 => Some(CodecName::Binary),
164        3 => Some(CodecName::RaBitQ),
165        4 => Some(CodecName::Bbq),
166        _ => None,
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173    use crate::rerank::codec::CodecName;
174    use crate::test_support::test_memory;
175
176    struct StubCodec;
177
178    impl RerankCodec for StubCodec {
179        fn encode(&self, v: &[f32]) -> Result<Vec<u8>, RerankError> {
180            Ok(v.iter().flat_map(|x| x.to_le_bytes()).collect())
181        }
182
183        fn prepare_query(&self, q: &[f32]) -> Result<PreparedQuery, RerankError> {
184            Ok(PreparedQuery::Raw(q.to_vec()))
185        }
186
187        fn distance_prepared(
188            &self,
189            prepared: &PreparedQuery,
190            encoded: &[u8],
191        ) -> Result<f32, RerankError> {
192            let query = match prepared {
193                PreparedQuery::Raw(v) => v,
194                _ => {
195                    return Err(RerankError::BadInput(
196                        "StubCodec expects Raw prepared query".into(),
197                    ));
198                }
199            };
200            let encoded_floats: Vec<f32> = encoded
201                .as_chunks::<4>()
202                .0
203                .iter()
204                .map(|b| f32::from_le_bytes(*b))
205                .collect();
206            if query.len() != encoded_floats.len() {
207                return Err(RerankError::BadInput("dimension mismatch".into()));
208            }
209            let dist = query
210                .iter()
211                .zip(encoded_floats.iter())
212                .map(|(a, b)| (a - b) * (a - b))
213                .sum::<f32>()
214                .sqrt();
215            Ok(dist)
216        }
217
218        fn name(&self) -> CodecName {
219            CodecName::Binary
220        }
221
222        fn to_bytes(&self) -> Result<Vec<u8>, RerankError> {
223            Err(RerankError::BadInput(
224                "StubCodec does not support serialization".into(),
225            ))
226        }
227    }
228
229    fn make_sidecar() -> CodecSidecar {
230        CodecSidecar::new(Arc::new(StubCodec))
231    }
232
233    #[test]
234    fn insert_and_get() {
235        let mut s = make_sidecar();
236        assert!(s.is_empty());
237        s.encode_and_insert(1, &[1.0, 2.0]).unwrap();
238        s.encode_and_insert(2, &[3.0, 4.0]).unwrap();
239        s.encode_and_insert(3, &[5.0, 6.0]).unwrap();
240        assert_eq!(s.len(), 3);
241
242        let expected_1: Vec<u8> = [1.0f32, 2.0f32]
243            .iter()
244            .flat_map(|x| x.to_le_bytes())
245            .collect();
246        assert_eq!(s.get(1), Some(expected_1.as_slice()));
247    }
248
249    #[test]
250    fn remove_returns_none() {
251        let mut s = make_sidecar();
252        s.encode_and_insert(10, &[1.0]).unwrap();
253        s.remove(10);
254        assert_eq!(s.get(10), None);
255        let prepared = s.prepare_query(&[1.0]).unwrap();
256        assert_eq!(s.distance_prepared(&prepared, 10).unwrap(), None);
257    }
258
259    #[test]
260    fn distance_prepared_correct() {
261        let mut s = make_sidecar();
262        s.encode_and_insert(5, &[0.0, 0.0]).unwrap();
263        let prepared = s.prepare_query(&[3.0, 4.0]).unwrap();
264        let dist = s.distance_prepared(&prepared, 5).unwrap().unwrap();
265        assert!((dist - 5.0).abs() < 1e-5, "expected L2=5.0, got {dist}");
266    }
267
268    #[test]
269    fn codec_name_passthrough() {
270        let s = make_sidecar();
271        assert_eq!(s.codec_name(), CodecName::Binary);
272        assert_eq!(s.codec_name().as_str(), "binary");
273    }
274
275    #[test]
276    fn len_and_is_empty() {
277        let mut s = make_sidecar();
278        assert!(s.is_empty());
279        s.encode_and_insert(1, &[1.0]).unwrap();
280        assert!(!s.is_empty());
281        assert_eq!(s.len(), 1);
282        s.remove(1);
283        assert!(s.is_empty());
284    }
285
286    // ── Sidecar serialization tests ────────────────────────────────────────────
287
288    fn det_vec(i: usize, dim: usize) -> Vec<f32> {
289        (0..dim)
290            .map(|j| ((i * 31 + j) % 100) as f32 / 100.0)
291            .collect()
292    }
293
294    #[test]
295    fn sidecar_roundtrip_sq8() {
296        use crate::rerank::codecs::Sq8Rerank;
297        let dim = 16;
298        let mut codec = Sq8Rerank::new(dim);
299        let samples: Vec<Vec<f32>> = (0..20).map(|i| det_vec(i, dim)).collect();
300        let refs: Vec<&[f32]> = samples.iter().map(|v| v.as_slice()).collect();
301        codec.train(&refs).unwrap();
302
303        let mut s = CodecSidecar::new(Arc::new(codec));
304        for i in 0..5u32 {
305            s.encode_and_insert(i, &det_vec(i as usize, dim)).unwrap();
306        }
307        let bytes = s.to_bytes().expect("to_bytes");
308        let s2 = CodecSidecar::from_bytes(&bytes, test_memory()).expect("from_bytes");
309        assert_eq!(s2.codec_name(), CodecName::Sq8);
310        for i in 0..5u32 {
311            assert_eq!(s.get(i), s2.get(i), "encoded bytes differ for id {i}");
312        }
313    }
314
315    #[test]
316    fn sidecar_roundtrip_binary() {
317        use crate::rerank::codecs::BinaryRerank;
318        let dim = 16;
319        let mut s = CodecSidecar::new(Arc::new(BinaryRerank::new(dim)));
320        for i in 0..5u32 {
321            s.encode_and_insert(i, &det_vec(i as usize, dim)).unwrap();
322        }
323        let bytes = s.to_bytes().expect("to_bytes");
324        let s2 = CodecSidecar::from_bytes(&bytes, test_memory()).expect("from_bytes");
325        assert_eq!(s2.codec_name(), CodecName::Binary);
326        for i in 0..5u32 {
327            assert_eq!(s.get(i), s2.get(i), "encoded bytes differ for id {i}");
328        }
329    }
330
331    #[test]
332    fn sidecar_roundtrip_pq() {
333        use crate::rerank::codecs::PqRerank;
334        let dim = 16;
335        let m = 4;
336        let k = 8;
337        let mut codec = PqRerank::new(dim, m, k, test_memory());
338        let samples: Vec<Vec<f32>> = (0..32).map(|i| det_vec(i, dim)).collect();
339        let refs: Vec<&[f32]> = samples.iter().map(|v| v.as_slice()).collect();
340        codec.train(&refs).unwrap();
341
342        let mut s = CodecSidecar::new(Arc::new(codec));
343        for i in 0..5u32 {
344            s.encode_and_insert(i, &det_vec(i as usize, dim)).unwrap();
345        }
346        let bytes = s.to_bytes().expect("to_bytes");
347        let s2 = CodecSidecar::from_bytes(&bytes, test_memory()).expect("from_bytes");
348        assert_eq!(s2.codec_name(), CodecName::Pq);
349        for i in 0..5u32 {
350            assert_eq!(s.get(i), s2.get(i), "encoded bytes differ for id {i}");
351        }
352    }
353
354    #[test]
355    fn sidecar_roundtrip_rabitq() {
356        use crate::rerank::codecs::RaBitQRerank;
357        let dim = 16;
358        let mut codec = RaBitQRerank::new(dim, 0xDEADBEEF_C0FFEE42);
359        let samples: Vec<Vec<f32>> = (0..20).map(|i| det_vec(i, dim)).collect();
360        let refs: Vec<&[f32]> = samples.iter().map(|v| v.as_slice()).collect();
361        codec.train(&refs).unwrap();
362
363        let mut s = CodecSidecar::new(Arc::new(codec));
364        for i in 0..5u32 {
365            s.encode_and_insert(i, &det_vec(i as usize, dim)).unwrap();
366        }
367        let bytes = s.to_bytes().expect("to_bytes");
368        let s2 = CodecSidecar::from_bytes(&bytes, test_memory()).expect("from_bytes");
369        assert_eq!(s2.codec_name(), CodecName::RaBitQ);
370        for i in 0..5u32 {
371            assert_eq!(s.get(i), s2.get(i), "encoded bytes differ for id {i}");
372        }
373    }
374
375    #[test]
376    fn sidecar_roundtrip_bbq() {
377        use crate::rerank::codecs::BbqRerank;
378        let dim = 16;
379        let mut codec = BbqRerank::new(dim, 4);
380        let samples: Vec<Vec<f32>> = (0..20).map(|i| det_vec(i, dim)).collect();
381        let refs: Vec<&[f32]> = samples.iter().map(|v| v.as_slice()).collect();
382        codec.train(&refs).unwrap();
383
384        let mut s = CodecSidecar::new(Arc::new(codec));
385        for i in 0..5u32 {
386            s.encode_and_insert(i, &det_vec(i as usize, dim)).unwrap();
387        }
388        let bytes = s.to_bytes().expect("to_bytes");
389        let s2 = CodecSidecar::from_bytes(&bytes, test_memory()).expect("from_bytes");
390        assert_eq!(s2.codec_name(), CodecName::Bbq);
391        for i in 0..5u32 {
392            assert_eq!(s.get(i), s2.get(i), "encoded bytes differ for id {i}");
393        }
394    }
395
396    #[test]
397    fn sidecar_bad_magic_returns_error() {
398        use crate::rerank::codecs::BinaryRerank;
399        let s = CodecSidecar::new(Arc::new(BinaryRerank::new(4)));
400        let mut bytes = s.to_bytes().unwrap();
401        bytes[0] = b'X';
402        assert!(CodecSidecar::from_bytes(&bytes, test_memory()).is_err());
403    }
404
405    #[test]
406    fn sidecar_bad_version_returns_error() {
407        use crate::rerank::codecs::BinaryRerank;
408        let s = CodecSidecar::new(Arc::new(BinaryRerank::new(4)));
409        let mut bytes = s.to_bytes().unwrap();
410        bytes[4] = 99;
411        assert!(CodecSidecar::from_bytes(&bytes, test_memory()).is_err());
412    }
413
414    #[test]
415    fn sidecar_distance_matches_after_roundtrip() {
416        use crate::rerank::codecs::Sq8Rerank;
417        let dim = 16;
418        let mut codec = Sq8Rerank::new(dim);
419        let samples: Vec<Vec<f32>> = (0..20).map(|i| det_vec(i, dim)).collect();
420        let refs: Vec<&[f32]> = samples.iter().map(|v| v.as_slice()).collect();
421        codec.train(&refs).unwrap();
422
423        let mut s = CodecSidecar::new(Arc::new(codec));
424        s.encode_and_insert(1, &det_vec(3, dim)).unwrap();
425        let query_vec = det_vec(7, dim);
426        let prepared_orig = s.prepare_query(&query_vec).unwrap();
427        let d_orig = s.distance_prepared(&prepared_orig, 1).unwrap().unwrap();
428
429        let bytes = s.to_bytes().unwrap();
430        let s2 = CodecSidecar::from_bytes(&bytes, test_memory()).unwrap();
431        let prepared_rest = s2.prepare_query(&query_vec).unwrap();
432        let d_rest = s2.distance_prepared(&prepared_rest, 1).unwrap().unwrap();
433
434        assert!(
435            (d_orig - d_rest).abs() < 1e-5,
436            "distance diverged: {d_orig} vs {d_rest}"
437        );
438    }
439}
nodedb_vector/rerank/sidecar.rs

nodedb_vector/rerank/
sidecar.rs
