nodedb_vector/
index_config.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Vector index configuration: unified config for HNSW, HNSW+PQ, and IVF-PQ.
4
5use crate::hnsw::HnswParams;
6
7/// Index type selection for vector collections.
8#[derive(
9    Debug,
10    Clone,
11    Default,
12    PartialEq,
13    Eq,
14    serde::Serialize,
15    serde::Deserialize,
16    zerompk::ToMessagePack,
17    zerompk::FromMessagePack,
18)]
19#[non_exhaustive]
20pub enum IndexType {
21    /// Pure HNSW with FP32 vectors. Best recall (~99%), highest memory.
22    #[default]
23    Hnsw,
24    /// HNSW graph with PQ-compressed storage for traversal.
25    HnswPq,
26    /// IVF-PQ flat index. Lowest memory (~16 bytes/vector), best for >10M vectors.
27    IvfPq,
28}
29
30impl IndexType {
31    pub fn parse(s: &str) -> Option<Self> {
32        match s.to_lowercase().as_str() {
33            "hnsw" | "" => Some(Self::Hnsw),
34            "hnsw_pq" => Some(Self::HnswPq),
35            "ivf_pq" => Some(Self::IvfPq),
36            _ => None,
37        }
38    }
39}
40
41/// Unified vector index configuration.
42#[derive(
43    Debug,
44    Clone,
45    serde::Serialize,
46    serde::Deserialize,
47    zerompk::ToMessagePack,
48    zerompk::FromMessagePack,
49)]
50pub struct IndexConfig {
51    /// HNSW parameters (used for Hnsw and HnswPq types).
52    pub hnsw: HnswParams,
53    /// Index type.
54    pub index_type: IndexType,
55    /// PQ subvectors (for HnswPq and IvfPq). Must divide dim evenly.
56    pub pq_m: usize,
57    /// IVF cells (for IvfPq only).
58    pub ivf_cells: usize,
59    /// IVF probe count (for IvfPq only).
60    pub ivf_nprobe: usize,
61}
62
63/// Default PQ subquantizer count (segments per vector).
64pub const DEFAULT_PQ_M: usize = 8;
65/// Default IVF cell count (Voronoi partitions).
66pub const DEFAULT_IVF_CELLS: usize = 256;
67/// Default IVF probe count (cells searched per query).
68pub const DEFAULT_IVF_NPROBE: usize = 16;
69
70impl Default for IndexConfig {
71    fn default() -> Self {
72        Self {
73            hnsw: HnswParams::default(),
74            index_type: IndexType::Hnsw,
75            pq_m: DEFAULT_PQ_M,
76            ivf_cells: DEFAULT_IVF_CELLS,
77            ivf_nprobe: DEFAULT_IVF_NPROBE,
78        }
79    }
80}
81
82impl IndexConfig {
83    /// Build IVF-PQ params from this config.
84    pub fn to_ivf_params(&self) -> crate::ivf::IvfPqParams {
85        crate::ivf::IvfPqParams {
86            n_cells: self.ivf_cells,
87            pq_m: self.pq_m,
88            pq_k: 256,
89            nprobe: self.ivf_nprobe,
90            metric: self.hnsw.metric,
91        }
92    }
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98
99    #[test]
100    fn parse_index_type() {
101        assert_eq!(IndexType::parse("hnsw"), Some(IndexType::Hnsw));
102        assert_eq!(IndexType::parse(""), Some(IndexType::Hnsw));
103        assert_eq!(IndexType::parse("hnsw_pq"), Some(IndexType::HnswPq));
104        assert_eq!(IndexType::parse("ivf_pq"), Some(IndexType::IvfPq));
105        assert_eq!(IndexType::parse("ivfpq"), None);
106        assert_eq!(IndexType::parse("unknown"), None);
107    }
108
109    #[test]
110    fn default_is_hnsw() {
111        let cfg = IndexConfig::default();
112        assert_eq!(cfg.index_type, IndexType::Hnsw);
113    }
114}
nodedb_vector/index_config.rs

nodedb_vector/
index_config.rs
