uqa_storage/diskann_index/build/
generation.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7use sha2::{Digest, Sha256};
8use uqa_core::memory::BudgetedVec;
9
10use super::{invalid, DiskANNBuildInput, DiskANNMergedGraph};
11use crate::diskann_index::format::{
12    encode_page, DiskANNArtifactDigests, DiskANNBuildProvenance, DiskANNGeneration,
13    DiskANNManifest, DiskANNManifestInput, DiskANNNodeInput, DiskANNNodeLayout, DiskANNSideEntry,
14    DiskANNSideLayout, PAGE_HEADER_BYTES, PAGE_PAYLOAD_BYTES,
15};
16use crate::diskann_index::pages::{DiskANNMemoryBuilder, DiskANNRecordKey};
17use crate::diskann_index::PQTrainingOptions;
18use crate::{
19    key_value::KeyValueDiskANNStage, read_control::StorageReadControl, StorageBackendResult,
20};
21
22#[cfg(test)]
23mod tests;
24mod vectors;
25
26/// Effective bounded metadata batches. Graph pages keep their codec-defined size.
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub struct DiskANNGenerationOptions {
29    pub training: PQTrainingOptions,
30    pub code_batch_nodes: usize,
31    pub side_batch_entries: usize,
32    pub max_record_bytes: usize,
33}
34
35/// A caller-retained unpublished artifact owner. A failed build leaves physical recovery and pending transaction outcomes with this owner; the algorithm never publishes or discards it implicitly.
36pub trait DiskANNBuildSink {
37    fn generation(&self) -> DiskANNGeneration;
38    fn write_record(
39        &mut self,
40        key: DiskANNRecordKey,
41        bytes: &[u8],
42        maximum: usize,
43        control: &StorageReadControl,
44    ) -> StorageBackendResult<()>;
45    fn write_graph_page(
46        &mut self,
47        id: u64,
48        bytes: &[u8],
49        control: &StorageReadControl,
50    ) -> StorageBackendResult<()>;
51}
52
53impl DiskANNBuildSink for DiskANNMemoryBuilder {
54    fn generation(&self) -> DiskANNGeneration {
55        self.generation()
56    }
57    fn write_record(
58        &mut self,
59        key: DiskANNRecordKey,
60        bytes: &[u8],
61        maximum: usize,
62        control: &StorageReadControl,
63    ) -> StorageBackendResult<()> {
64        control.check_value_size(bytes.len(), maximum)?;
65        DiskANNMemoryBuilder::write_record(self, key, bytes, control)
66    }
67    fn write_graph_page(
68        &mut self,
69        id: u64,
70        bytes: &[u8],
71        control: &StorageReadControl,
72    ) -> StorageBackendResult<()> {
73        DiskANNMemoryBuilder::write_graph_page(self, id, bytes, control)
74    }
75}
76
77impl DiskANNBuildSink for KeyValueDiskANNStage {
78    fn generation(&self) -> DiskANNGeneration {
79        self.generation()
80    }
81    fn write_record(
82        &mut self,
83        key: DiskANNRecordKey,
84        bytes: &[u8],
85        maximum: usize,
86        control: &StorageReadControl,
87    ) -> StorageBackendResult<()> {
88        KeyValueDiskANNStage::write_record(self, key, bytes, maximum, control)
89    }
90    fn write_graph_page(
91        &mut self,
92        id: u64,
93        bytes: &[u8],
94        control: &StorageReadControl,
95    ) -> StorageBackendResult<()> {
96        KeyValueDiskANNStage::write_graph_page(self, id, bytes, control)
97    }
98}
99
100impl DiskANNBuildInput {
101    /// Stream immutable records to the selected private generation. The returned manifest must pass the sink owner's physical seal before opening; it grants no SQL/MVCC publication authority.
102    pub fn write_generation(
103        &self,
104        graph: &DiskANNMergedGraph,
105        options: DiskANNGenerationOptions,
106        sink: &mut dyn DiskANNBuildSink,
107    ) -> StorageBackendResult<DiskANNManifest> {
108        graph.check_input(self)?;
109        let generation = self.coverage.generation();
110        let parameters = graph.summary().partitions.parameters;
111        if sink.generation() != generation {
112            return Err(invalid("artifact sink generation differs"));
113        }
114        self.control.check_value_size(
115            DiskANNManifest::ENCODED_BYTES + DiskANNBuildProvenance::ENCODED_BYTES,
116            options.max_record_bytes,
117        )?;
118        let provenance = DiskANNBuildProvenance::from_merge(
119            graph.summary(),
120            self.node_count(),
121            options.training,
122            options.code_batch_nodes,
123            options.side_batch_entries,
124        )?;
125        let mut artifacts = DiskANNArtifactDigests::empty();
126        vectors::quantize(self, parameters.pq_bytes, options, sink, &mut artifacts)?;
127        let entry_node = vectors::entry(self, parameters.seed)?;
128        artifacts.side = self.write_sides(options, sink)?;
129        artifacts.graph = self.write_graph(graph, sink)?;
130        self.control.check()?;
131        DiskANNManifest::new(DiskANNManifestInput {
132            generation,
133            dimensions: self.dimensions,
134            parameters,
135            nodes: self.node_count(),
136            side_vectors: self.side_count(),
137            entry_node,
138            coverage: self.coverage,
139            artifacts,
140        })?
141        .with_build_provenance(provenance)
142    }
143
144    fn write_sides(
145        &self,
146        options: DiskANNGenerationOptions,
147        sink: &mut dyn DiskANNBuildSink,
148    ) -> StorageBackendResult<[u8; 32]> {
149        let layout = DiskANNSideLayout::new(
150            self.coverage.generation(),
151            self.dimensions,
152            self.side_count(),
153        )?;
154        let capacity = self.side_count().min(options.side_batch_entries as u64) as usize;
155        let mut entries = BudgetedVec::new(self.control.memory());
156        entries.reserve(capacity)?;
157        let mut first = 0;
158        let mut hash = Sha256::new();
159        for id in 0..self.side_count() {
160            let record = self.read_side(id)?;
161            entries.push(DiskANNSideEntry::from_raw(
162                self.dimensions,
163                record.doc_id(),
164                record.ordinal(),
165                record.version(),
166                record.raw(),
167                &self.control,
168            )?)?;
169            if entries.len() == capacity || id + 1 == self.side_count() {
170                let bytes = layout.encode(first, &entries, &self.control)?;
171                self.control
172                    .check_value_size(bytes.len(), options.max_record_bytes)?;
173                let batch = layout.decode(first, &bytes, &self.control)?;
174                for chunk in batch.bytes().chunks(4096) {
175                    self.control.check()?;
176                    hash.update(chunk);
177                }
178                sink.write_record(
179                    DiskANNRecordKey::Side(first),
180                    &bytes,
181                    options.max_record_bytes,
182                    &self.control,
183                )?;
184                first = id + 1;
185                entries.clear();
186            }
187        }
188        Ok(hash.finalize().into())
189    }
190
191    fn write_graph(
192        &self,
193        graph: &DiskANNMergedGraph,
194        sink: &mut dyn DiskANNBuildSink,
195    ) -> StorageBackendResult<[u8; 32]> {
196        let layout = DiskANNNodeLayout::new(
197            self.dimensions,
198            graph.summary().partitions.parameters.max_degree,
199            self.node_count(),
200        )?;
201        let generation = self.coverage.generation();
202        let mut payload = BudgetedVec::new(self.control.memory());
203        let mut page = 0;
204        let mut hash = Sha256::new();
205        graph.visit_neighbors(&mut |node_id, neighbors| {
206            let record = self.read_node(node_id)?;
207            let encoded = layout.encode_node(
208                &DiskANNNodeInput {
209                    node_id,
210                    doc_id: record.doc_id(),
211                    ordinal: record.ordinal(),
212                    version: record.version(),
213                    vector: record.raw(),
214                    neighbors,
215                },
216                &self.control,
217            )?;
218            let expected = layout.page_shape(page)?.payload_bytes as usize;
219            let mut emit = |payload: &[u8]| -> StorageBackendResult<()> {
220                let bytes = encode_page(generation, layout, page, payload, &self.control)?;
221                hash.update(&bytes[PAGE_HEADER_BYTES - 32..PAGE_HEADER_BYTES]);
222                sink.write_graph_page(page, &bytes, &self.control)?;
223                page += 1;
224                Ok(())
225            };
226            if layout.node_address(node_id)?.fragments > 1 {
227                for fragment in encoded.chunks(PAGE_PAYLOAD_BYTES) {
228                    self.control.check()?;
229                    emit(fragment)?;
230                }
231            } else {
232                if payload.is_empty() {
233                    payload.reserve(expected)?;
234                }
235                payload.extend_from_slice(&encoded)?;
236                if payload.len() == expected {
237                    emit(&payload)?;
238                    payload.clear();
239                }
240            }
241            Ok(())
242        })?;
243        if page != layout.page_count() || !payload.is_empty() {
244            return Err(invalid("incomplete graph page output"));
245        }
246        Ok(hash.finalize().into())
247    }
248}
uqa_storage/diskann_index/build/generation.rs

uqa_storage/diskann_index/build/
generation.rs
