Add benchmarks to instant-distance-py

Author: pierd

.github/workflows/rust.yml

@@ -53,7 +53,7 @@ jobs:
       - name: Test Python bindings
         run: |
           sudo apt-get install -y wamerican
-          cp target/release/libinstant_distance.so instant-distance-py/test/instant_distance.so
+          cp target/release/libinstant_distance_py.so instant-distance-py/test/instant_distance.so
           PYTHONPATH=instant-distance-py/test/ python3 -m test
 
   lint:

Makefile

@@ -1,6 +1,9 @@
-test-python:
-	cargo build --release
-	cp target/release/libinstant_distance.dylib instant-distance-py/test/instant_distance.so
+instant-distance-py/test/instant_distance.so: instant-distance-py/src/lib.rs
+	RUSTFLAGS="-C target-cpu=native" cargo build --release
+	([ -f target/release/libinstant_distance_py.dylib ] && cp target/release/libinstant_distance_py.dylib instant-distance-py/test/instant_distance.so) || \
+	([ -f target/release/libinstant_distance_py.so ] && cp target/release/libinstant_distance_py.so instant-distance-py/test/instant_distance.so)
+
+test-python: instant-distance-py/test/instant_distance.so
 	PYTHONPATH=instant-distance-py/test/ python3 -m test
 
 clean:

instant-distance-py/Cargo.toml

@@ -11,12 +11,23 @@ repository = "https://github.com/InstantDomain/instant-distance"
 readme = "../README.md"
 
 [lib]
+name = "instant_distance_py"
+crate-type = ["cdylib", "lib"]
+
+[package.metadata.maturin]
 name = "instant_distance"
-crate-type = ["cdylib"]
 
 [dependencies]
 bincode = "1.3.1"
 instant-distance = { version = "0.6", path = "../instant-distance", features = ["with-serde"] }
 pyo3 = { version = "0.18.0", features = ["extension-module"] }
 serde = { version = "1", features = ["derive"] }
 serde-big-array = "0.5.0"
+
+[dev-dependencies]
+bencher = "0.1.5"
+rand = { version = "0.8", features = ["small_rng"] }
+
+[[bench]]
+name = "all"
+harness = false

instant-distance-py/benches/all.rs

@@ -0,0 +1,43 @@
+use bencher::{benchmark_group, benchmark_main, Bencher};
+
+use instant_distance::{Builder, Point, Search};
+use instant_distance_py::FloatArray;
+use rand::{rngs::StdRng, Rng, SeedableRng};
+
+benchmark_main!(benches);
+benchmark_group!(benches, distance, build, query);
+
+fn distance(bench: &mut Bencher) {
+    let mut rng = StdRng::seed_from_u64(SEED);
+    let point_a = FloatArray([rng.gen(); 300]);
+    let point_b = FloatArray([rng.gen(); 300]);
+
+    bench.iter(|| point_a.distance(&point_b));
+}
+
+fn build(bench: &mut Bencher) {
+    let mut rng = StdRng::seed_from_u64(SEED);
+    let points = (0..1024)
+        .into_iter()
+        .map(|_| FloatArray([rng.gen(); 300]))
+        .collect::<Vec<_>>();
+
+    bench.iter(|| Builder::default().seed(SEED).build_hnsw(points.clone()));
+}
+
+fn query(bench: &mut Bencher) {
+    let mut rng = StdRng::seed_from_u64(SEED);
+    let points = (0..1024)
+        .into_iter()
+        .map(|_| FloatArray([rng.gen(); 300]))
+        .collect::<Vec<_>>();
+    let (hnsw, _) = Builder::default().seed(SEED).build_hnsw(points);
+    let point = FloatArray([rng.gen(); 300]);
+
+    bench.iter(|| {
+        let mut search = Search::default();
+        let _ = hnsw.search(&point, &mut search);
+    });
+}
+
+const SEED: u64 = 123456789;

instant-distance-py/pyproject.toml

@@ -4,3 +4,6 @@ name = "instant-distance"
 [build-system]
 requires = ["maturin >= 0.14, < 0.15"]
 build-backend = "maturin"
+
+[tool.maturin]
+python-source = "python"

instant-distance-py/python/instant_distance/__init__.py

@@ -0,0 +1,5 @@
+from .instant_distance import *
+
+__doc__ = instant_distance.__doc__
+if hasattr(instant_distance, "__all__"):
+    __all__ = instant_distance.__all__

instant-distance-py/src/lib.rs

@@ -348,7 +348,7 @@ impl Neighbor {
 
 #[repr(align(32))]
 #[derive(Clone, Deserialize, Serialize)]
-struct FloatArray(#[serde(with = "BigArray")] [f32; DIMENSIONS]);
+pub struct FloatArray(#[serde(with = "BigArray")] pub [f32; DIMENSIONS]);
 
 impl TryFrom<&PyAny> for FloatArray {
     type Error = PyErr;