From 8cb8a91e5282eb29285283cefb9a495e5da37f64 Mon Sep 17 00:00:00 2001 From: Nelson Spence Date: Thu, 28 May 2026 10:31:56 -0500 Subject: [PATCH 1/2] feat: add dense SignBitmap scoring Signed-off-by: Nelson Spence --- ordvec-python/python/ordvec/__init__.py | 14 +-- ordvec-python/src/lib.rs | 52 ++++++++++ ordvec-python/tests/test_sign_bitmap.py | 44 +++++++++ src/sign_bitmap.rs | 120 ++++++++++++++++++++++++ 4 files changed, 223 insertions(+), 7 deletions(-) diff --git a/ordvec-python/python/ordvec/__init__.py b/ordvec-python/python/ordvec/__init__.py index 6010dcb..8282058 100644 --- a/ordvec-python/python/ordvec/__init__.py +++ b/ordvec-python/python/ordvec/__init__.py @@ -24,13 +24,13 @@ caller-supplied data, validate or sandbox the path first, exactly as you would before a bare ``open()``. -Threading: ``search`` / ``search_asymmetric`` / ``add`` and the candidate -generators release the GIL during the heavy Rust scan, so other Python threads -run concurrently. The input arrays are *read in place* (not copied) for that -window — do not mutate an array from another thread while a call that received -it is in progress, or the scan races the write and may return inconsistent -results. This is the standard contract for GIL-releasing numeric extensions -(NumPy itself behaves this way). +Threading: ``search`` / ``search_asymmetric`` / ``add`` and the dense scoring / +candidate generator methods release the GIL during the heavy Rust scan, so other +Python threads run concurrently. The input arrays are *read in place* (not +copied) for that window — do not mutate an array from another thread while a +call that received it is in progress, or the scan races the write and may return +inconsistent results. This is the standard contract for GIL-releasing numeric +extensions (NumPy itself behaves this way). """ from ._ordvec import ( diff --git a/ordvec-python/src/lib.rs b/ordvec-python/src/lib.rs index a770cb8..4cde592 100644 --- a/ordvec-python/src/lib.rs +++ b/ordvec-python/src/lib.rs @@ -1199,6 +1199,58 @@ impl SignBitmap { .into_pyarray(py)) } + /// Dense full-corpus sign-agreement scores for a single query. Returns a + /// 1-D `uint32` array of length `len(index)`, aligned by document id. + fn score_all<'py>( + &self, + py: Python<'py>, + query: &Bound<'py, PyAny>, + ) -> PyResult>> { + let query = as_f32_1d(query, Some(self.inner.dim()))?; + let arr = query.as_array(); + let slice = arr.as_slice().ok_or_else(|| { + pyo3::exceptions::PyValueError::new_err( + "array must be C-contiguous; call np.ascontiguousarray() first", + ) + })?; + let scores = py.detach(|| self.inner.score_all(slice)); + Ok(scores.into_pyarray(py)) + } + + /// Batched dense full-corpus sign-agreement scores. Returns a 2-D `uint32` + /// array of shape `(batch, len(index))`, aligned by query row and document id. + fn score_all_batched<'py>( + &self, + py: Python<'py>, + queries: &Bound<'py, PyAny>, + ) -> PyResult>> { + let queries = as_f32_2d(queries, self.inner.dim())?; + let arr = queries.as_array(); + let batch = arr.nrows(); + let slice = arr.as_slice().ok_or_else(|| { + pyo3::exceptions::PyValueError::new_err( + "array must be C-contiguous; call np.ascontiguousarray() first", + ) + })?; + let n = self.inner.len(); + let qpv = self.inner.dim() / 64; + batch.checked_mul(n.max(qpv)).ok_or_else(|| { + pyo3::exceptions::PyValueError::new_err("batch * index size overflows usize") + })?; + let result = py.detach(|| self.inner.score_all_batched(slice)); + let total = batch.checked_mul(n).ok_or_else(|| { + pyo3::exceptions::PyValueError::new_err("result size (batch * n) overflows usize") + })?; + let mut flat: Vec = Vec::with_capacity(total); + for row in &result { + debug_assert_eq!(row.len(), n); + flat.extend_from_slice(row); + } + Ok(numpy::ndarray::Array2::from_shape_vec((batch, n), flat) + .expect("internal: batched dense score flatten shape invariant") + .into_pyarray(py)) + } + /// Build the query-side sign bitmap from an FP32 query, returned as a 1-D /// `uint64` array of `dim / 64` words (`bit j` set iff `q[j] > 0`). fn build_query_bitmap<'py>( diff --git a/ordvec-python/tests/test_sign_bitmap.py b/ordvec-python/tests/test_sign_bitmap.py index 69a0552..000378b 100644 --- a/ordvec-python/tests/test_sign_bitmap.py +++ b/ordvec-python/tests/test_sign_bitmap.py @@ -22,6 +22,10 @@ def unit_vectors(n: int, dim: int, seed: int = 0) -> np.ndarray: return v +def sign_agreement_reference(vectors: np.ndarray, query: np.ndarray) -> np.ndarray: + return ((vectors > 0.0) == (query > 0.0)).sum(axis=1).astype(np.uint32) + + def test_new_reports_dim_and_is_empty(): idx = SignBitmap(dim=128) assert idx.dim == 128 @@ -118,6 +122,46 @@ def test_empty_batch_against_empty_index_yields_zero_columns(): assert batched.shape == (0, 0) +def test_score_all_shape_dtype_and_reference_values(): + vectors = unit_vectors(25, 128, seed=21) + query = unit_vectors(1, 128, seed=22)[0] + idx = SignBitmap(dim=128) + idx.add(vectors) + + scores = idx.score_all(query) + + assert scores.shape == (25,) + assert scores.dtype == np.uint32 + np.testing.assert_array_equal(scores, sign_agreement_reference(vectors, query)) + + +def test_score_all_batched_shape_and_matches_single_query(): + vectors = unit_vectors(30, 128, seed=23) + queries = unit_vectors(4, 128, seed=24) + idx = SignBitmap(dim=128) + idx.add(vectors) + + batched = idx.score_all_batched(queries) + + assert batched.shape == (4, 30) + assert batched.dtype == np.uint32 + expected = np.vstack([idx.score_all(q) for q in queries]) + np.testing.assert_array_equal(batched, expected) + + +def test_score_all_empty_shapes(): + idx = SignBitmap(dim=64) + q = unit_vectors(1, 64, seed=25)[0] + assert idx.score_all(q).shape == (0,) + + queries = unit_vectors(3, 64, seed=26) + assert idx.score_all_batched(queries).shape == (3, 0) + + idx.add(unit_vectors(5, 64, seed=27)) + empty_q = np.empty((0, 64), dtype=np.float32) + assert idx.score_all_batched(empty_q).shape == (0, 5) + + def test_dim_not_multiple_of_64_rejected(): # The binding validates that dim is a positive multiple of 64 -> ValueError. with pytest.raises(ValueError, match="multiple of 64"): diff --git a/src/sign_bitmap.rs b/src/sign_bitmap.rs index 650d54b..62f796f 100644 --- a/src/sign_bitmap.rs +++ b/src/sign_bitmap.rs @@ -206,6 +206,67 @@ impl SignBitmap { .collect() } + /// Score every indexed document against one query and return dense + /// sign-agreement counts aligned by document id. + /// + /// `scores[di] = dim - popcount(q_bits ^ doc_bits[di])`, so higher is + /// better. This is a full-corpus scoring primitive, not a retrieval helper: + /// it performs no top-k selection and no sorting. + #[must_use = "this scans the corpus to score every document; dropping the result discards that work"] + pub fn score_all(&self, q: &[f32]) -> Vec { + let qb = self.build_query_bitmap(q); + let mut scores = vec![0u32; self.n_vectors]; // Hamming distance first. + sign_scan_collect( + &self.bitmaps, + self.n_vectors, + self.qwords_per_vec, + &qb, + &mut scores, + ); + let dim = u32::try_from(self.dim).expect("sign bitmap dim fits u32"); + scores.par_iter_mut().for_each(|h| *h = dim - *h); + scores + } + + /// Batched dense scoring. Returns one full-corpus sign-agreement row per + /// query, with columns aligned by document id and no sorting. + #[must_use = "this scans the corpus to score every document per query; dropping the result discards that work"] + pub fn score_all_batched(&self, queries: &[f32]) -> Vec> { + let dim = self.dim; + let batch = queries.len() / dim; + assert_eq!(queries.len(), batch * dim); + if batch == 0 { + return Vec::new(); + } + let n = self.n_vectors; + let qpv = self.qwords_per_vec; + + let q_batch_len = batch + .checked_mul(qpv) + .expect("batched query-bitmap buffer length (batch * qpv) overflows usize"); + let mut q_batch = vec![0u64; q_batch_len]; + for bi in 0..batch { + let qb = self.build_query_bitmap(&queries[bi * dim..(bi + 1) * dim]); + q_batch[bi * qpv..(bi + 1) * qpv].copy_from_slice(&qb); + } + + if n == 0 { + return vec![Vec::new(); batch]; + } + + let scores_len = batch + .checked_mul(n) + .expect("batched dense score buffer length (batch * n) overflows usize"); + let mut scores = vec![0u32; scores_len]; // Hamming distance first. + sign_scan_collect_batched(&self.bitmaps, n, qpv, &q_batch, batch, &mut scores); + + let dim = u32::try_from(dim).expect("sign bitmap dim fits u32"); + scores + .par_chunks_mut(n) + .for_each(|row| row.iter_mut().for_each(|h| *h = dim - *h)); + scores.chunks(n).map(|row| row.to_vec()).collect() + } + pub fn len(&self) -> usize { self.n_vectors } @@ -562,6 +623,65 @@ mod tests { } } + #[test] + fn score_all_returns_sign_agreement_by_doc_id() { + let n = 37; + let corpus = make_corpus(27, n); + let mut idx = SignBitmap::new(D); + idx.add(&corpus); + let mut rng = ChaCha8Rng::seed_from_u64(28); + let query: Vec = (0..D).map(|_| rng.random_range(-1.0..1.0)).collect(); + + let scores = idx.score_all(&query); + assert_eq!(scores.len(), n); + let qbm = idx.build_query_bitmap(&query); + for (di, &score) in scores.iter().enumerate() { + let off = di * idx.qwords_per_vec; + let dbm = &idx.bitmaps[off..off + idx.qwords_per_vec]; + assert_eq!( + score, + D as u32 - scalar_hamming(&qbm, dbm), + "score_all must return sign agreement for doc {di}", + ); + } + } + + #[test] + fn score_all_batched_matches_single_query() { + let n = 75; + let corpus = make_corpus(29, n); + let mut idx = SignBitmap::new(D); + idx.add(&corpus); + let mut rng = ChaCha8Rng::seed_from_u64(30); + let batch = 6; + let queries: Vec = (0..batch * D) + .map(|_| rng.random_range(-1.0..1.0)) + .collect(); + + let batched = idx.score_all_batched(&queries); + assert_eq!(batched.len(), batch); + for bi in 0..batch { + assert_eq!( + batched[bi], + idx.score_all(&queries[bi * D..(bi + 1) * D]), + "batched dense scoring diverged at batch idx {bi}", + ); + } + } + + #[test] + fn score_all_empty_shapes() { + let idx = SignBitmap::new(D); + let query = vec![1.0f32; D]; + assert!(idx.score_all(&query).is_empty()); + + let queries = vec![1.0f32; 2 * D]; + assert_eq!(idx.score_all_batched(&queries), vec![Vec::::new(); 2]); + + let empty_queries: Vec = Vec::new(); + assert!(idx.score_all_batched(&empty_queries).is_empty()); + } + #[test] fn large_dim_above_u16_max_roundtrips() { // Regression for the Codex stop-time finding: SignBitmap::new From aafc591d2c3997af21dee364f72bf35a383abadf Mon Sep 17 00:00:00 2001 From: Nelson Spence Date: Thu, 28 May 2026 11:19:29 -0500 Subject: [PATCH 2/2] fix sign bitmap batched dense scoring copies Signed-off-by: Nelson Spence --- ordvec-python/src/lib.rs | 12 ++------- src/sign_bitmap.rs | 56 ++++++++++++++++++++++++++++++++++++---- 2 files changed, 53 insertions(+), 15 deletions(-) diff --git a/ordvec-python/src/lib.rs b/ordvec-python/src/lib.rs index 4cde592..25a5557 100644 --- a/ordvec-python/src/lib.rs +++ b/ordvec-python/src/lib.rs @@ -1237,16 +1237,8 @@ impl SignBitmap { batch.checked_mul(n.max(qpv)).ok_or_else(|| { pyo3::exceptions::PyValueError::new_err("batch * index size overflows usize") })?; - let result = py.detach(|| self.inner.score_all_batched(slice)); - let total = batch.checked_mul(n).ok_or_else(|| { - pyo3::exceptions::PyValueError::new_err("result size (batch * n) overflows usize") - })?; - let mut flat: Vec = Vec::with_capacity(total); - for row in &result { - debug_assert_eq!(row.len(), n); - flat.extend_from_slice(row); - } - Ok(numpy::ndarray::Array2::from_shape_vec((batch, n), flat) + let scores = py.detach(|| self.inner.score_all_batched_flat(slice)); + Ok(numpy::ndarray::Array2::from_shape_vec((batch, n), scores) .expect("internal: batched dense score flatten shape invariant") .into_pyarray(py)) } diff --git a/src/sign_bitmap.rs b/src/sign_bitmap.rs index 62f796f..2f5435a 100644 --- a/src/sign_bitmap.rs +++ b/src/sign_bitmap.rs @@ -228,10 +228,11 @@ impl SignBitmap { scores } - /// Batched dense scoring. Returns one full-corpus sign-agreement row per - /// query, with columns aligned by document id and no sorting. + /// Batched dense scoring. Returns a flat row-major buffer of full-corpus + /// sign-agreement scores of length `batch * len(index)`, with columns + /// aligned by document id and no sorting. #[must_use = "this scans the corpus to score every document per query; dropping the result discards that work"] - pub fn score_all_batched(&self, queries: &[f32]) -> Vec> { + pub fn score_all_batched_flat(&self, queries: &[f32]) -> Vec { let dim = self.dim; let batch = queries.len() / dim; assert_eq!(queries.len(), batch * dim); @@ -251,7 +252,7 @@ impl SignBitmap { } if n == 0 { - return vec![Vec::new(); batch]; + return Vec::new(); } let scores_len = batch @@ -264,7 +265,22 @@ impl SignBitmap { scores .par_chunks_mut(n) .for_each(|row| row.iter_mut().for_each(|h| *h = dim - *h)); - scores.chunks(n).map(|row| row.to_vec()).collect() + scores + } + + /// Batched dense scoring. Returns one full-corpus sign-agreement row per + /// query, with columns aligned by document id and no sorting. + #[must_use = "this scans the corpus to score every document per query; dropping the result discards that work"] + pub fn score_all_batched(&self, queries: &[f32]) -> Vec> { + let dim = self.dim; + let batch = queries.len() / dim; + assert_eq!(queries.len(), batch * dim); + let n = self.n_vectors; + let flat = self.score_all_batched_flat(queries); + if n == 0 { + return vec![Vec::new(); batch]; + } + flat.chunks(n).map(|row| row.to_vec()).collect() } pub fn len(&self) -> usize { @@ -669,6 +685,29 @@ mod tests { } } + #[test] + fn score_all_batched_flat_matches_single_query() { + let n = 75; + let corpus = make_corpus(31, n); + let mut idx = SignBitmap::new(D); + idx.add(&corpus); + let mut rng = ChaCha8Rng::seed_from_u64(32); + let batch = 6; + let queries: Vec = (0..batch * D) + .map(|_| rng.random_range(-1.0..1.0)) + .collect(); + + let batched = idx.score_all_batched_flat(&queries); + assert_eq!(batched.len(), batch * n); + for bi in 0..batch { + assert_eq!( + &batched[bi * n..(bi + 1) * n], + idx.score_all(&queries[bi * D..(bi + 1) * D]), + "flat batched dense scoring diverged at batch idx {bi}", + ); + } + } + #[test] fn score_all_empty_shapes() { let idx = SignBitmap::new(D); @@ -676,9 +715,16 @@ mod tests { assert!(idx.score_all(&query).is_empty()); let queries = vec![1.0f32; 2 * D]; + assert!(idx.score_all_batched_flat(&queries).is_empty()); assert_eq!(idx.score_all_batched(&queries), vec![Vec::::new(); 2]); let empty_queries: Vec = Vec::new(); + assert!(idx.score_all_batched_flat(&empty_queries).is_empty()); + assert!(idx.score_all_batched(&empty_queries).is_empty()); + + let mut idx = SignBitmap::new(D); + idx.add(&make_corpus(33, 5)); + assert!(idx.score_all_batched_flat(&empty_queries).is_empty()); assert!(idx.score_all_batched(&empty_queries).is_empty()); }