docling_rag/
math.rs

1//! Small, dependency-free vector maths used by embedders and stores.
2
3/// Dot product of two equal-length slices (unchecked length; extra tail ignored).
4pub fn dot(a: &[f32], b: &[f32]) -> f32 {
5    a.iter().zip(b).map(|(x, y)| x * y).sum()
6}
7
8/// Euclidean (L2) norm.
9pub fn norm(a: &[f32]) -> f32 {
10    dot(a, a).sqrt()
11}
12
13/// Cosine similarity in `[-1, 1]`; returns 0 if either vector is zero-length.
14pub fn cosine(a: &[f32], b: &[f32]) -> f32 {
15    let na = norm(a);
16    let nb = norm(b);
17    if na == 0.0 || nb == 0.0 {
18        0.0
19    } else {
20        dot(a, b) / (na * nb)
21    }
22}
23
24/// L2-normalize a vector in place. A zero vector is left unchanged.
25pub fn normalize(v: &mut [f32]) {
26    let n = norm(v);
27    if n > 0.0 {
28        for x in v.iter_mut() {
29            *x /= n;
30        }
31    }
32}
33
34/// Encode a vector as a little-endian `f32` byte blob (for DB storage).
35pub fn to_bytes(v: &[f32]) -> Vec<u8> {
36    let mut out = Vec::with_capacity(v.len() * 4);
37    for x in v {
38        out.extend_from_slice(&x.to_le_bytes());
39    }
40    out
41}
42
43/// Decode a little-endian `f32` byte blob back into a vector.
44pub fn from_bytes(bytes: &[u8]) -> Vec<f32> {
45    bytes
46        .chunks_exact(4)
47        .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
48        .collect()
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54
55    #[test]
56    fn cosine_basics() {
57        assert!((cosine(&[1.0, 0.0], &[1.0, 0.0]) - 1.0).abs() < 1e-6);
58        assert!(cosine(&[1.0, 0.0], &[0.0, 1.0]).abs() < 1e-6);
59        assert_eq!(cosine(&[0.0, 0.0], &[1.0, 1.0]), 0.0);
60    }
61
62    #[test]
63    fn normalize_makes_unit_length() {
64        let mut v = vec![3.0, 4.0];
65        normalize(&mut v);
66        assert!((norm(&v) - 1.0).abs() < 1e-6);
67    }
68
69    #[test]
70    fn bytes_roundtrip() {
71        let v = vec![1.5f32, -2.25, 0.0, 1024.5];
72        assert_eq!(from_bytes(&to_bytes(&v)), v);
73    }
74}
docling_rag/math.rs

docling_rag/
math.rs
