Skip to main content

uqa_storage/vector_index/config/
diskann.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Resolved `DiskANN` parameters, independent of public access-method routing.
8
9use 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;