Introduce Metric trait

Author: pierd

instant-distance-py/benches/all.rs

@@ -1,7 +1,7 @@
 use bencher::{benchmark_group, benchmark_main, Bencher};
 
-use instant_distance::{Builder, Point, Search};
-use instant_distance_py::FloatArray;
+use instant_distance::{Builder, Metric, Search};
+use instant_distance_py::{EuclidMetric, FloatArray};
 use rand::{rngs::StdRng, Rng, SeedableRng};
 
 benchmark_main!(benches);
@@ -12,7 +12,7 @@ fn distance(bench: &mut Bencher) {
     let point_a = FloatArray([rng.gen(); 300]);
     let point_b = FloatArray([rng.gen(); 300]);
 
-    bench.iter(|| point_a.distance(&point_b));
+    bench.iter(|| EuclidMetric::distance(&point_a, &point_b));
 }
 
 fn build(bench: &mut Bencher) {
@@ -25,7 +25,7 @@ fn build(bench: &mut Bencher) {
     bench.iter(|| {
         Builder::default()
             .seed(SEED)
-            .build_hnsw::<_, _, Vec<FloatArray>>(points.clone())
+            .build_hnsw::<_, _, EuclidMetric, Vec<FloatArray>>(points.clone())
     });
 }
 
@@ -37,7 +37,7 @@ fn query(bench: &mut Bencher) {
         .collect::<Vec<_>>();
     let (hnsw, _) = Builder::default()
         .seed(SEED)
-        .build_hnsw::<_, _, Vec<FloatArray>>(points);
+        .build_hnsw::<_, _, EuclidMetric, Vec<FloatArray>>(points);
     let point = FloatArray([rng.gen(); 300]);
 
     bench.iter(|| {

instant-distance-py/src/lib.rs

@@ -6,7 +6,7 @@ use std::fs::File;
 use std::io::{BufReader, BufWriter};
 use std::iter::FromIterator;
 
-use instant_distance::Point;
+use instant_distance::Metric;
 use pyo3::conversion::IntoPy;
 use pyo3::exceptions::{PyTypeError, PyValueError};
 use pyo3::types::{PyList, PyModule, PyString};
@@ -29,7 +29,7 @@ fn instant_distance_py(_: Python, m: &PyModule) -> PyResult<()> {
 
 #[pyclass]
 struct HnswMap {
-    inner: instant_distance::HnswMap<FloatArray, MapValue, Vec<FloatArray>>,
+    inner: instant_distance::HnswMap<FloatArray, EuclidMetric, MapValue, Vec<FloatArray>>,
 }
 
 #[pymethods]
@@ -54,7 +54,7 @@ impl HnswMap {
     /// Load an index from the given file name
     #[staticmethod]
     fn load(fname: &str) -> PyResult<Self> {
-        let hnsw_map = bincode::deserialize_from::<_, instant_distance::HnswMap<_, _, _>>(
+        let hnsw_map = bincode::deserialize_from::<_, instant_distance::HnswMap<_, _, _, _>>(
             BufReader::with_capacity(32 * 1024 * 1024, File::open(fname)?),
         )
         .map_err(|e| PyValueError::new_err(format!("deserialization error: {e:?}")))?;
@@ -90,7 +90,7 @@ impl HnswMap {
 /// with a squared Euclidean distance metric.
 #[pyclass]
 struct Hnsw {
-    inner: instant_distance::Hnsw<FloatArray, Vec<FloatArray>>,
+    inner: instant_distance::Hnsw<FloatArray, EuclidMetric, Vec<FloatArray>>,
 }
 
 #[pymethods]
@@ -111,7 +111,7 @@ impl Hnsw {
     /// Load an index from the given file name
     #[staticmethod]
     fn load(fname: &str) -> PyResult<Self> {
-        let hnsw = bincode::deserialize_from::<_, instant_distance::Hnsw<_, _>>(
+        let hnsw = bincode::deserialize_from::<_, instant_distance::Hnsw<_, _, _>>(
             BufReader::with_capacity(32 * 1024 * 1024, File::open(fname)?),
         )
         .map_err(|e| PyValueError::new_err(format!("deserialization error: {e:?}")))?;
@@ -144,7 +144,7 @@ impl Hnsw {
 /// Search buffer and result set
 #[pyclass]
 struct Search {
-    inner: instant_distance::Search<FloatArray>,
+    inner: instant_distance::Search<FloatArray, EuclidMetric>,
     cur: Option<(HnswType, usize)>,
 }
 
@@ -364,20 +364,23 @@ impl TryFrom<&PyAny> for FloatArray {
     }
 }
 
-impl Point for FloatArray {
-    fn distance(&self, rhs: &Self) -> f32 {
+#[derive(Clone, Copy, Deserialize, Serialize)]
+pub struct EuclidMetric;
+
+impl Metric<FloatArray> for EuclidMetric {
+    fn distance(lhs: &FloatArray, rhs: &FloatArray) -> f32 {
         #[cfg(target_arch = "x86_64")]
         {
             use std::arch::x86_64::{
                 _mm256_castps256_ps128, _mm256_extractf128_ps, _mm256_fmadd_ps, _mm256_load_ps,
                 _mm256_setzero_ps, _mm256_sub_ps, _mm_add_ps, _mm_add_ss, _mm_cvtss_f32,
                 _mm_fmadd_ps, _mm_load_ps, _mm_movehl_ps, _mm_shuffle_ps, _mm_sub_ps,
             };
-            debug_assert_eq!(self.0.len() % 8, 4);
+            debug_assert_eq!(lhs.0.len() % 8, 4);
 
             unsafe {
                 let mut acc_8x = _mm256_setzero_ps();
-                for (lh_slice, rh_slice) in self.0.chunks_exact(8).zip(rhs.0.chunks_exact(8)) {
+                for (lh_slice, rh_slice) in lhs.0.chunks_exact(8).zip(rhs.0.chunks_exact(8)) {
                     let lh_8x = _mm256_load_ps(lh_slice.as_ptr());
                     let rh_8x = _mm256_load_ps(rh_slice.as_ptr());
                     let diff = _mm256_sub_ps(lh_8x, rh_8x);
@@ -388,7 +391,7 @@ impl Point for FloatArray {
                 let right = _mm256_castps256_ps128(acc_8x); // lower half
                 acc_4x = _mm_add_ps(acc_4x, right); // sum halves
 
-                let lh_4x = _mm_load_ps(self.0[DIMENSIONS - 4..].as_ptr());
+                let lh_4x = _mm_load_ps(lhs.0[DIMENSIONS - 4..].as_ptr());
                 let rh_4x = _mm_load_ps(rhs.0[DIMENSIONS - 4..].as_ptr());
                 let diff = _mm_sub_ps(lh_4x, rh_4x);
                 acc_4x = _mm_fmadd_ps(diff, diff, acc_4x);
@@ -401,7 +404,7 @@ impl Point for FloatArray {
             }
         }
         #[cfg(not(target_arch = "x86_64"))]
-        self.0
+        lhs.0
             .iter()
             .zip(rhs.0.iter())
             .map(|(&a, &b)| (a - b).powi(2))

instant-distance/benches/all.rs

@@ -2,7 +2,7 @@ use bencher::{benchmark_group, benchmark_main, Bencher};
 use rand::rngs::StdRng;
 use rand::{Rng, SeedableRng};
 
-use instant_distance::Builder;
+use instant_distance::{Builder, Metric};
 
 benchmark_main!(benches);
 benchmark_group!(benches, build_heuristic);
@@ -11,13 +11,13 @@ fn build_heuristic(bench: &mut Bencher) {
     let mut rng = StdRng::seed_from_u64(SEED);
     let points = (0..1024)
         .into_iter()
-        .map(|_| Point(rng.gen(), rng.gen()))
+        .map(|_| [rng.gen(), rng.gen()])
         .collect::<Vec<_>>();
 
     bench.iter(|| {
         Builder::default()
             .seed(SEED)
-            .build_hnsw::<Point, Point, Vec<Point>>(points.clone())
+            .build_hnsw::<[f32; 2], [f32; 2], EuclidMetric, Vec<[f32; 2]>>(points.clone())
     })
 }
 
@@ -51,12 +51,15 @@ fn randomized(builder: Builder) -> (u64, usize) {
 }
 */
 
-#[derive(Clone, Copy, Debug)]
-struct Point(f32, f32);
+struct EuclidMetric;
 
-impl instant_distance::Point for Point {
-    fn distance(&self, other: &Self) -> f32 {
+impl Metric<[f32; 2]> for EuclidMetric {
+    fn distance(a: &[f32; 2], b: &[f32; 2]) -> f32 {
         // Euclidean distance metric
-        ((self.0 - other.0).powi(2) + (self.1 - other.1).powi(2)).sqrt()
+        a.iter()
+            .zip(b.iter())
+            .map(|(&a, &b)| (a - b).powi(2))
+            .sum::<f32>()
+            .sqrt()
     }
 }

instant-distance/examples/colors.rs

@@ -1,10 +1,11 @@
-use instant_distance::{Builder, Search};
+use instant_distance::{Builder, Metric, Search};
 
 fn main() {
     let points = vec![Point(255, 0, 0), Point(0, 255, 0), Point(0, 0, 255)];
     let values = vec!["red", "green", "blue"];
 
-    let map = Builder::default().build::<Point, Point, &str, Vec<Point>>(points, values);
+    let map =
+        Builder::default().build::<Point, Point, EuclidMetric, &str, Vec<Point>>(points, values);
     let mut search = Search::default();
 
     let burnt_orange = Point(204, 85, 0);
@@ -17,10 +18,11 @@ fn main() {
 #[derive(Clone, Copy, Debug)]
 struct Point(isize, isize, isize);
 
-impl instant_distance::Point for Point {
-    fn distance(&self, other: &Self) -> f32 {
+struct EuclidMetric;
+
+impl Metric<Point> for EuclidMetric {
+    fn distance(a: &Point, b: &Point) -> f32 {
         // Euclidean distance metric
-        (((self.0 - other.0).pow(2) + (self.1 - other.1).pow(2) + (self.2 - other.2).pow(2)) as f32)
-            .sqrt()
+        (((a.0 - b.0).pow(2) + (a.1 - b.1).pow(2) + (a.2 - b.2).pow(2)) as f32).sqrt()
     }
 }