uqa_storage/diskann_index/build/
merge.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7use sha2::{Digest, Sha256};
8use std::path::Path;
9use uqa_core::memory::BudgetedVec;
10
11use super::runs::{RunReader, RunWriter};
12use super::temporary::TemporaryRun;
13use super::{invalid, DiskANNBuildInput, DiskANNPartitionRuns, DiskANNPartitionSummary};
14use crate::diskann_index::{vamana::Selection, NavigationInput, NavigationVector};
15use crate::{read_control::StorageReadControl, StorageBackendResult};
16
17mod sort;
18#[cfg(test)]
19mod tests;
20
21/// An explicit record capacity for initial sorted chunks. Every merge pass uses two readers; the configured capacity is never silently reduced.
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub struct DiskANNMergeOptions {
24    pub sort_buffer_records: usize,
25}
26
27/// Global adjacency construction provenance, without physical artifact sealing or SQL publication authority.
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub struct DiskANNMergeSummary {
30    pub work_order_revision: u32,
31    pub partitions: DiskANNPartitionSummary,
32    pub options: DiskANNMergeOptions,
33    pub merge_passes: u8,
34    pub edges: u64,
35    pub adjacency_digest: [u8; 32],
36}
37
38/// Encrypted fixed-width adjacency, with unique sorted neighbors and a reserved global successor for every non-singleton node.
39pub struct DiskANNMergedGraph {
40    run: TemporaryRun,
41    nodes: u64,
42    degree: usize,
43    summary: DiskANNMergeSummary,
44    control: StorageReadControl,
45}
46
47impl DiskANNMergedGraph {
48    pub(super) fn check_input(&self, input: &DiskANNBuildInput) -> StorageBackendResult<()> {
49        input.control.check()?;
50        if self.nodes != input.node_count()
51            || self.summary.partitions.coverage != input.coverage
52            || !self.control.shares_context(&input.control)
53            || !self.run.uses_allowance(&input.temporary)
54        {
55            return Err(invalid("merged graph source or build allowance differs"));
56        }
57        Ok(())
58    }
59
60    pub fn summary(&self) -> &DiskANNMergeSummary {
61        &self.summary
62    }
63
64    /// Visit every dense node in order, including empty/singleton neighborhoods. Buffers and errors retain the original build control.
65    pub fn visit_neighbors(
66        &self,
67        visitor: &mut dyn FnMut(u64, &[u64]) -> StorageBackendResult<()>,
68    ) -> StorageBackendResult<()> {
69        self.run.read(&self.control, |file| {
70            let mut hash = adjacency_hash(&self.summary.partitions, self.nodes, self.degree);
71            let mut edges = 0_u64;
72            if self.nodes != 0 {
73                let mut reader = RunReader::new(file, &self.control)?;
74                let mut neighbors = BudgetedVec::new(self.control.memory());
75                neighbors.reserve(self.degree)?;
76                for node in 0..self.nodes {
77                    self.control.check()?;
78                    let bytes = reader.record::<8>()?;
79                    hash.update(bytes);
80                    let count = u64::from_le_bytes(bytes);
81                    if count > self.degree as u64 || (self.nodes > 1 && count == 0) {
82                        return Err(invalid("invalid merged neighborhood length"));
83                    }
84                    neighbors.clear();
85                    for slot in 0..self.degree {
86                        let bytes = reader.record::<8>()?;
87                        hash.update(bytes);
88                        let neighbor = u64::from_le_bytes(bytes);
89                        if (slot as u64) < count {
90                            if neighbor >= self.nodes
91                                || neighbor == node
92                                || neighbors.last().is_some_and(|&last| last >= neighbor)
93                            {
94                                return Err(invalid("invalid merged neighbor"));
95                            }
96                            neighbors.push(neighbor)?;
97                        } else if neighbor != 0 {
98                            return Err(invalid("nonzero merged neighborhood padding"));
99                        }
100                    }
101                    if self.nodes > 1 && !neighbors.contains(&((node + 1) % self.nodes)) {
102                        return Err(invalid("merged neighborhood lacks global successor"));
103                    }
104                    edges = edges
105                        .checked_add(count)
106                        .ok_or_else(|| invalid("merged edge overflow"))?;
107                    visitor(node, &neighbors)?;
108                }
109            }
110            if edges != self.summary.edges
111                || <[u8; 32]>::from(hash.finalize()) != self.summary.adjacency_digest
112            {
113                return Err(invalid(
114                    "merged adjacency differs from its construction digest",
115                ));
116            }
117            Ok(())
118        })
119    }
120}
121
122impl DiskANNBuildInput {
123    /// Consume partition candidates, externally order them, and prune using full navigation vectors. Capture and candidates must share source coverage and both resource allowances.
124    pub fn merge_partitions(
125        &self,
126        partitions: DiskANNPartitionRuns,
127        directory: &Path,
128        options: DiskANNMergeOptions,
129    ) -> StorageBackendResult<DiskANNMergedGraph> {
130        self.control.check()?;
131        if options.sort_buffer_records == 0 {
132            return Err(invalid("global sort capacity must be positive"));
133        }
134        let (run, summary) = partitions.into_source(self)?;
135        merge(self, run, &summary, directory, options)
136    }
137}
138
139fn merge(
140    input: &DiskANNBuildInput,
141    candidates: TemporaryRun,
142    partitions: &DiskANNPartitionSummary,
143    directory: &Path,
144    options: DiskANNMergeOptions,
145) -> StorageBackendResult<DiskANNMergedGraph> {
146    let sorted = sort::sort(
147        input,
148        candidates,
149        partitions.edges,
150        directory,
151        options.sort_buffer_records,
152    )?;
153    let control = &input.control;
154    let nodes = input.node_count();
155    let degree = (partitions.parameters.max_degree as u64).min(nodes.saturating_sub(1)) as usize;
156    let mut writer = RunWriter::new(directory, &input.temporary, control)?;
157    if nodes != 0 {
158        writer.prepare()?;
159    }
160    let mut hash = adjacency_hash(partitions, nodes, degree);
161    let mut edges = 0_u64;
162    sorted.run.read(control, |file| {
163        if nodes == 0 {
164            return Ok(());
165        }
166        let mut candidates = sort::Range::new(file, 0, sorted.count, nodes, control)?;
167        let mut head = candidates.next()?;
168        for node in 0..nodes {
169            control.check()?;
170            let successor = (node + 1) % nodes;
171            let mut selected = Selection::new(
172                partitions.parameters.alpha,
173                degree.saturating_sub(1),
174                control,
175            )?;
176            let mut vectors = BudgetedVec::new(control.memory());
177            vectors.reserve(degree.saturating_sub(1))?;
178            let mut previous = None;
179            while let Some(candidate) = head.filter(|candidate| candidate.source == node) {
180                control.check()?;
181                if previous != Some(candidate.neighbor)
182                    && candidate.neighbor != successor
183                    && !selected.is_full()
184                {
185                    let vector = navigation(input, candidate.neighbor)?;
186                    if selected.consider(
187                        candidate.neighbor,
188                        candidate.distance,
189                        control,
190                        |position, _| {
191                            let neighbor: &NavigationVector = &vectors[position];
192                            neighbor.squared_distance(&vector, control)
193                        },
194                    )? {
195                        vectors.push(vector)?;
196                    }
197                }
198                previous = Some(candidate.neighbor);
199                head = candidates.next()?;
200            }
201            drop(vectors);
202            let mut neighbors = selected.finish();
203            if nodes > 1 {
204                neighbors.push(successor)?;
205            }
206            neighbors.sort_unstable();
207            edges = edges
208                .checked_add(neighbors.len() as u64)
209                .ok_or_else(|| invalid("merged edge overflow"))?;
210            write_word(neighbors.len() as u64, &mut writer, &mut hash)?;
211            for slot in 0..degree {
212                control.check()?;
213                write_word(
214                    neighbors.get(slot).copied().unwrap_or(0),
215                    &mut writer,
216                    &mut hash,
217                )?;
218            }
219        }
220        if head.is_some() {
221            return Err(invalid("global candidate source was not consumed"));
222        }
223        Ok(())
224    })?;
225    let summary = DiskANNMergeSummary {
226        work_order_revision: 1,
227        partitions: *partitions,
228        options,
229        merge_passes: sorted.passes,
230        edges,
231        adjacency_digest: hash.finalize().into(),
232    };
233    drop(sorted);
234    let run = writer.finish()?;
235    control.check()?;
236    Ok(DiskANNMergedGraph {
237        run,
238        nodes,
239        degree,
240        summary,
241        control: control.clone(),
242    })
243}
244
245fn navigation(input: &DiskANNBuildInput, node: u64) -> StorageBackendResult<NavigationVector> {
246    let raw = input.read_node(node)?;
247    match NavigationInput::from_raw(input.dimensions, raw.raw(), &input.control)? {
248        NavigationInput::Navigable(vector) => Ok(vector),
249        NavigationInput::Exact(_) => Err(invalid(
250            "global candidate changed navigation classification",
251        )),
252    }
253}
254
255fn adjacency_hash(partitions: &DiskANNPartitionSummary, nodes: u64, degree: usize) -> Sha256 {
256    crate::diskann_index::format::adjacency_hash(partitions.coverage, nodes, degree)
257}
258
259fn write_word(value: u64, writer: &mut RunWriter, hash: &mut Sha256) -> StorageBackendResult<()> {
260    let bytes = value.to_le_bytes();
261    writer.append(&bytes)?;
262    hash.update(bytes);
263    Ok(())
264}
uqa_storage/diskann_index/build/merge.rs

uqa_storage/diskann_index/build/
merge.rs
