Skip to main content

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}