uqa_storage/diskann_index/format/
codes.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7use uqa_core::memory::BudgetedVec;
8
9use super::{field, invalid, record, DiskANNQuantizationIdentity};
10use crate::diskann_index::{metric::checkpoint, pq::PQCodebook};
11use crate::{read_control::StorageReadControl, StorageBackendResult};
12
13const MAGIC: [u8; 8] = *b"UQADNCD\0";
14const METADATA_BYTES: usize = 72;
15
16#[derive(Debug)]
17pub struct DiskANNCodeBatch<'a> {
18    start: u64,
19    count: u64,
20    chunks: usize,
21    bytes: &'a [u8],
22}
23
24impl DiskANNCodeBatch<'_> {
25    pub fn first_node(&self) -> u64 {
26        self.start
27    }
28    pub fn node_count(&self) -> u64 {
29        self.count
30    }
31    pub fn bytes(&self) -> &[u8] {
32        self.bytes
33    }
34    pub fn code(&self, node: u64) -> Option<&[u8]> {
35        let offset = node.checked_sub(self.start)?;
36        if offset >= self.count {
37            return None;
38        }
39        let start = offset as usize * self.chunks;
40        Some(&self.bytes[start..start + self.chunks])
41    }
42}
43
44impl DiskANNQuantizationIdentity {
45    pub fn encode_codes(
46        self,
47        start: u64,
48        codes: &[u8],
49        control: &StorageReadControl,
50    ) -> StorageBackendResult<BudgetedVec<u8>> {
51        let count = self.validate_codes(start, codes, control)?;
52        let size = METADATA_BYTES
53            .checked_add(codes.len())
54            .ok_or_else(|| invalid("code record size overflow"))?;
55        let mut bytes = record::begin(MAGIC, self.generation, size, control)?;
56        for value in [
57            self.dimensions,
58            self.chunks as u32,
59            u32::from(self.centroids),
60            PQCodebook::CODEC_REVISION,
61        ] {
62            bytes.extend_from_slice(&value.to_le_bytes())?;
63        }
64        for value in [self.nodes, start, count] {
65            bytes.extend_from_slice(&value.to_le_bytes())?;
66        }
67        bytes.extend_from_slice(&self.digest)?;
68        for part in codes.chunks(4096) {
69            control.check()?;
70            bytes.extend_from_slice(part)?;
71        }
72        record::finish(bytes, control)
73    }
74
75    pub fn decode_codes<'a>(
76        self,
77        start: u64,
78        bytes: &'a [u8],
79        control: &StorageReadControl,
80    ) -> StorageBackendResult<DiskANNCodeBatch<'a>> {
81        let body = record::open(MAGIC, self.generation, bytes, control)?;
82        if body.len() < METADATA_BYTES {
83            return Err(invalid("truncated code batch"));
84        }
85        for (offset, expected) in [
86            (0, self.dimensions),
87            (4, self.chunks as u32),
88            (8, u32::from(self.centroids)),
89            (12, PQCodebook::CODEC_REVISION),
90        ] {
91            if record::u32_at(body, offset)? != expected {
92                return Err(invalid("code batch layout differs from codebook"));
93            }
94        }
95        if record::u64_at(body, 16)? != self.nodes
96            || record::u64_at(body, 24)? != start
97            || field::<32>(body, 40)? != self.digest
98        {
99            return Err(invalid("code batch corpus, address or codebook differs"));
100        }
101        let codes = &body[METADATA_BYTES..];
102        let count = self.validate_codes(start, codes, control)?;
103        if record::u64_at(body, 32)? != count {
104            return Err(invalid("code batch count differs"));
105        }
106        Ok(DiskANNCodeBatch {
107            start,
108            count,
109            chunks: self.chunks,
110            bytes: codes,
111        })
112    }
113
114    fn validate_codes(
115        self,
116        start: u64,
117        codes: &[u8],
118        control: &StorageReadControl,
119    ) -> StorageBackendResult<u64> {
120        control.check()?;
121        if codes.is_empty() || !codes.len().is_multiple_of(self.chunks) {
122            return Err(invalid("codes must contain complete nonempty node entries"));
123        }
124        let count = (codes.len() / self.chunks) as u64;
125        if start.checked_add(count).is_none_or(|end| end > self.nodes) {
126            return Err(invalid("code batch exceeds the generation"));
127        }
128        for (offset, &label) in codes.iter().enumerate() {
129            checkpoint(offset, control)?;
130            if u16::from(label) >= self.centroids {
131                return Err(invalid("code label exceeds actual centroid count"));
132            }
133        }
134        control.check()?;
135        Ok(count)
136    }
137}
uqa_storage/diskann_index/format/codes.rs

uqa_storage/diskann_index/format/
codes.rs
