Skip to main content

uqa_storage/diskann_index/
build.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Encrypted canonical build input and retained source membership. A fingerprint alone does not establish visibility or authorize generation publication.
8
9use 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/// One captured canonical vector. Its original bits and version are independent of navigation, payloads and scores; its raw buffer retains the build allowance.
43#[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
70/// Repeatable encrypted input, with fixed-width navigation and numeric-side records. Dense navigation IDs follow canonical logical order without a resident global directory.
71pub 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    /// Capture a retained provider source without fabricating mutation origins or collecting its corpus. The resulting fingerprint is still not a visibility token or publication authority.
82    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    /// Enumerate every selected canonical ordinal once in document order. Any source or consumer failure discards the entire capture, including failures the source suppresses; temporary data is never a staged or published generation.
100    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    /// Train the existing bounded reservoir by replaying captured navigation values. Empty and all-side inputs have no codebook; invalid settings still fail validation.
172    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}