Skip to main content

uqa_storage/vector_index/config/
ivf.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! IVF parameters and persisted-catalog decoding.
8
9use std::collections::BTreeMap;
10
11use super::parsing::{read_positive_usize, reject_unknown_parameters};
12use crate::{StorageBackendError, StorageBackendResult};
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub struct IVFIndexParams {
16    pub nlist: usize,
17    pub nprobe: usize,
18    pub train_threshold: usize,
19}
20
21impl Default for IVFIndexParams {
22    fn default() -> Self {
23        Self {
24            nlist: 100,
25            nprobe: 10,
26            train_threshold: 256,
27        }
28    }
29}
30
31impl IVFIndexParams {
32    pub fn validate(self) -> StorageBackendResult<Self> {
33        for (name, value) in [
34            ("nlist", self.nlist),
35            ("nprobe", self.nprobe),
36            ("train_threshold", self.train_threshold),
37        ] {
38            if value == 0 {
39                return Err(StorageBackendError::Other(format!(
40                    "IVF parameter `{name}` must be greater than zero"
41                )));
42            }
43        }
44        Ok(self)
45    }
46
47    pub fn from_catalog_map(parameters: &BTreeMap<String, String>) -> StorageBackendResult<Self> {
48        reject_unknown_parameters(
49            parameters,
50            &[
51                "lists",
52                "nlist",
53                "probes",
54                "nprobe",
55                "train_threshold",
56                "train-threshold",
57                "min_train",
58            ],
59            "IVF",
60        )?;
61        let defaults = Self::default();
62        Self {
63            nlist: read_positive_usize(parameters, &["lists", "nlist"], defaults.nlist, "IVF")?,
64            nprobe: read_positive_usize(parameters, &["probes", "nprobe"], defaults.nprobe, "IVF")?,
65            train_threshold: read_positive_usize(
66                parameters,
67                &["train_threshold", "train-threshold", "min_train"],
68                defaults.train_threshold,
69                "IVF",
70            )?,
71        }
72        .validate()
73    }
74}
75
76#[cfg(test)]
77mod tests;