nodedb_vector/
index_config.rs

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

nodedb_vector/
index_config.rs
