Implement Hamming distance for binary strings

proptest-regressions/unaligned_vector/binary_test.txt

@@ -0,0 +1,7 @@
+# Seeds for failure cases proptest has generated in the past. It is
+# automatically read and these particular cases re-run before any
+# novel cases are generated.
+#
+# It is recommended to check this file in to source control so that
+# everyone who runs the test benefits from these saved cases.
+cc 4044d46b46fbadeb98b1c37a1b7ec57f00fe21008401479d30f2f39aa83e3195 # shrinks to original = [0.0]

src/distance/hamming.rs

@@ -0,0 +1,148 @@
+use std::fmt;
+
+use crate::distance::Distance;
+use crate::node::Leaf;
+use crate::parallel::ImmutableSubsetLeafs;
+use crate::unaligned_vector::{Binary, UnalignedVector};
+use bytemuck::{Pod, Zeroable};
+use rand::Rng;
+
+/// The Hamming distance between two vectors is the number of positions at
+/// which the corresponding symbols are different.
+///
+/// `d(u,v) = ||u ^ v||₁`
+///
+/// /!\ This distance function is binary, which means it loses all its precision
+///     and their scalar values are converted to `0` or `1` under the rule
+///     `x > 0.0 => 1`, otherwise `0`
+#[derive(Debug, Clone)]
+pub enum Hamming {}
+
+/// The header of BinaryEuclidean leaf nodes.
+#[repr(C)]
+#[derive(Pod, Zeroable, Clone, Copy)]
+pub struct NodeHeaderHamming {
+    idx: usize,
+}
+impl fmt::Debug for NodeHeaderHamming {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("NodeHeaderHamming ").field("idx", &format!("{}", self.idx)).finish()
+    }
+}
+
+impl Distance for Hamming {
+    const DEFAULT_OVERSAMPLING: usize = 3;
+
+    type Header = NodeHeaderHamming;
+    type VectorCodec = Binary;
+
+    fn name() -> &'static str {
+        "hamming"
+    }
+
+    fn new_header(_vector: &UnalignedVector<Self::VectorCodec>) -> Self::Header {
+        NodeHeaderHamming { idx: 0 }
+    }
+
+    fn built_distance(p: &Leaf<Self>, q: &Leaf<Self>) -> f32 {
+        hamming_bitwise_fast(p.vector.as_bytes(), q.vector.as_bytes())
+    }
+
+    fn normalized_distance(d: f32, _: usize) -> f32 {
+        d
+    }
+
+    fn norm_no_header(v: &UnalignedVector<Self::VectorCodec>) -> f32 {
+        v.as_bytes().iter().map(|b| b.count_ones() as i32).sum::<i32>() as f32
+    }
+
+    fn init(_node: &mut Leaf<Self>) {}
+
+    fn create_split<'a, R: Rng>(
+        children: &'a ImmutableSubsetLeafs<Self>,
+        rng: &mut R,
+    ) -> heed::Result<Leaf<'a, Self>> {
+        // unlike other distances which build a seperating hyperplane we
+        // construct an LSH by bit sampling and storing the splitting index
+        // in the node header.
+        // https://en.wikipedia.org/wiki/Locality-sensitive_hashing#Bit_sampling_for_Hamming_distance
+
+        const ITERATION_STEPS: usize = 200;
+
+        let is_valid_split = |n: &Leaf<'a, Self>, rng: &mut R| {
+            let mut count = 0;
+            for _ in 0..ITERATION_STEPS {
+                let u = children.choose(rng)?.unwrap();
+                if <Self as Distance>::margin(n, &u).is_sign_positive() {
+                    count += 1;
+                }
+            }
+            Ok::<bool, heed::Error>(count > 0 && count < ITERATION_STEPS)
+        };
+
+        // first try random index
+        let dim = children.choose(rng)?.unwrap().vector.len();
+        let idx = rng.gen_range(0..dim);
+        let mut normal =
+            Leaf { header: NodeHeaderHamming { idx }, vector: UnalignedVector::from_vec(vec![]) };
+
+        if is_valid_split(&normal, rng)? {
+            return Ok(normal);
+        }
+
+        // otherwise brute-force search for a splitting coordinate
+        for j in 0..dim {
+            normal.header.idx = j;
+            if is_valid_split(&normal, rng)? {
+                return Ok(normal);
+            }
+        }
+
+        // fallback
+        Ok(normal)
+    }
+
+    fn margin(n: &Leaf<Self>, q: &Leaf<Self>) -> f32 {
+        let v = q.vector.as_bytes();
+        let byte = n.header.idx / 8;
+        let bit = n.header.idx % 8;
+        if (v[byte] >> bit) & 1 == 1 {
+            1.0
+        } else {
+            -1.0
+        }
+    }
+}
+
+#[inline]
+pub fn hamming_bitwise_fast(u: &[u8], v: &[u8]) -> f32 {
+    // based on : https://github.com/emschwartz/hamming-bitwise-fast
+    // Explicitly structuring the code as below lends itself to SIMD optimizations by
+    // the compiler -> https://matklad.github.io/2023/04/09/can-you-trust-a-compiler-to-optimize-your-code.html
+    assert_eq!(u.len(), v.len());
+
+    type BitPackedWord = u64;
+    const CHUNK_SIZE: usize = std::mem::size_of::<BitPackedWord>();
+
+    let mut distance = u
+        .chunks_exact(CHUNK_SIZE)
+        .zip(v.chunks_exact(CHUNK_SIZE))
+        .map(|(u_chunk, v_chunk)| {
+            let u_val = BitPackedWord::from_ne_bytes(u_chunk.try_into().unwrap());
+            let v_val = BitPackedWord::from_ne_bytes(v_chunk.try_into().unwrap());
+            (u_val ^ v_val).count_ones()
+        })
+        .sum::<u32>();
+
+    if u.len() % CHUNK_SIZE != 0 {
+        distance += u
+            .chunks_exact(CHUNK_SIZE)
+            .remainder()
+            .iter()
+            .zip(v.chunks_exact(CHUNK_SIZE).remainder())
+            .map(|(u_byte, v_byte)| (u_byte ^ v_byte).count_ones())
+            .sum::<u32>();
+    }
+
+    distance as f32
+}

src/distance/mod.rs

@@ -11,6 +11,7 @@ use bytemuck::{Pod, Zeroable};
 pub use cosine::{Cosine, NodeHeaderCosine};
 pub use dot_product::{DotProduct, NodeHeaderDotProduct};
 pub use euclidean::{Euclidean, NodeHeaderEuclidean};
+pub use hamming::{Hamming, NodeHeaderHamming};
 use heed::{RwPrefix, RwTxn};
 pub use manhattan::{Manhattan, NodeHeaderManhattan};
 use rand::Rng;
@@ -27,6 +28,7 @@ mod binary_quantized_manhattan;
 mod cosine;
 mod dot_product;
 mod euclidean;
+mod hamming;
 mod manhattan;
 
 fn new_leaf<D: Distance>(vec: Vec<f32>) -> Leaf<'static, D> {

src/lib.rs

@@ -114,7 +114,7 @@ pub mod internals {
     pub use crate::distance::{
         NodeHeaderBinaryQuantizedCosine, NodeHeaderBinaryQuantizedEuclidean,
         NodeHeaderBinaryQuantizedManhattan, NodeHeaderCosine, NodeHeaderDotProduct,
-        NodeHeaderEuclidean, NodeHeaderManhattan,
+        NodeHeaderEuclidean, NodeHeaderHamming, NodeHeaderManhattan,
     };
     pub use crate::key::KeyCodec;
     pub use crate::node::{Leaf, NodeCodec};
@@ -145,7 +145,7 @@ pub mod internals {
 pub mod distances {
     pub use crate::distance::{
         BinaryQuantizedCosine, BinaryQuantizedEuclidean, BinaryQuantizedManhattan, Cosine,
-        DotProduct, Euclidean, Manhattan,
+        DotProduct, Euclidean, Hamming, Manhattan,
     };
 }
 

src/tests/binary.rs

@@ -0,0 +1,55 @@
+use crate::{
+    distance::Hamming,
+    tests::{create_database, rng},
+    Writer,
+};
+
+#[test]
+fn write_and_retrieve_binary_vector() {
+    let handle = create_database::<Hamming>();
+    let mut wtxn = handle.env.write_txn().unwrap();
+    let writer = Writer::new(handle.database, 0, 16);
+    writer
+        .add_item(
+            &mut wtxn,
+            0,
+            &[
+                -2.0, -1.0, 0.0, -0.1, 2.0, 2.0, -12.4, 21.2, -2.0, -1.0, 0.0, 1.0, 2.0, 2.0,
+                -12.4, 21.2,
+            ],
+        )
+        .unwrap();
+    let vec = writer.item_vector(&wtxn, 0).unwrap().unwrap();
+    insta::assert_debug_snapshot!(vec, @r###"
+    [
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        1.0,
+        1.0,
+        0.0,
+        1.0,
+        0.0,
+        0.0,
+        0.0,
+        1.0,
+        1.0,
+        1.0,
+        0.0,
+        1.0,
+    ]
+    "###);
+
+    writer.builder(&mut rng()).n_trees(1).build(&mut wtxn).unwrap();
+    wtxn.commit().unwrap();
+
+    insta::assert_snapshot!(handle, @r#"
+    ==================
+    Dumping index 0
+    Root: Metadata { dimensions: 16, items: RoaringBitmap<[0]>, roots: [0], distance: "hamming" }
+    Version: Version { major: 0, minor: 7, patch: 0 }
+    Tree 0: Descendants(Descendants { descendants: [0] })
+    Item 0: Leaf(Leaf { header: NodeHeaderHamming  { idx: "0" }, vector: [0.0000, 0.0000, 0.0000, 0.0000, 1.0000, 1.0000, 0.0000, 1.0000, 0.0000, 0.0000, "other ..."] })
+    "#);
+}

src/tests/mod.rs

@@ -11,6 +11,7 @@ use tempfile::TempDir;
 use crate::version::VersionCodec;
 use crate::{Database, Distance, MetadataCodec, NodeCodec, NodeMode, Reader};
 
+mod binary;
 mod binary_quantized;
 mod reader;
 mod upgrade;