Add custom storage support

Author: pierd

instant-distance-py/benches/all.rs

@@ -21,15 +21,21 @@ fn build(bench: &mut Bencher) {
         .map(|_| FloatArray([rng.gen(); 300]))
         .collect::<Vec<_>>();
 
-    bench.iter(|| Builder::default().seed(SEED).build_hnsw(points.clone()));
+    bench.iter(|| {
+        Builder::default()
+            .seed(SEED)
+            .build_hnsw::<_, _, Vec<FloatArray>>(points.clone())
+    });
 }
 
 fn query(bench: &mut Bencher) {
     let mut rng = StdRng::seed_from_u64(SEED);
     let points = (0..1024)
         .map(|_| FloatArray([rng.gen(); 300]))
         .collect::<Vec<_>>();
-    let (hnsw, _) = Builder::default().seed(SEED).build_hnsw(points);
+    let (hnsw, _) = Builder::default()
+        .seed(SEED)
+        .build_hnsw::<_, _, Vec<FloatArray>>(points);
     let point = FloatArray([rng.gen(); 300]);
 
     bench.iter(|| {

instant-distance-py/src/lib.rs

@@ -29,7 +29,7 @@ fn instant_distance_py(_: Python, m: &PyModule) -> PyResult<()> {
 
 #[pyclass]
 struct HnswMap {
-    inner: instant_distance::HnswMap<FloatArray, MapValue>,
+    inner: instant_distance::HnswMap<FloatArray, MapValue, Vec<FloatArray>>,
 }
 
 #[pymethods]
@@ -54,11 +54,10 @@ 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<FloatArray, MapValue>>(
-                BufReader::with_capacity(32 * 1024 * 1024, File::open(fname)?),
-            )
-            .map_err(|e| PyValueError::new_err(format!("deserialization error: {e:?}")))?;
+        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:?}")))?;
         Ok(Self { inner: hnsw_map })
     }
 
@@ -91,7 +90,7 @@ impl HnswMap {
 /// with a squared Euclidean distance metric.
 #[pyclass]
 struct Hnsw {
-    inner: instant_distance::Hnsw<FloatArray>,
+    inner: instant_distance::Hnsw<FloatArray, Vec<FloatArray>>,
 }
 
 #[pymethods]
@@ -112,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<FloatArray>>(
+        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:?}")))?;
@@ -145,7 +144,7 @@ impl Hnsw {
 /// Search buffer and result set
 #[pyclass]
 struct Search {
-    inner: instant_distance::Search,
+    inner: instant_distance::Search<FloatArray>,
     cur: Option<(HnswType, usize)>,
 }
 

instant-distance/benches/all.rs

@@ -14,7 +14,11 @@ fn build_heuristic(bench: &mut Bencher) {
         .map(|_| Point(rng.gen(), rng.gen()))
         .collect::<Vec<_>>();
 
-    bench.iter(|| Builder::default().seed(SEED).build_hnsw(points.clone()))
+    bench.iter(|| {
+        Builder::default()
+            .seed(SEED)
+            .build_hnsw::<Point, Point, Vec<Point>>(points.clone())
+    })
 }
 
 const SEED: u64 = 123456789;

instant-distance/examples/colors.rs

@@ -4,7 +4,7 @@ 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(points, values);
+    let map = Builder::default().build::<Point, Point, &str, Vec<Point>>(points, values);
     let mut search = Search::default();
 
     let burnt_orange = Point(204, 85, 0);