Skip to main content

nodedb_vector/collection/
codec_build.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Methods on `VectorCollection` for building the collection-level
4//! codec-dispatch index (RaBitQ, BBQ).
5//!
6//! All `impl VectorCollection` blocks in `collection/` extend the same type.
7
8use super::codec_dispatch::{CollectionCodec, build_collection_codec};
9use super::lifecycle::VectorCollection;
10
11impl VectorCollection {
12    /// Collect all live FP32 vectors from every segment (growing, building,
13    /// and sealed) in insertion order. Used to train the collection-level
14    /// codec-dispatch index.
15    pub(crate) fn gather_all_vectors_fp32(&self) -> Vec<Vec<f32>> {
16        let total = self.len();
17        let mut out = Vec::with_capacity(total);
18
19        for i in 0..self.growing.len() as u32 {
20            if let Some(v) = self.growing.get_vector(i) {
21                out.push(v.to_vec());
22            }
23        }
24
25        for seg in &self.building {
26            for i in 0..seg.flat.len() as u32 {
27                if let Some(v) = seg.flat.get_vector(i) {
28                    out.push(v.to_vec());
29                }
30            }
31        }
32
33        for seg in &self.sealed {
34            let n = seg.index.len();
35            for i in 0..n as u32 {
36                if !seg.index.is_deleted(i)
37                    && let Some(v) = seg.index.get_vector(i)
38                {
39                    out.push(v.to_vec());
40                }
41            }
42        }
43
44        out
45    }
46
47    /// Build a codec-dispatched index over all current vectors using the
48    /// requested quantization. Replaces any existing dispatch index for
49    /// this collection. Idempotent.
50    ///
51    /// Returns a reference to the new index, or `None` if the quantization
52    /// tag is not supported (falls back to per-segment Sq8/PQ paths) or there
53    /// are no vectors to train on.
54    pub fn build_codec_dispatch(&mut self, quantization: &str) -> Option<&CollectionCodec> {
55        let vectors = self.gather_all_vectors_fp32();
56        let dim = self.dim;
57        let m = self.params.m;
58        let ef_construction = self.params.ef_construction;
59        let seed = 42_u64;
60        self.codec_dispatch =
61            build_collection_codec(quantization, &vectors, dim, m, ef_construction, seed);
62        self.codec_dispatch.as_ref()
63    }
64}