nodedb_vector/ivf/params.rs
1// SPDX-License-Identifier: Apache-2.0
2
3//! IVF-PQ index parameters.
4
5use crate::distance::DistanceMetric;
6
7/// IVF-PQ index configuration.
8#[derive(Debug, Clone, zerompk::ToMessagePack, zerompk::FromMessagePack)]
9pub struct IvfPqParams {
10 /// Number of Voronoi cells (partitions). Typical: sqrt(N).
11 pub n_cells: usize,
12 /// Number of PQ subvectors. Must divide dimension evenly.
13 pub pq_m: usize,
14 /// Centroids per PQ subvector (at most 256 for u8 codes).
15 pub pq_k: usize,
16 /// Number of cells to probe at query time. Higher = better recall.
17 pub nprobe: usize,
18 /// Distance metric.
19 pub metric: DistanceMetric,
20}
21
22impl IvfPqParams {
23 /// Vectors an index needs before it can train: one per coarse cell and
24 /// one per PQ centroid, since both k-means runs need at least `k` points.
25 pub fn training_threshold(&self) -> usize {
26 self.n_cells.max(self.pq_k).max(1)
27 }
28}
29
30impl Default for IvfPqParams {
31 fn default() -> Self {
32 Self {
33 n_cells: 256,
34 pq_m: 8,
35 pq_k: 256,
36 nprobe: 16,
37 metric: DistanceMetric::L2,
38 }
39 }
40}