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

nodedb_vector/
index_config.rs
