Binary quantization

examples/compare_with_hnsw.rs

@@ -1,9 +1,8 @@
-use std::borrow::Cow;
 use std::num::NonZeroUsize;
 use std::time::Instant;
 
 use arroy::distances::Euclidean;
-use arroy::internals::{Leaf, UnalignedF32Slice};
+use arroy::internals::{Leaf, UnalignedVector};
 use arroy::{Database, Distance, ItemId, Reader, Result, Writer};
 use heed::{EnvOpenOptions, RwTxn};
 use instant_distance::{Builder, HnswMap, MapItem};
@@ -94,10 +93,10 @@ struct Point(Vec<f32>);
 
 impl instant_distance::Point for Point {
     fn distance(&self, other: &Self) -> f32 {
-        let this = UnalignedF32Slice::from_slice(&self.0);
-        let other = UnalignedF32Slice::from_slice(&other.0);
-        let p = Leaf { header: Euclidean::new_header(this), vector: Cow::Borrowed(this) };
-        let q = Leaf { header: Euclidean::new_header(other), vector: Cow::Borrowed(other) };
+        let this = UnalignedVector::from_slice(&self.0);
+        let other = UnalignedVector::from_slice(&other.0);
+        let p = Leaf { header: Euclidean::new_header(&this), vector: this };
+        let q = Leaf { header: Euclidean::new_header(&other), vector: other };
         arroy::distances::Euclidean::built_distance(&p, &q).sqrt()
     }
 }

src/distance/angular.rs

@@ -5,9 +5,10 @@ use rand::Rng;
 
 use super::two_means;
 use crate::distance::Distance;
-use crate::node::{Leaf, UnalignedF32Slice};
+use crate::node::Leaf;
 use crate::parallel::ImmutableSubsetLeafs;
 use crate::spaces::simple::dot_product;
+use crate::unaligned_vector::UnalignedVector;
 
 /// The Cosine similarity is a measure of similarity between two
 /// non-zero vectors defined in an inner product space. Cosine similarity
@@ -24,12 +25,13 @@ pub struct NodeHeaderAngular {
 
 impl Distance for Angular {
     type Header = NodeHeaderAngular;
+    type VectorFormat = f32;
 
     fn name() -> &'static str {
         "angular"
     }
 
-    fn new_header(vector: &UnalignedF32Slice) -> Self::Header {
+    fn new_header(vector: &UnalignedVector<Self::VectorFormat>) -> Self::Header {
         NodeHeaderAngular { norm: Self::norm_no_header(vector) }
     }
 
@@ -54,24 +56,32 @@ impl Distance for Angular {
         d
     }
 
+    fn norm_no_header(v: &UnalignedVector<Self::VectorFormat>) -> f32 {
+        dot_product(v, v).sqrt()
+    }
+
     fn init(node: &mut Leaf<Self>) {
         node.header.norm = dot_product(&node.vector, &node.vector).sqrt();
     }
 
-    fn create_split<R: Rng>(
-        children: &ImmutableSubsetLeafs<Self>,
+    fn create_split<'a, R: Rng>(
+        children: &'a ImmutableSubsetLeafs<Self>,
         rng: &mut R,
-    ) -> heed::Result<Vec<f32>> {
+    ) -> heed::Result<Cow<'a, UnalignedVector<Self::VectorFormat>>> {
         let [node_p, node_q] = two_means(rng, children, true)?;
-        let vector = node_p.vector.iter().zip(node_q.vector.iter()).map(|(p, q)| p - q).collect();
-        let mut normal =
-            Leaf { header: NodeHeaderAngular { norm: 0.0 }, vector: Cow::Owned(vector) };
+        let vector: Vec<f32> =
+            node_p.vector.iter().zip(node_q.vector.iter()).map(|(p, q)| p - q).collect();
+        let unaligned_vector = Self::craft_owned_unaligned_vector_from_f32(vector);
+        let mut normal = Leaf { header: NodeHeaderAngular { norm: 0.0 }, vector: unaligned_vector };
         Self::normalize(&mut normal);
 
-        Ok(normal.vector.into_owned())
+        Ok(normal.vector)
     }
 
-    fn margin_no_header(p: &UnalignedF32Slice, q: &UnalignedF32Slice) -> f32 {
+    fn margin_no_header(
+        p: &UnalignedVector<Self::VectorFormat>,
+        q: &UnalignedVector<Self::VectorFormat>,
+    ) -> f32 {
         dot_product(p, q)
     }
 }

src/distance/binary_quantized_euclidean.rs

@@ -0,0 +1,79 @@
+use std::borrow::Cow;
+
+use bytemuck::{Pod, Zeroable};
+use rand::Rng;
+
+use super::two_means;
+use crate::distance::Distance;
+use crate::node::Leaf;
+use crate::parallel::ImmutableSubsetLeafs;
+use crate::unaligned_vector::{self, BinaryQuantized, UnalignedVector};
+
+/// The Euclidean distance between two points in Euclidean space
+/// is the length of the line segment between them.
+///
+/// `d(p, q) = sqrt((p - q)²)`
+#[derive(Debug, Clone)]
+pub enum BinaryQuantizedEuclidean {}
+
+/// The header of BinaryQuantizedEuclidean leaf nodes.
+#[repr(C)]
+#[derive(Pod, Zeroable, Debug, Clone, Copy)]
+pub struct NodeHeaderBinaryQuantizedEuclidean {
+    /// An extra constant term to determine the offset of the plane
+    bias: f32,
+}
+
+impl Distance for BinaryQuantizedEuclidean {
+    type Header = NodeHeaderBinaryQuantizedEuclidean;
+    type VectorFormat = unaligned_vector::BinaryQuantized;
+
+    fn name() -> &'static str {
+        "binary quantized euclidean"
+    }
+
+    fn new_header(_vector: &UnalignedVector<Self::VectorFormat>) -> Self::Header {
+        NodeHeaderBinaryQuantizedEuclidean { bias: 0.0 }
+    }
+
+    fn built_distance(p: &Leaf<Self>, q: &Leaf<Self>) -> f32 {
+        dot_product(&p.vector, &q.vector)
+    }
+
+    fn norm_no_header(v: &UnalignedVector<Self::VectorFormat>) -> f32 {
+        dot_product(v, v).sqrt()
+    }
+
+    fn init(_node: &mut Leaf<Self>) {}
+
+    fn create_split<'a, R: Rng>(
+        children: &'a ImmutableSubsetLeafs<Self>,
+        rng: &mut R,
+    ) -> heed::Result<Cow<'a, UnalignedVector<Self::VectorFormat>>> {
+        let [node_p, node_q] = two_means(rng, children, false)?;
+        let vector: Vec<f32> =
+            node_p.vector.iter().zip(node_q.vector.iter()).map(|(p, q)| p - q).collect();
+        let mut normal = Leaf {
+            header: NodeHeaderBinaryQuantizedEuclidean { bias: 0.0 },
+            vector: UnalignedVector::from_slice(&vector),
+        };
+        Self::normalize(&mut normal);
+
+        Ok(Cow::Owned(normal.vector.into_owned()))
+    }
+
+    fn margin(p: &Leaf<Self>, q: &Leaf<Self>) -> f32 {
+        p.header.bias + dot_product(&p.vector, &q.vector)
+    }
+
+    fn margin_no_header(
+        p: &UnalignedVector<Self::VectorFormat>,
+        q: &UnalignedVector<Self::VectorFormat>,
+    ) -> f32 {
+        dot_product(p, q)
+    }
+}
+
+fn dot_product(u: &UnalignedVector<BinaryQuantized>, v: &UnalignedVector<BinaryQuantized>) -> f32 {
+    u.as_bytes().iter().zip(v.as_bytes()).map(|(u, v)| (u ^ v).count_ones()).sum::<u32>() as f32
+}

src/distance/dot_product.rs

@@ -7,9 +7,10 @@ use rand::Rng;
 use super::two_means;
 use crate::distance::Distance;
 use crate::internals::KeyCodec;
-use crate::node::{Leaf, UnalignedF32Slice};
+use crate::node::Leaf;
 use crate::parallel::ImmutableSubsetLeafs;
 use crate::spaces::simple::dot_product;
+use crate::unaligned_vector::UnalignedVector;
 use crate::{Node, NodeCodec};
 
 /// In mathematics, the dot product or scalar product is an algebraic
@@ -29,12 +30,13 @@ pub struct NodeHeaderDotProduct {
 
 impl Distance for DotProduct {
     type Header = NodeHeaderDotProduct;
+    type VectorFormat = f32;
 
     fn name() -> &'static str {
         "dot-product"
     }
 
-    fn new_header(_vector: &UnalignedF32Slice) -> Self::Header {
+    fn new_header(_vector: &UnalignedVector<Self::VectorFormat>) -> Self::Header {
         // We compute the norm when we preprocess the vector, before generating the tree nodes.
         NodeHeaderDotProduct { extra_dim: 0.0, norm: 0.0 }
     }
@@ -64,14 +66,19 @@ impl Distance for DotProduct {
         (dot + leaf.header.extra_dim * leaf.header.extra_dim).sqrt()
     }
 
+    fn norm_no_header(v: &UnalignedVector<Self::VectorFormat>) -> f32 {
+        dot_product(v, v).sqrt()
+    }
+
     fn normalized_distance(d: f32) -> f32 {
         -d
     }
 
     fn normalize(node: &mut Leaf<Self>) {
         let norm = Self::norm(node);
         if norm > 0.0 {
-            node.vector.to_mut().iter_mut().for_each(|x| *x /= norm);
+            let vec: Vec<_> = node.vector.iter().map(|x| x / norm).collect();
+            node.vector = UnalignedVector::from_vec(vec);
             node.header.extra_dim /= norm;
         }
     }
@@ -80,27 +87,31 @@ impl Distance for DotProduct {
         node.header.norm = dot_product(&node.vector, &node.vector);
     }
 
-    fn create_split<R: Rng>(
-        children: &ImmutableSubsetLeafs<Self>,
+    fn create_split<'a, R: Rng>(
+        children: &'a ImmutableSubsetLeafs<Self>,
         rng: &mut R,
-    ) -> heed::Result<Vec<f32>> {
+    ) -> heed::Result<Cow<'a, UnalignedVector<Self::VectorFormat>>> {
         let [node_p, node_q] = two_means(rng, children, true)?;
-        let vector = node_p.vector.iter().zip(node_q.vector.iter()).map(|(p, q)| p - q).collect();
+        let vector: Vec<f32> =
+            node_p.vector.iter().zip(node_q.vector.iter()).map(|(p, q)| p - q).collect();
         let mut normal = Leaf::<Self> {
             header: NodeHeaderDotProduct { norm: 0.0, extra_dim: 0.0 },
-            vector: Cow::Owned(vector),
+            vector: Self::craft_owned_unaligned_vector_from_f32(vector),
         };
         normal.header.extra_dim = node_p.header.extra_dim - node_q.header.extra_dim;
         Self::normalize(&mut normal);
 
-        Ok(normal.vector.into_owned())
+        Ok(normal.vector)
     }
 
     fn margin(p: &Leaf<Self>, q: &Leaf<Self>) -> f32 {
         dot_product(&p.vector, &q.vector) + p.header.extra_dim * q.header.extra_dim
     }
 
-    fn margin_no_header(p: &UnalignedF32Slice, q: &UnalignedF32Slice) -> f32 {
+    fn margin_no_header(
+        p: &UnalignedVector<Self::VectorFormat>,
+        q: &UnalignedVector<Self::VectorFormat>,
+    ) -> f32 {
         dot_product(p, q)
     }
 

src/distance/euclidean.rs

@@ -5,9 +5,10 @@ use rand::Rng;
 
 use super::two_means;
 use crate::distance::Distance;
-use crate::node::{Leaf, UnalignedF32Slice};
+use crate::node::Leaf;
 use crate::parallel::ImmutableSubsetLeafs;
 use crate::spaces::simple::{dot_product, euclidean_distance};
+use crate::unaligned_vector::UnalignedVector;
 
 /// The Euclidean distance between two points in Euclidean space
 /// is the length of the line segment between them.
@@ -26,29 +27,37 @@ pub struct NodeHeaderEuclidean {
 
 impl Distance for Euclidean {
     type Header = NodeHeaderEuclidean;
+    type VectorFormat = f32;
 
     fn name() -> &'static str {
         "euclidean"
     }
 
-    fn new_header(_vector: &UnalignedF32Slice) -> Self::Header {
+    fn new_header(_vector: &UnalignedVector<Self::VectorFormat>) -> Self::Header {
         NodeHeaderEuclidean { bias: 0.0 }
     }
 
     fn built_distance(p: &Leaf<Self>, q: &Leaf<Self>) -> f32 {
         euclidean_distance(&p.vector, &q.vector)
     }
 
+    fn norm_no_header(v: &UnalignedVector<Self::VectorFormat>) -> f32 {
+        dot_product(v, v).sqrt()
+    }
+
     fn init(_node: &mut Leaf<Self>) {}
 
-    fn create_split<R: Rng>(
-        children: &ImmutableSubsetLeafs<Self>,
+    fn create_split<'a, R: Rng>(
+        children: &'a ImmutableSubsetLeafs<Self>,
         rng: &mut R,
-    ) -> heed::Result<Vec<f32>> {
+    ) -> heed::Result<Cow<'a, UnalignedVector<Self::VectorFormat>>> {
         let [node_p, node_q] = two_means(rng, children, false)?;
-        let vector = node_p.vector.iter().zip(node_q.vector.iter()).map(|(p, q)| p - q).collect();
-        let mut normal =
-            Leaf { header: NodeHeaderEuclidean { bias: 0.0 }, vector: Cow::Owned(vector) };
+        let vector: Vec<_> =
+            node_p.vector.iter().zip(node_q.vector.iter()).map(|(p, q)| p - q).collect();
+        let mut normal = Leaf {
+            header: NodeHeaderEuclidean { bias: 0.0 },
+            vector: Self::craft_owned_unaligned_vector_from_f32(vector),
+        };
         Self::normalize(&mut normal);
 
         normal.header.bias = normal
@@ -59,14 +68,17 @@ impl Distance for Euclidean {
             .map(|((n, p), q)| -n * (p + q) / 2.0)
             .sum();
 
-        Ok(normal.vector.into_owned())
+        Ok(normal.vector)
     }
 
     fn margin(p: &Leaf<Self>, q: &Leaf<Self>) -> f32 {
         p.header.bias + dot_product(&p.vector, &q.vector)
     }
 
-    fn margin_no_header(p: &UnalignedF32Slice, q: &UnalignedF32Slice) -> f32 {
+    fn margin_no_header(
+        p: &UnalignedVector<Self::VectorFormat>,
+        q: &UnalignedVector<Self::VectorFormat>,
+    ) -> f32 {
         dot_product(p, q)
     }
 }