uqa_storage/diskann_index/build/
partitions.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7use sha2::{Digest, Sha256};
8use std::path::Path;
9
10use super::runs::{RunReader, RunWriter};
11use super::temporary::TemporaryRun;
12use super::{invalid, DiskANNBuildCoverage, DiskANNBuildInput};
13use crate::diskann_index::PQTrainingOptions;
14use crate::{
15    read_control::StorageReadControl, vector_index::DiskANNIndexParams, StorageBackendResult,
16};
17
18mod leaf;
19mod source;
20#[cfg(test)]
21mod tests;
22
23use source::Ids;
24
25/// Effective construction settings, retained with the resulting provenance. Limits do not silently change sample sizes or partition capacity.
26#[derive(Debug, Clone, Copy, PartialEq, Eq)]
27pub struct DiskANNPartitionOptions {
28    pub max_partition_points: usize,
29    pub coarse_training: PQTrainingOptions,
30    pub max_depth: u8,
31}
32
33impl DiskANNPartitionOptions {
34    pub const WORK_ORDER_REVISION: u32 = 1;
35
36    pub(in crate::diskann_index) fn validate(self) -> StorageBackendResult<Self> {
37        self.coarse_training.validate()?;
38        if self.max_partition_points < 2
39            || self.max_depth > 32
40            || self.coarse_training.max_centroids < 3
41            || self.coarse_training.max_samples < u32::from(self.coarse_training.max_centroids)
42        {
43            return Err(invalid("partition capacity must be at least two, depth at most 32 and coarse samples must cover at least three centroids"));
44        }
45        Ok(self)
46    }
47}
48
49/// Compact construction provenance, not a physical seal or publication permission. The final generation builder must encode it with its manifest.
50#[derive(Debug, Clone, Copy, PartialEq, Eq)]
51pub struct DiskANNPartitionSummary {
52    pub work_order_revision: u32,
53    pub options: DiskANNPartitionOptions,
54    pub parameters: DiskANNIndexParams,
55    pub coverage: DiskANNBuildCoverage,
56    pub partitions: u64,
57    pub memberships: u64,
58    pub edges: u64,
59    pub maximum_depth: u8,
60    pub maximum_partition_points: usize,
61    pub assignment_digest: [u8; 32],
62}
63
64/// Encrypted, deterministic global edge candidates from overlapping local graphs. Their union still requires external ordering, final pruning and global connectivity.
65pub struct DiskANNPartitionRuns {
66    edges: TemporaryRun,
67    summary: DiskANNPartitionSummary,
68    nodes: u64,
69    control: StorageReadControl,
70}
71
72impl DiskANNPartitionRuns {
73    pub fn summary(&self) -> &DiskANNPartitionSummary {
74        &self.summary
75    }
76
77    pub(super) fn into_source(
78        self,
79        input: &DiskANNBuildInput,
80    ) -> StorageBackendResult<(TemporaryRun, DiskANNPartitionSummary)> {
81        input.control.check()?;
82        if self.nodes != input.node_count()
83            || self.summary.coverage != input.coverage
84            || !self.control.shares_context(&input.control)
85            || !self.edges.uses_allowance(&input.temporary)
86        {
87            return Err(invalid(
88                "partition source, coverage or build allowance differs",
89            ));
90        }
91        Ok((self.edges, self.summary))
92    }
93
94    pub fn visit_edges(
95        &self,
96        visitor: &mut dyn FnMut(u64, u64) -> StorageBackendResult<()>,
97    ) -> StorageBackendResult<()> {
98        self.edges.read(&self.control, |file| {
99            if self.summary.edges == 0 {
100                return Ok(());
101            }
102            let mut reader = RunReader::new(file, &self.control)?;
103            for _ in 0..self.summary.edges {
104                let bytes = reader.record::<16>()?;
105                let source = u64::from_le_bytes(bytes[..8].try_into().expect("source ID"));
106                let neighbor = u64::from_le_bytes(bytes[8..].try_into().expect("neighbor ID"));
107                if source >= self.nodes || neighbor >= self.nodes || source == neighbor {
108                    return Err(invalid("invalid global partition edge"));
109                }
110                visitor(source, neighbor)?;
111            }
112            Ok(())
113        })
114    }
115}
116
117impl DiskANNBuildInput {
118    /// Construct one admitted Vamana graph at a time. All temporary runs share the capture's existing disk allowance, and all buffers retain its original memory/cancellation control.
119    pub fn build_partitions(
120        &self,
121        directory: &Path,
122        parameters: DiskANNIndexParams,
123        options: DiskANNPartitionOptions,
124    ) -> StorageBackendResult<DiskANNPartitionRuns> {
125        self.control.check()?;
126        let parameters = parameters.validate(self.dimensions)?;
127        let options = options.validate()?;
128        let writer = RunWriter::new(directory, &self.temporary, &self.control)?;
129        let mut builder = Builder {
130            input: self,
131            directory,
132            writer,
133            summary: DiskANNPartitionSummary {
134                work_order_revision: DiskANNPartitionOptions::WORK_ORDER_REVISION,
135                options,
136                parameters,
137                coverage: self.coverage,
138                partitions: 0,
139                memberships: 0,
140                edges: 0,
141                maximum_depth: 0,
142                maximum_partition_points: 0,
143                assignment_digest: [0; 32],
144            },
145            hash: provenance(self.coverage, parameters, options),
146        };
147        builder.partition(Ids::All(self.node_count()), 0, options.coarse_training.seed)?;
148        self.control.check()?;
149        builder
150            .hash
151            .update(builder.summary.partitions.to_le_bytes());
152        builder
153            .hash
154            .update(builder.summary.memberships.to_le_bytes());
155        builder.summary.assignment_digest = builder.hash.finalize().into();
156        Ok(DiskANNPartitionRuns {
157            edges: builder.writer.finish()?,
158            summary: builder.summary,
159            nodes: self.node_count(),
160            control: self.control.clone(),
161        })
162    }
163}
164
165struct Builder<'a> {
166    input: &'a DiskANNBuildInput,
167    directory: &'a Path,
168    writer: RunWriter,
169    summary: DiskANNPartitionSummary,
170    hash: Sha256,
171}
172
173impl Builder<'_> {
174    fn partition(&mut self, source: Ids<'_>, depth: u8, seed: u64) -> StorageBackendResult<()> {
175        self.input.control.check()?;
176        let count = source.len();
177        if count == 0 {
178            return Ok(());
179        }
180        self.summary.maximum_depth = self.summary.maximum_depth.max(depth);
181        let options = self.summary.options;
182        if count <= options.max_partition_points as u64 || depth == options.max_depth {
183            return self.capacity(source);
184        }
185        let split = source::split(
186            self.input,
187            source,
188            self.directory,
189            options.coarse_training,
190            seed,
191        )?;
192        let reduces = split
193            .counts
194            .iter()
195            .all(|&child| u128::from(child) * 4 <= u128::from(count) * 3);
196        self.hash.update([b'S', depth, u8::from(reduces)]);
197        self.hash.update(seed.to_le_bytes());
198        self.hash.update(split.digest);
199        if !reduces {
200            drop(split);
201            return self.capacity(source);
202        }
203        for label in 0..split.counts.len() {
204            self.partition(
205                Ids::Child(&split, label as u8),
206                depth + 1,
207                source::child_seed(seed, label as u8),
208            )?;
209        }
210        Ok(())
211    }
212
213    fn capacity(&mut self, source: Ids<'_>) -> StorageBackendResult<()> {
214        leaf::windows(
215            source,
216            self.summary.options.max_partition_points,
217            &self.input.control,
218            &mut |ids| self.leaf(ids),
219        )
220    }
221}
222
223fn provenance(
224    coverage: DiskANNBuildCoverage,
225    parameters: DiskANNIndexParams,
226    options: DiskANNPartitionOptions,
227) -> Sha256 {
228    let mut hash = Sha256::new();
229    hash.update(b"UQA DiskANN partitions\0\x01");
230    hash.update(coverage.digest());
231    for value in [
232        parameters.max_degree as u64,
233        parameters.build_list_size as u64,
234        parameters.search_list_size as u64,
235        parameters.alpha.get().to_bits(),
236        parameters.beam_width as u64,
237        parameters.pq_bytes as u64,
238        parameters.seed,
239        u64::from(parameters.format_revision),
240        u64::from(parameters.algorithm_revision),
241        options.max_partition_points as u64,
242        u64::from(options.max_depth),
243        u64::from(options.coarse_training.max_samples),
244        u64::from(options.coarse_training.max_iterations),
245        u64::from(options.coarse_training.max_centroids),
246        options.coarse_training.seed,
247    ] {
248        hash.update(value.to_le_bytes());
249    }
250    hash
251}
uqa_storage/diskann_index/build/partitions.rs

uqa_storage/diskann_index/build/
partitions.rs
