nodedb_vector/collection/
codec_dispatch.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Per-collection codec selection. Wraps the generic `HnswCodecIndex<C>`
4//! for codecs other than Sq8 (which retains its specialised fast path
5//! in `quantize.rs` / `search.rs`).
6
7use nodedb_codec::vector_quant::bbq::BbqCodec;
8use nodedb_codec::vector_quant::rabitq::RaBitQCodec;
9
10use crate::codec_index::HnswCodecIndex;
11use crate::error::VectorError;
12
13/// One built codec-index per collection (other than Sq8). Variants match
14/// the publicly-selectable quantization choices that route through
15/// `HnswCodecIndex`.
16#[non_exhaustive]
17pub enum CollectionCodec {
18    RaBitQ(HnswCodecIndex<RaBitQCodec>),
19    Bbq(HnswCodecIndex<BbqCodec>),
20}
21
22impl CollectionCodec {
23    /// Forwarding `search` so the collection layer doesn't have to match
24    /// on the variant for the common case.
25    pub fn search(
26        &self,
27        query: &[f32],
28        k: usize,
29        ef_search: usize,
30    ) -> Result<Vec<crate::codec_index::CodecSearchResult>, VectorError> {
31        match self {
32            Self::RaBitQ(idx) => idx.search(query, k, ef_search),
33            Self::Bbq(idx) => idx.search(query, k, ef_search),
34        }
35    }
36
37    /// Forwarding `insert`.
38    pub fn insert(&mut self, id: u32, v: &[f32]) -> Result<(), VectorError> {
39        match self {
40            Self::RaBitQ(idx) => idx.insert(id, v),
41            Self::Bbq(idx) => idx.insert(id, v),
42        }
43    }
44
45    /// Total nodes (including deleted).
46    pub fn len(&self) -> usize {
47        match self {
48            Self::RaBitQ(idx) => idx.len(),
49            Self::Bbq(idx) => idx.len(),
50        }
51    }
52
53    pub fn is_empty(&self) -> bool {
54        self.len() == 0
55    }
56
57    /// Quantization tag for stats reporting.
58    pub fn quantization(&self) -> &'static str {
59        match self {
60            Self::RaBitQ(_) => "rabitq",
61            Self::Bbq(_) => "bbq",
62        }
63    }
64}
65
66/// Build a `CollectionCodec` from a quantization tag and training vectors.
67///
68/// Returns `Ok(None)` for unsupported or unrecognised tags (e.g. "sq8",
69/// "pq", "none") — those variants use separate per-segment code paths — and
70/// for an empty vector set. A vector without `dim` components fails with
71/// [`VectorError::DimensionMismatch`].
72///
73/// Each entry is `(id, vector)`; the index reports `id` in its results, so
74/// it must be the collection's global vector id.
75pub fn build_collection_codec(
76    quantization: &str,
77    vectors: &[(u32, Vec<f32>)],
78    dim: usize,
79    m: usize,
80    ef_construction: usize,
81    seed: u64,
82) -> Result<Option<CollectionCodec>, VectorError> {
83    if vectors.is_empty() {
84        return Ok(None);
85    }
86    // Calibration reads every vector as `dim` components; check before it.
87    for (_, v) in vectors {
88        crate::error::check_dim(dim, v.len())?;
89    }
90    let refs: Vec<&[f32]> = vectors.iter().map(|(_, v)| v.as_slice()).collect();
91    match quantization {
92        "rabitq" => {
93            let codec = RaBitQCodec::calibrate(&refs, dim, seed);
94            let mut idx = HnswCodecIndex::new(dim, m, ef_construction, codec, seed);
95            for (id, v) in vectors {
96                idx.insert(*id, v)?;
97            }
98            Ok(Some(CollectionCodec::RaBitQ(idx)))
99        }
100        "bbq" => {
101            let codec = BbqCodec::calibrate(&refs, dim, 3);
102            let mut idx = HnswCodecIndex::new(dim, m, ef_construction, codec, seed);
103            for (id, v) in vectors {
104                idx.insert(*id, v)?;
105            }
106            Ok(Some(CollectionCodec::Bbq(idx)))
107        }
108        _ => Ok(None),
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115
116    fn make_vectors(n: usize, dim: usize) -> Vec<(u32, Vec<f32>)> {
117        (0..n)
118            .map(|i| {
119                let v = (0..dim).map(|d| (i * dim + d) as f32 * 0.01).collect();
120                (i as u32, v)
121            })
122            .collect()
123    }
124
125    #[test]
126    fn build_rabitq_returns_some() {
127        let vecs = make_vectors(50, 8);
128        let result = build_collection_codec("rabitq", &vecs, 8, 16, 100, 42).unwrap();
129        assert!(
130            matches!(result, Some(CollectionCodec::RaBitQ(_))),
131            "expected RaBitQ variant"
132        );
133    }
134
135    #[test]
136    fn build_bbq_returns_some() {
137        let vecs = make_vectors(50, 8);
138        let result = build_collection_codec("bbq", &vecs, 8, 16, 100, 42).unwrap();
139        assert!(
140            matches!(result, Some(CollectionCodec::Bbq(_))),
141            "expected Bbq variant"
142        );
143    }
144
145    #[test]
146    fn unknown_codec_returns_none() {
147        let vecs = make_vectors(50, 8);
148        let result = build_collection_codec("unknown_codec", &vecs, 8, 16, 100, 42).unwrap();
149        assert!(result.is_none(), "unknown codec should return None");
150    }
151
152    #[test]
153    fn sq8_tag_returns_none() {
154        let vecs = make_vectors(50, 8);
155        let result = build_collection_codec("sq8", &vecs, 8, 16, 100, 42).unwrap();
156        assert!(
157            result.is_none(),
158            "sq8 tag should fall through to per-segment path"
159        );
160    }
161
162    #[test]
163    fn empty_vectors_returns_none() {
164        let result = build_collection_codec("rabitq", &[], 8, 16, 100, 42).unwrap();
165        assert!(result.is_none(), "empty vectors should return None");
166    }
167
168    #[test]
169    fn len_and_is_empty() {
170        let vecs = make_vectors(20, 4);
171        let codec = build_collection_codec("bbq", &vecs, 4, 8, 50, 1)
172            .unwrap()
173            .unwrap();
174        assert_eq!(codec.len(), 20);
175        assert!(!codec.is_empty());
176    }
177
178    #[test]
179    fn quantization_tag() {
180        let vecs = make_vectors(10, 4);
181        let rabitq = build_collection_codec("rabitq", &vecs, 4, 8, 50, 1)
182            .unwrap()
183            .unwrap();
184        assert_eq!(rabitq.quantization(), "rabitq");
185        let bbq = build_collection_codec("bbq", &vecs, 4, 8, 50, 1)
186            .unwrap()
187            .unwrap();
188        assert_eq!(bbq.quantization(), "bbq");
189    }
190
191    #[test]
192    fn wrong_dimension_query_is_a_typed_error() {
193        let vecs = make_vectors(20, 4);
194        for tag in ["rabitq", "bbq"] {
195            let mut codec = build_collection_codec(tag, &vecs, 4, 8, 50, 1)
196                .unwrap()
197                .unwrap();
198            assert!(matches!(
199                codec.search(&[0.0; 3], 5, 20),
200                Err(VectorError::DimensionMismatch {
201                    expected: 4,
202                    got: 3
203                })
204            ));
205            assert!(matches!(
206                codec.insert(99, &[0.0; 5]),
207                Err(VectorError::DimensionMismatch {
208                    expected: 4,
209                    got: 5
210                })
211            ));
212        }
213        let short = vec![(0_u32, vec![0.0_f32; 3])];
214        assert!(matches!(
215            build_collection_codec("rabitq", &short, 4, 8, 50, 1),
216            Err(VectorError::DimensionMismatch {
217                expected: 4,
218                got: 3
219            })
220        ));
221    }
222}
nodedb_vector/collection/codec_dispatch.rs

nodedb_vector/collection/
codec_dispatch.rs
