uqa_storage/diskann_index/format/
coverage.rs1use sha2::{Digest, Sha256};
8use uqa_core::DocId;
9
10use super::{invalid, DiskANNGeneration, DiskANNVectorVersion};
11use crate::diskann_index::metric::checkpoint;
12use crate::vector_index::validate_vector_values_controlled;
13use crate::{read_control::StorageReadControl, StorageBackendResult};
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub struct DiskANNBuildCoverage {
18 pub(super) generation: DiskANNGeneration,
19 pub(super) dimensions: u32,
20 pub(super) vectors: u64,
21 pub(super) digest: [u8; 32],
22}
23
24impl DiskANNBuildCoverage {
25 pub fn generation(self) -> DiskANNGeneration {
26 self.generation
27 }
28
29 pub fn vector_count(self) -> u64 {
30 self.vectors
31 }
32 pub fn digest(self) -> [u8; 32] {
33 self.digest
34 }
35}
36
37pub struct DiskANNCoverageBuilder {
38 generation: DiskANNGeneration,
39 dimensions: u32,
40 vectors: u64,
41 previous: Option<(DocId, u32)>,
42 hash: Sha256,
43}
44
45impl DiskANNCoverageBuilder {
46 pub fn new(generation: DiskANNGeneration, dimensions: u32) -> StorageBackendResult<Self> {
47 if dimensions == 0 {
48 return Err(invalid("coverage dimensions must be positive"));
49 }
50 let mut hash = Sha256::new();
51 hash.update(b"UQA DiskANN canonical coverage\0\x01");
52 hash.update(generation.database);
53 for value in [generation.table, generation.index, generation.generation] {
54 hash.update(value.to_le_bytes());
55 }
56 hash.update(dimensions.to_le_bytes());
57 Ok(Self {
58 generation,
59 dimensions,
60 vectors: 0,
61 previous: None,
62 hash,
63 })
64 }
65
66 pub fn push(
68 &mut self,
69 doc: DocId,
70 ordinal: u32,
71 version: DiskANNVectorVersion,
72 raw: &[f32],
73 control: &StorageReadControl,
74 ) -> StorageBackendResult<()> {
75 control.check()?;
76 let ordered = match self.previous {
77 None => ordinal == 0,
78 Some((previous, last)) if previous == doc => last.checked_add(1) == Some(ordinal),
79 Some((previous, _)) => doc > previous && ordinal == 0,
80 };
81 if !ordered {
82 return Err(invalid(
83 "coverage input must contain ordered complete tensor ordinals",
84 ));
85 }
86 validate_vector_values_controlled(self.dimensions, raw, Some(control))?;
87 let vectors = self
88 .vectors
89 .checked_add(1)
90 .ok_or_else(|| invalid("coverage count overflow"))?;
91 let mut hash = self.hash.clone();
92 hash.update(doc.to_le_bytes());
93 hash.update(ordinal.to_le_bytes());
94 hash.update(version.writer.database().as_bytes());
95 hash.update(version.writer.allocation().to_le_bytes());
96 hash.update(version.revision.to_le_bytes());
97 for (offset, value) in raw.iter().enumerate() {
98 checkpoint(offset, control)?;
99 hash.update(value.to_bits().to_le_bytes());
100 }
101 control.check()?;
102 self.hash = hash;
103 self.vectors = vectors;
104 self.previous = Some((doc, ordinal));
105 Ok(())
106 }
107
108 pub fn finish(mut self) -> DiskANNBuildCoverage {
109 self.hash.update(self.vectors.to_le_bytes());
110 DiskANNBuildCoverage {
111 generation: self.generation,
112 dimensions: self.dimensions,
113 vectors: self.vectors,
114 digest: self.hash.finalize().into(),
115 }
116 }
117}