uqa_storage/diskann_index/format/
quantization.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, zeros, DiskANNGeneration, DiskANNManifest};
10use crate::diskann_index::{
11    metric::checkpoint,
12    pq::{PQCodebook, PQTrainingOptions, PQTrainingSummary},
13};
14use crate::{read_control::StorageReadControl, StorageBackendResult};
15
16const MAGIC: [u8; 8] = *b"UQADNPQ\0";
17const METADATA_BYTES: usize = 64;
18const SCALAR_REVISION: u32 = 1;
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
21pub struct DiskANNQuantizationIdentity {
22    pub(super) generation: DiskANNGeneration,
23    pub(super) dimensions: u32,
24    pub(super) chunks: usize,
25    pub(super) centroids: u16,
26    pub(super) nodes: u64,
27    pub(super) digest: [u8; 32],
28}
29
30impl DiskANNQuantizationIdentity {
31    pub fn codebook_digest(self) -> [u8; 32] {
32        self.digest
33    }
34    pub fn node_count(self) -> u64 {
35        self.nodes
36    }
37    pub fn pq_bytes(self) -> usize {
38        self.chunks
39    }
40
41    fn new(generation: DiskANNGeneration, codebook: &PQCodebook, digest: [u8; 32]) -> Self {
42        Self {
43            generation,
44            dimensions: codebook.dimensions() as u32,
45            chunks: codebook.pq_bytes(),
46            centroids: codebook.centroid_count(),
47            nodes: codebook.training().observed_vectors,
48            digest,
49        }
50    }
51}
52
53fn size(chunks: usize, scalars: usize) -> StorageBackendResult<usize> {
54    chunks
55        .checked_add(1)
56        .and_then(|count| count.checked_mul(4))
57        .and_then(|offsets| {
58            scalars
59                .checked_mul(8)
60                .and_then(|coordinates| offsets.checked_add(coordinates))
61        })
62        .and_then(|body| body.checked_add(METADATA_BYTES))
63        .ok_or_else(|| invalid("codebook body size overflow"))
64}
65
66pub fn encode_codebook(
67    generation: DiskANNGeneration,
68    codebook: &PQCodebook,
69    control: &StorageReadControl,
70) -> StorageBackendResult<(BudgetedVec<u8>, DiskANNQuantizationIdentity)> {
71    let summary = codebook.training();
72    let options = summary.options;
73    let dimensions =
74        u32::try_from(codebook.dimensions()).map_err(|_| invalid("codebook dimension range"))?;
75    let chunks = codebook.pq_bytes();
76    let count = codebook.centroid_count();
77    let scalars = summary.validate(dimensions, chunks, count)?;
78    let mut bytes = record::begin(MAGIC, generation, size(chunks, scalars)?, control)?;
79    bytes.extend_from_slice(&dimensions.to_le_bytes())?;
80    bytes.extend_from_slice(&(chunks as u32).to_le_bytes())?;
81    bytes.extend_from_slice(&count.to_le_bytes())?;
82    bytes.extend_from_slice(&[0; 2])?;
83    for value in [
84        SCALAR_REVISION,
85        PQCodebook::CODEC_REVISION,
86        PQCodebook::TRAINING_REVISION,
87        options.max_samples,
88        options.max_iterations,
89    ] {
90        bytes.extend_from_slice(&value.to_le_bytes())?;
91    }
92    bytes.extend_from_slice(&options.max_centroids.to_le_bytes())?;
93    bytes.extend_from_slice(&[0; 2])?;
94    bytes.extend_from_slice(&summary.sampled_vectors.to_le_bytes())?;
95    for value in [options.seed, summary.observed_vectors, scalars as u64] {
96        bytes.extend_from_slice(&value.to_le_bytes())?;
97    }
98    for chunk in 0..chunks {
99        checkpoint(chunk, control)?;
100        let start = codebook.chunk_range(chunk).expect("validated chunk").start;
101        bytes.extend_from_slice(&(start as u32).to_le_bytes())?;
102    }
103    bytes.extend_from_slice(&dimensions.to_le_bytes())?;
104    for chunk in 0..chunks {
105        for (offset, coordinate) in codebook
106            .chunk_centroids(chunk)
107            .expect("validated chunk")
108            .iter()
109            .enumerate()
110        {
111            checkpoint(offset, control)?;
112            bytes.extend_from_slice(&coordinate.to_bits().to_le_bytes())?;
113        }
114    }
115    let bytes = record::finish(bytes, control)?;
116    let digest = record::artifact_digest(&bytes, control)?;
117    Ok((
118        bytes,
119        DiskANNQuantizationIdentity::new(generation, codebook, digest),
120    ))
121}
122
123pub fn decode_codebook(
124    manifest: &DiskANNManifest,
125    bytes: &[u8],
126    control: &StorageReadControl,
127) -> StorageBackendResult<(PQCodebook, DiskANNQuantizationIdentity)> {
128    let input = manifest.input();
129    let body = record::open(MAGIC, input.generation, bytes, control)?;
130    if input.nodes == 0
131        || body.len() < METADATA_BYTES
132        || record::artifact_digest(bytes, control)? != input.artifacts.codebook
133    {
134        return Err(invalid(
135            "codebook is absent or differs from manifest digest",
136        ));
137    }
138    let dimensions = record::u32_at(body, 0)?;
139    let chunks = record::u32_at(body, 4)? as usize;
140    let count = u16::from_le_bytes(field(body, 8)?);
141    zeros(&body[10..12], control)?;
142    zeros(&body[34..36], control)?;
143    for (offset, expected) in [
144        (12, SCALAR_REVISION),
145        (16, PQCodebook::CODEC_REVISION),
146        (20, PQCodebook::TRAINING_REVISION),
147    ] {
148        if record::u32_at(body, offset)? != expected {
149            return Err(invalid(
150                "unrecognized codebook scalar, codec or training revision",
151            ));
152        }
153    }
154    let summary = PQTrainingSummary {
155        options: PQTrainingOptions {
156            max_samples: record::u32_at(body, 24)?,
157            max_iterations: record::u32_at(body, 28)?,
158            max_centroids: u16::from_le_bytes(field(body, 32)?),
159            seed: record::u64_at(body, 40)?,
160        },
161        sampled_vectors: record::u32_at(body, 36)?,
162        observed_vectors: record::u64_at(body, 48)?,
163    };
164    if dimensions != input.dimensions
165        || chunks != input.parameters.pq_bytes
166        || summary.observed_vectors != input.nodes
167        || summary.options.seed != input.parameters.seed
168        || manifest
169            .build_provenance()
170            .is_some_and(|build| build.training() != summary.options)
171    {
172        return Err(invalid(
173            "codebook shape, corpus or seed differs from manifest",
174        ));
175    }
176    let scalars = summary.validate(dimensions, chunks, count)?;
177    if record::u64_at(body, 56)? != scalars as u64 || body.len() != size(chunks, scalars)? {
178        return Err(invalid("codebook scalar count or length differs"));
179    }
180    for chunk in 0..=chunks {
181        checkpoint(chunk, control)?;
182        let expected = if chunk == chunks {
183            dimensions as usize
184        } else {
185            crate::diskann_index::pq::chunk_range(dimensions as usize, chunks, chunk).start
186        };
187        if record::u32_at(body, METADATA_BYTES + chunk * 4)? as usize != expected {
188            return Err(invalid("codebook chunk offsets differ"));
189        }
190    }
191    let mut centroids = BudgetedVec::new(control.memory());
192    centroids.reserve(scalars)?;
193    let start = METADATA_BYTES + (chunks + 1) * 4;
194    for (offset, scalar) in body[start..].chunks_exact(8).enumerate() {
195        checkpoint(offset, control)?;
196        centroids.push(f64::from_bits(u64::from_le_bytes(field(scalar, 0)?)))?;
197    }
198    let codebook = PQCodebook::restore(dimensions, chunks, count, summary, centroids, control)?;
199    let identity =
200        DiskANNQuantizationIdentity::new(input.generation, &codebook, input.artifacts.codebook);
201    Ok((codebook, identity))
202}
uqa_storage/diskann_index/format/quantization.rs

uqa_storage/diskann_index/format/
quantization.rs
