nodedb_vector/collection/
quantize.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Quantizer training helpers for `VectorCollection`.
4//!
5//! All methods here are `impl VectorCollection` blocks — Rust allows a
6//! type's impl to be split across files.
7
8use crate::hnsw::{HnswIndex, HnswParams};
9use crate::index_config::{IndexConfig, IndexType};
10use crate::quantize::pq::PqCodec;
11use crate::quantize::sq8::Sq8Codec;
12
13use super::lifecycle::VectorCollection;
14use super::segment::DEFAULT_SEAL_THRESHOLD;
15
16impl VectorCollection {
17    /// Convenience constructor for PQ-configured collections.
18    ///
19    /// Equivalent to building a full `IndexConfig` with
20    /// `index_type = HnswPq` and the given `pq_m`.
21    pub fn with_pq_config(dim: usize, hnsw: HnswParams, pq_m: usize) -> Self {
22        let config = IndexConfig {
23            hnsw,
24            index_type: IndexType::HnswPq,
25            pq_m,
26            ..IndexConfig::default()
27        };
28        Self::with_index_config(dim, config)
29    }
30
31    /// Convenience constructor for PQ-configured collections with a custom
32    /// seal threshold.
33    pub fn with_seal_threshold_and_pq_config(
34        dim: usize,
35        hnsw: HnswParams,
36        pq_m: usize,
37        seal_threshold: usize,
38    ) -> Self {
39        let config = IndexConfig {
40            hnsw,
41            index_type: IndexType::HnswPq,
42            pq_m,
43            ..IndexConfig::default()
44        };
45        Self::with_seal_threshold_and_config(dim, config, seal_threshold)
46    }
47
48    /// Build SQ8 quantized data for an HNSW index.
49    ///
50    /// Returns `None` when there are too few live vectors for stable
51    /// min/max calibration.
52    pub fn build_sq8_for_index(index: &HnswIndex) -> Option<(Sq8Codec, Vec<u8>)> {
53        if index.live_count() < 1000 {
54            return None;
55        }
56        let dim = index.dim();
57        let n = index.len();
58
59        // no-governor: cold SQ8 training; slices borrow already-live vector data, no new heap
60        let mut refs: Vec<&[f32]> = Vec::with_capacity(n);
61        for i in 0..n {
62            if !index.is_deleted(i as u32)
63                && let Some(v) = index.get_vector(i as u32)
64            {
65                refs.push(v);
66            }
67        }
68        if refs.is_empty() {
69            return None;
70        }
71
72        let codec = Sq8Codec::calibrate(&refs, dim);
73
74        // no-governor: cold SQ8 quantize-all; governed at segment build call site
75        let mut data = Vec::with_capacity(dim * n);
76        for i in 0..n {
77            if let Some(v) = index.get_vector(i as u32) {
78                data.extend(codec.quantize(v));
79            } else {
80                data.extend(vec![0u8; dim]);
81            }
82        }
83
84        Some((codec, data))
85    }
86
87    /// Train a PQ codec from a built HNSW index's live vectors.
88    pub fn build_pq_for_index(index: &HnswIndex, pq_m: usize) -> Option<(PqCodec, Vec<u8>)> {
89        let dim = index.dim();
90        if pq_m == 0 || !dim.is_multiple_of(pq_m) {
91            return None;
92        }
93        let n = index.len();
94        // no-governor: cold PQ training; vectors copied for k-means, governed at segment build call site
95        let mut refs: Vec<Vec<f32>> = Vec::with_capacity(n);
96        for i in 0..n {
97            if !index.is_deleted(i as u32)
98                && let Some(v) = index.get_vector(i as u32)
99            {
100                refs.push(v.to_vec());
101            }
102        }
103        if refs.is_empty() {
104            return None;
105        }
106        let refs_slices: Vec<&[f32]> = refs.iter().map(|v| v.as_slice()).collect();
107        let k = 256usize.min(refs.len());
108        let codec = PqCodec::train(&refs_slices, dim, pq_m, k, 20);
109        let codes = codec.encode_batch(&refs_slices).ok()?;
110        Some((codec, codes))
111    }
112}
113
114// Keep the DEFAULT_SEAL_THRESHOLD import live when future refactors move
115// additional ctors into this file; explicitly referenced to suppress
116// an otherwise-unused warning.
117const _: usize = DEFAULT_SEAL_THRESHOLD;
nodedb_vector/collection/quantize.rs

nodedb_vector/collection/
quantize.rs
