uqa_storage/diskann_index/format/
origins.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, DocId};
9
10use super::{invalid, record, DiskANNCanonicalOrigin, DiskANNGeneration, CANONICAL_ORIGIN_BYTES};
11use crate::{read_control::StorageReadControl, StorageBackendResult};
12
13#[cfg(test)]
14mod tests;
15
16const MAGIC: [u8; 8] = *b"UQADNOR\0";
17const METADATA_BYTES: usize = 32;
18pub const ORIGIN_ENTRY_BYTES: usize = 8 + CANONICAL_ORIGIN_BYTES;
19pub const ORIGIN_BATCH_DOCUMENTS: usize = 64;
20
21/// Complete ordered origin artifact metadata. This is an integrity descriptor, not authority to publish or retire changes.
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub struct DiskANNOriginSummary {
24    documents: u64,
25    digest: [u8; 32],
26}
27
28impl DiskANNOriginSummary {
29    pub const ENCODED_BYTES: usize = 40;
30
31    pub(crate) fn new(documents: u64, digest: [u8; 32]) -> StorageBackendResult<Self> {
32        if documents == 0 && digest != <[u8; 32]>::from(Sha256::digest([])) {
33            return Err(invalid("empty origin stream requires the empty digest"));
34        }
35        Ok(Self { documents, digest })
36    }
37
38    pub fn documents(self) -> u64 {
39        self.documents
40    }
41    pub fn digest(self) -> [u8; 32] {
42        self.digest
43    }
44}
45
46/// One complete selected tensor, including an explicit empty replacement.
47#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub struct DiskANNOriginEntry {
49    document: DocId,
50    origin: DiskANNCanonicalOrigin,
51}
52
53impl DiskANNOriginEntry {
54    pub fn new(document: DocId, origin: DiskANNCanonicalOrigin) -> Self {
55        Self { document, origin }
56    }
57    pub fn document(self) -> DocId {
58        self.document
59    }
60    pub fn origin(self) -> DiskANNCanonicalOrigin {
61        self.origin
62    }
63
64    pub(crate) fn encode(self) -> [u8; ORIGIN_ENTRY_BYTES] {
65        let mut bytes = [0; ORIGIN_ENTRY_BYTES];
66        bytes[..8].copy_from_slice(&self.document.to_le_bytes());
67        bytes[8..].copy_from_slice(&self.origin.encode());
68        bytes
69    }
70
71    pub(crate) fn decode(bytes: &[u8], dimensions: u32) -> StorageBackendResult<Self> {
72        if bytes.len() != ORIGIN_ENTRY_BYTES {
73            return Err(invalid("origin entry length differs"));
74        }
75        Ok(Self::new(
76            record::u64_at(bytes, 0)?,
77            DiskANNCanonicalOrigin::decode(&bytes[8..], dimensions)?,
78        ))
79    }
80}
81
82/// Fixed-capacity batches use dense document positions while preserving sparse document IDs. Point lookup needs no resident directory.
83#[derive(Debug, Clone, Copy)]
84pub struct DiskANNOriginLayout {
85    generation: DiskANNGeneration,
86    dimensions: u32,
87    documents: u64,
88}
89
90pub struct DiskANNOriginBatch<'a> {
91    dimensions: u32,
92    entries: &'a [u8],
93}
94
95impl DiskANNOriginBatch<'_> {
96    pub fn len(&self) -> usize {
97        self.entries.len() / ORIGIN_ENTRY_BYTES
98    }
99    pub fn is_empty(&self) -> bool {
100        self.entries.is_empty()
101    }
102    pub fn bytes(&self) -> &[u8] {
103        self.entries
104    }
105    pub fn entry(&self, index: usize) -> Option<DiskANNOriginEntry> {
106        let start = index.checked_mul(ORIGIN_ENTRY_BYTES)?;
107        let end = start.checked_add(ORIGIN_ENTRY_BYTES)?;
108        Some(
109            DiskANNOriginEntry::decode(self.entries.get(start..end)?, self.dimensions)
110                .expect("validated origin entry"),
111        )
112    }
113}
114
115impl DiskANNOriginLayout {
116    pub const MAX_ENCODED_BYTES: usize =
117        record::HEADER_BYTES + METADATA_BYTES + ORIGIN_BATCH_DOCUMENTS * ORIGIN_ENTRY_BYTES;
118
119    pub fn new(
120        generation: DiskANNGeneration,
121        dimensions: u32,
122        documents: u64,
123    ) -> StorageBackendResult<Self> {
124        if dimensions == 0 {
125            return Err(invalid("origin dimensions must be positive"));
126        }
127        Ok(Self {
128            generation,
129            dimensions,
130            documents,
131        })
132    }
133
134    pub fn encode(
135        self,
136        first: u64,
137        entries: &[DiskANNOriginEntry],
138        control: &StorageReadControl,
139    ) -> StorageBackendResult<BudgetedVec<u8>> {
140        self.check_extent(first, entries.len())?;
141        let mut bytes = record::begin(
142            MAGIC,
143            self.generation,
144            self.encoded_bytes(first)? - record::HEADER_BYTES,
145            control,
146        )?;
147        bytes.extend_from_slice(&self.dimensions.to_le_bytes())?;
148        bytes.extend_from_slice(&(ORIGIN_BATCH_DOCUMENTS as u32).to_le_bytes())?;
149        for value in [self.documents, first, entries.len() as u64] {
150            bytes.extend_from_slice(&value.to_le_bytes())?;
151        }
152        let mut previous = None;
153        for &entry in entries {
154            control.check()?;
155            let encoded = entry.encode();
156            DiskANNOriginEntry::decode(&encoded, self.dimensions)?;
157            ordered(&mut previous, entry)?;
158            bytes.extend_from_slice(&encoded)?;
159        }
160        record::finish(bytes, control)
161    }
162
163    pub fn decode<'a>(
164        self,
165        first: u64,
166        bytes: &'a [u8],
167        control: &StorageReadControl,
168    ) -> StorageBackendResult<DiskANNOriginBatch<'a>> {
169        let body = record::open(MAGIC, self.generation, bytes, control)?;
170        if body.len() < METADATA_BYTES
171            || !(body.len() - METADATA_BYTES).is_multiple_of(ORIGIN_ENTRY_BYTES)
172        {
173            return Err(invalid("origin batch length differs"));
174        }
175        let entries = &body[METADATA_BYTES..];
176        let count = entries.len() / ORIGIN_ENTRY_BYTES;
177        self.check_extent(first, count)?;
178        if record::u32_at(body, 0)? != self.dimensions
179            || record::u32_at(body, 4)? != ORIGIN_BATCH_DOCUMENTS as u32
180            || record::u64_at(body, 8)? != self.documents
181            || record::u64_at(body, 16)? != first
182            || record::u64_at(body, 24)? != count as u64
183        {
184            return Err(invalid("origin batch shape or address differs"));
185        }
186        let mut previous = None;
187        for bytes in entries.as_chunks::<ORIGIN_ENTRY_BYTES>().0 {
188            control.check()?;
189            ordered(
190                &mut previous,
191                DiskANNOriginEntry::decode(bytes, self.dimensions)?,
192            )?;
193        }
194        control.check()?;
195        Ok(DiskANNOriginBatch {
196            dimensions: self.dimensions,
197            entries,
198        })
199    }
200
201    fn check_extent(self, first: u64, count: usize) -> StorageBackendResult<()> {
202        if first >= self.documents
203            || !first.is_multiple_of(ORIGIN_BATCH_DOCUMENTS as u64)
204            || count as u64 != (self.documents - first).min(ORIGIN_BATCH_DOCUMENTS as u64)
205        {
206            return Err(invalid(
207                "origin batch does not cover its exact fixed extent",
208            ));
209        }
210        Ok(())
211    }
212
213    pub fn encoded_bytes(self, first: u64) -> StorageBackendResult<usize> {
214        let count = self
215            .documents
216            .saturating_sub(first)
217            .min(ORIGIN_BATCH_DOCUMENTS as u64) as usize;
218        self.check_extent(first, count)?;
219        Ok(record::HEADER_BYTES + METADATA_BYTES + count * ORIGIN_ENTRY_BYTES)
220    }
221}
222
223fn ordered(previous: &mut Option<DocId>, entry: DiskANNOriginEntry) -> StorageBackendResult<()> {
224    if previous.is_some_and(|last| last >= entry.document) {
225        return Err(invalid("origin documents must be unique and ordered"));
226    }
227    *previous = Some(entry.document);
228    Ok(())
229}
uqa_storage/diskann_index/format/origins.rs

uqa_storage/diskann_index/format/
origins.rs
