rig_core/embeddings/
distance.rs

1//! Distance and similarity helpers for embedding vectors.
2//!
3//! The [`VectorDistance`] implementation for [`Embedding`](crate::embeddings::Embedding)
4//! uses iterator-based calculations by default and switches to Rayon-backed
5//! parallel iterators when the `rayon` feature is enabled.
6
7/// Distance and similarity metrics for embedding vectors.
8pub trait VectorDistance {
9    /// Get dot product of two embedding vectors
10    fn dot_product(&self, other: &Self) -> f64;
11
12    /// Get cosine similarity of two embedding vectors.
13    /// If `normalized` is true, the dot product is returned.
14    fn cosine_similarity(&self, other: &Self, normalized: bool) -> f64;
15
16    /// Get angular distance of two embedding vectors.
17    fn angular_distance(&self, other: &Self, normalized: bool) -> f64;
18
19    /// Get euclidean distance of two embedding vectors.
20    fn euclidean_distance(&self, other: &Self) -> f64;
21
22    /// Get manhattan distance of two embedding vectors.
23    fn manhattan_distance(&self, other: &Self) -> f64;
24
25    /// Get chebyshev distance of two embedding vectors.
26    fn chebyshev_distance(&self, other: &Self) -> f64;
27}
28
29/// Generates the [`VectorDistance`] method bodies for [`Embedding`](crate::embeddings::Embedding).
30///
31/// The math is identical between the sequential and Rayon-backed implementations;
32/// only the iterator constructor (`iter` vs `par_iter`) and the max-reduction
33/// differ (`Iterator::fold` vs `ParallelIterator::reduce`), so both are supplied
34/// by the caller. Keeping one source prevents the two copies from drifting.
35macro_rules! impl_vector_distance {
36    ($iter:ident, $($max_reduce:tt)+) => {
37        fn dot_product(&self, other: &Self) -> f64 {
38            self.vec
39                .$iter()
40                .zip(other.vec.$iter())
41                .map(|(x, y)| x * y)
42                .sum()
43        }
44
45        fn cosine_similarity(&self, other: &Self, normalized: bool) -> f64 {
46            let dot_product = self.dot_product(other);
47
48            if normalized {
49                dot_product
50            } else {
51                let magnitude1: f64 = self.vec.$iter().map(|x| x.powi(2)).sum::<f64>().sqrt();
52                let magnitude2: f64 = other.vec.$iter().map(|x| x.powi(2)).sum::<f64>().sqrt();
53
54                dot_product / (magnitude1 * magnitude2)
55            }
56        }
57
58        fn angular_distance(&self, other: &Self, normalized: bool) -> f64 {
59            let cosine_sim = self.cosine_similarity(other, normalized);
60            cosine_sim.acos() / std::f64::consts::PI
61        }
62
63        fn euclidean_distance(&self, other: &Self) -> f64 {
64            self.vec
65                .$iter()
66                .zip(other.vec.$iter())
67                .map(|(x, y)| (x - y).powi(2))
68                .sum::<f64>()
69                .sqrt()
70        }
71
72        fn manhattan_distance(&self, other: &Self) -> f64 {
73            self.vec
74                .$iter()
75                .zip(other.vec.$iter())
76                .map(|(x, y)| (x - y).abs())
77                .sum()
78        }
79
80        fn chebyshev_distance(&self, other: &Self) -> f64 {
81            self.vec
82                .$iter()
83                .zip(other.vec.$iter())
84                .map(|(x, y)| (x - y).abs())
85                .$($max_reduce)+
86        }
87    };
88}
89
90#[cfg(not(feature = "rayon"))]
91impl VectorDistance for crate::embeddings::Embedding {
92    impl_vector_distance!(iter, fold(0.0, f64::max));
93}
94
95#[cfg(feature = "rayon")]
96mod rayon {
97    use crate::embeddings::{Embedding, distance::VectorDistance};
98    use rayon::prelude::*;
99
100    impl VectorDistance for Embedding {
101        // `ParallelIterator` has no scalar `fold`; use `reduce` (0.0 is a valid
102        // identity for `max` since every mapped value is a non-negative abs diff).
103        impl_vector_distance!(par_iter, reduce(|| 0.0, f64::max));
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::VectorDistance;
110    use crate::embeddings::Embedding;
111
112    fn embeddings() -> (Embedding, Embedding) {
113        let embedding_1 = Embedding {
114            document: "test".to_string(),
115            vec: vec![1.0, 2.0, 3.0],
116        };
117
118        let embedding_2 = Embedding {
119            document: "test".to_string(),
120            vec: vec![1.0, 5.0, 7.0],
121        };
122
123        (embedding_1, embedding_2)
124    }
125
126    #[test]
127    fn test_dot_product() {
128        let (embedding_1, embedding_2) = embeddings();
129
130        assert_eq!(embedding_1.dot_product(&embedding_2), 32.0)
131    }
132
133    #[test]
134    fn test_cosine_similarity() {
135        let (embedding_1, embedding_2) = embeddings();
136
137        assert_eq!(
138            embedding_1.cosine_similarity(&embedding_2, false),
139            0.9875414397573881
140        )
141    }
142
143    #[test]
144    fn test_angular_distance() {
145        let (embedding_1, embedding_2) = embeddings();
146
147        assert_eq!(
148            embedding_1.angular_distance(&embedding_2, false),
149            0.0502980301830343
150        )
151    }
152
153    #[test]
154    fn test_euclidean_distance() {
155        let (embedding_1, embedding_2) = embeddings();
156
157        assert_eq!(embedding_1.euclidean_distance(&embedding_2), 5.0)
158    }
159
160    #[test]
161    fn test_manhattan_distance() {
162        let (embedding_1, embedding_2) = embeddings();
163
164        assert_eq!(embedding_1.manhattan_distance(&embedding_2), 7.0)
165    }
166
167    #[test]
168    fn test_chebyshev_distance() {
169        let (embedding_1, embedding_2) = embeddings();
170
171        assert_eq!(embedding_1.chebyshev_distance(&embedding_2), 4.0)
172    }
173}
rig_core/embeddings/distance.rs

rig_core/embeddings/
distance.rs
