uqa_storage/vector_index/config/
ivf.rs1use 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;