uqa_storage/vector_index/
canonical.rs1use crate::{read_control::StorageReadControl, StorageBackendResult};
10use std::sync::Arc;
11use uqa_core::{
12 memory::{Budgeted, BudgetedVec},
13 DocId,
14};
15
16mod fingerprint;
17mod ordered;
18mod selection;
19pub(crate) use fingerprint::selected_fingerprint;
20pub(crate) use ordered::{copy_vector, ordinal_count};
21pub use selection::SelectedVectorRead;
22
23pub trait VectorRead: Send + Sync {
25 fn check_control(&self, control: &StorageReadControl) -> StorageBackendResult<()>;
26 fn dimensions(&self) -> u32;
27 fn corpus_fingerprint(
29 &self,
30 control: &StorageReadControl,
31 ) -> StorageBackendResult<Option<[u8; 32]>> {
32 self.check_control(control)?;
33 Ok(None)
34 }
35 fn next_document_after(
36 &self,
37 after: Option<DocId>,
38 control: &StorageReadControl,
39 ) -> StorageBackendResult<Option<DocId>>;
40 fn document_vector_count(
42 &self,
43 document: DocId,
44 control: &StorageReadControl,
45 ) -> StorageBackendResult<u64>;
46 fn read_vector(
48 &self,
49 document: DocId,
50 ordinal: u32,
51 control: &StorageReadControl,
52 ) -> StorageBackendResult<Option<BudgetedVec<f32>>>;
53}
54
55pub type VectorReadSnapshot = Arc<dyn VectorRead>;
56
57pub fn decode_vector_bytes(
59 bytes: &[u8],
60 dimensions: u32,
61 control: &StorageReadControl,
62) -> StorageBackendResult<BudgetedVec<f32>> {
63 control.check()?;
64 let expected = usize::try_from(dimensions)
65 .map_err(|_| uqa_core::memory::MemoryError::SizeOverflow)?
66 .checked_mul(4)
67 .ok_or(uqa_core::memory::MemoryError::SizeOverflow)?;
68 if bytes.len() != expected {
69 return Err(
70 crate::mvcc::VersionError::InvalidEncoding("canonical vector width mismatch")
71 .into_storage_error(),
72 );
73 }
74 let mut vector = BudgetedVec::new(control.memory());
75 vector.reserve(bytes.len() / 4)?;
76 for (position, bytes) in bytes.chunks_exact(4).enumerate() {
77 if position.is_multiple_of(1024) {
78 control.check()?;
79 }
80 vector.push(f32::from_le_bytes(
81 bytes.try_into().expect("four-byte float"),
82 ))?;
83 }
84 super::validate_vector_values_controlled(dimensions, &vector, Some(control))?;
85 Ok(vector)
86}
87
88impl<T: VectorRead> VectorRead for Budgeted<T> {
89 fn check_control(&self, control: &StorageReadControl) -> StorageBackendResult<()> {
90 (**self).check_control(control)
91 }
92 fn dimensions(&self) -> u32 {
93 (**self).dimensions()
94 }
95 fn corpus_fingerprint(
96 &self,
97 control: &StorageReadControl,
98 ) -> StorageBackendResult<Option<[u8; 32]>> {
99 (**self).corpus_fingerprint(control)
100 }
101 fn next_document_after(
102 &self,
103 after: Option<DocId>,
104 control: &StorageReadControl,
105 ) -> StorageBackendResult<Option<DocId>> {
106 (**self).next_document_after(after, control)
107 }
108 fn document_vector_count(
109 &self,
110 document: DocId,
111 control: &StorageReadControl,
112 ) -> StorageBackendResult<u64> {
113 (**self).document_vector_count(document, control)
114 }
115 fn read_vector(
116 &self,
117 document: DocId,
118 ordinal: u32,
119 control: &StorageReadControl,
120 ) -> StorageBackendResult<Option<BudgetedVec<f32>>> {
121 (**self).read_vector(document, ordinal, control)
122 }
123}