diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index fbd07b9..a20b9ef 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -9,14 +9,14 @@ permissions: pull-requests: write jobs: - benchmark-base: - name: Benchmark base branch + benchmark: + name: Benchmark base and PR runs-on: ubuntu-latest steps: - - name: Checkout base branch + - name: Checkout repository uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: - ref: ${{ github.event.pull_request.base.sha }} + fetch-depth: 0 - name: Set up Python 3.11 uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 @@ -27,51 +27,47 @@ jobs: run: | python -m pip install --upgrade pip pip install pytest pytest-benchmark - pip install . - - name: Run benchmarks + - name: Record benchmark host run: | - pytest test/ -k bench --benchmark-only --benchmark-json=results.json --benchmark-disable-gc + uname -a + lscpu - - name: Upload benchmark results + - name: Benchmark base branch + env: + BASE_SHA: ${{ github.event.pull_request.base.sha }} + run: | + git checkout --detach "$BASE_SHA" + pip install --force-reinstall --no-deps . + mkdir -p base + pytest test/ -k bench --benchmark-only --benchmark-json=base/results.json --benchmark-disable-gc + + - name: Benchmark PR branch + env: + HEAD_SHA: ${{ github.event.pull_request.head.sha }} + run: | + git checkout --detach "$HEAD_SHA" + pip install --force-reinstall --no-deps . + mkdir -p pr + pytest test/ -k bench --benchmark-only --benchmark-json=pr/results.json --benchmark-disable-gc + + - name: Upload base benchmark results uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: benchmark-base - path: results.json + path: base/results.json retention-days: 1 - benchmark-pr: - name: Benchmark PR branch - runs-on: ubuntu-latest - steps: - - name: Checkout PR branch - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - - - name: Set up Python 3.11 - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 - with: - python-version: "3.11" - - - name: Install dependencies - run: | - python -m pip install --upgrade pip - pip install pytest pytest-benchmark - pip install . - - - name: Run benchmarks - run: | - pytest test/ -k bench --benchmark-only --benchmark-json=results.json --benchmark-disable-gc - - - name: Upload benchmark results + - name: Upload PR benchmark results uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 with: name: benchmark-pr - path: results.json + path: pr/results.json retention-days: 1 compare: name: Compare benchmarks - needs: [benchmark-base, benchmark-pr] + needs: benchmark runs-on: ubuntu-latest steps: - name: Download base benchmark results diff --git a/README.rst b/README.rst index 06f14c7..218d5b4 100644 --- a/README.rst +++ b/README.rst @@ -179,6 +179,89 @@ distance less than ``max_dist``, or ``[]`` if none do. Tip: When you're assembling the long array of records to compare against, don't concatenate the different ``bytes`` together. As they're immutable that is a very slow operation. Use a ``bytearray`` instead, and cast it to ``bytes`` at the end. See https://www.guyrutenberg.com/2020/04/04/fast-bytes-concatenation-in-python/ for more info and tests. +Batch APIs +~~~~~~~~~~ + +The per-call APIs above are still the right choice for one-off distances, but +computing many distances in Python for-loops pays repeated FFI overhead. +The batch APIs below fold that overhead into a single call by taking +contiguous buffers. + +Pairwise distances between two equal-length contiguous buffers of fixed-width +records: + +:: + + >>> from hexhamming import ( + ... hamming_distances_bytes, + ... hamming_distances_bytes_packed, + ... hamming_distances_bytes_into, + ... ) + >>> a = b"\xde\xad\xbe\xef" * 4 + >>> b = b"\x00" * 16 + >>> hamming_distances_bytes(a, b, 4) # list[int] + [24, 24, 24, 24] + >>> hamming_distances_bytes_packed(a, b, 4).hex() # little-endian u64 bytes + '1800000000000000180000000000000018000000000000001800000000000000' + >>> out = bytearray(4 * 8) + >>> hamming_distances_bytes_into(a, b, 4, out) # writes u64 LE into `out` + 4 + +``hamming_distances_bytes_into`` requires ``out`` to be a writable, +C-contiguous byte buffer of exactly ``count * 8`` bytes; read-only, +non-contiguous, or wrong-size outputs raise ``ValueError``. Writable ``_into`` +APIs are unavailable on free-threaded Python because the buffer protocol does +not provide exclusive access; use the corresponding ``_packed`` API there. +On standard Python builds, ``_into`` keeps the GIL while writing; use +``_packed`` when detached computation is more important than buffer reuse. + +Multi-query catalog scans run one catalog against many contiguous queries in +one call, mirroring the shape of repeated single-query calls: + +:: + + >>> from hexhamming import ( + ... check_bytes_arrays_first_many_within_dist, + ... check_bytes_arrays_best_many_within_dist, + ... check_bytes_arrays_all_many_within_dist, + ... ) + >>> catalog = b"\xaa\xaa\xbb\xbb\xcc\xcc\xdd\xdd\xee\xee\xff\xff" + >>> queries = b"\xff\xff\xef\xfe" + >>> check_bytes_arrays_first_many_within_dist(catalog, queries, 2, 4) + [1, 4] + >>> check_bytes_arrays_best_many_within_dist(catalog, queries, 2, 4) + [(0, 5), (2, 4)] + >>> check_bytes_arrays_all_many_within_dist(catalog, queries, 2, 4) + [[(4, 1), (4, 3), (4, 4), (0, 5)], [(2, 4), (2, 5)]] + +Semantics match the single-query calls exactly: ``-1`` and ``(-1, -1)`` +sentinels for no-match, lowest-index tie-breaking for ``best_many``, exact-match +short-circuiting, and ascending index order for ``all_many``. + +Dense/compact match transport for ``all_within_dist`` uses ``u16`` +distances and ``u32`` indices instead of Python tuples: + +:: + + >>> from hexhamming import ( + ... check_bytes_arrays_all_within_dist_packed, + ... check_bytes_arrays_all_within_dist_into, + ... ) + >>> dbytes, ibytes = check_bytes_arrays_all_within_dist_packed(catalog, b"\xff\xff", 4) + >>> [int.from_bytes(dbytes[i:i+2], "little") for i in range(0, len(dbytes), 2)] + [4, 4, 4, 0] + >>> d_out = bytearray(len(catalog) // 2 * 2) # worst case: num_records * 2 + >>> i_out = bytearray(len(catalog) // 2 * 4) # worst case: num_records * 4 + >>> check_bytes_arrays_all_within_dist_into(catalog, b"\xff\xff", 4, d_out, i_out) + 4 + +The ``_packed`` variant returns two ``bytes`` objects; ``_into`` writes into +caller-provided writable buffers and returns the match count. Element widths +whose maximum possible distance exceeds ``u16::MAX`` bits, and catalogs with +more than ``u32::MAX`` records, are rejected. On free-threaded Python, use +``_packed`` because writable ``_into`` buffers cannot be made exclusive through +the Python buffer protocol. + Benchmark --------- @@ -371,3 +454,57 @@ at 64 KiB; generic scans parallelize with Rayon at 5 MiB, while the optimized 16/32-byte NEON scanners use a measured 16 MiB crossover. The ``first`` variant additionally short-circuits on the first hit, so a match near the start is much faster than one near the end. + +Batch APIs vs. Python for-loops +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +These numbers use ``.benchmarks/batch_measure.py`` (``timeit.repeat`` with 50 +calls per sample, three independent runs, median of medians) on the same M4 +Max. The "loop" columns run the equivalent single-call API inside a Python +for-loop. Speedups are relative to the loop baseline. + +Pairwise distances between two contiguous buffers of ``count`` records: + +============================================ ========== ========== ========== ========== +Case loop (ns) list (ns) packed (ns) into (ns) +============================================ ========== ========== ========== ========== +pairwise 100×16 15,796.7 598.3 287.5 211.7 +pairwise 1,000×16 147,985.0 4,494.2 1,528.3 1,245.8 +pairwise 10,000×16 1,424,663.3 43,612.5 14,660.8 11,720.8 +pairwise 100×32 16,091.7 821.7 511.7 438.3 +pairwise 1,000×32 160,530.0 6,787.5 3,855.0 3,545.8 +pairwise 10,000×32 1,542,458.3 66,835.8 37,700.8 34,721.7 +============================================ ========== ========== ========== ========== + +Multi-query catalog scans against a 1,024×16-byte catalog with 100 queries: + +================================================== ========== ========== +Case loop (ns) batch (ns) +================================================== ========== ========== +first_many 100×1024×16 (permissive threshold) 10,948.3 742.5 +best_many 100×1024×16 (max_dist=128) 76,889.2 66,529.2 +================================================== ========== ========== + +Dense-match transport for a single query against a 1,024×16-byte catalog: + +========================================== ========== ========== ========== +Case list (ns) packed (ns) into (ns) +========================================== ========== ========== ========== +all 1024×16 (max_dist=128, all match) 30,681.7 3,067.5 2,022.5 +========================================== ========== ========== ========== + +Interpretation: + +* Pairwise: the ``list`` API is 26–33× faster than the Python for-loop and is + the recommended default. ``packed`` and ``into`` skip the per-distance + Python ``int`` allocation for another 2–3× on top; use them when the caller + can consume little-endian ``u64`` bytes directly. +* Multi-query ``first_many`` is a very large win (≈15×) because each inner + scan short-circuits on the first hit and Python-loop overhead dominates. + ``best_many`` and ``all_many`` are more modest wins (≈1.15–1.2×) because + their inner scans always traverse the whole catalog and the per-call FFI + overhead is proportionally smaller. +* Dense ``all_within_dist``: ``packed`` avoids allocating ``num_records`` + Python 2-tuples (≈10×); ``into`` additionally reuses caller-owned + buffers (≈15×) and matches the throughput of Rust code that never + touches the Python heap. diff --git a/src/api.rs b/src/api.rs index 09d4799..0b866ae 100644 --- a/src/api.rs +++ b/src/api.rs @@ -11,34 +11,35 @@ use rayon::prelude::*; use std::sync::atomic::Ordering; /// Minimum total byte size of big_array before we use rayon parallel paths. -const PAR_THRESHOLD_BYTES: usize = 5 * 1024 * 1024; +pub(crate) const PAR_THRESHOLD_BYTES: usize = 5 * 1024 * 1024; /// The fixed-width NEON scanners make serial scans substantially cheaper, so /// use a larger crossover before paying Rayon scheduling and partition costs. -const FIXED_WIDTH_PAR_THRESHOLD_BYTES: usize = 16 * 1024 * 1024; +pub(crate) const FIXED_WIDTH_PAR_THRESHOLD_BYTES: usize = 16 * 1024 * 1024; /// Keep byte-array scans to a small number of coarse jobs. More workers spend /// more time scheduling these very small per-record calculations than running /// them on current many-core CPUs. const PAR_JOBS: usize = 4; -type ArrayFirstScanner = fn(&[u8], &[u8], i64) -> Option; -type ArrayBestScanner = fn(&[u8], &[u8], i64) -> Option<(u64, usize)>; -type ArrayAllScanner = fn(&[u8], &[u8], i64) -> Vec<(u64, usize)>; +pub(crate) type ArrayFirstScanner = fn(&[u8], &[u8], i64) -> Option; +pub(crate) type ArrayBestScanner = fn(&[u8], &[u8], i64) -> Option<(u64, usize)>; +pub(crate) type ArrayAllScanner = fn(&[u8], &[u8], i64) -> Vec<(u64, usize)>; #[derive(Clone, Copy)] -struct ArrayScanner { - first: ArrayFirstScanner, - best: ArrayBestScanner, - all: ArrayAllScanner, +pub(crate) struct ArrayScanner { + pub(crate) first: ArrayFirstScanner, + pub(crate) best: ArrayBestScanner, + pub(crate) all: ArrayAllScanner, } #[inline] -fn select_array_scanner_for_width(width: usize) -> Option { +pub(crate) fn select_array_scanner_for_width(width: usize) -> Option { #[cfg(target_arch = "aarch64")] { let algo = CURRENT_ALGO.load(Ordering::Relaxed); if algo != ALGO_NATIVE && algo != ALGO_NEON { return None; } + return match width { 16 => Some(ArrayScanner { first: crate::neon_simd::array_first_neon_16, @@ -92,6 +93,16 @@ fn select_array_scanner_for_width(width: usize) -> Option { } } +#[inline] +pub(crate) fn should_parallel_array_scan(byte_len: usize, width: usize) -> bool { + let threshold = if select_array_scanner_for_width(width).is_some() { + FIXED_WIDTH_PAR_THRESHOLD_BYTES + } else { + PAR_THRESHOLD_BYTES + }; + byte_len >= threshold +} + #[inline] fn partition_element_ranges(num_elements: usize) -> [(usize, usize); PAR_JOBS] { let base = num_elements / PAR_JOBS; @@ -197,7 +208,7 @@ pub fn bytes_array_first_within_dist( } #[inline] -fn serial_first_within_dist( +pub(crate) fn serial_first_within_dist( big_array: &[u8], small_array: &[u8], max_dist: i64, @@ -274,7 +285,7 @@ fn merge_best(a: Option<(u64, usize)>, b: Option<(u64, usize)>) -> Option<(u64, } #[inline] -fn serial_best_within_dist( +pub(crate) fn serial_best_within_dist( big_array: &[u8], small_array: &[u8], max_dist: i64, @@ -364,7 +375,7 @@ pub fn bytes_array_all_within_dist( } #[inline] -fn serial_all_within_dist( +pub(crate) fn serial_all_within_dist( big_array: &[u8], small_array: &[u8], max_dist: i64, diff --git a/src/batch.rs b/src/batch.rs new file mode 100644 index 0000000..9a76c36 --- /dev/null +++ b/src/batch.rs @@ -0,0 +1,527 @@ +//! Batch APIs that amortize the Python↔Rust boundary across many distance +//! calculations. +//! +//! Callers pay a single validation + dispatch resolution once per batch, then +//! run a tight loop that reuses the already-resolved kernel or fixed-width +//! scanner. These APIs preserve the semantics of repeated single-call use: +//! +//! * Ordering of results matches the caller's element order. +//! * Ties in `best_*` break on lowest index, exact matches short-circuit. +//! * Threshold semantics for `max_dist` are unchanged (see `api.rs`). +//! * `set_algo` still controls the backend at call time; resolution happens +//! once per batch call. + +use crate::api::{ + select_array_scanner_for_width, serial_all_within_dist, serial_best_within_dist, + serial_first_within_dist, should_parallel_array_scan, ArrayScanner, +}; +use crate::{select_bytes_kernel_for_width, BytesKernel}; + +/// Compute Hamming distances between corresponding fixed-width records in two +/// contiguous buffers. Returns one `u64` per record in original order. +/// +/// # Errors +/// * `element_size` is zero. +/// * `a.len() != b.len()`. +/// * Buffer length is not a multiple of `element_size`. +pub fn bytes_pairwise_distances( + a: &[u8], + b: &[u8], + element_size: usize, +) -> Result, &'static str> { + if element_size == 0 { + return Err("`element_size` must be >0"); + } + if a.len() != b.len() { + return Err("bytes are NOT the same length"); + } + if a.len() % element_size != 0 { + return Err("length must be a multiple of `element_size`"); + } + let count = a.len() / element_size; + if count == 0 { + return Ok(Vec::new()); + } + let kernel = select_bytes_kernel_for_width(element_size); + Ok(a.chunks_exact(element_size) + .zip(b.chunks_exact(element_size)) + .map(|(a_chunk, b_chunk)| kernel(a_chunk, b_chunk, -1)) + .collect()) +} + +/// Compute Hamming distances and write them as little-endian `u64` values into +/// `out`. `out` must be exactly `count * 8` bytes, where `count = a.len() / +/// element_size`. +/// +/// Returns the number of distances written. +/// +/// # Errors +/// * Same input validation as [`bytes_pairwise_distances`]. +/// * `out.len() != count * 8`. +pub fn bytes_pairwise_distances_into( + a: &[u8], + b: &[u8], + element_size: usize, + out: &mut [u8], +) -> Result { + if element_size == 0 { + return Err("`element_size` must be >0"); + } + if a.len() != b.len() { + return Err("bytes are NOT the same length"); + } + if a.len() % element_size != 0 { + return Err("length must be a multiple of `element_size`"); + } + let count = a.len() / element_size; + let expected = count.checked_mul(8).ok_or("output capacity overflows")?; + if out.len() != expected { + return Err("`out` must be exactly count*8 bytes"); + } + if count == 0 { + return Ok(0); + } + let kernel = select_bytes_kernel_for_width(element_size); + for ((a_chunk, b_chunk), out_chunk) in a + .chunks_exact(element_size) + .zip(b.chunks_exact(element_size)) + .zip(out.chunks_exact_mut(8)) + { + out_chunk.copy_from_slice(&kernel(a_chunk, b_chunk, -1).to_le_bytes()); + } + Ok(count) +} + +/// Validate common catalog + queries inputs and return `(query_count, +/// element_size, kernel, scanner)` for the multi-query APIs. +#[inline] +fn resolve_multi_scan<'a>( + catalog: &'a [u8], + queries: &'a [u8], + query_width: usize, +) -> Result<(usize, BytesKernel, Option), &'static str> { + if query_width == 0 { + return Err("`query_width` must be >0"); + } + if catalog.len() % query_width != 0 { + return Err("catalog length must be a multiple of `query_width`"); + } + if queries.len() % query_width != 0 { + return Err("queries length must be a multiple of `query_width`"); + } + let kernel = select_bytes_kernel_for_width(query_width); + let scanner = select_array_scanner_for_width(query_width); + let query_count = queries.len() / query_width; + Ok((query_count, kernel, scanner)) +} + +/// Multi-query variant of [`bytes_array_first_within_dist`]: +/// runs the same scan for every fixed-width slice of `queries` against the +/// same `catalog`, returning one `Option` per query in query order. +pub fn bytes_array_first_many_within_dist( + catalog: &[u8], + queries: &[u8], + query_width: usize, + max_dist: i64, +) -> Result>, &'static str> { + let (query_count, kernel, scanner) = resolve_multi_scan(catalog, queries, query_width)?; + let mut out = Vec::with_capacity(query_count); + // `first` is intentionally serial: the early-exit dominates parallel setup. + for q in 0..query_count { + let query = &queries[q * query_width..(q + 1) * query_width]; + let result = match scanner { + Some(sc) => (sc.first)(catalog, query, max_dist), + None => serial_first_within_dist(catalog, query, max_dist, kernel), + }; + out.push(result); + } + Ok(out) +} + +/// Multi-query variant of [`bytes_array_best_within_dist`]. +pub fn bytes_array_best_many_within_dist( + catalog: &[u8], + queries: &[u8], + query_width: usize, + max_dist: i64, +) -> Result>, &'static str> { + let (query_count, kernel, scanner) = resolve_multi_scan(catalog, queries, query_width)?; + let mut out = Vec::with_capacity(query_count); + for q in 0..query_count { + let query = &queries[q * query_width..(q + 1) * query_width]; + let result = if should_parallel_array_scan(catalog.len(), query_width) { + crate::bytes_array_best_within_dist(catalog, query, max_dist)? + } else { + match scanner { + Some(sc) => (sc.best)(catalog, query, max_dist), + None => serial_best_within_dist(catalog, query, max_dist, kernel), + } + }; + out.push(result); + } + Ok(out) +} + +/// Multi-query variant of [`bytes_array_all_within_dist`]. +pub fn bytes_array_all_many_within_dist( + catalog: &[u8], + queries: &[u8], + query_width: usize, + max_dist: i64, +) -> Result>, &'static str> { + let (query_count, kernel, scanner) = resolve_multi_scan(catalog, queries, query_width)?; + let mut out = Vec::with_capacity(query_count); + for q in 0..query_count { + let query = &queries[q * query_width..(q + 1) * query_width]; + let result = if should_parallel_array_scan(catalog.len(), query_width) { + crate::bytes_array_all_within_dist(catalog, query, max_dist)? + } else { + match scanner { + Some(sc) => (sc.all)(catalog, query, max_dist), + None => serial_all_within_dist(catalog, query, max_dist, kernel), + } + }; + out.push(result); + } + Ok(out) +} + +/// Dense-transport variant of [`bytes_array_all_within_dist`]: returns matched +/// distances as `Vec` and matched indices as `Vec` in ascending +/// index order. +/// +/// # Errors +/// * `small_array` is empty. +/// * `big_array.len() % small_array.len() != 0`. +/// * Element width would allow distances exceeding `u16::MAX` bits. +/// * Catalog would produce indices exceeding `u32::MAX`. +pub fn bytes_array_all_within_dist_packed( + big_array: &[u8], + small_array: &[u8], + max_dist: i64, +) -> Result<(Vec, Vec), &'static str> { + let width = small_array.len(); + if width == 0 { + return Err("`elem_to_compare` size must be >0"); + } + if big_array.len() % width != 0 { + return Err("`array_of_elems` size must be multiplier of `elem_to_compare`"); + } + let max_bits = (width as u64).saturating_mul(8); + if max_bits > u16::MAX as u64 { + return Err("element width too large for u16 packed distances"); + } + let num_records = big_array.len() / width; + if num_records > u32::MAX as usize { + return Err("catalog record count exceeds u32::MAX"); + } + let matches = crate::bytes_array_all_within_dist(big_array, small_array, max_dist)?; + let mut distances = Vec::with_capacity(matches.len()); + let mut indices = Vec::with_capacity(matches.len()); + for (d, i) in matches { + distances.push(d as u16); + indices.push(i as u32); + } + Ok((distances, indices)) +} + +/// Write `all_within_dist` results into caller-provided u16 distance and u32 +/// index buffers. Returns the number of matches written. +/// +/// The buffers must be able to hold the worst case (all records match); the +/// caller is responsible for sizing them appropriately. +/// +/// # Errors +/// * Same as [`bytes_array_all_within_dist_packed`]. +/// * `out_distances_u16.len() < num_records`. +/// * `out_indices_u32.len() < num_records`. +pub fn bytes_array_all_within_dist_into( + big_array: &[u8], + small_array: &[u8], + max_dist: i64, + out_distances_u16: &mut [u16], + out_indices_u32: &mut [u32], +) -> Result { + let width = small_array.len(); + if width == 0 { + return Err("`elem_to_compare` size must be >0"); + } + if big_array.len() % width != 0 { + return Err("`array_of_elems` size must be multiplier of `elem_to_compare`"); + } + let max_bits = (width as u64).saturating_mul(8); + if max_bits > u16::MAX as u64 { + return Err("element width too large for u16 packed distances"); + } + let num_records = big_array.len() / width; + if num_records > u32::MAX as usize { + return Err("catalog record count exceeds u32::MAX"); + } + if out_distances_u16.len() < num_records || out_indices_u32.len() < num_records { + return Err("output buffers must have capacity for at least num_records entries"); + } + let matches = crate::bytes_array_all_within_dist(big_array, small_array, max_dist)?; + for (i, (d, idx)) in matches.iter().enumerate() { + out_distances_u16[i] = *d as u16; + out_indices_u32[i] = *idx as u32; + } + Ok(matches.len()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn make_records(width: usize, count: usize, seed: u64) -> Vec { + let mut state = seed; + let mut out = Vec::with_capacity(width * count); + for _ in 0..(width * count) { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + out.push((state >> 33) as u8); + } + out + } + + fn oracle_distance(a: &[u8], b: &[u8]) -> u64 { + assert_eq!(a.len(), b.len()); + a.iter() + .zip(b.iter()) + .map(|(&x, &y)| (x ^ y).count_ones() as u64) + .sum() + } + + #[test] + fn pairwise_matches_oracle_various_widths() { + for &width in &[1usize, 16, 24, 32, 33] { + let count = 41; + let a = make_records(width, count, 0xDEAD_BEEF); + let b = make_records(width, count, 0xCAFE_F00D); + let expected: Vec = (0..count) + .map(|i| { + oracle_distance( + &a[i * width..(i + 1) * width], + &b[i * width..(i + 1) * width], + ) + }) + .collect(); + let got = bytes_pairwise_distances(&a, &b, width).unwrap(); + assert_eq!(got, expected, "mismatch at width {width}"); + } + } + + #[test] + fn pairwise_empty_batch() { + let got = bytes_pairwise_distances(&[], &[], 16).unwrap(); + assert!(got.is_empty()); + } + + #[test] + fn pairwise_error_shapes() { + assert!(bytes_pairwise_distances(b"aa", b"bb", 0).is_err()); + assert!(bytes_pairwise_distances(b"aaa", b"bbb", 2).is_err()); + assert!(bytes_pairwise_distances(b"aa", b"bbb", 1).is_err()); + } + + #[test] + fn pairwise_into_writes_le_and_returns_count() { + let width = 16; + let count = 5; + let a = make_records(width, count, 1); + let b = make_records(width, count, 2); + let mut out = vec![0u8; count * 8]; + let n = bytes_pairwise_distances_into(&a, &b, width, &mut out).unwrap(); + assert_eq!(n, count); + let list = bytes_pairwise_distances(&a, &b, width).unwrap(); + for (i, d) in list.iter().enumerate() { + let bytes: [u8; 8] = out[i * 8..(i + 1) * 8].try_into().unwrap(); + assert_eq!(u64::from_le_bytes(bytes), *d); + } + } + + #[test] + fn pairwise_into_rejects_wrong_size() { + let width = 16; + let count = 5; + let a = make_records(width, count, 1); + let b = make_records(width, count, 2); + let mut ok = vec![0u8; count * 8]; + assert!(bytes_pairwise_distances_into(&a, &b, width, &mut ok).is_ok()); + let mut short = vec![0u8; count * 8 - 1]; + assert!(bytes_pairwise_distances_into(&a, &b, width, &mut short).is_err()); + let mut long = vec![0u8; count * 8 + 1]; + assert!(bytes_pairwise_distances_into(&a, &b, width, &mut long).is_err()); + } + + #[test] + fn multi_query_first_matches_repeated_calls() { + let width = 16; + let catalog = make_records(width, 100, 11); + let queries = make_records(width, 7, 12); + let batch = bytes_array_first_many_within_dist(&catalog, &queries, width, 8).unwrap(); + assert_eq!(batch.len(), 7); + for (q_index, want) in batch.iter().enumerate() { + let query = &queries[q_index * width..(q_index + 1) * width]; + let got = crate::bytes_array_first_within_dist(&catalog, query, 8).unwrap(); + assert_eq!(*want, got, "query {q_index}"); + } + } + + #[test] + fn multi_query_best_matches_repeated_calls() { + let width = 16; + let catalog = make_records(width, 200, 21); + let queries = make_records(width, 5, 22); + let batch = bytes_array_best_many_within_dist(&catalog, &queries, width, 64).unwrap(); + for (q_index, want) in batch.iter().enumerate() { + let query = &queries[q_index * width..(q_index + 1) * width]; + let got = crate::bytes_array_best_within_dist(&catalog, query, 64).unwrap(); + assert_eq!(*want, got, "query {q_index}"); + } + } + + #[test] + fn multi_query_all_matches_repeated_calls() { + let width = 16; + let catalog = make_records(width, 40, 31); + let queries = make_records(width, 3, 32); + let batch = bytes_array_all_many_within_dist(&catalog, &queries, width, 62).unwrap(); + for (q_index, want) in batch.iter().enumerate() { + let query = &queries[q_index * width..(q_index + 1) * width]; + let got = crate::bytes_array_all_within_dist(&catalog, query, 62).unwrap(); + assert_eq!(*want, got, "query {q_index}"); + } + } + + #[test] + fn multi_query_empty_queries() { + let catalog = make_records(16, 5, 1); + let queries: Vec = Vec::new(); + assert!( + bytes_array_first_many_within_dist(&catalog, &queries, 16, 4) + .unwrap() + .is_empty() + ); + assert!(bytes_array_best_many_within_dist(&catalog, &queries, 16, 4) + .unwrap() + .is_empty()); + assert!(bytes_array_all_many_within_dist(&catalog, &queries, 16, 4) + .unwrap() + .is_empty()); + } + + #[test] + fn packed_all_matches_list_and_dense_case() { + let width = 16; + let catalog = make_records(width, 128, 41); + let query = &catalog[3 * width..4 * width]; // dense: at least one exact match + let list = crate::bytes_array_all_within_dist(&catalog, query, 128).unwrap(); + let (dists, idxs) = bytes_array_all_within_dist_packed(&catalog, query, 128).unwrap(); + assert_eq!(dists.len(), list.len()); + assert_eq!(idxs.len(), list.len()); + for (i, (d, idx)) in list.iter().enumerate() { + assert_eq!(dists[i] as u64, *d); + assert_eq!(idxs[i] as usize, *idx); + } + // include an exact match at least once + assert!(list.iter().any(|(d, _)| *d == 0)); + } + + #[test] + fn packed_all_sparse_case() { + let width = 16; + let catalog = make_records(width, 200, 51); + let query = &catalog[10 * width..11 * width]; + // Very tight max_dist: only exact matches. + let (dists, idxs) = bytes_array_all_within_dist_packed(&catalog, query, 0).unwrap(); + assert!(!dists.is_empty()); + for (i, &d) in dists.iter().enumerate() { + assert_eq!(d, 0); + let idx = idxs[i] as usize; + assert_eq!(&catalog[idx * width..(idx + 1) * width], query); + } + } + + #[test] + fn packed_into_matches_packed() { + let width = 16; + let catalog = make_records(width, 128, 61); + let query = &catalog[5 * width..6 * width]; + let num_records = catalog.len() / width; + let (dists, idxs) = bytes_array_all_within_dist_packed(&catalog, query, 128).unwrap(); + let mut out_d = vec![0u16; num_records]; + let mut out_i = vec![0u32; num_records]; + let n = + bytes_array_all_within_dist_into(&catalog, query, 128, &mut out_d, &mut out_i).unwrap(); + assert_eq!(n, dists.len()); + assert_eq!(&out_d[..n], &dists[..]); + assert_eq!(&out_i[..n], &idxs[..]); + } + + #[test] + fn packed_into_rejects_short_buffers() { + let width = 16; + let catalog = make_records(width, 8, 71); + let query = &catalog[0..width]; + let mut short_d = vec![0u16; 4]; + let mut ok_i = vec![0u32; 8]; + assert!( + bytes_array_all_within_dist_into(&catalog, query, 128, &mut short_d, &mut ok_i,) + .is_err() + ); + let mut ok_d = vec![0u16; 8]; + let mut short_i = vec![0u32; 4]; + assert!( + bytes_array_all_within_dist_into(&catalog, query, 128, &mut ok_d, &mut short_i,) + .is_err() + ); + } + + #[test] + fn algorithm_invariance_pairwise() { + let width = 16; + let a = make_records(width, 50, 81); + let b = make_records(width, 50, 82); + let baseline = bytes_pairwise_distances(&a, &b, width).unwrap(); + for algo in ["classic", "native"] { + crate::api::set_algorithm(algo).unwrap(); + let got = bytes_pairwise_distances(&a, &b, width).unwrap(); + assert_eq!(got, baseline, "mismatch under algo {algo}"); + } + crate::api::set_algorithm("native").unwrap(); + } + + #[test] + fn best_many_tiebreak_lowest_index() { + let width = 16; + let mut catalog = vec![0xFFu8; width * 20]; + // Two exact-match entries at indices 4 and 9. + for &idx in &[4usize, 9] { + for byte in &mut catalog[idx * width..(idx + 1) * width] { + *byte = 0; + } + } + let query = vec![0u8; width]; + let batch = bytes_array_best_many_within_dist(&catalog, &query, width, 128).unwrap(); + assert_eq!(batch, vec![Some((0u64, 4usize))]); + } + + #[test] + fn all_many_ordering_preserved() { + let width = 16; + let mut catalog = vec![0xFFu8; width * 12]; + for &idx in &[1usize, 5, 8, 11] { + for byte in &mut catalog[idx * width..(idx + 1) * width] { + *byte = 0; + } + } + let query = vec![0u8; width]; + // 0xFF vs 0 distance = 128 bits over 16 bytes; use max_dist 4 so only + // the exact matches survive. + let batch = bytes_array_all_many_within_dist(&catalog, &query, width, 4).unwrap(); + assert_eq!(batch.len(), 1); + let matches: Vec = batch[0].iter().map(|&(_, i)| i).collect(); + assert_eq!(matches, vec![1, 5, 8, 11]); + } +} diff --git a/src/lib.rs b/src/lib.rs index 193b6b0..9b9fe65 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -20,6 +20,7 @@ use std::sync::atomic::{AtomicU8, Ordering}; // enabled, `python` registers the PyO3 module and maps Python errors/GIL policy // onto the same dispatch functions used by Rust callers. mod api; +mod batch; mod classic; mod hex; mod native; @@ -33,6 +34,11 @@ mod tests; mod x86_simd; pub use api::*; +pub use batch::{ + bytes_array_all_many_within_dist, bytes_array_all_within_dist_into, + bytes_array_all_within_dist_packed, bytes_array_best_many_within_dist, + bytes_array_first_many_within_dist, bytes_pairwise_distances, bytes_pairwise_distances_into, +}; /// Lookup table for popcount of 4-bit values (0-15). /// Hex string distance is computed one nibble at a time, so this avoids a diff --git a/src/python.rs b/src/python.rs index 7ee6072..3525aff 100644 --- a/src/python.rs +++ b/src/python.rs @@ -57,17 +57,35 @@ impl SimpleByteBuffer { } fn acquire(mut self: Pin<&mut Self>, obj: &Bound<'_, PyAny>) -> PyResult<()> { + self.as_mut().acquire_mode(obj, false) + } + + fn acquire_writable(mut self: Pin<&mut Self>, obj: &Bound<'_, PyAny>) -> PyResult<()> { + self.as_mut().acquire_mode(obj, true) + } + + fn acquire_mode( + mut self: Pin<&mut Self>, + obj: &Bound<'_, PyAny>, + writable: bool, + ) -> PyResult<()> { let this = unsafe { self.as_mut().get_unchecked_mut() }; - let result = unsafe { - ffi::PyObject_GetBuffer( - obj.as_ptr(), - &mut this.raw, - ffi::PyBUF_ND | ffi::PyBUF_FORMAT, - ) + let flags = if writable { + ffi::PyBUF_C_CONTIGUOUS | ffi::PyBUF_FORMAT | ffi::PyBUF_WRITABLE + } else { + ffi::PyBUF_ND | ffi::PyBUF_FORMAT }; + let result = unsafe { ffi::PyObject_GetBuffer(obj.as_ptr(), &mut this.raw, flags) }; if result != 0 { let _ = PyErr::fetch(obj.py()); - let message = if unsafe { ffi::PyObject_CheckBuffer(obj.as_ptr()) } != 0 { + let supports_buffer = unsafe { ffi::PyObject_CheckBuffer(obj.as_ptr()) } != 0; + let message = if writable { + if supports_buffer { + "output must be a writable, C-contiguous buffer" + } else { + "output must support the buffer protocol" + } + } else if supports_buffer { "input must be contiguous" } else { "error occurred while parsing arguments" @@ -76,19 +94,30 @@ impl SimpleByteBuffer { } this.acquired = true; + if writable && this.raw.readonly != 0 { + return Err(PyValueError::new_err("output must be writable")); + } let compatible_format = if this.raw.format.is_null() { - false + writable } else { let format = unsafe { CStr::from_ptr(this.raw.format) }; ::is_compatible_format(format) }; if this.raw.itemsize != 1 || !compatible_format { - return Err(PyValueError::new_err( - "error occurred while parsing arguments", - )); + let message = if writable { + "output must be a byte buffer (itemsize 1)" + } else { + "error occurred while parsing arguments" + }; + return Err(PyValueError::new_err(message)); } if this.raw.len < 0 || (this.raw.buf.is_null() && this.raw.len != 0) { - return Err(PyValueError::new_err("invalid buffer view")); + let message = if writable { + "invalid output buffer" + } else { + "invalid buffer view" + }; + return Err(PyValueError::new_err(message)); } Ok(()) } @@ -98,8 +127,28 @@ impl SimpleByteBuffer { #[inline] unsafe fn as_slice(self: Pin<&Self>) -> &[u8] { let raw = &self.get_ref().raw; + if raw.len == 0 { + return &[]; + } std::slice::from_raw_parts(raw.buf as *const u8, raw.len as usize) } + + #[inline] + fn len(&self) -> usize { + self.raw.len as usize + } + + /// The pointer is valid while the pinned guard is alive. Dereferencing it + /// or constructing references from it remains the caller's responsibility. + #[inline] + fn raw_mut_ptr(mut self: Pin<&mut Self>) -> *mut u8 { + let this = unsafe { self.as_mut().get_unchecked_mut() }; + if this.raw.len == 0 { + std::ptr::NonNull::::dangling().as_ptr() + } else { + this.raw.buf as *mut u8 + } + } } impl Drop for SimpleByteBuffer { @@ -176,8 +225,7 @@ fn hamming_distance_bytes( let mut buf_b = std::pin::pin!(SimpleByteBuffer::new()); buf_a.as_mut().acquire(a)?; buf_b.as_mut().acquire(b)?; - let a_slice = unsafe { buf_a.as_ref().as_slice() }; - let b_slice = unsafe { buf_b.as_ref().as_slice() }; + let (a_slice, b_slice) = unsafe { (buf_a.as_ref().as_slice(), buf_b.as_ref().as_slice()) }; if a_slice.len() != b_slice.len() { return Err(PyValueError::new_err("bytes are NOT the same length")); @@ -320,8 +368,7 @@ fn check_bytes_within_dist( let mut buf_b = std::pin::pin!(SimpleByteBuffer::new()); buf_a.as_mut().acquire(a)?; buf_b.as_mut().acquire(b)?; - let a_slice = unsafe { buf_a.as_ref().as_slice() }; - let b_slice = unsafe { buf_b.as_ref().as_slice() }; + let (a_slice, b_slice) = unsafe { (buf_a.as_ref().as_slice(), buf_b.as_ref().as_slice()) }; if a_slice.is_empty() || b_slice.is_empty() { return Err(PyValueError::new_err("array size must be >0")); @@ -579,6 +626,463 @@ fn check_bytes_arrays_all_within_dist( Ok(results) } +// --------------------------------------------------------------------------- +// Batch APIs — amortize the Python↔Rust boundary across many distance calls. +// --------------------------------------------------------------------------- + +/// Acquire read-only buffers for `a` and `b`, invoking the closure with their +/// slices. Both buffers live for the whole call (dropped on return). +#[inline] +fn with_two_readonly_buffers(a: &Bound<'_, PyAny>, b: &Bound<'_, PyAny>, f: F) -> PyResult +where + F: FnOnce(&[u8], &[u8], bool) -> PyResult, +{ + if let (Some(a_slice), Some(b_slice)) = (exact_bytes(a), exact_bytes(b)) { + return f(a_slice, b_slice, true); + } + let mut buf_a = std::pin::pin!(SimpleByteBuffer::new()); + let mut buf_b = std::pin::pin!(SimpleByteBuffer::new()); + buf_a.as_mut().acquire(a)?; + buf_b.as_mut().acquire(b)?; + let (a_slice, b_slice) = unsafe { (buf_a.as_ref().as_slice(), buf_b.as_ref().as_slice()) }; + // General buffer-protocol inputs may be writable. Keep the GIL attached + // while reading them so another Python thread cannot mutate the storage. + f(a_slice, b_slice, false) +} + +#[inline] +fn buffer_ranges_overlap(a_ptr: *const u8, a_len: usize, b_ptr: *const u8, b_len: usize) -> bool { + if a_len == 0 || b_len == 0 { + return false; + } + let a_start = a_ptr as usize; + let b_start = b_ptr as usize; + let a_end = a_start.checked_add(a_len).unwrap_or(usize::MAX); + let b_end = b_start.checked_add(b_len).unwrap_or(usize::MAX); + a_start < b_end && b_start < a_end +} + +#[inline] +fn ensure_writable_into_supported() -> PyResult<()> { + #[cfg(Py_GIL_DISABLED)] + { + return Err(PyValueError::new_err( + "writable `_into` APIs are unavailable on free-threaded Python; use the packed API", + )); + } + #[cfg(not(Py_GIL_DISABLED))] + { + Ok(()) + } +} + +/// Compute Hamming distances between corresponding fixed-width records in `a` +/// and `b`. Returns a list of `int` distances, one per record. +/// +/// `a` and `b` must be equal-length buffer-protocol objects whose length is a +/// multiple of `element_size`. +#[pyfunction] +#[pyo3(signature = (a, b, element_size))] +fn hamming_distances_bytes( + py: Python<'_>, + a: &Bound<'_, PyAny>, + b: &Bound<'_, PyAny>, + element_size: usize, +) -> PyResult> { + with_two_readonly_buffers(a, b, |a_slice, b_slice, can_detach| { + // Small batches stay attached; large batches detach the GIL. The raw + // slices remain valid because the Py_buffer guards outlive `f`. + let compute = || { + crate::bytes_pairwise_distances(a_slice, b_slice, element_size) + .map_err(PyValueError::new_err) + }; + if can_detach && a_slice.len() >= ARRAY_GIL_RELEASE_THRESHOLD { + py.detach(compute) + } else { + compute() + } + }) +} + +/// Compute Hamming distances and return them as `bytes` of little-endian u64 +/// values (8 bytes per distance). +#[pyfunction] +#[pyo3(signature = (a, b, element_size))] +fn hamming_distances_bytes_packed<'py>( + py: Python<'py>, + a: &Bound<'_, PyAny>, + b: &Bound<'_, PyAny>, + element_size: usize, +) -> PyResult> { + let distances = with_two_readonly_buffers(a, b, |a_slice, b_slice, can_detach| { + let compute = || { + crate::bytes_pairwise_distances(a_slice, b_slice, element_size) + .map_err(PyValueError::new_err) + }; + if can_detach && a_slice.len() >= ARRAY_GIL_RELEASE_THRESHOLD { + py.detach(compute) + } else { + compute() + } + })?; + let byte_len = distances + .len() + .checked_mul(8) + .ok_or_else(|| PyValueError::new_err("distance buffer size overflows platform usize"))?; + PyBytes::new_with(py, byte_len, |dst| { + for (chunk, d) in dst.chunks_exact_mut(8).zip(&distances) { + chunk.copy_from_slice(&d.to_le_bytes()); + } + Ok(()) + }) +} + +/// Compute Hamming distances and write them as little-endian u64 values into +/// `output`. Returns the number of distances written (== `len(a)//element_size`). +/// +/// `output` must be a writable, C-contiguous byte buffer of exactly +/// `count * 8` bytes. Non-contiguous, read-only, or wrong-size outputs are +/// rejected with `ValueError`. +#[pyfunction] +#[pyo3(signature = (a, b, element_size, output))] +fn hamming_distances_bytes_into( + _py: Python<'_>, + a: &Bound<'_, PyAny>, + b: &Bound<'_, PyAny>, + element_size: usize, + output: &Bound<'_, PyAny>, +) -> PyResult { + ensure_writable_into_supported()?; + let mut buf_out = std::pin::pin!(SimpleByteBuffer::new()); + buf_out.as_mut().acquire_writable(output)?; + let out_len = buf_out.len(); + let out_ptr_addr = buf_out.as_mut().raw_mut_ptr() as usize; + + with_two_readonly_buffers(a, b, |a_slice, b_slice, _can_detach| { + if buffer_ranges_overlap( + a_slice.as_ptr(), + a_slice.len(), + out_ptr_addr as *const u8, + out_len, + ) || buffer_ranges_overlap( + b_slice.as_ptr(), + b_slice.len(), + out_ptr_addr as *const u8, + out_len, + ) { + return Err(PyValueError::new_err( + "output buffer must not overlap input buffers", + )); + } + if element_size == 0 { + return Err(PyValueError::new_err("`element_size` must be >0")); + } + if a_slice.len() != b_slice.len() { + return Err(PyValueError::new_err("bytes are NOT the same length")); + } + if a_slice.len() % element_size != 0 { + return Err(PyValueError::new_err( + "length must be a multiple of `element_size`", + )); + } + let count = a_slice.len() / element_size; + let expected = count + .checked_mul(8) + .ok_or_else(|| PyValueError::new_err("output capacity overflows"))?; + if out_len != expected { + return Err(PyValueError::new_err("`out` must be exactly count*8 bytes")); + } + + // Writable buffer exports are not exclusive. Keep the GIL attached + // while constructing and using the mutable Rust slice. + let out_slice = unsafe { std::slice::from_raw_parts_mut(out_ptr_addr as *mut u8, out_len) }; + crate::bytes_pairwise_distances_into(a_slice, b_slice, element_size, out_slice) + .map_err(PyValueError::new_err) + }) +} + +// --------------------------------------------------------------------------- +// Multi-query catalog scans +// --------------------------------------------------------------------------- + +#[inline] +fn resolve_and_dispatch_multi( + py: Python<'_>, + catalog: &Bound<'_, PyAny>, + queries: &Bound<'_, PyAny>, + query_width: usize, + max_dist: i64, + compute: F, +) -> PyResult +where + F: FnOnce(&[u8], &[u8], usize, i64) -> Result + Send, + R: Send, +{ + if max_dist < 0 { + return Err(PyValueError::new_err("`max_dist` must be >=0")); + } + with_two_readonly_buffers( + catalog, + queries, + |catalog_slice, queries_slice, can_detach| { + // Detach when total work (catalog * queries scan) is non-trivial. + // catalog.len() is a proxy since queries usually << catalog. + let compute_call = || compute(catalog_slice, queries_slice, query_width, max_dist); + let total_work = catalog_slice.len().saturating_mul(queries_slice.len()); + let result = if can_detach && total_work >= 4 * ARRAY_GIL_RELEASE_THRESHOLD { + py.detach(compute_call) + } else { + compute_call() + }; + result.map_err(PyValueError::new_err) + }, + ) +} + +/// Run [`check_bytes_arrays_first_within_dist`] against `catalog` for every +/// fixed-width slice of `queries`. Returns a list of `int` indices (or `-1` +/// when no record matches), one per query. +#[pyfunction] +#[pyo3(signature = (catalog, queries, query_width, max_dist))] +fn check_bytes_arrays_first_many_within_dist( + py: Python<'_>, + catalog: &Bound<'_, PyAny>, + queries: &Bound<'_, PyAny>, + query_width: usize, + max_dist: i64, +) -> PyResult> { + resolve_and_dispatch_multi(py, catalog, queries, query_width, max_dist, |c, q, w, m| { + let results = crate::bytes_array_first_many_within_dist(c, q, w, m)?; + Ok(results + .into_iter() + .map(|r| r.map(|i| i as i64).unwrap_or(-1)) + .collect()) + }) +} + +/// Run [`check_bytes_arrays_best_within_dist`] against `catalog` for every +/// fixed-width slice of `queries`. Returns a list of `(distance, index)` +/// tuples, using `(-1, -1)` for queries with no match. +#[pyfunction] +#[pyo3(signature = (catalog, queries, query_width, max_dist))] +fn check_bytes_arrays_best_many_within_dist( + py: Python<'_>, + catalog: &Bound<'_, PyAny>, + queries: &Bound<'_, PyAny>, + query_width: usize, + max_dist: i64, +) -> PyResult> { + resolve_and_dispatch_multi(py, catalog, queries, query_width, max_dist, |c, q, w, m| { + let results = crate::bytes_array_best_many_within_dist(c, q, w, m)?; + Ok(results + .into_iter() + .map(|r| r.map(|(d, i)| (d as i64, i as i64)).unwrap_or((-1, -1))) + .collect()) + }) +} + +/// Run [`check_bytes_arrays_all_within_dist`] against `catalog` for every +/// fixed-width slice of `queries`. Returns a list of lists of +/// `(distance, index)` tuples, one inner list per query. +#[pyfunction] +#[pyo3(signature = (catalog, queries, query_width, max_dist))] +fn check_bytes_arrays_all_many_within_dist( + py: Python<'_>, + catalog: &Bound<'_, PyAny>, + queries: &Bound<'_, PyAny>, + query_width: usize, + max_dist: i64, +) -> PyResult>> { + resolve_and_dispatch_multi(py, catalog, queries, query_width, max_dist, |c, q, w, m| { + let results = crate::bytes_array_all_many_within_dist(c, q, w, m)?; + Ok(results + .into_iter() + .map(|matches| { + matches + .into_iter() + .map(|(d, i)| (d, i as u64)) + .collect::>() + }) + .collect()) + }) +} + +// --------------------------------------------------------------------------- +// Packed/into transport for dense all-results +// --------------------------------------------------------------------------- + +/// Dense-transport variant of `check_bytes_arrays_all_within_dist`. +/// +/// Returns a `(distance_bytes, index_bytes)` tuple. `distance_bytes` is a +/// contiguous little-endian u16 buffer (2 bytes per match); `index_bytes` is a +/// contiguous little-endian u32 buffer (4 bytes per match). Matches are +/// ordered by ascending catalog index. Fails if the element width would allow +/// distances exceeding `u16::MAX` bits, or if the catalog exceeds +/// `u32::MAX` records. +#[pyfunction] +#[pyo3(signature = (array_of_elems, elem_to_compare, max_dist))] +fn check_bytes_arrays_all_within_dist_packed<'py>( + py: Python<'py>, + array_of_elems: &Bound<'_, PyAny>, + elem_to_compare: &Bound<'_, PyAny>, + max_dist: i64, +) -> PyResult<(Bound<'py, PyBytes>, Bound<'py, PyBytes>)> { + if max_dist < 0 { + return Err(PyValueError::new_err("`max_dist` must be >=0")); + } + let (dists, idxs) = + with_two_readonly_buffers(array_of_elems, elem_to_compare, |big, small, can_detach| { + let compute = || { + crate::bytes_array_all_within_dist_packed(big, small, max_dist) + .map_err(PyValueError::new_err) + }; + if can_detach && big.len() >= ARRAY_GIL_RELEASE_THRESHOLD { + py.detach(compute) + } else { + compute() + } + })?; + let d_byte_len = dists + .len() + .checked_mul(2) + .ok_or_else(|| PyValueError::new_err("distance buffer size overflows platform usize"))?; + let i_byte_len = idxs + .len() + .checked_mul(4) + .ok_or_else(|| PyValueError::new_err("index buffer size overflows platform usize"))?; + let d_bytes = PyBytes::new_with(py, d_byte_len, |dst| { + for (chunk, value) in dst.chunks_exact_mut(2).zip(&dists) { + chunk.copy_from_slice(&value.to_le_bytes()); + } + Ok(()) + })?; + let i_bytes = PyBytes::new_with(py, i_byte_len, |dst| { + for (chunk, value) in dst.chunks_exact_mut(4).zip(&idxs) { + chunk.copy_from_slice(&value.to_le_bytes()); + } + Ok(()) + })?; + Ok((d_bytes, i_bytes)) +} + +/// Write `check_bytes_arrays_all_within_dist` results into caller-provided +/// writable buffers. `out_distances` receives little-endian u16 distances (2 +/// bytes per match); `out_indices` receives little-endian u32 indices (4 +/// bytes per match). Returns the number of matches written. +/// +/// Both output buffers must be sized for the worst case (`num_records * 2` +/// and `num_records * 4` bytes respectively); if the resolved match count +/// exceeds the capacity, `ValueError` is raised. +#[pyfunction] +#[pyo3(signature = (array_of_elems, elem_to_compare, max_dist, out_distances, out_indices))] +fn check_bytes_arrays_all_within_dist_into( + _py: Python<'_>, + array_of_elems: &Bound<'_, PyAny>, + elem_to_compare: &Bound<'_, PyAny>, + max_dist: i64, + out_distances: &Bound<'_, PyAny>, + out_indices: &Bound<'_, PyAny>, +) -> PyResult { + ensure_writable_into_supported()?; + if max_dist < 0 { + return Err(PyValueError::new_err("`max_dist` must be >=0")); + } + let mut buf_d = std::pin::pin!(SimpleByteBuffer::new()); + let mut buf_i = std::pin::pin!(SimpleByteBuffer::new()); + buf_d.as_mut().acquire_writable(out_distances)?; + buf_i.as_mut().acquire_writable(out_indices)?; + let d_len = buf_d.len(); + let i_len = buf_i.len(); + let d_ptr_addr = buf_d.as_mut().raw_mut_ptr() as usize; + let i_ptr_addr = buf_i.as_mut().raw_mut_ptr() as usize; + if buffer_ranges_overlap( + d_ptr_addr as *const u8, + d_len, + i_ptr_addr as *const u8, + i_len, + ) { + return Err(PyValueError::new_err( + "output buffers must not overlap each other", + )); + } + + with_two_readonly_buffers( + array_of_elems, + elem_to_compare, + |big, small, _can_detach| { + for (ptr, len) in [ + (d_ptr_addr as *const u8, d_len), + (i_ptr_addr as *const u8, i_len), + ] { + if buffer_ranges_overlap(big.as_ptr(), big.len(), ptr, len) + || buffer_ranges_overlap(small.as_ptr(), small.len(), ptr, len) + { + return Err(PyValueError::new_err( + "output buffers must not overlap input buffers", + )); + } + } + let width = small.len(); + if width == 0 { + return Err(PyValueError::new_err("`elem_to_compare` size must be >0")); + } + if big.len() % width != 0 { + return Err(PyValueError::new_err( + "`array_of_elems` size must be multiplier of `elem_to_compare`", + )); + } + let num_records = big.len() / width; + let max_bits = (width as u64).saturating_mul(8); + if max_bits > u16::MAX as u64 { + return Err(PyValueError::new_err( + "element width too large for u16 packed distances", + )); + } + if num_records > u32::MAX as usize { + return Err(PyValueError::new_err( + "catalog record count exceeds u32::MAX", + )); + } + // Worst-case capacity check up front so partial writes are never + // observable. + let required_distances = num_records.checked_mul(2).ok_or_else(|| { + PyValueError::new_err("distance buffer size overflows platform usize") + })?; + let required_indices = num_records.checked_mul(4).ok_or_else(|| { + PyValueError::new_err("index buffer size overflows platform usize") + })?; + if d_len < required_distances { + return Err(PyValueError::new_err( + "out_distances must have capacity for num_records * 2 bytes", + )); + } + if i_len < required_indices { + return Err(PyValueError::new_err( + "out_indices must have capacity for num_records * 4 bytes", + )); + } + + let matches = crate::bytes_array_all_within_dist(big, small, max_dist) + .map_err(PyValueError::new_err)?; + // Writable buffer exports are not exclusive. Keep the GIL attached + // while serializing into the caller's Python-owned buffers. + for (k, (distance, index)) in matches.iter().enumerate() { + unsafe { + std::ptr::write_unaligned( + (d_ptr_addr as *mut u8).add(k * 2) as *mut u16, + (*distance as u16).to_le(), + ); + std::ptr::write_unaligned( + (i_ptr_addr as *mut u8).add(k * 4) as *mut u32, + (*index as u32).to_le(), + ); + } + } + Ok(matches.len()) + }, + ) +} + // --------------------------------------------------------------------------- // §9 + §14: set_algo — delegate to api::set_algorithm // --------------------------------------------------------------------------- @@ -613,6 +1117,29 @@ fn hexhamming(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(check_bytes_arrays_first_within_dist, m)?)?; m.add_function(wrap_pyfunction!(check_bytes_arrays_best_within_dist, m)?)?; m.add_function(wrap_pyfunction!(check_bytes_arrays_all_within_dist, m)?)?; + m.add_function(wrap_pyfunction!(hamming_distances_bytes, m)?)?; + m.add_function(wrap_pyfunction!(hamming_distances_bytes_packed, m)?)?; + m.add_function(wrap_pyfunction!(hamming_distances_bytes_into, m)?)?; + m.add_function(wrap_pyfunction!( + check_bytes_arrays_first_many_within_dist, + m + )?)?; + m.add_function(wrap_pyfunction!( + check_bytes_arrays_best_many_within_dist, + m + )?)?; + m.add_function(wrap_pyfunction!( + check_bytes_arrays_all_many_within_dist, + m + )?)?; + m.add_function(wrap_pyfunction!( + check_bytes_arrays_all_within_dist_packed, + m + )?)?; + m.add_function(wrap_pyfunction!( + check_bytes_arrays_all_within_dist_into, + m + )?)?; m.add_function(wrap_pyfunction!(set_algo, m)?)?; // Auto-detect best algorithm on module load diff --git a/test/test_batch.py b/test/test_batch.py new file mode 100644 index 0000000..6c6fa52 --- /dev/null +++ b/test/test_batch.py @@ -0,0 +1,452 @@ +"""Tests for the batch APIs added in perf/python-batch-apis. + +Covers: +* Pairwise list/packed/into against a hand-oracle and each other. +* Multi-query first/best/all_many parity with repeated single-call use. +* Packed and _into transport for `all_within_dist` (dense + sparse). +* Error shapes for element_size / length mismatches, readonly outputs, + non-contiguous outputs, wrong-size outputs, and unaligned writable memoryview + targets. +* Algorithm invariance and ordering / tie behavior preserved. +""" + +import array +import random +import sys + +import pytest +from hexhamming import ( + check_bytes_arrays_all_many_within_dist, + check_bytes_arrays_all_within_dist, + check_bytes_arrays_all_within_dist_into, + check_bytes_arrays_all_within_dist_packed, + check_bytes_arrays_best_many_within_dist, + check_bytes_arrays_best_within_dist, + check_bytes_arrays_first_many_within_dist, + check_bytes_arrays_first_within_dist, + hamming_distances_bytes, + hamming_distances_bytes_into, + hamming_distances_bytes_packed, + set_algo, +) + + +def _random_bytes(length: int, seed: int) -> bytes: + rng = random.Random(seed) + return bytes(rng.randrange(256) for _ in range(length)) + + +def _oracle_distance(a: bytes, b: bytes) -> int: + assert len(a) == len(b) + total = 0 + for x, y in zip(a, b): + total += (x ^ y).bit_count() + return total + + +# --------------------------------------------------------------------------- +# Pairwise batch API +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("width", (1, 16, 24, 32, 33)) +def test_pairwise_matches_oracle_random(width): + count = 41 + a = _random_bytes(width * count, 0x51_0000 + width) + b = _random_bytes(width * count, 0x52_0000 + width) + expected = [ + _oracle_distance(a[i * width : (i + 1) * width], b[i * width : (i + 1) * width]) + for i in range(count) + ] + assert hamming_distances_bytes(a, b, width) == expected + + +def test_pairwise_empty_batch(): + assert hamming_distances_bytes(b"", b"", 16) == [] + assert hamming_distances_bytes_packed(b"", b"", 16) == b"" + out = bytearray(0) + assert hamming_distances_bytes_into(b"", b"", 16, out) == 0 + assert bytes(out) == b"" + + +def test_pairwise_error_shapes(): + with pytest.raises(ValueError): + hamming_distances_bytes(b"aa", b"bb", 0) + with pytest.raises(ValueError): + hamming_distances_bytes(b"aa", b"bbb", 1) + with pytest.raises(ValueError): + hamming_distances_bytes(b"aaa", b"bbb", 2) + + +def test_pairwise_packed_matches_list(): + width = 16 + count = 5 + a = _random_bytes(width * count, 1) + b = _random_bytes(width * count, 2) + dists = hamming_distances_bytes(a, b, width) + packed = hamming_distances_bytes_packed(a, b, width) + assert len(packed) == count * 8 + parsed = [ + int.from_bytes(packed[i * 8 : (i + 1) * 8], "little") for i in range(count) + ] + assert parsed == dists + + +def test_pairwise_into_matches_packed(): + width = 32 + count = 7 + a = _random_bytes(width * count, 10) + b = _random_bytes(width * count, 11) + dists = hamming_distances_bytes(a, b, width) + out = bytearray(count * 8) + n = hamming_distances_bytes_into(a, b, width, out) + assert n == count + parsed = [int.from_bytes(out[i * 8 : (i + 1) * 8], "little") for i in range(count)] + assert parsed == dists + + +def test_pairwise_into_rejects_readonly(): + width = 16 + count = 3 + a = _random_bytes(width * count, 1) + b = _random_bytes(width * count, 2) + with pytest.raises(ValueError): + hamming_distances_bytes_into(a, b, width, bytes(count * 8)) + + +def test_pairwise_into_rejects_wrong_size(): + width = 16 + count = 3 + a = _random_bytes(width * count, 1) + b = _random_bytes(width * count, 2) + with pytest.raises(ValueError): + hamming_distances_bytes_into(a, b, width, bytearray(count * 8 - 1)) + with pytest.raises(ValueError): + hamming_distances_bytes_into(a, b, width, bytearray(count * 8 + 1)) + + +def test_pairwise_into_rejects_noncontiguous_output(): + width = 16 + count = 4 + a = _random_bytes(width * count, 1) + b = _random_bytes(width * count, 2) + backing = bytearray(count * 8 * 2) + # Every-other-byte stride is not C-contiguous. + mv = memoryview(backing)[::2] + with pytest.raises(ValueError): + hamming_distances_bytes_into(a, b, width, mv) + + +def test_pairwise_into_accepts_unaligned_memoryview(): + # Provide an intentionally unaligned writable byte view: slice a byte off + # the front of a larger backing buffer. Writes must be safe regardless of + # alignment because we use write_unaligned. + width = 16 + count = 5 + a = _random_bytes(width * count, 1) + b = _random_bytes(width * count, 2) + backing = bytearray(count * 8 + 3) + view = memoryview(backing)[3 : 3 + count * 8] + n = hamming_distances_bytes_into(a, b, width, view) + assert n == count + dists = hamming_distances_bytes(a, b, width) + parsed = [ + int.from_bytes(bytes(view[i * 8 : (i + 1) * 8]), "little") for i in range(count) + ] + assert parsed == dists + + +def test_pairwise_into_accepts_live_writable_alias(): + width = 16 + count = 5 + a = _random_bytes(width * count, 1) + b = _random_bytes(width * count, 2) + out = bytearray(count * 8) + alias = memoryview(out) + assert hamming_distances_bytes_into(a, b, width, out) == count + assert bytes(alias) == hamming_distances_bytes_packed(a, b, width) + + +def test_pairwise_into_rejects_overlapping_input(): + backing = bytearray(range(20)) + a = memoryview(backing)[:16] + b = bytes(16) + output = memoryview(backing)[4:20] + with pytest.raises(ValueError, match="must not overlap input"): + hamming_distances_bytes_into(a, b, 8, output) + + +def test_pairwise_algorithm_invariance(): + width = 16 + count = 20 + a = _random_bytes(width * count, 1) + b = _random_bytes(width * count, 2) + baseline = hamming_distances_bytes(a, b, width) + for algo in ("classic", "native"): + assert set_algo(algo) == "" + assert hamming_distances_bytes(a, b, width) == baseline + set_algo("native") + + +# --------------------------------------------------------------------------- +# Multi-query catalog scans +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("width", (16, 24, 32)) +def test_multi_query_first_matches_repeated(width): + catalog = _random_bytes(width * 50, 0x100 + width) + queries = _random_bytes(width * 4, 0x200 + width) + got = check_bytes_arrays_first_many_within_dist(catalog, queries, width, 8) + expected = [ + check_bytes_arrays_first_within_dist( + catalog, queries[i * width : (i + 1) * width], 8 + ) + for i in range(len(queries) // width) + ] + assert got == expected + + +@pytest.mark.parametrize("width", (16, 24, 32)) +def test_multi_query_best_matches_repeated(width): + catalog = _random_bytes(width * 50, 0x300 + width) + queries = _random_bytes(width * 4, 0x400 + width) + got = check_bytes_arrays_best_many_within_dist(catalog, queries, width, 32) + expected = [ + check_bytes_arrays_best_within_dist( + catalog, queries[i * width : (i + 1) * width], 32 + ) + for i in range(len(queries) // width) + ] + assert got == expected + + +@pytest.mark.parametrize("width", (16, 24, 32)) +def test_multi_query_all_matches_repeated(width): + catalog = _random_bytes(width * 30, 0x500 + width) + queries = _random_bytes(width * 3, 0x600 + width) + max_dist = 30 + got = check_bytes_arrays_all_many_within_dist(catalog, queries, width, max_dist) + expected = [ + check_bytes_arrays_all_within_dist( + catalog, queries[i * width : (i + 1) * width], max_dist + ) + for i in range(len(queries) // width) + ] + assert got == expected + + +def test_multi_query_empty_queries_produces_empty_list(): + catalog = _random_bytes(16 * 5, 1) + assert check_bytes_arrays_first_many_within_dist(catalog, b"", 16, 8) == [] + assert check_bytes_arrays_best_many_within_dist(catalog, b"", 16, 8) == [] + assert check_bytes_arrays_all_many_within_dist(catalog, b"", 16, 8) == [] + + +def test_multi_query_missing_uses_minus_one_sentinels(): + catalog = b"\xff" * 32 + query = b"\x00" * 16 # distance 128, always > max_dist 4 + got_first = check_bytes_arrays_first_many_within_dist(catalog, query, 16, 4) + assert got_first == [-1] + got_best = check_bytes_arrays_best_many_within_dist(catalog, query, 16, 4) + assert got_best == [(-1, -1)] + got_all = check_bytes_arrays_all_many_within_dist(catalog, query, 16, 4) + assert got_all == [[]] + + +def test_multi_query_error_shapes(): + catalog = _random_bytes(16 * 3, 1) + with pytest.raises(ValueError): + check_bytes_arrays_first_many_within_dist(catalog, b"", 0, 0) + with pytest.raises(ValueError): + check_bytes_arrays_first_many_within_dist(b"\x00\x00\x00", b"\x00\x00", 2, 0) + with pytest.raises(ValueError): + check_bytes_arrays_first_many_within_dist(catalog, b"\x00\x00\x00", 2, 0) + with pytest.raises(ValueError): + check_bytes_arrays_first_many_within_dist(catalog, b"\x00" * 16, 16, -1) + + +def test_multi_query_best_tiebreak_lowest_index(): + width = 16 + records = [b"\xff" * width] * 20 + records[4] = b"\x00" * width + records[9] = b"\x00" * width # duplicate exact match at higher index + catalog = b"".join(records) + got = check_bytes_arrays_best_many_within_dist(catalog, b"\x00" * width, width, 4) + assert got == [(0, 4)] + + +def test_multi_query_all_ordering_preserved(): + width = 16 + records = [b"\xff" * width] * 12 + for idx in (1, 5, 8, 11): + records[idx] = b"\x00" * width + catalog = b"".join(records) + got = check_bytes_arrays_all_many_within_dist(catalog, b"\x00" * width, width, 4) + assert got == [[(0, 1), (0, 5), (0, 8), (0, 11)]] + + +# --------------------------------------------------------------------------- +# Packed / into transport for all-results +# --------------------------------------------------------------------------- + + +def test_packed_all_dense_parity_with_list(): + width = 16 + catalog = _random_bytes(width * 64, 71) + query = catalog[3 * width : 4 * width] + list_result = check_bytes_arrays_all_within_dist(catalog, query, 128) + d_bytes, i_bytes = check_bytes_arrays_all_within_dist_packed(catalog, query, 128) + assert len(d_bytes) == len(list_result) * 2 + assert len(i_bytes) == len(list_result) * 4 + for k, (d, i) in enumerate(list_result): + assert int.from_bytes(d_bytes[k * 2 : (k + 1) * 2], "little") == d + assert int.from_bytes(i_bytes[k * 4 : (k + 1) * 4], "little") == i + + +def test_packed_all_sparse_parity_with_list(): + width = 16 + catalog = _random_bytes(width * 200, 81) + query = catalog[7 * width : 8 * width] + list_result = check_bytes_arrays_all_within_dist(catalog, query, 0) + d_bytes, i_bytes = check_bytes_arrays_all_within_dist_packed(catalog, query, 0) + assert len(d_bytes) // 2 == len(list_result) + parsed = [ + ( + int.from_bytes(d_bytes[k * 2 : (k + 1) * 2], "little"), + int.from_bytes(i_bytes[k * 4 : (k + 1) * 4], "little"), + ) + for k in range(len(list_result)) + ] + assert parsed == list_result + + +def test_into_all_matches_packed(): + width = 16 + count = 64 + catalog = _random_bytes(width * count, 91) + query = catalog[2 * width : 3 * width] + d_bytes, i_bytes = check_bytes_arrays_all_within_dist_packed(catalog, query, 128) + matches = len(d_bytes) // 2 + d_out = bytearray(count * 2) + i_out = bytearray(count * 4) + n = check_bytes_arrays_all_within_dist_into(catalog, query, 128, d_out, i_out) + assert n == matches + assert bytes(d_out[: n * 2]) == bytes(d_bytes) + assert bytes(i_out[: n * 4]) == bytes(i_bytes) + + +def test_into_all_rejects_readonly_buffers(): + width = 16 + catalog = _random_bytes(width * 8, 101) + query = catalog[:width] + with pytest.raises(ValueError): + check_bytes_arrays_all_within_dist_into( + catalog, query, 128, bytes(16), bytearray(32) + ) + with pytest.raises(ValueError): + check_bytes_arrays_all_within_dist_into( + catalog, query, 128, bytearray(16), bytes(32) + ) + + +def test_into_all_rejects_short_buffers(): + width = 16 + catalog = _random_bytes(width * 8, 111) + query = catalog[:width] + with pytest.raises(ValueError): + check_bytes_arrays_all_within_dist_into( + catalog, query, 128, bytearray(4), bytearray(32) + ) + with pytest.raises(ValueError): + check_bytes_arrays_all_within_dist_into( + catalog, query, 128, bytearray(16), bytearray(4) + ) + + +def test_into_all_accepts_unaligned_memoryview(): + width = 16 + count = 16 + catalog = _random_bytes(width * count, 121) + query = catalog[3 * width : 4 * width] + d_backing = bytearray(count * 2 + 5) + i_backing = bytearray(count * 4 + 7) + d_view = memoryview(d_backing)[5 : 5 + count * 2] + i_view = memoryview(i_backing)[7 : 7 + count * 4] + n = check_bytes_arrays_all_within_dist_into(catalog, query, 128, d_view, i_view) + d_bytes, i_bytes = check_bytes_arrays_all_within_dist_packed(catalog, query, 128) + assert n == len(d_bytes) // 2 + assert bytes(d_view[: n * 2]) == bytes(d_bytes) + assert bytes(i_view[: n * 4]) == bytes(i_bytes) + + +def test_into_all_rejects_overlapping_outputs(): + catalog = bytes(range(64)) + query = bytes(16) + backing = bytearray(20) + with pytest.raises(ValueError, match="must not overlap each other"): + check_bytes_arrays_all_within_dist_into( + catalog, + query, + 128, + memoryview(backing)[:8], + memoryview(backing)[4:20], + ) + + +def test_into_all_accepts_live_writable_aliases(): + width = 16 + count = 16 + catalog = _random_bytes(width * count, 131) + query = catalog[:width] + d_out = bytearray(count * 2) + i_out = bytearray(count * 4) + d_alias = memoryview(d_out) + i_alias = memoryview(i_out) + n = check_bytes_arrays_all_within_dist_into(catalog, query, 128, d_out, i_out) + d_bytes, i_bytes = check_bytes_arrays_all_within_dist_packed(catalog, query, 128) + assert bytes(d_alias[: n * 2]) == bytes(d_bytes) + assert bytes(i_alias[: n * 4]) == bytes(i_bytes) + + +def test_into_all_rejects_overlapping_input(): + catalog = bytearray(range(64)) + query = bytes(16) + with pytest.raises(ValueError, match="must not overlap input"): + check_bytes_arrays_all_within_dist_into( + catalog, + query, + 128, + memoryview(catalog)[:8], + bytearray(16), + ) + + +@pytest.mark.skipif( + not hasattr(sys, "_is_gil_enabled") or sys._is_gil_enabled(), + reason="requires free-threaded Python", +) +def test_writable_into_apis_rejected_without_gil(): + with pytest.raises(ValueError, match="unavailable on free-threaded Python"): + hamming_distances_bytes_into(b"\x00" * 16, b"\x00" * 16, 16, bytearray(8)) + with pytest.raises(ValueError, match="unavailable on free-threaded Python"): + check_bytes_arrays_all_within_dist_into( + b"\x00" * 16, + b"\x00" * 16, + 0, + bytearray(2), + bytearray(4), + ) + + +def test_pairwise_accepts_bytearray_and_memoryview_and_array(): + width = 16 + count = 4 + a = _random_bytes(width * count, 1) + b = _random_bytes(width * count, 2) + expected = hamming_distances_bytes(a, b, width) + assert hamming_distances_bytes(bytearray(a), bytearray(b), width) == expected + assert hamming_distances_bytes(memoryview(a), memoryview(b), width) == expected + aa = array.array("B", a) + bb = array.array("B", b) + assert hamming_distances_bytes(aa, bb, width) == expected