1use crate::{HawDBError, Result, ValueRef, VectorDimension};
9
10#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
11pub struct Vector {
12 encoded: Vec<u8>,
13}
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
16pub struct VectorRef<'a> {
17 encoded: &'a [u8],
18 dimension: VectorDimension,
19}
20
21impl Vector {
22 pub fn new(coordinates: &[f32]) -> Result<Self> {
23 let dimension = u32::try_from(coordinates.len())
24 .ok()
25 .and_then(VectorDimension::new)
26 .ok_or_else(|| invalid("dimension must be a positive u32"))?;
27 if coordinates.iter().any(|value| !value.is_finite()) {
28 return Err(invalid("coordinates must be finite"));
29 }
30 let bytes = coordinates
31 .len()
32 .checked_mul(4)
33 .and_then(|n| n.checked_add(4))
34 .ok_or_else(|| invalid("payload length overflow"))?;
35 let mut encoded = Vec::with_capacity(bytes);
36 encoded.extend_from_slice(&dimension.get().to_le_bytes());
37 for value in coordinates {
38 encoded.extend_from_slice(&value.to_le_bytes());
39 }
40 Ok(Self { encoded })
41 }
42
43 pub fn from_encoded(encoded: Vec<u8>) -> Result<Self> {
45 Self::from_encoded_with_work(encoded, || Ok(()), |_| Ok(()), std::convert::identity)
46 }
47
48 pub fn from_encoded_with_work<T, E>(
50 encoded: Vec<u8>,
51 start: impl FnMut() -> std::result::Result<T, E>,
52 finish: impl FnMut(T) -> std::result::Result<(), E>,
53 invalid: impl Fn(HawDBError) -> E,
54 ) -> std::result::Result<Self, E> {
55 VectorRef::from_encoded_with_work(&encoded, start, finish, invalid)?;
56 Ok(Self { encoded })
57 }
58
59 pub fn dimension(&self) -> VectorDimension {
60 VectorDimension::new(u32::from_le_bytes(
61 self.encoded[..4].try_into().expect("validated dimension"),
62 ))
63 .expect("nonzero dimension")
64 }
65 pub fn retained_capacity(&self) -> usize {
66 self.encoded.capacity()
67 }
68
69 pub fn as_ref(&self) -> VectorRef<'_> {
70 VectorRef {
71 dimension: self.dimension(),
72 encoded: &self.encoded,
73 }
74 }
75
76 pub fn encoded(&self) -> &[u8] {
77 &self.encoded
78 }
79}
80
81impl<'a> VectorRef<'a> {
82 pub fn from_encoded(encoded: &'a [u8]) -> Result<Self> {
84 Self::from_encoded_with_work(encoded, || Ok(()), |_| Ok(()), std::convert::identity)
85 }
86
87 pub fn from_encoded_with_work<T, E>(
89 encoded: &'a [u8],
90 mut start: impl FnMut() -> std::result::Result<T, E>,
91 mut finish: impl FnMut(T) -> std::result::Result<(), E>,
92 map_invalid: impl Fn(HawDBError) -> E,
93 ) -> std::result::Result<Self, E> {
94 let shape = || -> Result<VectorDimension> {
95 let prefix: [u8; 4] = encoded
96 .get(..4)
97 .and_then(|bytes| bytes.try_into().ok())
98 .ok_or_else(|| invalid("truncated dimension"))?;
99 let dimension = VectorDimension::new(u32::from_le_bytes(prefix))
100 .ok_or_else(|| invalid("dimension must be positive"))?;
101 let expected = usize::try_from(dimension.get())
102 .ok()
103 .and_then(|n| n.checked_mul(4))
104 .and_then(|n| n.checked_add(4))
105 .ok_or_else(|| invalid("payload length overflow"))?;
106 if encoded.len() != expected {
107 return Err(invalid("payload length does not match dimension"));
108 }
109 Ok(dimension)
110 };
111 let dimension = shape().map_err(&map_invalid)?;
112 for block in encoded[4..].chunks(64 * 1024) {
113 let unit = start()?;
114 let finite = block.chunks_exact(4).all(|bytes| {
115 f32::from_le_bytes(bytes.try_into().expect("fixed-width coordinate")).is_finite()
116 });
117 finish(unit)?;
118 if !finite {
119 return Err(map_invalid(invalid("coordinates must be finite")));
120 }
121 }
122 Ok(Self { dimension, encoded })
123 }
124
125 pub const fn dimension(self) -> VectorDimension {
126 self.dimension
127 }
128 pub const fn encoded(self) -> &'a [u8] {
129 self.encoded
130 }
131
132 pub fn coordinates(self) -> impl ExactSizeIterator<Item = f32> + 'a {
133 self.encoded[4..]
134 .chunks_exact(4)
135 .map(|bytes| f32::from_le_bytes(bytes.try_into().expect("fixed-width coordinate")))
136 }
137
138 pub fn to_owned(self) -> Vector {
140 Vector {
141 encoded: self.encoded.to_vec(),
142 }
143 }
144}
145
146#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
147pub enum VectorDistance {
148 Cosine,
149 Euclidean,
150}
151
152impl VectorDistance {
153 pub fn from_function_name(name: &str) -> Option<Self> {
154 if name.eq_ignore_ascii_case("cosine_distance") {
155 Some(Self::Cosine)
156 } else if name.eq_ignore_ascii_case("l2_distance") {
157 Some(Self::Euclidean)
158 } else {
159 None
160 }
161 }
162
163 pub const fn function_name(self) -> &'static str {
164 match self {
165 Self::Cosine => "cosine_distance",
166 Self::Euclidean => "l2_distance",
167 }
168 }
169
170 pub fn distance(self, left: VectorRef<'_>, right: VectorRef<'_>) -> Result<f64> {
171 self.distance_with_work(left, right, || Ok(()))
172 }
173
174 pub fn distance_with_work(
176 self,
177 left: VectorRef<'_>,
178 right: VectorRef<'_>,
179 checkpoint: impl FnMut() -> Result<()>,
180 ) -> Result<f64> {
181 if left.dimension != right.dimension {
182 return Err(invalid("distance dimensions must match"));
183 }
184 self.distance_coordinates(left.coordinates(), right.coordinates(), checkpoint)
185 }
186
187 pub fn distance_slices_with_work(
190 self,
191 left: &[f32],
192 right: &[f32],
193 checkpoint: impl FnMut() -> Result<()>,
194 ) -> Result<f64> {
195 if left.len() != right.len() {
196 return Err(invalid("distance dimensions must match"));
197 }
198 if left.is_empty() || u32::try_from(left.len()).is_err() {
199 return Err(invalid("dimension must be a positive u32"));
200 }
201 self.distance_coordinates(left.iter().copied(), right.iter().copied(), checkpoint)
202 }
203
204 pub fn distance_slice_to_vector_with_work(
207 self,
208 left: &[f32],
209 right: VectorRef<'_>,
210 checkpoint: impl FnMut() -> Result<()>,
211 ) -> Result<f64> {
212 if left.len() != right.dimension.get() as usize {
213 return Err(invalid("distance dimensions must match"));
214 }
215 self.distance_coordinates(left.iter().copied(), right.coordinates(), checkpoint)
216 }
217
218 fn distance_coordinates(
219 self,
220 left: impl Iterator<Item = f32>,
221 right: impl Iterator<Item = f32>,
222 mut checkpoint: impl FnMut() -> Result<()>,
223 ) -> Result<f64> {
224 let (mut dot, mut left_norm, mut right_norm, mut squared) =
225 (0.0f64, 0.0f64, 0.0f64, 0.0f64);
226 for (index, (left, right)) in left.zip(right).enumerate() {
227 if index % (64 * 1024 / (2 * std::mem::size_of::<f32>())) == 0 {
229 checkpoint()?;
230 }
231 if !left.is_finite() || !right.is_finite() {
232 return Err(invalid("coordinates must be finite"));
233 }
234 let (left, right) = (f64::from(left), f64::from(right));
235 match self {
236 Self::Euclidean => squared += (left - right).powi(2),
237 Self::Cosine => {
238 dot += left * right;
239 left_norm += left * left;
240 right_norm += right * right;
241 }
242 }
243 }
244 checkpoint()?;
245 match self {
246 Self::Euclidean => Ok(squared.sqrt()),
247 Self::Cosine if left_norm == 0.0 || right_norm == 0.0 => {
248 Err(invalid("cosine distance requires nonzero vectors"))
249 }
250 Self::Cosine => {
251 Ok(1.0 - (dot / (left_norm.sqrt() * right_norm.sqrt())).clamp(-1.0, 1.0))
252 }
253 }
254 }
255
256 pub fn evaluate(self, left: ValueRef<'_>, right: ValueRef<'_>) -> Result<Option<f64>> {
257 match (left, right) {
258 (ValueRef::Null, _) | (_, ValueRef::Null) => Ok(None),
259 (ValueRef::Vector(left), ValueRef::Vector(right)) => {
260 self.distance(left, right).map(Some)
261 }
262 _ => Err(invalid("distance arguments must be vectors")),
263 }
264 }
265}
266
267fn invalid(reason: &str) -> HawDBError {
268 HawDBError::Semantic(format!("invalid vector: {reason}"))
269}
270
271#[cfg(test)]
272mod tests {
273 use super::*;
274
275 #[test]
276 fn projection_distance_preserves_signed_cosine_and_canonical_parity() {
277 let query = [1.0, 0.0];
278 let values = [[1.0, 0.0], [0.0, 1.0], [-0.6, 0.8], [-1.0, 0.0]];
279 for metric in [VectorDistance::Cosine, VectorDistance::Euclidean] {
280 for coordinates in values {
281 let canonical = metric
282 .distance(
283 Vector::new(&query).unwrap().as_ref(),
284 Vector::new(&coordinates).unwrap().as_ref(),
285 )
286 .unwrap();
287 let projection = metric
288 .distance_slices_with_work(&query, &coordinates, || Ok(()))
289 .unwrap();
290 assert_eq!(projection.to_bits(), canonical.to_bits());
291 let encoded = Vector::new(&coordinates).unwrap();
292 let mut unaligned = vec![0];
293 unaligned.extend_from_slice(encoded.encoded());
294 let mapped = metric
295 .distance_slice_to_vector_with_work(
296 &query,
297 VectorRef::from_encoded(&unaligned[1..]).unwrap(),
298 || Ok(()),
299 )
300 .unwrap();
301 assert_eq!(mapped.to_bits(), canonical.to_bits());
302 }
303 }
304 let distance = |coordinates: &[f32]| {
305 VectorDistance::Cosine
306 .distance_slices_with_work(&query, coordinates, || Ok(()))
307 .unwrap()
308 };
309 assert!(distance(&values[2]) > 1.0);
310 assert!(distance(&values[2]) < distance(&values[3]));
311 assert_eq!(distance(&values[3]), 2.0);
312 for coordinates in [&[][..], &[0.0][..], &[f32::NAN, 0.0][..], &[0.0, 0.0][..]] {
313 assert!(VectorDistance::Cosine
314 .distance_slices_with_work(&query, coordinates, || Ok(()))
315 .is_err());
316 }
317 }
318
319 #[test]
320 fn projection_distance_preserves_cooperative_failure() {
321 let coordinates = vec![1.0; 20_000];
322 let mut checkpoints = 0;
323 let error = VectorDistance::Cosine
324 .distance_slices_with_work(&coordinates, &coordinates, || {
325 checkpoints += 1;
326 if checkpoints == 2 {
327 Err(HawDBError::ProjectionBusy)
328 } else {
329 Ok(())
330 }
331 })
332 .unwrap_err();
333 assert!(matches!(error, HawDBError::ProjectionBusy));
334 assert_eq!(checkpoints, 2);
335 let encoded = Vector::new(&coordinates).unwrap();
336 let error = VectorDistance::Cosine
337 .distance_slice_to_vector_with_work(&coordinates, encoded.as_ref(), || {
338 Err(HawDBError::ProjectionBusy)
339 })
340 .unwrap_err();
341 assert!(matches!(error, HawDBError::ProjectionBusy));
342 assert!(VectorDistance::Cosine
343 .distance_slice_to_vector_with_work(&[1.0], encoded.as_ref(), || Ok(()))
344 .is_err());
345 }
346
347 #[test]
348 fn vector_encoding_rejects_malformed_and_nonfinite_payloads() {
349 let value = Vector::new(&[1.0, -2.5]).unwrap();
350 assert_eq!(value.encoded(), &[2, 0, 0, 0, 0, 0, 128, 63, 0, 0, 32, 192]);
351 assert_eq!(
352 Vector::from_encoded(value.encoded().to_vec()).unwrap(),
353 value
354 );
355 let mut unaligned = vec![0];
356 unaligned.extend_from_slice(value.encoded());
357 assert_eq!(
358 VectorRef::from_encoded(&unaligned[1..]).unwrap(),
359 value.as_ref()
360 );
361 for size in 0..value.encoded().len() {
362 assert!(VectorRef::from_encoded(&value.encoded()[..size]).is_err());
363 }
364 assert!(VectorRef::from_encoded(&[0, 0, 0, 0]).is_err());
365 assert!(VectorRef::from_encoded(&u32::MAX.to_le_bytes()).is_err());
366 for component in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] {
367 assert!(Vector::new(&[component]).is_err());
368 let mut encoded = 1u32.to_le_bytes().to_vec();
369 encoded.extend_from_slice(&component.to_le_bytes());
370 assert!(VectorRef::from_encoded(&encoded).is_err());
371 }
372 assert!(Vector::new(&[]).is_err());
373 }
374
375 #[test]
376 fn vector_distance_is_mathematical_and_shared() {
377 let left = Vector::new(&[1.0, 0.0]).unwrap();
378 let opposite = Vector::new(&[-1.0, 0.0]).unwrap();
379 let zero = Vector::new(&[0.0, 0.0]).unwrap();
380 assert_eq!(
381 VectorDistance::Cosine
382 .distance(left.as_ref(), opposite.as_ref())
383 .unwrap(),
384 2.0
385 );
386 assert_eq!(
387 VectorDistance::Euclidean
388 .distance(left.as_ref(), zero.as_ref())
389 .unwrap(),
390 1.0
391 );
392 assert!(VectorDistance::Cosine
393 .distance(left.as_ref(), zero.as_ref())
394 .is_err());
395 assert!(VectorDistance::Euclidean
396 .distance(left.as_ref(), Vector::new(&[1.0]).unwrap().as_ref())
397 .is_err());
398 assert_eq!(
399 VectorDistance::Cosine
400 .evaluate(ValueRef::Null, ValueRef::Vector(zero.as_ref()))
401 .unwrap(),
402 None
403 );
404 let large = Vector::new(&[f32::MAX, f32::MAX]).unwrap();
405 assert!(VectorDistance::Euclidean
406 .distance(large.as_ref(), zero.as_ref())
407 .unwrap()
408 .is_finite());
409 assert!(
410 VectorDistance::Cosine
411 .distance(large.as_ref(), large.as_ref())
412 .unwrap()
413 .abs()
414 < 1e-12
415 );
416 }
417}
418
419impl std::hash::Hash for VectorRef<'_> {
420 fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
421 self.encoded.hash(state);
422 }
423}
424
425#[cfg(test)]
426mod work_tests {
427 use super::*;
428
429 #[test]
430 fn shared_vector_validation_stops_between_bounded_coordinate_blocks() {
431 let vector = Vector::new(&vec![1.0; 32769]).unwrap();
432 let calls = std::cell::Cell::new(0);
433 let error = Vector::from_encoded_with_work(
434 vector.encoded().to_vec(),
435 || {
436 calls.set(calls.get() + 1);
437 if calls.get() == 2 {
438 Err("cancelled".to_owned())
439 } else {
440 Ok(())
441 }
442 },
443 |_| Ok(()),
444 |error| error.to_string(),
445 )
446 .unwrap_err();
447 assert_eq!(error, "cancelled");
448 assert_eq!(calls.get(), 2);
449 let mut encoded = vector.encoded().to_vec();
450 let last = encoded.len() - 4;
451 encoded[last..].copy_from_slice(&f32::NAN.to_le_bytes());
452 assert!(Vector::from_encoded(encoded).is_err());
453 }
454}