difflib-fast 0.4.0__tar.gz → 0.5.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (23) hide show
  1. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/.github/workflows/CI.yml +3 -4
  2. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/Cargo.lock +1 -1
  3. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/Cargo.toml +1 -1
  4. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/PKG-INFO +11 -14
  5. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/README.md +10 -13
  6. difflib_fast-0.5.0/examples/simjoin_double.rs +110 -0
  7. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/src/gestalt.rs +1 -1
  8. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/src/simjoin.rs +590 -105
  9. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/src/simjoin_gpu.rs +8 -1
  10. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/.gitignore +0 -0
  11. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/benchmarks.md +0 -0
  12. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/examples/simjoin_bench.rs +0 -0
  13. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/examples/simjoin_gpu_bench.rs +0 -0
  14. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/examples/simjoin_pypi.rs +0 -0
  15. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/pyproject.toml +0 -0
  16. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/python/difflib_fast/__init__.py +0 -0
  17. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/python/difflib_fast/__init__.pyi +0 -0
  18. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/python/difflib_fast/py.typed +0 -0
  19. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/src/bin/bench.rs +0 -0
  20. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/src/gpu.rs +0 -0
  21. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/src/lib.rs +0 -0
  22. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/src/rationer.rs +0 -0
  23. {difflib_fast-0.4.0 → difflib_fast-0.5.0}/tests/approx_ro.rs +0 -0
@@ -1,8 +1,7 @@
1
1
  # Build + test, then on a version tag (`vX.Y.Z`) attach wheels + sdist to a GitHub Release.
2
- # crates.io is published manually (`cargo publish`). No PyPI — install the Python package from a
3
- # release wheel or via `pip install git+https://github.com/prostomarkeloff/difflib-fast`.
4
- #
5
- # No secrets needed — the GitHub Release uses the built-in GITHUB_TOKEN.
2
+ # crates.io is published manually (`cargo publish`), and so is PyPI: download the release's
3
+ # wheels + sdist and `uv publish` them. Nothing here needs a secret — the GitHub Release uses the
4
+ # built-in GITHUB_TOKEN.
6
5
  #
7
6
  # abi3 (pyo3 abi3-py39) → one wheel per platform/arch works on CPython 3.9+, so no per-Python matrix.
8
7
  name: CI
@@ -78,7 +78,7 @@ checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28"
78
78
 
79
79
  [[package]]
80
80
  name = "difflib-fast"
81
- version = "0.4.0"
81
+ version = "0.5.0"
82
82
  dependencies = [
83
83
  "metal",
84
84
  "mimalloc",
@@ -1,6 +1,6 @@
1
1
  [package]
2
2
  name = "difflib-fast"
3
- version = "0.4.0"
3
+ version = "0.5.0"
4
4
  description = "Fast, byte-for-byte exact difflib Ratcliff–Obershelp (gestalt) similarity ratio + single-linkage clustering (suffix automaton), plus an exact all-pairs weighted-cosine similarity join (L2AP, CPU+GPU)."
5
5
  keywords = ["difflib", "similarity", "ratcliff-obershelp", "suffix-automaton", "fuzzy"]
6
6
  categories = ["algorithms", "text-processing"]
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: difflib-fast
3
- Version: 0.4.0
3
+ Version: 0.5.0
4
4
  Classifier: Programming Language :: Rust
5
5
  Classifier: Programming Language :: Python :: Implementation :: CPython
6
6
  Classifier: Topic :: Text Processing :: General
@@ -58,7 +58,7 @@ assert_eq!(ratio("the quick brown fox", "the quick brown dog"), 0.89473684210526
58
58
  ```
59
59
 
60
60
  ```toml
61
- difflib-fast = "0.3"
61
+ difflib-fast = "0.5"
62
62
  ```
63
63
 
64
64
  ---
@@ -269,27 +269,24 @@ Swap `difflib.SequenceMatcher(...).ratio()` for `difflib_fast.ratio(...)` and th
269
269
  back — hundreds of times faster per call, and **thousands of times** faster when you score a whole
270
270
  corpus. That's the entire change: `import difflib_fast`.
271
271
 
272
- **No PyPI** — `pip install difflib-fast` won't work. Two ways to install:
273
-
274
- **1. Prebuilt wheel — no Rust toolchain needed.** The wheels are `cp39-abi3`, so one wheel per platform
275
- works on **every CPython ≥ 3.9, including 3.14**. GitHub Releases isn't a package index, so pip can't
276
- pick the wheel for you — grab the one for your platform from the
277
- [Releases page](https://github.com/prostomarkeloff/difflib-fast/releases/latest) and install it by URL:
278
-
279
272
  ```bash
280
- # macOS Apple Silicon — swap the filename for your platform (see below):
281
- pip install https://github.com/prostomarkeloff/difflib-fast/releases/download/v0.3.5/difflib_fast-0.3.5-cp39-abi3-macosx_11_0_arm64.whl
273
+ pip install difflib-fast
282
274
  ```
283
275
 
284
- | platform | wheel suffix |
276
+ The wheels on [PyPI](https://pypi.org/project/difflib-fast/) are `cp39-abi3`, so one wheel per
277
+ platform works on **every CPython ≥ 3.9, including 3.14** — no Rust toolchain needed:
278
+
279
+ | platform | wheel |
285
280
  |---|---|
286
- | macOS Apple Silicon | `…-macosx_11_0_arm64.whl` |
281
+ | macOS Apple Silicon | `…-macosx_11_0_arm64.whl` (ships the Metal GPU path) |
287
282
  | macOS Intel | `…-macosx_10_12_x86_64.whl` |
288
283
  | Linux x86_64 | `…-manylinux_2_17_x86_64.manylinux2014_x86_64.whl` |
289
284
  | Linux aarch64 | `…-manylinux_2_17_aarch64.manylinux2014_aarch64.whl` |
290
285
  | Windows x64 | `…-win_amd64.whl` |
291
286
 
292
- **2. From source — needs a Rust toolchain** (pip drives maturin automatically, no manual build):
287
+ The same wheels are attached to every
288
+ [GitHub release](https://github.com/prostomarkeloff/difflib-fast/releases/latest). Any other
289
+ platform builds from source — pip drives maturin, so it only needs a Rust toolchain:
293
290
 
294
291
  ```bash
295
292
  pip install git+https://github.com/prostomarkeloff/difflib-fast
@@ -42,7 +42,7 @@ assert_eq!(ratio("the quick brown fox", "the quick brown dog"), 0.89473684210526
42
42
  ```
43
43
 
44
44
  ```toml
45
- difflib-fast = "0.3"
45
+ difflib-fast = "0.5"
46
46
  ```
47
47
 
48
48
  ---
@@ -253,27 +253,24 @@ Swap `difflib.SequenceMatcher(...).ratio()` for `difflib_fast.ratio(...)` and th
253
253
  back — hundreds of times faster per call, and **thousands of times** faster when you score a whole
254
254
  corpus. That's the entire change: `import difflib_fast`.
255
255
 
256
- **No PyPI** — `pip install difflib-fast` won't work. Two ways to install:
257
-
258
- **1. Prebuilt wheel — no Rust toolchain needed.** The wheels are `cp39-abi3`, so one wheel per platform
259
- works on **every CPython ≥ 3.9, including 3.14**. GitHub Releases isn't a package index, so pip can't
260
- pick the wheel for you — grab the one for your platform from the
261
- [Releases page](https://github.com/prostomarkeloff/difflib-fast/releases/latest) and install it by URL:
262
-
263
256
  ```bash
264
- # macOS Apple Silicon — swap the filename for your platform (see below):
265
- pip install https://github.com/prostomarkeloff/difflib-fast/releases/download/v0.3.5/difflib_fast-0.3.5-cp39-abi3-macosx_11_0_arm64.whl
257
+ pip install difflib-fast
266
258
  ```
267
259
 
268
- | platform | wheel suffix |
260
+ The wheels on [PyPI](https://pypi.org/project/difflib-fast/) are `cp39-abi3`, so one wheel per
261
+ platform works on **every CPython ≥ 3.9, including 3.14** — no Rust toolchain needed:
262
+
263
+ | platform | wheel |
269
264
  |---|---|
270
- | macOS Apple Silicon | `…-macosx_11_0_arm64.whl` |
265
+ | macOS Apple Silicon | `…-macosx_11_0_arm64.whl` (ships the Metal GPU path) |
271
266
  | macOS Intel | `…-macosx_10_12_x86_64.whl` |
272
267
  | Linux x86_64 | `…-manylinux_2_17_x86_64.manylinux2014_x86_64.whl` |
273
268
  | Linux aarch64 | `…-manylinux_2_17_aarch64.manylinux2014_aarch64.whl` |
274
269
  | Windows x64 | `…-win_amd64.whl` |
275
270
 
276
- **2. From source — needs a Rust toolchain** (pip drives maturin automatically, no manual build):
271
+ The same wheels are attached to every
272
+ [GitHub release](https://github.com/prostomarkeloff/difflib-fast/releases/latest). Any other
273
+ platform builds from source — pip drives maturin, so it only needs a Rust toolchain:
277
274
 
278
275
  ```bash
279
276
  pip install git+https://github.com/prostomarkeloff/difflib-fast
@@ -0,0 +1,110 @@
1
+ //! Bench for the fused two-corpus join [`difflib_fast::simjoin::cosine_join_double`] — the typical
2
+ //! "two views of the same objects" shape (e.g. line-vectors and part-vectors of the same functions:
3
+ //! two corpora with the same row count). Times one fused join against two separate
4
+ //! [`cosine_join`]s and checks each half matches the separate result.
5
+ //!
6
+ //! `cargo run --release --example simjoin_double -- <a.bin> [b.bin] [t] [reps]`
7
+ //! defaults: b.bin = a.bin, t=0.8, reps=3. `SJ_NSUB=<n>` truncates both corpora; thread count via
8
+ //! `RAYON_NUM_THREADS`.
9
+
10
+ #![allow(
11
+ clippy::cast_possible_truncation,
12
+ clippy::cast_precision_loss,
13
+ clippy::many_single_char_names,
14
+ clippy::doc_markdown
15
+ )]
16
+
17
+ use std::time::Instant;
18
+
19
+ use difflib_fast::simjoin::{cosine_join, cosine_join_double, Corpus};
20
+
21
+ fn le_u32(b: &[u8], p: usize) -> u32 {
22
+ u32::from_le_bytes(b[p..p + 4].try_into().unwrap())
23
+ }
24
+
25
+ fn arg<T: std::str::FromStr>(i: usize, def: T) -> T {
26
+ std::env::args().nth(i).and_then(|s| s.parse().ok()).unwrap_or(def)
27
+ }
28
+
29
+ /// Parse the `SIMJOIN1` binary into raw `(dim, weight)` rows (same format as `examples/simjoin_pypi`).
30
+ fn load(path: &str) -> Vec<Vec<(u32, f64)>> {
31
+ let b = std::fs::read(path).expect("read corpus");
32
+ assert!(b.len() >= 16, "corpus too small");
33
+ let magic = u64::from_le_bytes(b[0..8].try_into().unwrap());
34
+ assert_eq!(magic, 0x5349_4D4A_4F49_4E31, "bad magic");
35
+ let n = le_u32(&b, 8) as usize;
36
+ let mut p = 16usize;
37
+ let mut rows = Vec::with_capacity(n);
38
+ for _ in 0..n {
39
+ let nnz = le_u32(&b, p) as usize;
40
+ p += 4;
41
+ let mut row = Vec::with_capacity(nnz);
42
+ for _ in 0..nnz {
43
+ let d = le_u32(&b, p);
44
+ p += 4;
45
+ let w = f64::from(f32::from_le_bytes(b[p..p + 4].try_into().unwrap()));
46
+ p += 4;
47
+ row.push((d, w));
48
+ }
49
+ rows.push(row);
50
+ }
51
+ rows
52
+ }
53
+
54
+ /// Canonical order for the parity check (join output order is arbitrary).
55
+ fn sort_pairs(mut v: Vec<(usize, usize, f64)>) -> Vec<(usize, usize, u64)> {
56
+ v.sort_by_key(|a| (a.0, a.1));
57
+ v.into_iter().map(|(a, b, s)| (a, b, s.to_bits())).collect()
58
+ }
59
+
60
+ fn main() {
61
+ let path_a: String = arg(1, "perf-local/pypi-type3.simjoin.bin".to_string());
62
+ let path_b: String = arg(2, path_a.clone());
63
+ let t: f64 = arg(3, 0.8);
64
+ let reps: usize = arg(4, 3);
65
+
66
+ let mut rows_a = load(&path_a);
67
+ let mut rows_b = load(&path_b);
68
+ if let Ok(nsub) = std::env::var("SJ_NSUB") {
69
+ let nsub: usize = nsub.parse().expect("SJ_NSUB");
70
+ rows_a.truncate(nsub);
71
+ rows_b.truncate(nsub);
72
+ }
73
+ assert_eq!(rows_a.len(), rows_b.len(), "corpora must share the row count");
74
+ let n = rows_a.len();
75
+
76
+ let corpus_a = Corpus::from_rows(&rows_a);
77
+ let corpus_b = Corpus::from_rows(&rows_b);
78
+
79
+ let mut fused: Vec<f64> = Vec::with_capacity(reps);
80
+ let mut separate: Vec<f64> = Vec::with_capacity(reps);
81
+ let (mut pa, mut pb) = (0usize, 0usize);
82
+ for r in 0..reps {
83
+ let (ga, gb) = {
84
+ let t0 = Instant::now();
85
+ let r = cosine_join_double(&corpus_a, &corpus_b, t);
86
+ fused.push(t0.elapsed().as_secs_f64() * 1000.0);
87
+ r
88
+ };
89
+ let (sa, sb) = {
90
+ let t0 = Instant::now();
91
+ let r = (cosine_join(&corpus_a, t), cosine_join(&corpus_b, t));
92
+ separate.push(t0.elapsed().as_secs_f64() * 1000.0);
93
+ r
94
+ };
95
+ if r == 0 {
96
+ pa = ga.len();
97
+ pb = gb.len();
98
+ // The fused halves must be bit-identical to the separate joins (real-data smoke).
99
+ assert_eq!(sort_pairs(ga), sort_pairs(sa), "half a differs from the separate join");
100
+ assert_eq!(sort_pairs(gb), sort_pairs(sb), "half b differs from the separate join");
101
+ }
102
+ }
103
+ fused.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
104
+ separate.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
105
+ let (f, s) = (fused[fused.len() / 2], separate[separate.len() / 2]);
106
+ eprintln!(
107
+ "double n={n} t={t} | pairs=({pa},{pb}) | fused={f:.1}ms separate={s:.1}ms speedup={:.2}x",
108
+ s / f
109
+ );
110
+ }
@@ -894,7 +894,7 @@ impl Builder {
894
894
  fn sort_big(&mut self, out: &mut Sam, nstates: usize, blen: usize) {
895
895
  // sorted copies for the big endpos sets (a handful of states: short, frequent substrings);
896
896
  // positions are below `blen`, so a radix sort needs one pass per byte of that
897
- let passes = ((usize::BITS - blen.leading_zeros()) as usize).div_ceil(8).max(1);
897
+ let passes = (blen.bit_width() as usize).div_ceil(8).max(1);
898
898
  let scratch = &mut self.child_arr;
899
899
  for s in 0..nstates {
900
900
  let cnt = out.chain_slot[s][7] as usize;
@@ -43,6 +43,9 @@ use rayon::prelude::*;
43
43
 
44
44
  use crate::Concurrency;
45
45
 
46
+ /// One join's result: `(row_a, row_b, cosine)` triples with `row_a < row_b`, arbitrary order.
47
+ pub type JoinPairs = Vec<(usize, usize, f64)>;
48
+
46
49
  /// A corpus of L2-normalised sparse non-negative vectors in CSR form, with dimensions relabelled to
47
50
  /// a global rank (decreasing max weight). Built once; joinable at any threshold.
48
51
  pub struct Corpus {
@@ -58,6 +61,41 @@ pub struct Corpus {
58
61
  maxw: Vec<f64>,
59
62
  }
60
63
 
64
+ /// Merge duplicate dims of one row (weights summed in input order within each dim group, dims in
65
+ /// ascending dim order — deterministic, unlike the per-call-random `HashMap` iteration order of the
66
+ /// previous implementation) and L2-normalise. A row whose norm is 0 (all weights 0) becomes empty.
67
+ #[must_use]
68
+ fn norm_row(r: &[(u32, f64)]) -> Vec<(u32, f64)> {
69
+ let mut v = r.to_vec();
70
+ v.sort_by_key(|&(d, _)| d); // stable: preserves input order within each duplicate group
71
+ let mut out: Vec<(u32, f64)> = Vec::with_capacity(v.len());
72
+ let mut sq = 0.0;
73
+ let mut iter = v.into_iter();
74
+ if let Some((mut d, mut w)) = iter.next() {
75
+ for (d2, w2) in iter {
76
+ if d2 == d {
77
+ w += w2;
78
+ } else {
79
+ sq += w * w;
80
+ out.push((d, w));
81
+ d = d2;
82
+ w = w2;
83
+ }
84
+ }
85
+ sq += w * w;
86
+ out.push((d, w));
87
+ }
88
+ let norm = sq.sqrt();
89
+ if norm > 0.0 {
90
+ for (_, w) in &mut out {
91
+ *w /= norm;
92
+ }
93
+ out
94
+ } else {
95
+ Vec::new()
96
+ }
97
+ }
98
+
61
99
  impl Corpus {
62
100
  /// Number of vectors.
63
101
  #[must_use]
@@ -131,61 +169,131 @@ impl Corpus {
131
169
  /// Build a corpus from raw `(dim, weight)` rows. Duplicate dims within a row are summed; each row
132
170
  /// is L2-normalised; dims are relabelled to a global rank by decreasing max weight. Weights are
133
171
  /// expected non-negative (the prefix-filter bound requires it — IDF weights satisfy this).
172
+ ///
173
+ /// The build is rayon-parallel (per-row normalise, global max-weight reduction, per-row relabel
174
+ /// into the CSR) and **deterministic**: rows are merged in dim order, the rank is a function of
175
+ /// the exact per-dim max weight (ties by dim id), and parallel reductions only combine
176
+ /// order-independent operations — the resulting corpus is a pure function of `rows`.
177
+ ///
178
+ /// Dims that occur in exactly one row (`df = 1`) are dropped from the rank and the rows **after**
179
+ /// normalisation: no pair shares them, so they contribute exactly `0.0` to every pairwise dot and
180
+ /// their postings are never walked — the join output is bit-identical with or without them, while
181
+ /// the index and the rows get smaller and the pruning bounds tighter. Normalising over the full
182
+ /// row first is what keeps the surviving weights — and every emitted score — unchanged.
134
183
  #[must_use]
135
184
  pub fn from_rows(rows: &[Vec<(u32, f64)>]) -> Corpus {
185
+ Self::from_rows_with(rows, true)
186
+ }
187
+
188
+ /// Build a corpus from normalized rows with optional **hapax stripping**. Hapax dims (present in
189
+ /// exactly one row) are dropped from the rank and the rows **after** L2 normalisation — no pair
190
+ /// shares them, so they contribute exactly `0.0` to every pairwise dot, their postings are never
191
+ /// walked (probes visit only postings of their own dims), and dropping their mass only tightens
192
+ /// the Cauchy–Schwarz bounds (which are filters). The join output is therefore bit-identical
193
+ /// with stripping on or off; the strip just shrinks the index/rows and prunes harder. Normalising
194
+ /// over the full row first keeps the surviving weights — and every emitted score — unchanged
195
+ /// (removing the mass *before* normalisation would rescale the weights and change the indexing).
196
+ /// `strip_hapax = false` is the un-stripped reference build, used by tests to pin the equality.
197
+ fn from_rows_with(rows: &[Vec<(u32, f64)>], strip_hapax: bool) -> Corpus {
136
198
  let n = rows.len();
137
- // 1. Merge duplicate dims + L2-normalise each row (kept as (orig_dim, weight)).
138
- let normed: Vec<Vec<(u32, f64)>> = rows
139
- .iter()
140
- .map(|r| {
141
- let mut m: HashMap<u32, f64> = HashMap::new();
142
- for &(d, w) in r {
143
- *m.entry(d).or_insert(0.0) += w;
199
+ // 1. Merge duplicate dims + L2-normalise each row (kept as (orig_dim, weight)), in parallel.
200
+ let normed: Vec<Vec<(u32, f64)>> = rows.par_iter().map(|r| norm_row(r)).collect();
201
+ // 2. Document frequency per original dim (over the normalised rows) — a parallel
202
+ // fold+reduce of exact counts, so the chunking cannot change the result.
203
+ let df: HashMap<u32, u32> = normed
204
+ .par_iter()
205
+ .fold(HashMap::new, |mut m, v| {
206
+ for &(d, _) in v {
207
+ *m.entry(d).or_insert(0) += 1;
144
208
  }
145
- let norm = m.values().map(|w| w * w).sum::<f64>().sqrt();
146
- if norm > 0.0 {
147
- m.into_iter().map(|(d, w)| (d, w / norm)).collect()
148
- } else {
149
- Vec::new()
209
+ m
210
+ })
211
+ .reduce(HashMap::new, |mut a, mut b| {
212
+ if a.len() < b.len() {
213
+ std::mem::swap(&mut a, &mut b);
214
+ }
215
+ for (d, c) in b {
216
+ *a.entry(d).or_insert(0) += c;
150
217
  }
218
+ a
219
+ });
220
+ // A dim earns a rank only if some pair could share it. `strip_hapax = false` keeps every
221
+ // dim (test-only reference).
222
+ let shared = |d: u32| !strip_hapax || df.get(&d).is_some_and(|&c| c >= 2);
223
+ // 3. Max normalised weight per original dim — a parallel fold+reduce. `max` is exact and
224
+ // order-independent, so the chunking cannot change any value.
225
+ let maxw_orig: HashMap<u32, f64> = normed
226
+ .par_iter()
227
+ .fold(HashMap::new, |mut m, v| {
228
+ for &(d, w) in v {
229
+ if shared(d) {
230
+ let e = m.entry(d).or_insert(0.0);
231
+ if w > *e {
232
+ *e = w;
233
+ }
234
+ }
235
+ }
236
+ m
151
237
  })
152
- .collect();
153
- // 2. Max normalised weight per original dim.
154
- let mut maxw_orig: HashMap<u32, f64> = HashMap::new();
155
- for v in &normed {
156
- for &(d, w) in v {
157
- let e = maxw_orig.entry(d).or_insert(0.0);
158
- if w > *e {
159
- *e = w;
238
+ .reduce(HashMap::new, |mut a, mut b| {
239
+ if a.len() < b.len() {
240
+ std::mem::swap(&mut a, &mut b);
160
241
  }
161
- }
162
- }
163
- // 3. Rank dims by (max weight ASC, dim asc) → dense global order. Ascending so the common,
242
+ for (d, w) in b {
243
+ let e = a.entry(d).or_insert(0.0);
244
+ if w > *e {
245
+ *e = w;
246
+ }
247
+ }
248
+ a
249
+ });
250
+ // 4. Rank dims by (max weight ASC, dim asc) → dense global order. Ascending so the common,
164
251
  // low-weight dims land at the FRONT: they fill the un-indexed prefix (their tiny
165
252
  // `w·maxw` keeps the cumulative bound under `t` for many of them), and only the rare,
166
253
  // high-weight tail is indexed — short postings. (Reversing this indexes the common dims
167
254
  // and their huge postings, which is correct but orders of magnitude slower.)
255
+ // The comparator is a total order on distinct dims, so the unstable parallel sort yields
256
+ // the same unique sorted sequence the serial stable sort it replaces would.
168
257
  let mut dims_sorted: Vec<(u32, f64)> = maxw_orig.into_iter().collect();
169
- dims_sorted.sort_by(|a, b| {
258
+ dims_sorted.par_sort_unstable_by(|a, b| {
170
259
  a.1.partial_cmp(&b.1).unwrap_or(Ordering::Equal).then(a.0.cmp(&b.0))
171
260
  });
172
261
  let rank: HashMap<u32, u32> =
173
262
  dims_sorted.iter().enumerate().map(|(i, &(d, _))| (d, i as u32)).collect();
174
263
  let ndims = dims_sorted.len();
175
264
  let maxw: Vec<f64> = dims_sorted.iter().map(|&(_, w)| w).collect();
176
- // 4. CSR with relabelled dims, ascending within each row.
265
+ // 5. CSR with relabelled dims, ascending within each row. Relabel + sort each row in
266
+ // parallel (rank lookups are the cost), then a serial flatten into the flat CSR arrays —
267
+ // a pure memcpy; a parallel fill into disjoint slices would need unsafe borrows for no
268
+ // measurable win. Hapax dims have no rank, so the `rank.get` filter drops them.
177
269
  let mut indptr = Vec::with_capacity(n + 1);
178
270
  indptr.push(0);
179
- let mut dims = Vec::new();
180
- let mut wts = Vec::new();
271
+ let mut off = 0usize;
181
272
  for v in &normed {
182
- let mut rv: Vec<(u32, f64)> = v.iter().map(|&(d, w)| (rank[&d], w)).collect();
183
- rv.sort_unstable_by_key(|&(d, _)| d);
273
+ let kept = if strip_hapax {
274
+ v.iter().filter(|&&(d, _)| rank.contains_key(&d)).count()
275
+ } else {
276
+ v.len()
277
+ };
278
+ off += kept;
279
+ indptr.push(off);
280
+ }
281
+ let relabelled: Vec<Vec<(u32, f64)>> = normed
282
+ .par_iter()
283
+ .map(|v| {
284
+ let mut rv: Vec<(u32, f64)> =
285
+ v.iter().filter_map(|&(d, w)| rank.get(&d).map(|&r| (r, w))).collect();
286
+ rv.sort_unstable_by_key(|&(d, _)| d);
287
+ rv
288
+ })
289
+ .collect();
290
+ let mut dims = Vec::with_capacity(off);
291
+ let mut wts = Vec::with_capacity(off);
292
+ for rv in relabelled {
184
293
  for (d, w) in rv {
185
294
  dims.push(d);
186
295
  wts.push(w);
187
296
  }
188
- indptr.push(dims.len());
189
297
  }
190
298
  Corpus { n, ndims, indptr, dims, wts, maxw }
191
299
  }
@@ -248,8 +356,53 @@ struct Scratch {
248
356
  touched: Vec<u32>,
249
357
  /// Probe prefix L2 norms for this probe: `xpn[k] = ‖wi[..k]‖₂`, length `nnz+1`.
250
358
  xpn: Vec<f64>,
359
+ /// Sparse-mode accumulator: `(candidate, partial dot)` pairs sorted by candidate id — the
360
+ /// touched list and the accumulator merged (classic L2AP). Used instead of `acc`/`touched` when
361
+ /// the corpus's average per-probe candidate count is low; the probe falls back to the dense pair
362
+ /// if the list overflows [`LIST_ACC_MAX`].
363
+ list: Vec<(u32, f64)>,
364
+ }
365
+
366
+ /// Accumulator strategy for the probe phase, chosen per join from the corpus's shape (and forced
367
+ /// per-mode in tests — the two must agree bit-for-bit).
368
+ #[derive(Clone, Copy, PartialEq, Eq, Debug)]
369
+ enum AccMode {
370
+ /// Dense `acc[n]` array: O(1) scatter per posting entry. Wins when probes touch many candidates.
371
+ Dense,
372
+ /// Sorted `(y, acc)` list: few candidates per probe ⇒ the list stays tiny and L1-resident, and
373
+ /// the n-wide dense array's footprint never competes with the postings for L2.
374
+ Sparse,
251
375
  }
252
376
 
377
+ /// Average per-probe posting touches at/below which [`cosine_join`] picks the [`AccMode::Sparse`]
378
+ /// accumulator. The touch volume is estimated from the built index as `Σ len²/2` (posting `d` is
379
+ /// walked by about half the probes, each walk visiting `len` entries) divided by `n` — the per-probe
380
+ /// mean walk volume, a corpus-shape signal. Per-probe outliers are caught by [`LIST_ACC_MAX`]'s
381
+ /// dense fallback.
382
+ const LIST_ACC_MAX_TOUCHES: usize = 32;
383
+
384
+ /// Pick the accumulator mode from a built inverted index: sparse list when the average probe walks
385
+ /// few posting entries, dense scatter otherwise.
386
+ #[must_use]
387
+ fn acc_mode_for(index: &[Vec<(u32, f64)>], n: usize) -> AccMode {
388
+ let mut walks = 0u64;
389
+ for p in index {
390
+ let len = p.len() as u64;
391
+ walks += len * len;
392
+ }
393
+ if walks / 2 / (n.max(1) as u64) < LIST_ACC_MAX_TOUCHES as u64 {
394
+ AccMode::Sparse
395
+ } else {
396
+ AccMode::Dense
397
+ }
398
+ }
399
+
400
+ /// Hard cap on one probe's sparse list: beyond this the per-insert memmove dominates the O(1)
401
+ /// dense scatter it replaces, so [`accumulate_list`] bails out and the probe is redone densely
402
+ /// (re-walking its postings — rare by construction; the additions repeat in the same order, so the
403
+ /// values are unchanged).
404
+ const LIST_ACC_MAX: usize = 512;
405
+
253
406
  /// Exact all-pairs cosine join via inverted index + **L2AP** prefix filtering and accumulation-time
254
407
  /// pruning. Returns `(j, i, cos)` with `j < i` for every pair with `cos ≥ t`. Bit-identical pair set
255
408
  /// to [`cosine_join_bruteforce`].
@@ -276,19 +429,32 @@ pub fn cosine_join(c: &Corpus, t: f64) -> Vec<(usize, usize, f64)> {
276
429
  let (di, wi) = c.row(i);
277
430
  index_suffix(c, i, (di, wi), t, &mut index, &mut cached);
278
431
  }
279
- // Probe vectors in parallel; each worker keeps one reusable `Scratch` (an `n`-wide accumulator).
280
- // `with_min_len` batches many probes per rayon task so the per-probe work (tiny) isn't dwarfed by
281
- // task-splitting / scheduling overhead (`swtch_pri` in the profile).
432
+ // Accumulator strategy from the corpus shape: few candidates per probe on average (sparse /
433
+ // duplicate-heavy corpora) ⇒ the sorted list wins; many (dense rows) ⇒ the O(1) dense scatter.
434
+ let mode = acc_mode_for(&index, n);
435
+ cosine_join_mode(c, t, &index, &cached, mode)
436
+ }
437
+
438
+ /// The join body under an explicit [`AccMode`]: probe vectors in parallel; each worker keeps one
439
+ /// reusable `Scratch`. `with_min_len` batches many probes per rayon task so the per-probe work
440
+ /// (tiny) isn't dwarfed by task-splitting / scheduling overhead (`swtch_pri` in the profile).
441
+ #[must_use]
442
+ fn cosine_join_mode(
443
+ c: &Corpus,
444
+ t: f64,
445
+ index: &[Vec<(u32, f64)>],
446
+ cached: &Cached,
447
+ mode: AccMode,
448
+ ) -> Vec<(usize, usize, f64)> {
449
+ let n = c.n;
282
450
  (0..n)
283
451
  .into_par_iter()
284
452
  .with_min_len(256)
285
453
  .map_init(
286
- || Scratch { acc: vec![-1.0; n], touched: Vec::new(), xpn: Vec::new() },
454
+ || Scratch { acc: vec![-1.0; n], touched: Vec::new(), xpn: Vec::new(), list: Vec::new() },
287
455
  |scratch, i| {
288
- let (di, wi) = c.row(i);
289
- accumulate(&index, (di, wi), i as u32, scratch);
290
456
  let mut out = Vec::new();
291
- verify_pruned(c, i, t, scratch, &cached, &mut out);
457
+ probe_half(c, i, t, index, cached, mode, scratch, &mut out);
292
458
  out
293
459
  },
294
460
  )
@@ -296,6 +462,118 @@ pub fn cosine_join(c: &Corpus, t: f64) -> Vec<(usize, usize, f64)> {
296
462
  .collect()
297
463
  }
298
464
 
465
+ /// One probe's accumulate + prune + score against `c`'s index, under `mode`, appending to `out`.
466
+ /// Shared by the single join and the fused two-corpus join so the per-half work is structurally
467
+ /// identical (the fused join relies on that for its bit-identical halves).
468
+ #[allow(clippy::too_many_arguments, clippy::inline_always)]
469
+ #[inline(always)] // per-probe dispatch only — must fold into the rayon closure, not add a call
470
+ fn probe_half(
471
+ c: &Corpus,
472
+ i: usize,
473
+ t: f64,
474
+ index: &[Vec<(u32, f64)>],
475
+ cached: &Cached,
476
+ mode: AccMode,
477
+ s: &mut Scratch,
478
+ out: &mut Vec<(usize, usize, f64)>,
479
+ ) {
480
+ let (di, wi) = c.row(i);
481
+ match mode {
482
+ AccMode::Dense => {
483
+ accumulate(index, (di, wi), i as u32, s);
484
+ verify_pruned(c, i, t, s, cached, out);
485
+ }
486
+ AccMode::Sparse => {
487
+ if accumulate_list(index, (di, wi), i as u32, s) {
488
+ verify_list(c, i, t, s, cached, out);
489
+ } else {
490
+ // The sparse list overflowed — redo this probe with the dense pass (identical
491
+ // additions, identical order, so the values are unchanged).
492
+ accumulate(index, (di, wi), i as u32, s);
493
+ verify_pruned(c, i, t, s, cached, out);
494
+ }
495
+ }
496
+ }
497
+ }
498
+
499
+ /// [`cosine_join`] under an explicitly forced [`AccMode`]. Test-only: the fuzz gates must cover both
500
+ /// modes, and the shape-based pick exercises only one of them per corpus.
501
+ #[cfg(test)]
502
+ #[must_use]
503
+ fn cosine_join_for_mode(c: &Corpus, t: f64, mode: AccMode) -> Vec<(usize, usize, f64)> {
504
+ let n = c.n;
505
+ let mut index: Vec<Vec<(u32, f64)>> = vec![Vec::new(); c.ndims];
506
+ let mut cached = Cached { bound: vec![Bound { pnorm: 0.0, split: u32::MAX }; n] };
507
+ for i in 0..n {
508
+ let (di, wi) = c.row(i);
509
+ index_suffix(c, i, (di, wi), t, &mut index, &mut cached);
510
+ }
511
+ cosine_join_mode(c, t, &index, &cached, mode)
512
+ }
513
+
514
+ /// Fused two-corpus cosine join: one probe pass over `a` and `b` — two corpora with the **same row
515
+ /// count**, the typical "two views of the same objects" shape (e.g. line-vectors and part-vectors of
516
+ /// the same functions). Returns `(pairs_a, pairs_b)`; each half is **bit-identical** to a separate
517
+ /// [`cosine_join`] on its corpus — the per-half phases ([`accumulate`]/[`accumulate_list`],
518
+ /// [`verify_pruned`]/[`verify_list`]) run on the same inputs in the same order, only the rayon
519
+ /// schedule and the per-worker [`Scratch`] are shared (one `acc` buffer serves both halves: half A's
520
+ /// entries are reset by its verify pass before half B accumulates).
521
+ ///
522
+ /// Wins over two separate joins by paying the per-probe fixed costs once: one parallel schedule, one
523
+ /// `Scratch` per worker, and the probe iteration bookkeeping is shared. Pair order within each half
524
+ /// is arbitrary (as in [`cosine_join`]). The accumulator mode is chosen per corpus from its shape
525
+ /// (see [`cosine_join`]), so the halves may use different modes.
526
+ ///
527
+ /// # Panics
528
+ ///
529
+ /// Panics if `a.len() != b.len()`.
530
+ #[must_use]
531
+ pub fn cosine_join_double(a: &Corpus, b: &Corpus, t: f64) -> (JoinPairs, JoinPairs) {
532
+ assert_eq!(a.n, b.n, "cosine_join_double: corpora must share the row count");
533
+ let n = a.n;
534
+ let mut index_a: Vec<Vec<(u32, f64)>> = vec![Vec::new(); a.ndims];
535
+ let mut index_b: Vec<Vec<(u32, f64)>> = vec![Vec::new(); b.ndims];
536
+ let mut cached_a = Cached { bound: vec![Bound { pnorm: 0.0, split: u32::MAX }; n] };
537
+ let mut cached_b = Cached { bound: vec![Bound { pnorm: 0.0, split: u32::MAX }; n] };
538
+ for i in 0..n {
539
+ let (di, wi) = a.row(i);
540
+ index_suffix(a, i, (di, wi), t, &mut index_a, &mut cached_a);
541
+ let (di, wi) = b.row(i);
542
+ index_suffix(b, i, (di, wi), t, &mut index_b, &mut cached_b);
543
+ }
544
+ let mode_a = acc_mode_for(&index_a, n);
545
+ let mode_b = acc_mode_for(&index_b, n);
546
+ (0..n)
547
+ .into_par_iter()
548
+ .with_min_len(256)
549
+ .map_init(
550
+ || Scratch { acc: vec![-1.0; n], touched: Vec::new(), xpn: Vec::new(), list: Vec::new() },
551
+ |scratch, i| {
552
+ let mut out_a = Vec::new();
553
+ probe_half(a, i, t, &index_a, &cached_a, mode_a, scratch, &mut out_a);
554
+ let mut out_b = Vec::new();
555
+ probe_half(b, i, t, &index_b, &cached_b, mode_b, scratch, &mut out_b);
556
+ (out_a, out_b)
557
+ },
558
+ )
559
+ .fold(
560
+ || (Vec::new(), Vec::new()),
561
+ |(mut va, mut vb), (oa, ob)| {
562
+ va.extend(oa);
563
+ vb.extend(ob);
564
+ (va, vb)
565
+ },
566
+ )
567
+ .reduce(
568
+ || (Vec::new(), Vec::new()),
569
+ |(mut a1, mut b1), (mut a2, mut b2)| {
570
+ a1.append(&mut a2);
571
+ b1.append(&mut b2);
572
+ (a1, b1)
573
+ },
574
+ )
575
+ }
576
+
299
577
  /// Run the cosine join under a chosen [`Concurrency`] backend. Returns `(j, i, cos)` pairs with
300
578
  /// `j < i` and `cos ≥ t`, scores as `f64` (the `Gpu` mode's f32 cosines are widened losslessly).
301
579
  ///
@@ -478,9 +756,60 @@ fn accumulate(index: &[Vec<(u32, f64)>], (di, wi): (&[u32], &[f64]), cutoff: u32
478
756
  unsafe { touched.set_len(tlen) };
479
757
  }
480
758
 
481
- /// Phase 2 — prune + verify. For each touched candidate `y`, reset its accumulator and test the
482
- /// **L2AP `l2` bound**: the dot mass missing from `acc[y]` (the dims in `prefix(y)`) is at most
483
- /// `‖x_{rank<split[y]}‖ · ‖prefix(y)‖` by Cauchy–Schwarz, where `x_{rank<split[y]}` is the probe
759
+ /// Fill `xpn` with the probe's running prefix L2 norms (`xpn[k] = ‖wi[..k]‖₂`) and return the full
760
+ /// probe norm plus whether the cheap monotone pre-bound is worth it for this row length. One sqrt
761
+ /// per probe dim, reused across all its candidates.
762
+ ///
763
+ /// Since `xpn` is monotonic, `xpn[kstar] ≤ xnorm` for every candidate, so `a + xnorm·pnorm ≥ a +
764
+ /// xpn[kstar]·pnorm` (the exact bound). A candidate failing the cheap `a + xnorm·pnorm < need` test
765
+ /// therefore also fails the exact bound — prune it without the per-candidate `partition_point`. The
766
+ /// survivor set is bit-identical; only the binary search is skipped for the ~99% of touched
767
+ /// candidates that can't reach `t` on dense corpora.
768
+ #[must_use]
769
+ fn probe_prefix_norms(wi: &[f64], xpn: &mut Vec<f64>) -> (f64, bool) {
770
+ xpn.clear();
771
+ xpn.push(0.0);
772
+ let mut sq = 0.0f64;
773
+ for &w in wi {
774
+ sq += w * w;
775
+ xpn.push(sq.sqrt());
776
+ }
777
+ (sq.sqrt(), wi.len() >= PREBOUND_MIN_DIMS)
778
+ }
779
+
780
+ /// Per-candidate Cauchy–Schwarz prune test, shared by the dense verify pass, the sparse verify pass
781
+ /// and the GPU survivor collector: `true` iff candidate `y` (accumulated partial dot `a`) survives —
782
+ /// i.e. its exact [`cos_full`] score must be computed. The bound is a filter only; the emitted value
783
+ /// is always the exact dot, so the pair set is bit-identical to brute force whatever this decides.
784
+ #[must_use]
785
+ #[inline(always)]
786
+ #[allow(clippy::inline_always)] // this runs once per touched candidate — call overhead is the cost
787
+ fn candidate_survives(
788
+ di: &[u32],
789
+ xpn: &[f64],
790
+ xnorm: f64,
791
+ prebound: bool,
792
+ need: f64,
793
+ a: f64,
794
+ bd: Bound,
795
+ ) -> bool {
796
+ // Cheap monotone pre-bound (dense rows only) — skip the binary search when even `xnorm` can't
797
+ // clear `need`. `prebound` is loop-invariant, so sparse rows pay nothing.
798
+ if prebound && a + xnorm * bd.pnorm < need {
799
+ return false;
800
+ }
801
+ // Number of probe dims with rank < split[y] → index into xpn (di sorted ascending).
802
+ let kstar = di.partition_point(|&d| d < bd.split);
803
+ // SAFETY: kstar ≤ di.len() == wi.len() = xpn.len()-1.
804
+ // (An added maxweight cap `min(…, Σ wx·maxw)` was tried + reverted: never tighter than the
805
+ // L2 cap on either synthetic or real data — it sums over all probe-prefix dims, not just the
806
+ // shared ones — so it pruned nothing and cost ~16% on the real PyPI corpus.)
807
+ a + unsafe { xpn.get_unchecked(kstar) } * bd.pnorm >= need
808
+ }
809
+
810
+ /// Phase 2 — prune + verify (dense mode). For each touched candidate `y`, reset its accumulator and
811
+ /// test the **L2AP `l2` bound**: the dot mass missing from `acc[y]` (the dims in `prefix(y)`) is at
812
+ /// most `‖x_{rank<split[y]}‖ · ‖prefix(y)‖` by Cauchy–Schwarz, where `x_{rank<split[y]}` is the probe
484
813
  /// restricted to the rank range `prefix(y)` lives in. Since the probe's mass sits in its rare
485
814
  /// (high-rank) dims, that restricted norm is tiny — a far tighter cap than the whole-probe `‖x‖=1`.
486
815
  /// Only if `acc[y] + that bound ≥ t` (minus FP slack) do we score exactly with [`cos_full`]. Filter
@@ -495,24 +824,9 @@ fn verify_pruned(
495
824
  out: &mut Vec<(usize, usize, f64)>,
496
825
  ) {
497
826
  let (di, wi) = c.row(i);
498
- // Probe prefix L2 norms: xpn[k] = ‖wi[..k]‖₂ (di ascending by rank, so xpn[k] = norm over the
499
- // probe's k lowest-rank dims). One sqrt per probe dim, reused across all its candidates.
500
- s.xpn.clear();
501
- s.xpn.push(0.0);
502
- let mut sq = 0.0f64;
503
- for &w in wi {
504
- sq += w * w;
505
- s.xpn.push(sq.sqrt());
506
- }
507
- // `‖probe‖` = the full prefix norm. Since `xpn` is monotonic, `xpn[kstar] ≤ xnorm` for every
508
- // candidate, so `a + xnorm·pnorm ≥ a + xpn[kstar]·pnorm` (the exact bound). A candidate failing
509
- // the cheap `a + xnorm·pnorm < need` test therefore also fails the exact bound — prune it without
510
- // the per-candidate `partition_point`. The survivor set is bit-identical; only the binary search
511
- // is skipped for the ~99% of touched candidates that can't reach `t` on dense corpora.
512
- let xnorm = sq.sqrt();
513
- let prebound = di.len() >= PREBOUND_MIN_DIMS;
827
+ let (xnorm, prebound) = probe_prefix_norms(wi, &mut s.xpn);
514
828
  let need = t - PRUNE_SLACK;
515
- let Scratch { acc, touched, xpn } = s;
829
+ let Scratch { acc, touched, xpn, .. } = s;
516
830
  // (Software-prefetching the candidate row a few ahead was tried + reverted: no measurable change
517
831
  // — with all cores gathering at once the join is memory-*bandwidth*-bound, not per-access
518
832
  // latency-bound, so prefetch can't add throughput.)
@@ -522,19 +836,73 @@ fn verify_pruned(
522
836
  let a = unsafe { std::mem::replace(acc.get_unchecked_mut(yu), -1.0) };
523
837
  // SAFETY: `yu < n`. One scattered load fetches both prune fields.
524
838
  let bd = unsafe { *cached.bound.get_unchecked(yu) };
525
- // Cheap monotone pre-bound (dense rows only) — skip the binary search when even `xnorm`
526
- // can't clear `need`. `prebound` is loop-invariant, so sparse rows pay nothing.
527
- if prebound && a + xnorm * bd.pnorm < need {
528
- continue;
839
+ if candidate_survives(di, xpn, xnorm, prebound, need, a, bd) {
840
+ let cos = cos_full((di, wi), c.row(yu));
841
+ if cos >= t {
842
+ out.push((yu, i, cos));
843
+ }
844
+ }
845
+ }
846
+ }
847
+
848
+ /// Sparse-mode phase 1 — the same partial-dot accumulation as [`accumulate`], but into the sorted
849
+ /// `(y, acc)` list instead of the dense `acc[n]` array. Wins when a probe touches few candidates:
850
+ /// the list stays tiny and L1-resident, and the n-wide dense array's footprint never competes with
851
+ /// the postings for L2. The additions happen in the same per-dim/posting order, so the accumulated
852
+ /// values are bit-identical to the dense pass. Returns `false` on [`LIST_ACC_MAX`] overflow — the
853
+ /// caller must redo the probe densely (identical additions, identical order).
854
+ #[cfg_attr(feature = "profiling", inline(never))]
855
+ fn accumulate_list(
856
+ index: &[Vec<(u32, f64)>],
857
+ (di, wi): (&[u32], &[f64]),
858
+ cutoff: u32,
859
+ s: &mut Scratch,
860
+ ) -> bool {
861
+ let Scratch { list, .. } = s;
862
+ list.clear();
863
+ for (&d, &w) in di.iter().zip(wi) {
864
+ for &(y, wy) in &index[d as usize] {
865
+ if y >= cutoff {
866
+ break;
867
+ }
868
+ let term = w * wy;
869
+ match list.binary_search_by_key(&y, |&(yy, _)| yy) {
870
+ Ok(k) => list[k].1 += term,
871
+ Err(k) => {
872
+ if list.len() >= LIST_ACC_MAX {
873
+ list.clear();
874
+ return false;
875
+ }
876
+ list.insert(k, (y, term));
877
+ }
878
+ }
529
879
  }
530
- // Number of probe dims with rank < split[y] → index into xpn (di sorted ascending).
531
- let kstar = di.partition_point(|&d| d < bd.split);
532
- // SAFETY: kstar ≤ di.len() == wi.len() = xpn.len()-1.
533
- // (An added maxweight cap `min(…, Σ wx·maxw)` was tried + reverted: never tighter than the
534
- // L2 cap on either synthetic or real data — it sums over all probe-prefix dims, not just the
535
- // shared ones — so it pruned nothing and cost ~16% on the real PyPI corpus.)
536
- let bound = a + unsafe { xpn.get_unchecked(kstar) } * bd.pnorm;
537
- if bound >= need {
880
+ }
881
+ true
882
+ }
883
+
884
+ /// Sparse-mode phase 2 — the same prune + score as [`verify_pruned`], over the sorted accumulator
885
+ /// list. No reset needed: the next probe's [`accumulate_list`] clears the list. Candidate order is
886
+ /// by id (the dense pass visits first-touch order); both are "arbitrary order" for the caller, and
887
+ /// the per-candidate arithmetic is identical, so the emitted pair set + scores are bit-identical.
888
+ #[cfg_attr(feature = "profiling", inline(never))]
889
+ fn verify_list(
890
+ c: &Corpus,
891
+ i: usize,
892
+ t: f64,
893
+ s: &mut Scratch,
894
+ cached: &Cached,
895
+ out: &mut Vec<(usize, usize, f64)>,
896
+ ) {
897
+ let (di, wi) = c.row(i);
898
+ let (xnorm, prebound) = probe_prefix_norms(wi, &mut s.xpn);
899
+ let need = t - PRUNE_SLACK;
900
+ let Scratch { list, xpn, .. } = s;
901
+ for &(y, a) in list.iter() {
902
+ let yu = y as usize;
903
+ // SAFETY: `yu < n` (candidate ids come from the postings, same as in `accumulate`).
904
+ let bd = unsafe { *cached.bound.get_unchecked(yu) };
905
+ if candidate_survives(di, xpn, xnorm, prebound, need, a, bd) {
538
906
  let cos = cos_full((di, wi), c.row(yu));
539
907
  if cos >= t {
540
908
  out.push((yu, i, cos));
@@ -592,34 +960,21 @@ fn index_suffix(
592
960
  const GPU_FILTER_MARGIN: f64 = 1e-4;
593
961
 
594
962
  /// Like the verify half of [`verify_pruned`], but instead of scoring, **collects** each surviving
595
- /// `(candidate, probe)` pair (`candidate < probe`) for batch scoring elsewhere. The bound here MUST
596
- /// stay identical to `verify_pruned`'s so the survivor set matches exactly.
963
+ /// `(candidate, probe)` pair (`candidate < probe`) for batch scoring elsewhere. The bound test is
964
+ /// [`candidate_survives`] — the same function `verify_pruned` uses — so the survivor set is
965
+ /// identical by construction.
597
966
  #[cfg(all(target_os = "macos", feature = "gpu"))]
598
967
  fn collect_survivors(c: &Corpus, i: usize, t: f64, s: &mut Scratch, cached: &Cached, out: &mut Vec<(u32, u32)>) {
599
968
  let (di, wi) = c.row(i);
600
- s.xpn.clear();
601
- s.xpn.push(0.0);
602
- let mut sq = 0.0f64;
603
- for &w in wi {
604
- sq += w * w;
605
- s.xpn.push(sq.sqrt());
606
- }
607
- let xnorm = sq.sqrt();
608
- let prebound = di.len() >= PREBOUND_MIN_DIMS;
969
+ let (xnorm, prebound) = probe_prefix_norms(wi, &mut s.xpn);
609
970
  let need = t - PRUNE_SLACK;
610
- let Scratch { acc, touched, xpn } = s;
971
+ let Scratch { acc, touched, xpn, .. } = s;
611
972
  for &y in touched.iter() {
612
973
  let yu = y as usize;
613
974
  // SAFETY: `yu < n` (same provenance as in `accumulate`).
614
975
  let a = unsafe { std::mem::replace(acc.get_unchecked_mut(yu), -1.0) };
615
976
  let bd = unsafe { *cached.bound.get_unchecked(yu) };
616
- // Cheap monotone pre-bound (see `verify_pruned`) — same survivor set, skips the binary search.
617
- if prebound && a + xnorm * bd.pnorm < need {
618
- continue;
619
- }
620
- let kstar = di.partition_point(|&d| d < bd.split);
621
- let bound = a + unsafe { xpn.get_unchecked(kstar) } * bd.pnorm;
622
- if bound >= need {
977
+ if candidate_survives(di, xpn, xnorm, prebound, need, a, bd) {
623
978
  out.push((y, i as u32)); // (candidate, probe), candidate < probe
624
979
  }
625
980
  }
@@ -706,7 +1061,7 @@ fn survivor_pairs(c: &Corpus, t: f64) -> (Vec<u32>, Vec<u32>) {
706
1061
  .into_par_iter()
707
1062
  .with_min_len(256)
708
1063
  .map_init(
709
- || Scratch { acc: vec![-1.0; n], touched: Vec::new(), xpn: Vec::new() },
1064
+ || Scratch { acc: vec![-1.0; n], touched: Vec::new(), xpn: Vec::new(), list: Vec::new() },
710
1065
  |scratch, i| {
711
1066
  let (di, wi) = c.row(i);
712
1067
  accumulate(&index, (di, wi), i as u32, scratch);
@@ -734,31 +1089,19 @@ pub fn cosine_join_counts(c: &Corpus, t: f64) -> (u64, u64, u64) {
734
1089
  let (di, wi) = c.row(i);
735
1090
  index_suffix(c, i, (di, wi), t, &mut index, &mut cached);
736
1091
  }
737
- let mut s = Scratch { acc: vec![-1.0; n], touched: Vec::new(), xpn: Vec::new() };
1092
+ let mut s = Scratch { acc: vec![-1.0; n], touched: Vec::new(), xpn: Vec::new(), list: Vec::new() };
738
1093
  let (mut ncand, mut survivors, mut pairs) = (0u64, 0u64, 0u64);
739
1094
  let need = t - PRUNE_SLACK;
740
1095
  for i in 0..n {
741
1096
  let (di, wi) = c.row(i);
742
1097
  accumulate(&index, (di, wi), i as u32, &mut s);
743
1098
  ncand += s.touched.len() as u64;
744
- s.xpn.clear();
745
- s.xpn.push(0.0);
746
- let mut sq = 0.0f64;
747
- for &w in wi {
748
- sq += w * w;
749
- s.xpn.push(sq.sqrt());
750
- }
751
- let xnorm = sq.sqrt();
752
- let prebound = di.len() >= PREBOUND_MIN_DIMS;
1099
+ let (xnorm, prebound) = probe_prefix_norms(wi, &mut s.xpn);
753
1100
  for &y in &s.touched {
754
1101
  let yu = y as usize;
755
1102
  let a = std::mem::replace(&mut s.acc[yu], -1.0);
756
1103
  let bd = cached.bound[yu];
757
- if prebound && a + xnorm * bd.pnorm < need {
758
- continue;
759
- }
760
- let kstar = di.partition_point(|&d| d < bd.split);
761
- if a + s.xpn[kstar] * bd.pnorm >= need {
1104
+ if candidate_survives(di, &s.xpn, xnorm, prebound, need, a, bd) {
762
1105
  survivors += 1;
763
1106
  if cos_full((di, wi), c.row(yu)) >= t {
764
1107
  pairs += 1;
@@ -787,7 +1130,10 @@ pub fn cosine_join_bruteforce(c: &Corpus, t: f64) -> Vec<(usize, usize, f64)> {
787
1130
 
788
1131
  #[cfg(test)]
789
1132
  mod tests {
790
- use super::{cosine_join, cosine_join_bruteforce, Corpus};
1133
+ use super::{
1134
+ cosine_join, cosine_join_bruteforce, cosine_join_double, cosine_join_for_mode, AccMode,
1135
+ Corpus,
1136
+ };
791
1137
 
792
1138
  fn xorshift(seed: u64) -> impl FnMut() -> u64 {
793
1139
  let mut s = seed;
@@ -827,6 +1173,145 @@ mod tests {
827
1173
  }
828
1174
  }
829
1175
 
1176
+ /// `Corpus::from_rows` must be a deterministic function of its input: building the same rows
1177
+ /// twice yields the same corpus, so the join output is bit-identical. (The previous
1178
+ /// HashMap-based per-row merge summed weights in per-call-random hash order, so this did not
1179
+ /// hold; it is now a pinned contract.)
1180
+ #[test]
1181
+ fn from_rows_is_deterministic() {
1182
+ let mut next = xorshift(0x6f6c_6c61_6279_6e6f);
1183
+ for _ in 0..20 {
1184
+ let n = (next() % 30 + 2) as usize;
1185
+ let dim_space = next() % 12 + 1;
1186
+ let rows: Vec<Vec<(u32, f64)>> = (0..n)
1187
+ .map(|_| {
1188
+ let nnz = (next() % 9) as usize;
1189
+ let mut r: Vec<(u32, f64)> = (0..nnz)
1190
+ .map(|_| ((next() % dim_space) as u32, (next() % 10 + 1) as f64))
1191
+ .collect();
1192
+ // Force duplicate dims so the merge path runs and the sum order matters.
1193
+ if nnz > 1 {
1194
+ r[1].0 = r[0].0;
1195
+ }
1196
+ r
1197
+ })
1198
+ .collect();
1199
+ let a = Corpus::from_rows(&rows);
1200
+ let b = Corpus::from_rows(&rows);
1201
+ for &t in &[0.3_f64, 0.7] {
1202
+ assert_eq!(
1203
+ sort_pairs(cosine_join(&a, t)),
1204
+ sort_pairs(cosine_join(&b, t)),
1205
+ "n={n} t={t}"
1206
+ );
1207
+ }
1208
+ }
1209
+ }
1210
+
1211
+ /// Stripping `df = 1` dims from the corpus must be invisible to the join: a hapax dim is shared
1212
+ /// by no pair, so it contributes exactly `0.0` to every pairwise dot (the branchless `cos_full`
1213
+ /// accumulates `+0.0` no-ops on non-negative sums), its posting is never walked, and the tighter
1214
+ /// Cauchy–Schwarz bounds only prune more. The stripped and un-stripped builds must therefore
1215
+ /// return bit-identical pair sets, and each must match its own brute-force oracle.
1216
+ #[test]
1217
+ fn hapax_stripping_is_invisible() {
1218
+ let mut next = xorshift(0x51a7_4e5f_0c92_83d6);
1219
+ for _ in 0..200 {
1220
+ let n = (next() % 40 + 2) as usize;
1221
+ let dim_space = next() % 12 + 1;
1222
+ let rows: Vec<Vec<(u32, f64)>> = (0..n)
1223
+ .map(|i| {
1224
+ let nnz = (next() % 8) as usize;
1225
+ let mut r: Vec<(u32, f64)> = (0..nnz)
1226
+ .map(|_| ((next() % dim_space) as u32, (next() % 10 + 1) as f64))
1227
+ .collect();
1228
+ // Every row carries one unique dim — a guaranteed hapax the strip must remove
1229
+ // invisibly (its mass participates in the L2 normalisation of the full row).
1230
+ r.push((1000 + i as u32, 1.0));
1231
+ r
1232
+ })
1233
+ .collect();
1234
+ let full = Corpus::from_rows_with(&rows, false);
1235
+ let stripped = Corpus::from_rows_with(&rows, true);
1236
+ for &t in &[0.1_f64, 0.25, 0.5, 0.75, 0.9, 1.0] {
1237
+ let want = sort_pairs(cosine_join_bruteforce(&stripped, t));
1238
+ assert_eq!(
1239
+ sort_pairs(cosine_join(&stripped, t)),
1240
+ want,
1241
+ "stripped vs oracle, n={n} t={t}"
1242
+ );
1243
+ assert_eq!(
1244
+ sort_pairs(cosine_join(&full, t)),
1245
+ want,
1246
+ "stripped vs full, n={n} t={t}"
1247
+ );
1248
+ assert_eq!(
1249
+ sort_pairs(cosine_join(&full, t)),
1250
+ sort_pairs(cosine_join_bruteforce(&full, t)),
1251
+ "full vs its own oracle, n={n} t={t}"
1252
+ );
1253
+ }
1254
+ }
1255
+ }
1256
+
1257
+ /// Both accumulator strategies must agree bit-for-bit with brute force. The shape-based mode
1258
+ /// pick in `cosine_join` means any given corpus exercises only one strategy — force each here
1259
+ /// (the dense pass is what the large-corpus shape selects, the sparse list the duplicate-heavy
1260
+ /// consumer shape).
1261
+ #[test]
1262
+ fn both_acc_modes_match_bruteforce() {
1263
+ let mut next = xorshift(0x5eed_0f1a_2b3c_4d5e);
1264
+ for _ in 0..200 {
1265
+ let n = (next() % 40 + 2) as usize;
1266
+ let dim_space = next() % 15 + 1;
1267
+ let rows: Vec<Vec<(u32, f64)>> = (0..n)
1268
+ .map(|_| {
1269
+ let nnz = (next() % 8) as usize;
1270
+ (0..nnz)
1271
+ .map(|_| ((next() % dim_space) as u32, (next() % 10 + 1) as f64))
1272
+ .collect()
1273
+ })
1274
+ .collect();
1275
+ let c = Corpus::from_rows(&rows);
1276
+ for &t in &[0.1_f64, 0.25, 0.5, 0.75, 0.9, 1.0] {
1277
+ let want = sort_pairs(cosine_join_bruteforce(&c, t));
1278
+ for mode in [AccMode::Dense, AccMode::Sparse] {
1279
+ let got = sort_pairs(cosine_join_for_mode(&c, t, mode));
1280
+ assert_eq!(got, want, "n={n} t={t} mode={mode:?}");
1281
+ }
1282
+ }
1283
+ }
1284
+ }
1285
+
1286
+ /// The fused two-corpus join must return, per half, the bit-identical pair set of a separate
1287
+ /// `cosine_join` on that corpus — fusion shares only schedule and scratch, never arithmetic.
1288
+ #[test]
1289
+ fn double_join_matches_single() {
1290
+ let mut next = xorshift(0xd0a2_7b7e_9c31_5f84);
1291
+ for _ in 0..200 {
1292
+ let n = (next() % 40 + 2) as usize;
1293
+ let dim_space = next() % 15 + 1;
1294
+ let mut gen = || -> Vec<Vec<(u32, f64)>> {
1295
+ (0..n)
1296
+ .map(|_| {
1297
+ let nnz = (next() % 8) as usize;
1298
+ (0..nnz)
1299
+ .map(|_| ((next() % dim_space) as u32, (next() % 10 + 1) as f64))
1300
+ .collect()
1301
+ })
1302
+ .collect()
1303
+ };
1304
+ let (rows_a, rows_b) = (gen(), gen());
1305
+ let a = Corpus::from_rows(&rows_a);
1306
+ let b = Corpus::from_rows(&rows_b);
1307
+ for &t in &[0.1_f64, 0.3, 0.5, 0.7, 0.9, 1.0] {
1308
+ let (ga, gb) = cosine_join_double(&a, &b, t);
1309
+ assert_eq!(sort_pairs(ga), sort_pairs(cosine_join(&a, t)), "half a n={n} t={t}");
1310
+ assert_eq!(sort_pairs(gb), sort_pairs(cosine_join(&b, t)), "half b n={n} t={t}");
1311
+ }
1312
+ }
1313
+ }
1314
+
830
1315
  /// The CPU+GPU hybrid [`super::cosine_join_gpu`] must return bit-identical results to the pure-CPU
831
1316
  /// [`cosine_join`] on fuzzed corpora — the GPU is only a conservative filter; every emitted score
832
1317
  /// is the exact CPU `f64` value. Skips when no Metal device is present.
@@ -164,7 +164,14 @@ impl BatchCosineGpu {
164
164
  }
165
165
 
166
166
  /// Copy `data` into a fresh shared (UMA) buffer — one memcpy into unified memory (zero device copy).
167
+ /// An empty slice gets a 1-byte zeroed buffer instead: `new_buffer_with_data` would memcpy 1 byte
168
+ /// from `data.as_ptr()`, which is a dangling pointer for empty slices (a real SIGSEGV once hapax
169
+ /// stripping made fully-empty CSR arrays reachable). The kernel never reads the buffer when the
170
+ /// corpus has no entries (every row's `indptr` window is empty), so the 1-byte placeholder is inert.
167
171
  fn upload<T: Copy>(device: &Device, data: &[T]) -> Buffer {
168
- let bytes = (std::mem::size_of_val(data) as NSUInteger).max(1);
172
+ if data.is_empty() {
173
+ return device.new_buffer(1, MTLResourceOptions::StorageModeShared);
174
+ }
175
+ let bytes = std::mem::size_of_val(data) as NSUInteger;
169
176
  device.new_buffer_with_data(data.as_ptr().cast::<c_void>(), bytes, MTLResourceOptions::StorageModeShared)
170
177
  }
File without changes
File without changes
File without changes
File without changes