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: vectors buffer, searched exactly, until `max(ivf_cells, pq_k)`
27    /// are held, then train k-means cells and PQ codebooks. A search probes
28    /// `ivf_nprobe` cells by PQ distance and reranks by exact distance.
29    IvfPq,
30}
31
32impl IndexType {
33    pub fn parse(s: &str) -> Option<Self> {
34        match s.to_lowercase().as_str() {
35            "hnsw" | "" => Some(Self::Hnsw),
36            "hnsw_pq" => Some(Self::HnswPq),
37            "ivf_pq" => Some(Self::IvfPq),
38            _ => None,
39        }
40    }
41}
42
43/// Unified vector index configuration.
44#[derive(
45    Debug,
46    Clone,
47    serde::Serialize,
48    serde::Deserialize,
49    zerompk::ToMessagePack,
50    zerompk::FromMessagePack,
51)]
52pub struct IndexConfig {
53    /// HNSW parameters (used for Hnsw and HnswPq types).
54    pub hnsw: HnswParams,
55    /// Index type.
56    pub index_type: IndexType,
57    /// PQ subvectors (for HnswPq and IvfPq). Must divide dim evenly.
58    pub pq_m: usize,
59    /// IVF cells (for IvfPq only).
60    pub ivf_cells: usize,
61    /// IVF probe count (for IvfPq only).
62    pub ivf_nprobe: usize,
63    /// Vector width declared when the index was created. Every vector written
64    /// to the index is checked against it, so a producer emitting the wrong
65    /// embedding width is rejected at the write rather than discovered as poor
66    /// search results. `0` = never declared, in which case the index adopts
67    /// the width of the first vector it sees. Defaulted on decode so indexes
68    /// snapshotted before the field existed keep loading.
69    #[serde(default)]
70    pub declared_dim: usize,
71}
72
73/// Default PQ subquantizer count (segments per vector).
74pub const DEFAULT_PQ_M: usize = 8;
75/// Default IVF cell count (Voronoi partitions).
76pub const DEFAULT_IVF_CELLS: usize = 256;
77/// Default IVF probe count (cells searched per query).
78pub const DEFAULT_IVF_NPROBE: usize = 16;
79
80impl Default for IndexConfig {
81    fn default() -> Self {
82        Self {
83            hnsw: HnswParams::default(),
84            index_type: IndexType::Hnsw,
85            pq_m: DEFAULT_PQ_M,
86            ivf_cells: DEFAULT_IVF_CELLS,
87            ivf_nprobe: DEFAULT_IVF_NPROBE,
88            declared_dim: 0,
89        }
90    }
91}
92
93impl IndexConfig {
94    /// Build IVF-PQ params from this config.
95    pub fn to_ivf_params(&self) -> crate::ivf::IvfPqParams {
96        crate::ivf::IvfPqParams {
97            n_cells: self.ivf_cells,
98            pq_m: self.pq_m,
99            pq_k: 256,
100            nprobe: self.ivf_nprobe,
101            metric: self.hnsw.metric,
102        }
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109
110    #[test]
111    fn parse_index_type() {
112        assert_eq!(IndexType::parse("hnsw"), Some(IndexType::Hnsw));
113        assert_eq!(IndexType::parse(""), Some(IndexType::Hnsw));
114        assert_eq!(IndexType::parse("hnsw_pq"), Some(IndexType::HnswPq));
115        assert_eq!(IndexType::parse("ivf_pq"), Some(IndexType::IvfPq));
116        assert_eq!(IndexType::parse("ivfpq"), None);
117        assert_eq!(IndexType::parse("unknown"), None);
118    }
119
120    #[test]
121    fn default_is_hnsw() {
122        let cfg = IndexConfig::default();
123        assert_eq!(cfg.index_type, IndexType::Hnsw);
124    }
125}
nodedb_vector/index_config.rs

nodedb_vector/
index_config.rs
