Stream dists to reduce max memory

Cargo.toml

@@ -1,12 +1,13 @@
 [package]
 name = "sketchlib"
-version = "0.3.0"
+version = "0.4.0"
 authors = [
     "John Lees <jlees@ebi.ac.uk>",
     "Nicholas Croucher <n.croucher@imperial.ac.uk>",
     "Johanna von Wachsmann <wachsmannj@ebi.ac.uk>",
     "Victor Rodriguez Bouza <vrbouza@ebi.ac.uk>",
-    "Joel Hellewell <joel@ebi.ac.uk",
+    "Matthew Russell <marussel@ebi.ac.uk>",
+    "Tommi Maklin",
 ]
 edition = "2021"
 description = "Genome and amino-acid sketching"
@@ -31,8 +32,10 @@ crate-type = ["cdylib", "lib"]
 [features]
 3di = ["dep:pyo3", "dep:pyo3-build-config"]
 
+
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
 # data structures
+aligned-vec = "0.6.4"
 hashbrown = { version = "0.14", features = ["rayon", "serde"] }
 ahash = ">=0.8.7"
 roaring = { version = "0.10", features = ["serde"] }
@@ -63,6 +66,7 @@ num_cpus = "1.0"
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]
 # data structures
+aligned-vec = "0.6.4"
 hashbrown = { version = "0.14", features = ["rayon", "serde"] }
 # ahash = ">=0.8.7"
 roaring = { version = "0.10", features = ["serde"] }

README.md

@@ -15,8 +15,11 @@ This is a reimplementation and extension of [pp-sketchlib](https://github.com/ba
 in the rust language. This version is optimised for larger sample numbers, particularly
 allowing subsets of samples to be compared.
 
-v0.2.0 is the first stable release. We intend to keep the file format unchanged after this
-point so sketch libraries will not need to be rebuilt.
+### News
+
+- **Data format changed in v0.4.0**. If you have an older database (see `sketchlib info <db_prefix> | grep "sketch_version"`)
+you should resketch samples with >v0.4. Otherwise you will need to use an older release available on crates.io.
+- v0.2.0 was the first stable release.
 
 ## Documentation
 

src/cli.rs

@@ -85,6 +85,18 @@ pub fn check_and_set_threads(threads: usize) {
         .unwrap();
 }
 
+/// Sets the global rayon thread pool size for operations that spawn an internal
+/// writer thread (sketching, and appending to an existing sketch database).
+///
+/// Reserves one extra thread for the writer so that `threads` workers remain
+/// available for compute; equivalent to `check_and_set_threads(threads + 1)`.
+/// Callers embedding this crate as a library (rather than via the CLI) should call
+/// this before sketching/appending, and plain [`check_and_set_threads`] before
+/// distance/query/delete operations that don't spawn a writer.
+pub fn set_threads_with_writer(threads: usize) {
+    check_and_set_threads(threads + 1);
+}
+
 /// Options that apply to all subcommands
 #[derive(Parser)]
 #[command(author, version, about, long_about = None)]

src/distances/distance_matrix.rs

@@ -164,6 +164,25 @@
     pub fn dists_mut(&mut self) -> &mut Vec<f32> {
         &mut self.distances
     }
+
+    /// Iterates over the distances as `(primary, accessory)` pairs.
+    ///
+    /// The second element is `Some` only for [`DistType::CoreAcc`] matrices
+    /// (interleaved core/accessory pairs written by `self_dists_all`/`cross_dists_all`);
+    /// `None` for Jaccard/ANI matrices, where each entry is a single value. Iteration
+    /// order matches the row-major order used by `Display` (ref-outer, query- or
+    /// ref-inner — see `self_dists_all`/`cross_dists_all` docs).
+    pub fn dists_iter(&self) -> Box<dyn Iterator<Item = (f32, Option<f32>)> + '_> {
+        if self.jaccard == DistType::CoreAcc {
+            Box::new(
+                self.distances
+                    .chunks_exact(2)
+                    .map(|pair| (pair[0], Some(pair[1]))),
+            )
+        } else {
+            Box::new(self.distances.iter().copied().map(|d| (d, None)))
+        }
+    }
 }
 
 impl<'a> Distances<'a> for DistanceMatrix<'a> {

src/distances/jaccard.rs

@@ -1,6 +1,6 @@
 //! Implementation of Jaccard, core and accessory distance calculations
 use crate::sketch::multisketch::MultiSketch;
-use crate::sketch::BBITS;
+use crate::sketch::BIN_BITS;
 
 /// Returns the Jaccard index between two samples
 pub fn jaccard_index(
@@ -11,37 +11,180 @@
     c2: Option<f64>,
     completeness_cutoff: f64,
 ) -> f64 {
+    let samebits = jaccard_same_bits(sketch1, sketch2) as f64;
     let unionsize = (u64::BITS as u64 * sketchsize64) as f64;
-    let samebits: u32 = sketch1
-        .chunks_exact(BBITS as usize)
-        .zip(sketch2.chunks_exact(BBITS as usize))
+    // Correction for random matches
+    let expected_random = unionsize / (1u64 << BIN_BITS) as f64;
+    let mut jaccard_index =
+        ((samebits - expected_random) / (unionsize - expected_random)).clamp(0.0, 1.0);
+
+    log::trace!("samebits:{samebits} expected_random:{expected_random} jaccard:{jaccard_index}");
+
+    // Apply completeness correction if both completeness values are provided
+    if let (Some(c1_val), Some(c2_val)) = (c1, c2) {
+        if c1_val * c2_val >= completeness_cutoff {
+            // Cap the corrected Jaccard index at 1.0 to prevent negative distances
+            jaccard_index = completeness_correction(jaccard_index, c1_val, c2_val).min(1.0);
+        }
+    }
+
+    jaccard_index
+}
+
+#[inline(always)]
+fn jaccard_same_bits(sketch1: &[u64], sketch2: &[u64]) -> u32 {
+    #[cfg(target_arch = "aarch64")]
+    unsafe {
+        jaccard_neon_unroll2_inner(sketch1, sketch2)
+    }
+
+    #[cfg(not(target_arch = "aarch64"))]
+    {
+        jaccard_same_bits_general(sketch1, sketch2)
+    }
+}
+
+#[cfg_attr(target_arch = "aarch64", allow(dead_code))]
+#[inline(always)]
+pub(crate) fn jaccard_same_bits_general(sketch1: &[u64], sketch2: &[u64]) -> u32 {
+    debug_assert_eq!(sketch1.len(), sketch2.len());
+    debug_assert_eq!(sketch1.len() % BIN_BITS, 0);
+    sketch1
+        .chunks_exact(BIN_BITS)
+        .zip(sketch2.chunks_exact(BIN_BITS))
         .map(|(chunk1, chunk2)| {
             let mut bits: u64 = !0;
             chunk1.iter().zip(chunk2.iter()).for_each(|(&s1, &s2)| {
                 bits &= !(s1 ^ s2);
             });
             bits.count_ones()
         })
-        .sum();
-    let maxnbits = sketchsize64 as u32 * u64::BITS;
-    let expected_samebits = maxnbits >> BBITS;
+        .sum()
+}
 
-    log::trace!("samebits:{samebits} expected_samebits:{expected_samebits} maxnbits:{maxnbits}");
-    let diff = samebits.saturating_sub(expected_samebits);
-    let intersize = (diff as f64 * maxnbits as f64) / (maxnbits - expected_samebits) as f64;
-    log::trace!("intersize:{intersize} unionsize:{unionsize}");
-    let mut jaccard_index = intersize / unionsize;
+#[cfg(target_arch = "aarch64")]
+#[target_feature(enable = "neon")]
+unsafe fn jaccard_neon_unroll2_inner(a: &[u64], b: &[u64]) -> u32 {
+    use std::arch::aarch64::*;
 
-    // Apply completeness correction if both completeness values are provided
-    if let (Some(c1_val), Some(c2_val)) = (c1, c2) {
-        if c1_val * c2_val >= completeness_cutoff {
-            jaccard_index = completeness_correction(jaccard_index, c1_val, c2_val);
-            // Cap the corrected Jaccard index at 1.0 to prevent negative distances
-            jaccard_index = jaccard_index.min(1.0);
+    debug_assert_eq!(a.len(), b.len());
+    debug_assert_eq!(a.len() % BIN_BITS, 0);
+
+    let chunk = BIN_BITS;
+    let n_chunks = a.len() / chunk;
+    let n_pairs = n_chunks / 2;
+    let mut total = 0u32;
+
+    #[inline(always)]
+    unsafe fn one_chunk(ap: *const u8, bp: *const u8) -> u8 {
+        let a0 = vld1q_u8(ap);
+        let a1 = vld1q_u8(ap.add(16));
+        let a2 = vld1q_u8(ap.add(32));
+        let a3 = vld1q_u8(ap.add(48));
+        let a4 = vld1q_u8(ap.add(64));
+        let a5 = vld1q_u8(ap.add(80));
+        let a6 = vld1q_u8(ap.add(96));
+        let a7 = vld1q_u8(ap.add(112));
+        let b0 = vld1q_u8(bp);
+        let b1 = vld1q_u8(bp.add(16));
+        let b2 = vld1q_u8(bp.add(32));
+        let b3 = vld1q_u8(bp.add(48));
+        let b4 = vld1q_u8(bp.add(64));
+        let b5 = vld1q_u8(bp.add(80));
+        let b6 = vld1q_u8(bp.add(96));
+        let b7 = vld1q_u8(bp.add(112));
+        let x0 = veorq_u8(a0, b0);
+        let x1 = veorq_u8(a1, b1);
+        let x2 = veorq_u8(a2, b2);
+        let x3 = veorq_u8(a3, b3);
+        let x4 = veorq_u8(a4, b4);
+        let x5 = veorq_u8(a5, b5);
+        let x6 = veorq_u8(a6, b6);
+        let x7 = veorq_u8(a7, b7);
+        let or01 = vorrq_u8(x0, x1);
+        let or23 = vorrq_u8(x2, x3);
+        let or45 = vorrq_u8(x4, x5);
+        let or67 = vorrq_u8(x6, x7);
+        let or_all = vorrq_u8(vorrq_u8(or01, or23), vorrq_u8(or45, or67));
+        let not_or = vmvnq_u8(or_all);
+        vaddv_u8(vcnt_u8(vand_u8(vget_low_u8(not_or), vget_high_u8(not_or))))
+    }
+
+    for i in 0..n_pairs {
+        let base = i * 2 * chunk;
+        total += one_chunk(a.as_ptr().add(base) as _, b.as_ptr().add(base) as _) as u32;
+        total += one_chunk(
+            a.as_ptr().add(base + chunk) as _,
+            b.as_ptr().add(base + chunk) as _,
+        ) as u32;
+    }
+    if n_chunks % 2 != 0 {
+        let base = (n_chunks - 1) * chunk;
+        total += one_chunk(a.as_ptr().add(base) as _, b.as_ptr().add(base) as _) as u32;
+    }
+    total
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn scalar_same_bits_counts_matching_bins() {
+        let sketch1 = [u64::MAX; BIN_BITS * 2];
+        let mut sketch2 = [u64::MAX; BIN_BITS * 2];
+        sketch2[BIN_BITS] = 0;
+
+        assert_eq!(jaccard_same_bits_general(&sketch1, &sketch1), 128);
+        assert_eq!(jaccard_same_bits_general(&sketch1, &sketch2), 64);
+    }
+
+    #[cfg(target_arch = "aarch64")]
+    #[test]
+    fn neon_same_bits_matches_scalar() {
+        let mut sketch1 = [0u64; BIN_BITS * 3];
+        let mut sketch2 = [0u64; BIN_BITS * 3];
+        for i in 0..sketch1.len() {
+            sketch1[i] = (i as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15);
+            sketch2[i] = sketch1[i] ^ ((i as u64) << (i % 17));
         }
+
+        let scalar = jaccard_same_bits_general(&sketch1, &sketch2);
+        let neon = unsafe { jaccard_neon_unroll2_inner(&sketch1, &sketch2) };
+        assert_eq!(neon, scalar);
     }
 
-    jaccard_index
+    #[test]
+    fn random_match_correction_applied_for_large_sketch_size() {
+        // sketchsize64 chosen so that unionsize = sketchsize64 * 64 exceeds
+        // 2^BIN_BITS, the regime the random-match correction is meant for.
+        let sketchsize64: u64 = 2048;
+        let unionsize = (u64::BITS as u64 * sketchsize64) as f64;
+        assert!(unionsize > (1u64 << BIN_BITS) as f64);
+
+        let n_chunks = sketchsize64 as usize;
+        let sketch1 = vec![u64::MAX; BIN_BITS * n_chunks];
+        let mut sketch2 = sketch1.clone();
+        // Mismatch exactly one chunk (one 64-bit "column") out of n_chunks.
+        for word in sketch2[..BIN_BITS].iter_mut() {
+            *word = 0;
+        }
+
+        let samebits = jaccard_same_bits_general(&sketch1, &sketch2) as f64;
+        let expected_random = unionsize / (1u64 << BIN_BITS) as f64;
+        let expected_jaccard =
+            ((samebits - expected_random) / (unionsize - expected_random)).clamp(0.0, 1.0);
+        // The correction should be non-trivial (i.e. not equal to the raw,
+        // uncorrected ratio) for this sketch size.
+        assert!((expected_jaccard - samebits / unionsize).abs() > f64::EPSILON);
+
+        let jaccard = jaccard_index(&sketch1, &sketch2, sketchsize64, None, None, 0.0);
+        assert_eq!(jaccard, expected_jaccard);
+
+        // Identical sketches should give a jaccard index of exactly 1.0.
+        let identical = jaccard_index(&sketch1, &sketch1, sketchsize64, None, None, 0.0);
+        assert_eq!(identical, 1.0);
+    }
 }
 
 /// Converts between Jaccard distance and ANI, using a Poisson model of mutations