uqa_storage/vector_index/config/
diskann.rs1use std::alloc::Layout;
10use std::collections::BTreeMap;
11
12use crate::{StorageBackendError, StorageBackendResult};
13
14mod alpha;
15mod catalog;
16
17pub use alpha::DiskANNAlpha;
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub struct DiskANNIndexParams {
21 pub max_degree: usize,
22 pub build_list_size: usize,
23 pub search_list_size: usize,
24 pub alpha: DiskANNAlpha,
25 pub beam_width: usize,
26 pub pq_bytes: usize,
27 pub seed: u64,
28 pub format_revision: u32,
29 pub algorithm_revision: u32,
30}
31
32impl DiskANNIndexParams {
33 pub const FORMAT_REVISION: u32 = 1;
34 pub const ALGORITHM_REVISION: u32 = 1;
35
36 pub fn for_dimensions(dimensions: u32) -> StorageBackendResult<Self> {
37 Self {
38 max_degree: 64,
39 build_list_size: 128,
40 search_list_size: 64,
41 alpha: DiskANNAlpha::default(),
42 beam_width: 4,
43 pq_bytes: usize::try_from(dimensions.min(32))
44 .map_err(|_| invalid("pq_bytes", "exceeds the platform size range"))?,
45 seed: 42,
46 format_revision: Self::FORMAT_REVISION,
47 algorithm_revision: Self::ALGORITHM_REVISION,
48 }
49 .validate(dimensions)
50 }
51
52 pub fn validate(self, dimensions: u32) -> StorageBackendResult<Self> {
53 let dimensions = usize::try_from(dimensions)
54 .map_err(|_| invalid("dimensions", "exceeds the platform size range"))?;
55 if dimensions == 0 {
56 return Err(invalid("dimensions", "must be greater than zero"));
57 }
58 if self.max_degree < 2 {
59 return Err(invalid("max_degree", "must be at least 2"));
60 }
61 if self.build_list_size < self.max_degree {
62 return Err(invalid("build_list_size", "must be at least max_degree"));
63 }
64 if self.search_list_size == 0 {
65 return Err(invalid("search_list_size", "must be greater than zero"));
66 }
67 if self.beam_width == 0 || self.beam_width > self.search_list_size {
68 return Err(invalid(
69 "beam_width",
70 "must be between 1 and search_list_size",
71 ));
72 }
73 if self.pq_bytes == 0 || self.pq_bytes > dimensions {
74 return Err(invalid("pq_bytes", "must be between 1 and dimensions"));
75 }
76 if self.format_revision != Self::FORMAT_REVISION {
77 return Err(invalid("format_revision", "is not supported"));
78 }
79 if self.algorithm_revision != Self::ALGORITHM_REVISION {
80 return Err(invalid("algorithm_revision", "is not supported"));
81 }
82 Layout::array::<u64>(self.max_degree)
83 .map_err(|_| invalid("max_degree", "overflows the neighbor buffer layout"))?;
84 for (name, count) in [
85 ("build_list_size", self.build_list_size),
86 ("search_list_size", self.search_list_size),
87 ] {
88 Layout::array::<(u64, f64)>(count)
89 .map_err(|_| invalid(name, "overflows the candidate buffer layout"))?;
90 }
91 Layout::array::<f32>(dimensions)
92 .map_err(|_| invalid("dimensions", "overflows the vector buffer layout"))?;
93 let lookup_entries = self
94 .pq_bytes
95 .checked_mul(256)
96 .ok_or_else(|| invalid("pq_bytes", "overflows the lookup entry count"))?;
97 Layout::array::<f64>(lookup_entries)
98 .map_err(|_| invalid("pq_bytes", "overflows the lookup buffer layout"))?;
99 Ok(self)
100 }
101
102 pub fn from_catalog_map(
103 dimensions: u32,
104 parameters: &BTreeMap<String, String>,
105 ) -> StorageBackendResult<Self> {
106 catalog::decode(dimensions, parameters)
107 }
108
109 pub fn to_catalog_map(self, dimensions: u32) -> StorageBackendResult<BTreeMap<String, String>> {
110 Ok(catalog::encode(self.validate(dimensions)?))
111 }
112}
113
114fn invalid(name: &str, reason: &str) -> StorageBackendError {
115 StorageBackendError::Other(format!("DiskANN parameter `{name}` {reason}"))
116}
117
118#[cfg(test)]
119mod tests;