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        let mut refs: Vec<&[f32]> = Vec::with_capacity(n);
60        for i in 0..n {
61            if !index.is_deleted(i as u32)
62                && let Some(v) = index.get_vector(i as u32)
63            {
64                refs.push(v);
65            }
66        }
67        if refs.is_empty() {
68            return None;
69        }
70
71        let codec = Sq8Codec::calibrate(&refs, dim);
72
73        let mut data = Vec::with_capacity(dim * n);
74        for i in 0..n {
75            if let Some(v) = index.get_vector(i as u32) {
76                data.extend(codec.quantize(v));
77            } else {
78                data.extend(vec![0u8; dim]);
79            }
80        }
81
82        Some((codec, data))
83    }
84
85    /// Train a PQ codec from a built HNSW index's live vectors.
86    pub fn build_pq_for_index(index: &HnswIndex, pq_m: usize) -> Option<(PqCodec, Vec<u8>)> {
87        let dim = index.dim();
88        if pq_m == 0 || !dim.is_multiple_of(pq_m) {
89            return None;
90        }
91        let n = index.len();
92        let mut refs: Vec<Vec<f32>> = Vec::with_capacity(n);
93        for i in 0..n {
94            if !index.is_deleted(i as u32)
95                && let Some(v) = index.get_vector(i as u32)
96            {
97                refs.push(v.to_vec());
98            }
99        }
100        if refs.is_empty() {
101            return None;
102        }
103        let refs_slices: Vec<&[f32]> = refs.iter().map(|v| v.as_slice()).collect();
104        let k = 256usize.min(refs.len());
105        let codec = PqCodec::train(&refs_slices, dim, pq_m, k, 20);
106        let codes = codec.encode_batch(&refs_slices).ok()?;
107        Some((codec, codes))
108    }
109}
110
111// Keep the DEFAULT_SEAL_THRESHOLD import live when future refactors move
112// additional ctors into this file; explicitly referenced to suppress
113// an otherwise-unused warning.
114const _: usize = DEFAULT_SEAL_THRESHOLD;
nodedb_vector/collection/quantize.rs

nodedb_vector/collection/
quantize.rs
