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 nodedb_mem::ScopedMemory;
9
10use crate::hnsw::{HnswIndex, HnswParams};
11use crate::index_config::{IndexConfig, IndexType};
12use crate::quantize::pq::PqCodec;
13use crate::quantize::sq8::Sq8Codec;
14
15use super::lifecycle::VectorCollection;
16use super::segment::DEFAULT_SEAL_THRESHOLD;
17
18impl VectorCollection {
19    /// Convenience constructor for PQ-configured collections.
20    ///
21    /// Equivalent to building a full `IndexConfig` with
22    /// `index_type = HnswPq` and the given `pq_m`.
23    pub fn with_pq_config(dim: usize, hnsw: HnswParams, pq_m: usize) -> Self {
24        let config = IndexConfig {
25            hnsw,
26            index_type: IndexType::HnswPq,
27            pq_m,
28            ..IndexConfig::default()
29        };
30        Self::with_index_config(dim, config)
31    }
32
33    /// Convenience constructor for PQ-configured collections with a custom
34    /// seal threshold.
35    pub fn with_seal_threshold_and_pq_config(
36        dim: usize,
37        hnsw: HnswParams,
38        pq_m: usize,
39        seal_threshold: usize,
40    ) -> Self {
41        let config = IndexConfig {
42            hnsw,
43            index_type: IndexType::HnswPq,
44            pq_m,
45            ..IndexConfig::default()
46        };
47        Self::with_seal_threshold_and_config(dim, config, seal_threshold)
48    }
49
50    /// Build SQ8 quantized data for an HNSW index.
51    ///
52    /// Returns `None` when there are too few live vectors for stable
53    /// min/max calibration.
54    pub fn build_sq8_for_index(index: &HnswIndex) -> Option<(Sq8Codec, Vec<u8>)> {
55        if index.live_count() < 1000 {
56            return None;
57        }
58        let dim = index.dim();
59        let n = index.len();
60
61        let mut refs: Vec<&[f32]> = Vec::with_capacity(n);
62        for i in 0..n {
63            if !index.is_deleted(i as u32)
64                && let Some(v) = index.get_vector(i as u32)
65            {
66                refs.push(v);
67            }
68        }
69        if refs.is_empty() {
70            return None;
71        }
72
73        // The refs are live vectors of this index, so calibration fails only
74        // for a zero dimension, which has nothing to quantize.
75        let codec = Sq8Codec::calibrate(&refs, dim).ok()?;
76
77        let mut data = Vec::with_capacity(dim * n);
78        for i in 0..n {
79            if let Some(v) = index.get_vector(i as u32) {
80                data.extend(codec.quantize(v));
81            } else {
82                data.extend(vec![0u8; dim]);
83            }
84        }
85
86        Some((codec, data))
87    }
88
89    /// Train a PQ codec from a built HNSW index's live vectors, tracking
90    /// codebook allocations against `memory`, and encode every node.
91    pub fn build_pq_for_index(
92        index: &HnswIndex,
93        pq_m: usize,
94        memory: ScopedMemory,
95    ) -> Option<(PqCodec, Vec<u8>)> {
96        let dim = index.dim();
97        if pq_m == 0 || !dim.is_multiple_of(pq_m) {
98            return None;
99        }
100        let n = index.len();
101        let mut refs: Vec<Vec<f32>> = Vec::with_capacity(n);
102        for i in 0..n {
103            if !index.is_deleted(i as u32)
104                && let Some(v) = index.get_vector(i as u32)
105            {
106                refs.push(v.to_vec());
107            }
108        }
109        if refs.is_empty() {
110            return None;
111        }
112        let refs_slices: Vec<&[f32]> = refs.iter().map(|v| v.as_slice()).collect();
113        let k = 256usize.min(refs.len());
114        // A PQ shape the codec cannot train (a codebook over its byte limit,
115        // a dimension over its decode limit) leaves the segment on plain
116        // HNSW, which answers the same queries exactly.
117        let codec = match PqCodec::train(&refs_slices, dim, pq_m, k, 20, memory) {
118            Ok(codec) => codec,
119            Err(e) => {
120                tracing::warn!(error = %e, dim, pq_m, "PQ training refused; segment stays unquantized");
121                return None;
122            }
123        };
124        // One code per local node id, soft-deleted nodes included: search
125        // reads the code of node `id` at `id * pq_m`.
126        let zero = vec![0.0f32; dim];
127        let all: Vec<&[f32]> = (0..n as u32)
128            .map(|i| index.get_vector(i).unwrap_or(zero.as_slice()))
129            .collect();
130        let codes = match codec.encode_batch(&all) {
131            Ok(codes) => codes,
132            Err(e) => {
133                tracing::warn!(error = %e, dim, pq_m, "PQ encoding refused; segment stays unquantized");
134                return None;
135            }
136        };
137        Some((codec, codes))
138    }
139}
140
141// Keep the DEFAULT_SEAL_THRESHOLD import live when future refactors move
142// additional ctors into this file; explicitly referenced to suppress
143// an otherwise-unused warning.
144const _: usize = DEFAULT_SEAL_THRESHOLD;
nodedb_vector/collection/quantize.rs

nodedb_vector/collection/
quantize.rs
