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