uqa_storage/diskann_index/
build.rs1use std::path::Path;
10
11use uqa_core::{memory::BudgetedVec, DocId};
12
13use super::format::{
14 DiskANNBuildCoverage, DiskANNCoverageBuilder, DiskANNGeneration, DiskANNVectorVersion,
15};
16use super::{
17 DiskANNCanonicalCorpusVisitor, DiskANNCanonicalRead, ExactVectorReason, NavigationInput,
18 PQCodebook, PQTrainer, PQTrainingOptions,
19};
20use crate::{read_control::StorageReadControl, StorageBackendError, StorageBackendResult};
21
22mod capture;
23mod generation;
24mod merge;
25mod partitions;
26mod records;
27mod runs;
28mod temporary;
29#[cfg(test)]
30mod tests;
31
32pub use capture::{DiskANNBuildCapture, DiskANNCanonicalCoverage};
33pub use generation::{DiskANNBuildSink, DiskANNGenerationOptions};
34pub use merge::{DiskANNMergeOptions, DiskANNMergeSummary, DiskANNMergedGraph};
35pub use partitions::{DiskANNPartitionOptions, DiskANNPartitionRuns, DiskANNPartitionSummary};
36pub use temporary::{DiskANNTemporaryBudget, DiskANNTemporaryError};
37
38use records::Records;
39
40pub type DiskANNBuildVisitor<'a> = DiskANNCanonicalCorpusVisitor<'a>;
41
42#[derive(Debug)]
44pub struct DiskANNBuildVector {
45 doc: DocId,
46 ordinal: u32,
47 version: DiskANNVectorVersion,
48 raw: BudgetedVec<f32>,
49 exact: Option<ExactVectorReason>,
50}
51
52impl DiskANNBuildVector {
53 pub fn doc_id(&self) -> DocId {
54 self.doc
55 }
56 pub fn ordinal(&self) -> u32 {
57 self.ordinal
58 }
59 pub fn version(&self) -> DiskANNVectorVersion {
60 self.version
61 }
62 pub fn raw(&self) -> &[f32] {
63 &self.raw
64 }
65 pub fn exact_reason(&self) -> Option<ExactVectorReason> {
66 self.exact
67 }
68}
69
70pub struct DiskANNBuildInput {
72 dimensions: u32,
73 navigation: Records,
74 side: Records,
75 coverage: DiskANNBuildCoverage,
76 control: StorageReadControl,
77 temporary: DiskANNTemporaryBudget,
78}
79
80impl DiskANNBuildInput {
81 pub fn capture_source(
83 generation: DiskANNGeneration,
84 source: &dyn DiskANNCanonicalRead,
85 directory: &Path,
86 temporary: &DiskANNTemporaryBudget,
87 control: &StorageReadControl,
88 ) -> StorageBackendResult<Self> {
89 Self::capture(
90 generation,
91 source.dimensions(),
92 directory,
93 temporary,
94 control,
95 |visit| source.visit_all(control, visit),
96 )
97 }
98
99 pub fn capture(
101 generation: DiskANNGeneration,
102 dimensions: u32,
103 directory: &Path,
104 temporary: &DiskANNTemporaryBudget,
105 control: &StorageReadControl,
106 read: impl FnOnce(&mut DiskANNBuildVisitor<'_>) -> StorageBackendResult<()>,
107 ) -> StorageBackendResult<Self> {
108 control.check()?;
109 let mut coverage = DiskANNCoverageBuilder::new(generation, dimensions)?;
110 let mut navigation = Records::new(dimensions, false, directory, temporary, control)?;
111 let mut side = Records::new(dimensions, true, directory, temporary, control)?;
112 let mut first_error = None;
113 let result = read(&mut |doc, ordinal, version, raw| {
114 if first_error.is_none() {
115 let result = (|| {
116 coverage.push(doc, ordinal, version, raw, control)?;
117 let (norm, _) = super::metric::norms(dimensions, raw, control)?;
118 let destination = if super::metric::exact_reason(norm).is_some() {
119 &mut side
120 } else {
121 &mut navigation
122 };
123 destination.push(doc, ordinal, version, raw, control)
124 })();
125 if let Err(error) = result {
126 first_error = Some(error);
127 }
128 }
129 if first_error.is_some() {
130 Err(invalid("canonical input consumer rejected a vector"))
131 } else {
132 Ok(())
133 }
134 });
135 if let Some(error) = first_error {
136 return Err(error);
137 }
138 result?;
139 control.check()?;
140 Ok(Self {
141 dimensions,
142 navigation,
143 side,
144 coverage: coverage.finish(),
145 control: control.clone(),
146 temporary: temporary.clone(),
147 })
148 }
149
150 pub fn dimensions(&self) -> u32 {
151 self.dimensions
152 }
153 pub fn node_count(&self) -> u64 {
154 self.navigation.len()
155 }
156 pub fn side_count(&self) -> u64 {
157 self.side.len()
158 }
159 pub fn coverage(&self) -> DiskANNBuildCoverage {
160 self.coverage
161 }
162
163 pub fn read_node(&self, node: u64) -> StorageBackendResult<DiskANNBuildVector> {
164 self.navigation.read(node, &self.control)
165 }
166
167 pub fn read_side(&self, index: u64) -> StorageBackendResult<DiskANNBuildVector> {
168 self.side.read(index, &self.control)
169 }
170
171 pub fn train(
173 &self,
174 pq_bytes: usize,
175 options: PQTrainingOptions,
176 ) -> StorageBackendResult<Option<PQCodebook>> {
177 let mut trainer = PQTrainer::new(self.dimensions, pq_bytes, options, &self.control)?;
178 self.navigation.visit(&self.control, &mut |record| {
179 let NavigationInput::Navigable(vector) =
180 NavigationInput::from_raw(self.dimensions, record.raw(), &self.control)?
181 else {
182 return Err(invalid("navigation record changed classification"));
183 };
184 trainer.observe(&vector)
185 })?;
186 if self.node_count() == 0 {
187 return Ok(None);
188 }
189 trainer.finish().map(Some)
190 }
191}
192
193fn invalid(message: &str) -> StorageBackendError {
194 StorageBackendError::Other(format!("invalid DiskANN build input: {message}"))
195}