1use super::{
10 codec::decode_value,
11 hnsw_persistence::{
12 checked_level, metadata_from_graph, metadata_header, PersistedHNSWMetadata,
13 PersistedHNSWNode,
14 },
15 record_json,
16 vector_records::{ordinal, tail, vector_bytes, GUARD},
17 TAG_HNSW_METADATA, TAG_HNSW_NODE, TAG_VECTOR,
18};
19use crate::{
20 hnsw_index::HNSWNodeSnapshot,
21 mvcc::{
22 HNSWRecordHeader, HNSWRecordKey as Key, HNSWRecordLayout, HNSWRecordValue as Value,
23 VersionError, VersionResult,
24 },
25 read_control::StorageReadControl,
26};
27use serde_json::value::RawValue;
28use uqa_core::{
29 memory::{Budgeted, BudgetedVec},
30 DocId,
31};
32
33pub struct KeyValueHNSWRecords;
34impl HNSWRecordLayout for KeyValueHNSWRecords {
35 fn metadata_key(
36 &self,
37 key: &[u8],
38 control: &StorageReadControl,
39 ) -> VersionResult<Option<BudgetedVec<u8>>> {
40 control.cancellation().check()?;
41 let count = match key.first() {
42 Some(&TAG_HNSW_METADATA) => 0,
43 Some(&TAG_HNSW_NODE) => 1,
44 _ => return Ok(None),
45 };
46 let (end, _) = tail(key, key[0], count)?;
47 let mut output = BudgetedVec::new(control.memory());
48 output.extend_from_slice(&key[..end])?;
49 output[0] = TAG_HNSW_METADATA;
50 Ok(Some(output))
51 }
52 fn key(
53 &self,
54 metadata: &[u8],
55 address: Key,
56 control: &StorageReadControl,
57 ) -> VersionResult<BudgetedVec<u8>> {
58 control.cancellation().check()?;
59 tail(metadata, TAG_HNSW_METADATA, 0)?;
60 let mut key = BudgetedVec::new(control.memory());
61 key.extend_from_slice(metadata)?;
62 match address {
63 Key::Structure => {
64 key[0] = GUARD;
65 key.push(0)?;
66 }
67 Key::Document(document) => {
68 key[0] = GUARD;
69 key.push(1)?;
70 key.extend_from_slice(&document.to_be_bytes())?;
71 }
72 Key::Vectors => key[0] = TAG_VECTOR,
73 Key::Nodes => key[0] = TAG_HNSW_NODE,
74 Key::Node(node) => {
75 key[0] = TAG_HNSW_NODE;
76 key.extend_from_slice(&node.to_be_bytes())?;
77 }
78 Key::Edge { .. } => {
79 return Err(VersionError::InvalidEncoding(
80 "KeyValue HNSW adjacency is embedded in nodes",
81 ))
82 }
83 }
84 Ok(key)
85 }
86 fn edges_prefix(
87 &self,
88 metadata: &[u8],
89 _source: Option<u64>,
90 control: &StorageReadControl,
91 ) -> VersionResult<Option<BudgetedVec<u8>>> {
92 control.cancellation().check()?;
93 tail(metadata, TAG_HNSW_METADATA, 0)?;
94 Ok(None)
95 }
96 fn header(
97 &self,
98 key: &[u8],
99 value: &[u8],
100 control: &StorageReadControl,
101 ) -> VersionResult<HNSWRecordHeader> {
102 control.cancellation().check()?;
103 tail(key, TAG_HNSW_METADATA, 0)?;
104 let meta: PersistedHNSWMetadata = decode_value(value)?;
105 Ok(metadata_header(&meta)?)
106 }
107
108 fn vector_id(&self, key: &[u8], control: &StorageReadControl) -> VersionResult<(DocId, u32)> {
109 control.cancellation().check()?;
110 let (_, [document, order]) = tail(key, TAG_VECTOR, 2)?;
111 Ok((document, ordinal(order)?))
112 }
113 fn vector(
114 &self,
115 key: &[u8],
116 value: &[u8],
117 control: &StorageReadControl,
118 ) -> VersionResult<(DocId, u32, BudgetedVec<f32>)> {
119 let (document, order) = self.vector_id(key, control)?;
120 Ok((document, order, vector_bytes(value, control)?))
121 }
122 fn node(
123 &self,
124 key: &[u8],
125 value: &[u8],
126 control: &StorageReadControl,
127 ) -> VersionResult<Budgeted<HNSWNodeSnapshot>> {
128 control.cancellation().check()?;
129 let (_, [id, _]) = tail(key, TAG_HNSW_NODE, 1)?;
130 let mut memory = control.memory().reserve(size_of::<HNSWNodeSnapshot>())?;
131 let node: PersistedHNSWNode<&RawValue, &RawValue> =
132 serde_json::from_slice(value).map_err(crate::StorageBackendError::from)?;
133 if id != node.node_id {
134 return Err(VersionError::InvalidEncoding(
135 "HNSW node identity disagrees with its key",
136 ));
137 }
138 let level = checked_level(node.level, "HNSW node level")?;
139 let vector = record_json::array::<f32>(node.raw_vector.get(), control)?;
140 let layers = record_json::array::<&RawValue>(node.neighbors.get(), control)?;
141 if layers.len() != level + 1 {
142 return Err(VersionError::InvalidEncoding(
143 "HNSW adjacency does not match its level",
144 ));
145 }
146 let mut neighbors = BudgetedVec::new(control.memory());
147 for layer in layers.iter() {
148 neighbors.reserve(1)?;
149 let (edges, retained) = record_json::array::<u64>(layer.get(), control)?.into_parts();
150 memory.absorb(retained);
151 neighbors.push(edges)?;
152 }
153 let (raw_vector, retained) = vector.into_parts();
154 memory.absorb(retained);
155 let (neighbors, retained) = neighbors.into_parts();
156 memory.absorb(retained);
157 Ok(Budgeted::new(
158 HNSWNodeSnapshot {
159 node_id: id,
160 doc_id: node.doc_id,
161 vector_ordinal: node.vector_ordinal,
162 raw_vector,
163 level,
164 deleted: node.deleted,
165 neighbors,
166 },
167 memory,
168 ))
169 }
170 fn edge(
171 &self,
172 _key: &[u8],
173 _value: &[u8],
174 _control: &StorageReadControl,
175 ) -> VersionResult<(u64, usize, u64)> {
176 Err(VersionError::InvalidEncoding(
177 "KeyValue HNSW adjacency is embedded in nodes",
178 ))
179 }
180 fn encode(
181 &self,
182 key: &[u8],
183 template: &[u8],
184 value: Value<'_>,
185 control: &StorageReadControl,
186 ) -> VersionResult<BudgetedVec<u8>> {
187 control.cancellation().check()?;
188 match value {
189 Value::Header { meta, revision } => {
190 let header = self.header(key, template, control)?;
191 record_json::encode(
192 &metadata_from_graph(
193 header.dimensions,
194 header.params,
195 meta,
196 revision.ok_or(VersionError::InvalidEncoding(
197 "missing HNSW mutation counter",
198 ))?,
199 )?,
200 control,
201 )
202 }
203 Value::Node(node) => {
204 let (_, [id, _]) = tail(key, TAG_HNSW_NODE, 1)?;
205 if id != node.node_id {
206 return Err(VersionError::InvalidEncoding(
207 "HNSW node identity disagrees with its key",
208 ));
209 }
210 record_json::encode(&PersistedHNSWNode::try_from(node)?, control)
211 }
212 Value::Edge => Err(VersionError::InvalidEncoding(
213 "KeyValue HNSW adjacency is embedded in nodes",
214 )),
215 }
216 }
217}