From 95be03b5bf29c5fd9b1f7a9404b1c2a4c0a02f9d Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Sun, 26 Jul 2026 14:11:28 -0700 Subject: [PATCH 1/5] Scan selection candidates in place (ABI 42 -> 43) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `select_range` prunes by zone map before it scans, and then threw the saving away: the pruned branch gathered `x[candidates]` and `y[candidates]` into two fresh f64 columns purely so it could call `range_indices`, which only knows how to walk a whole column. Sixteen bytes per candidate, which produces the backwards result that selecting *part* of a trace costs far more than selecting all of it — on a 10M-row scatter, a half-domain box peaked at 134.6 MB while a full-domain box peaked at 38.1 MB. `xy_range_indices_rows` is the rectangular twin of the lasso's `xy_polygon_select`: same `(x, y, len, rows, n_rows, ..., out)` shape, same inclusive comparison chain as `range_scan_scalar`, same order-preserving gather-down across worker segments. The pruned branch now reads through the row ids and gets canonical ids straight back, so both the gathers and the `candidates[hits]` re-index are gone. The half-domain case drops to 57.6 MB and the full-domain case is untouched. ABI_VERSION goes 42 -> 43 in `src/lib.rs` and `python/xy/_native.py` together. Parity is the gate, not inspection: the new kernel is tested against `range_indices` over the gathered form for five windows (including empty, degenerate-point and whole-domain) crossed with five row subsets and five sizes up to 100k, plus a 2M-row case that exercises the threaded path, NaN and infinities on both axes, inclusive bounds, unsorted row order, invalid windows and an out-of-range row id. `test_selection_prunes_non_overlapping_zone_chunks` now also asserts that exactly one scan happens and that it reads the whole column, which is what proves the gather is gone. 118 Rust tests, 133 ABI smoke checks and the full Python suite pass; a 97-artifact output fingerprint sweep is identical to main. --- python/xy/_native.py | 58 ++++++++++++++ python/xy/interaction.py | 11 ++- python/xy/kernels.py | 2 + src/kernels.rs | 80 +++++++++++++++++++ src/lib.rs | 45 +++++++++++ tests/test_range_indices_rows.py | 129 +++++++++++++++++++++++++++++++ tests/test_scatter.py | 24 ++++-- 7 files changed, 341 insertions(+), 8 deletions(-) create mode 100644 tests/test_range_indices_rows.py diff --git a/python/xy/_native.py b/python/xy/_native.py index 9f514213..509ccea5 100644 --- a/python/xy/_native.py +++ b/python/xy/_native.py @@ -558,6 +558,19 @@ def _load() -> ctypes.CDLL: ctypes.c_double, ctypes.c_void_p, ] + lib.xy_range_indices_rows.restype = ctypes.c_size_t + lib.xy_range_indices_rows.argtypes = [ + ctypes.c_void_p, # x + ctypes.c_void_p, # y + ctypes.c_size_t, # len + ctypes.c_void_p, # candidate row ids + ctypes.c_size_t, # candidate count + ctypes.c_double, # lo_x + ctypes.c_double, # hi_x + ctypes.c_double, # lo_y + ctypes.c_double, # hi_y + ctypes.c_void_p, # output row IDs + ] lib.xy_polygon_select.restype = ctypes.c_size_t lib.xy_polygon_select.argtypes = [ ctypes.c_void_p, # x @@ -2623,6 +2636,51 @@ def range_indices( return out if written == len(out) else out[:written].copy() +def range_indices_rows( + x: npt.NDArray[np.float64], + y: npt.NDArray[np.float64], + rows: npt.NDArray[np.uint32], + lo_x: float, + hi_x: float, + lo_y: float, + hi_y: float, +) -> npt.NDArray[np.uint32]: + """Those of `rows` whose canonical point is in the inclusive window. + + The row-restricted form of `range_indices`, shaped like `polygon_select`. + A caller that already knows its candidate rows (zone-map pruning, a drill + window) reads them in place instead of gathering `x[rows]`/`y[rows]` into + two fresh f64 columns first. + """ + lo_x, hi_x = _finite_ordered(lo_x, hi_x, "x range") + lo_y, hi_y = _finite_ordered(lo_y, hi_y, "y range") + x = _as_f64(x, "x") + y = _as_f64(y, "y") + if len(x) != len(y): + raise ValueError("x and y must have equal length") + rows = np.ascontiguousarray(rows, dtype=np.uint32) + if rows.ndim != 1: + raise ValueError("rows must be 1-D") + out = np.empty(len(rows), dtype=np.uint32) + if len(rows) == 0: + return out + written = _lib.xy_range_indices_rows( + _ptr_f64(x), + _ptr_f64(y), + len(x), + rows.ctypes.data, + len(rows), + lo_x, + hi_x, + lo_y, + hi_y, + out.ctypes.data, + ) + if written == _USIZE_MAX: + raise ValueError("invalid range_indices_rows arguments") + return out if written == len(out) else out[:written].copy() + + def polygon_select( x: npt.NDArray[np.float64], y: npt.NDArray[np.float64], diff --git a/python/xy/interaction.py b/python/xy/interaction.py index 940eef02..b33eb75e 100644 --- a/python/xy/interaction.py +++ b/python/xy/interaction.py @@ -299,11 +299,16 @@ def select_range( elif len(candidate_chunks) == 0: out[t.id] = np.empty(0, dtype=np.uint32) else: + # Scan the candidates in place. Gathering `x[candidates]` and + # `y[candidates]` first cost two fresh f64 columns — 16 bytes per + # candidate — so a partial selection could cost several times a + # whole-domain one: half of a 10M-row trace peaked at 134.6 MB + # where selecting every row peaked at 38.1 MB. The kernel returns + # canonical row ids directly, so the re-index is gone too. candidates = _expand_zone_chunks(t.x, candidate_chunks) - out[t.id] = kernels.range_indices( - t.x.values[candidates], t.y.values[candidates], lo_x, hi_x, lo_y, hi_y + out[t.id] = kernels.range_indices_rows( + t.x.values, t.y.values, candidates, lo_x, hi_x, lo_y, hi_y ) - out[t.id] = candidates[out[t.id]] out[t.id] = _drop_hidden_rows(t, out[t.id]) return out diff --git a/python/xy/kernels.py b/python/xy/kernels.py index 1b50a2b1..fc40e861 100644 --- a/python/xy/kernels.py +++ b/python/xy/kernels.py @@ -63,6 +63,7 @@ valid_indices_f64 = _impl.valid_indices_f64 remap_u8 = _impl.remap_u8 range_indices = _impl.range_indices +range_indices_rows = _impl.range_indices_rows polygon_select = _impl.polygon_select sample_mask = _impl.sample_mask sample_range_indices = _impl.sample_range_indices @@ -134,6 +135,7 @@ "pyramid_free", "quad_mesh_triangles", "range_indices", + "range_indices_rows", "rasterize", "rasterize_png", "remap_u8", diff --git a/src/kernels.rs b/src/kernels.rs index ad7b63c6..ba7f2c50 100644 --- a/src/kernels.rs +++ b/src/kernels.rs @@ -5062,6 +5062,86 @@ fn range_indices_impl( write } +/// Which of `rows` land inside the rectangular window (§34 selection), the +/// row-restricted twin of [`range_indices`]. Writes the surviving canonical row +/// ids to `out` (capacity `rows.len()`, order preserved) and returns how many. +/// +/// The predicate is the same inclusive comparison chain `range_scan_scalar` +/// uses, so a subset scan agrees with a full scan row for row — including on +/// NaN, which fails every comparison and is therefore never selected. +/// +/// This exists because the zone-map-pruned selection path already knows its +/// candidate rows before it scans. Without it, that path had to gather +/// `x[rows]` and `y[rows]` into two fresh f64 columns just to call +/// `range_indices` — 16 bytes per candidate, which is why selecting *half* a +/// 10M-row trace cost 134.6 MB where selecting *all* of it cost 38.1 MB. +#[allow(clippy::too_many_arguments)] +pub fn range_indices_rows( + x: &[f64], + y: &[f64], + rows: &[u32], + lo_x: f64, + hi_x: f64, + lo_y: f64, + hi_y: f64, + out: &mut [u32], +) -> usize { + assert_eq!(x.len(), y.len()); + assert!(out.len() >= rows.len()); + let n = rows.len(); + let threads = par_threads(n); + if threads <= 1 || n < threads { + return range_scan_rows(x, y, rows, lo_x, hi_x, lo_y, hi_y, out); + } + let chunk = n.div_ceil(threads); + let counts: Vec = std::thread::scope(|s| { + let handles: Vec<_> = rows + .chunks(chunk) + .zip(out[..n].chunks_mut(chunk)) + .map(|(rseg, oseg)| { + s.spawn(move || range_scan_rows(x, y, rseg, lo_x, hi_x, lo_y, hi_y, oseg)) + }) + .collect(); + handles + .into_iter() + .map(|hd| hd.join().expect("range_indices_rows worker panicked")) + .collect() + }); + // Same gather-down as `range_indices_impl`: each worker packed its hits at + // the front of its own out-segment, so slide the later segments back. + let mut write = counts[0]; + for (t, &c) in counts.iter().enumerate().skip(1) { + let start = t * chunk; + out.copy_within(start..start + c, write); + write += c; + } + write +} + +#[allow(clippy::too_many_arguments)] +fn range_scan_rows( + x: &[f64], + y: &[f64], + rows: &[u32], + lo_x: f64, + hi_x: f64, + lo_y: f64, + hi_y: f64, + out: &mut [u32], +) -> usize { + let mut n = 0usize; + for &row in rows { + let i = row as usize; + let xv = x[i]; + let yv = y[i]; + if xv >= lo_x && xv <= hi_x && yv >= lo_y && yv <= hi_y { + out[n] = row; + n += 1; + } + } + n +} + /// Per-point log-normalized local density for a subset. This fuses the /// grid-bin + point-lookup pass used during drill handoff, avoiding Python-side /// integer temp arrays. diff --git a/src/lib.rs b/src/lib.rs index d6ecc734..9de2d61d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2872,6 +2872,51 @@ pub unsafe extern "C" fn xy_range_indices( }) } +/// Canonical row ids from `rows` that fall inside the rectangular window — +/// the row-restricted twin of `xy_range_indices`, shaped like +/// `xy_polygon_select`. Returns the count written; `out` must hold `n_rows` +/// u32s. Row ids must be < `len`; an out-of-range id is caught by the bounds +/// check inside and returned as the error sentinel. +/// +/// # Safety +/// `x`/`y` must point to `len` readable f64s, `rows` to `n_rows` readable +/// u32s, and `out` to `n_rows` writable u32s. +#[no_mangle] +#[allow(clippy::too_many_arguments)] +pub unsafe extern "C" fn xy_range_indices_rows( + x: *const f64, + y: *const f64, + len: usize, + rows: *const u32, + n_rows: usize, + lo_x: f64, + hi_x: f64, + lo_y: f64, + hi_y: f64, + out: *mut u32, +) -> usize { + if !finite_ordered(lo_x, hi_x) || !finite_ordered(lo_y, hi_y) { + return usize::MAX; + } + // u32 index ceiling — see xy_m4_indices. + if len > u32::MAX as usize { + return usize::MAX; + } + if n_rows == 0 { + return 0; + } + if x.is_null() || y.is_null() || rows.is_null() || out.is_null() { + return usize::MAX; + } + let x = std::slice::from_raw_parts(x, len); + let y = std::slice::from_raw_parts(y, len); + let rows = std::slice::from_raw_parts(rows, n_rows); + let out = std::slice::from_raw_parts_mut(out, n_rows); + ffi_guard(usize::MAX, || { + kernels::range_indices_rows(x, y, rows, lo_x, hi_x, lo_y, hi_y, out) + }) +} + /// Canonical row ids from `rows` that fall inside the lasso polygon, by /// even-odd ray casting. Returns the count written; `out` must hold /// `n_rows` u32s. A polygon of fewer than 3 vertices selects nothing. diff --git a/tests/test_range_indices_rows.py b/tests/test_range_indices_rows.py new file mode 100644 index 00000000..5c45f887 --- /dev/null +++ b/tests/test_range_indices_rows.py @@ -0,0 +1,129 @@ +"""Row-restricted rectangular selection: `range_indices_rows` must answer +exactly what `range_indices` answers over the same rows (design dossier §34). + +This is the kernel that lets the zone-map-pruned branch of `select_range` scan +candidates in place. Before it, that branch gathered `x[candidates]` and +`y[candidates]` into two fresh f64 columns — 16 bytes per candidate, so a +*partial* selection cost several times a whole-domain one. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from xy import interaction, kernels +from xy._figure import Figure + + +def _reference( + x: np.ndarray, y: np.ndarray, rows: np.ndarray, win: tuple[float, float, float, float] +) -> np.ndarray: + """What the gather-then-scan form produced, kept executable.""" + lo_x, hi_x, lo_y, hi_y = win + hit = kernels.range_indices(x[rows], y[rows], lo_x, hi_x, lo_y, hi_y) + return rows[hit] + + +WINDOWS = [ + (-0.5, 0.5, -0.5, 0.5), + (-10.0, 10.0, -10.0, 10.0), # everything + (100.0, 200.0, 100.0, 200.0), # nothing + (0.0, 0.0, 0.0, 0.0), # degenerate point window + (-1.0, 0.0, 0.0, 1.0), # one quadrant +] + + +@pytest.mark.parametrize("win", WINDOWS, ids=lambda w: f"{w[0]}_{w[1]}_{w[2]}_{w[3]}") +@pytest.mark.parametrize("n", [1, 2, 17, 1000, 100_000]) +def test_matches_the_gathered_form(n: int, win: tuple[float, float, float, float]) -> None: + rng = np.random.default_rng(4) + x = rng.normal(size=n) + y = rng.normal(size=n) + for rows in ( + np.arange(n, dtype=np.uint32), # every row + np.arange(0, n, 3, dtype=np.uint32), # a strided subset + np.array([n - 1], dtype=np.uint32), # the last row alone + np.empty(0, dtype=np.uint32), # nothing + ): + got = kernels.range_indices_rows(x, y, rows, *win) + want = _reference(x, y, rows, win) + np.testing.assert_array_equal(got, want) + assert got.dtype == np.uint32 + + +def test_non_finite_is_never_inside() -> None: + """NaN and infinities fail every comparison, matching `range_indices`.""" + x = np.array([0.0, np.nan, np.inf, -np.inf, 0.5, 0.0, 0.0]) + y = np.array([0.0, 0.0, 0.0, 0.0, 0.5, np.nan, np.inf]) + rows = np.arange(len(x), dtype=np.uint32) + got = kernels.range_indices_rows(x, y, rows, -1.0, 1.0, -1.0, 1.0) + np.testing.assert_array_equal(got, np.array([0, 4], dtype=np.uint32)) + np.testing.assert_array_equal(got, _reference(x, y, rows, (-1.0, 1.0, -1.0, 1.0))) + + +def test_bounds_are_inclusive() -> None: + x = np.array([-1.0, -1.0, 1.0, 1.0, 0.0]) + y = np.array([-1.0, 1.0, -1.0, 1.0, 0.0]) + rows = np.arange(5, dtype=np.uint32) + got = kernels.range_indices_rows(x, y, rows, -1.0, 1.0, -1.0, 1.0) + np.testing.assert_array_equal(got, rows) + + +def test_row_order_is_preserved_not_sorted() -> None: + """Output follows `rows`, so a caller's ordering survives the scan.""" + x = np.array([0.0, 0.1, 0.2, 0.3]) + y = np.zeros(4) + rows = np.array([3, 1, 2, 0], dtype=np.uint32) + got = kernels.range_indices_rows(x, y, rows, -1.0, 1.0, -1.0, 1.0) + np.testing.assert_array_equal(got, rows) + + +def test_parallel_and_serial_agree() -> None: + """The threaded gather-down must produce the serial answer exactly.""" + rng = np.random.default_rng(5) + n = 2_000_000 + x = rng.normal(size=n) + y = rng.normal(size=n) + rows = np.arange(0, n, 2, dtype=np.uint32) + got = kernels.range_indices_rows(x, y, rows, -0.25, 0.25, -0.25, 0.25) + want = _reference(x, y, rows, (-0.25, 0.25, -0.25, 0.25)) + np.testing.assert_array_equal(got, want) + assert 0 < got.size < rows.size + + +@pytest.mark.parametrize("bad", [(np.nan, 1.0, -1.0, 1.0), (1.0, -1.0, -1.0, 1.0)]) +def test_invalid_window_raises(bad: tuple[float, float, float, float]) -> None: + x = np.zeros(4) + rows = np.arange(4, dtype=np.uint32) + with pytest.raises(ValueError): + kernels.range_indices_rows(x, x, rows, *bad) + + +def test_out_of_range_row_is_an_error_not_a_read() -> None: + """A row id past the column length must come back as the error sentinel.""" + x = np.zeros(4) + rows = np.array([0, 9], dtype=np.uint32) + with pytest.raises(ValueError): + kernels.range_indices_rows(x, x, rows, -1.0, 1.0, -1.0, 1.0) + + +def test_select_range_agrees_across_the_pruning_branches() -> None: + """Zone-pruned and whole-column selection must return the same rows. + + `select_range` takes the pruned branch only when some zone chunks are + excluded, so a sorted-x column with a narrow window exercises the path this + kernel replaced, and a full-domain window exercises the other one. + """ + rng = np.random.default_rng(6) + n = 300_000 + x = np.sort(rng.normal(size=n)) + y = rng.normal(size=n) + fig = Figure() + fig.scatter(x, y) + fig.build_payload() + + for lo, hi in ((-10.0, 10.0), (-0.5, 0.0), (-3.0, 3.0), (0.9, 1.1)): + got = interaction.select_range(fig, lo, hi, -0.5, 0.5)[0] + want = np.flatnonzero((x >= lo) & (x <= hi) & (y >= -0.5) & (y <= 0.5)) + np.testing.assert_array_equal(got.astype(np.int64), want) diff --git a/tests/test_scatter.py b/tests/test_scatter.py index 8cc3cbdd..ec5b3f01 100644 --- a/tests/test_scatter.py +++ b/tests/test_scatter.py @@ -1158,24 +1158,38 @@ def test_selection_prunes_non_overlapping_zone_chunks(monkeypatch): y = np.concatenate([np.zeros(ZONE_CHUNK), np.full(ZONE_CHUNK, 100.0)]) fig = Figure().scatter(x, y) seen = [] + scanned_columns = [] expanded = [] - original = interaction.kernels.range_indices + original = interaction.kernels.range_indices_rows + original_full = interaction.kernels.range_indices original_expand = interaction._expand_zone_chunks - def wrapped(xv, yv, *args): - seen.append(len(xv)) - return original(xv, yv, *args) + def wrapped(xv, yv, rows, *args): + # The candidate count is what pruning controls; the column length the + # kernel reads is the *whole* column, because the pruned branch scans + # rows in place rather than gathering them into a fresh pair of f64 + # columns first. + seen.append(len(rows)) + scanned_columns.append(len(xv)) + return original(xv, yv, rows, *args) + + def wrapped_full(xv, yv, *args): + scanned_columns.append(("full", len(xv))) + return original_full(xv, yv, *args) def wrapped_expand(col, chunks): expanded.append(len(chunks)) return original_expand(col, chunks) - monkeypatch.setattr(interaction.kernels, "range_indices", wrapped) + monkeypatch.setattr(interaction.kernels, "range_indices_rows", wrapped) + monkeypatch.setattr(interaction.kernels, "range_indices", wrapped_full) monkeypatch.setattr(interaction, "_expand_zone_chunks", wrapped_expand) selected = fig.select_range(-1.0, 1.0, -1.0, 1.0)[0] np.testing.assert_array_equal(selected, np.arange(ZONE_CHUNK, dtype=np.uint32)) assert seen == [ZONE_CHUNK] assert expanded == [1] + # Exactly one scan, of the candidate rows, with no gathered copy. + assert scanned_columns == [2 * ZONE_CHUNK] def test_to_shipped_indices_translates_nan_drop_before_payload_build(): From 22568ac77a27f4e0c50980a2280364443698b637 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Sun, 26 Jul 2026 16:07:49 -0700 Subject: [PATCH 2/5] Answer an out-of-range row id instead of aborting on it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `range_indices_rows` and `polygon_select` index `x[row]` and left the report to Rust's bounds check, which panics. `ffi_guard` turns a panic into the entry point's sentinel only where panics unwind, and the PyEmscripten wheel is built `-C panic=abort` precisely so they cannot: there the same id aborted the Pyodide instance (exit 134) rather than raising ValueError. `xy.kernels` is public API, so these row arrays are caller data — the documented contract has to hold on every target, not just the ones that unwind. Both kernels now read through `get` and return `Option`, with the FFI layer mapping `None` to the existing sentinel. This is free: the same two bounds checks happen either way, so only the reporting changes. A separate validating pass over `rows` was not free and is deliberately not what this does — one serial sweep of 5M u32 cost more than the entire parallel scan (1.0 ms -> 2.4 ms), where reading through `get` measures 1.01 ms against 1.02 ms, and `polygon_select` 36.3 ms against 36.4 ms. --- src/kernels.rs | 83 ++++++++++++++++++++++++++------ src/lib.rs | 15 +++--- tests/test_range_indices_rows.py | 17 ++++++- tests/test_select_polygon.py | 4 ++ 4 files changed, 97 insertions(+), 22 deletions(-) diff --git a/src/kernels.rs b/src/kernels.rs index ba7f2c50..94275773 100644 --- a/src/kernels.rs +++ b/src/kernels.rs @@ -4807,6 +4807,10 @@ pub fn range_indices( /// /// A non-finite coordinate fails every comparison and so is never inside, /// matching both the vectorized predicate and `range_indices`. +/// +/// `None` means a row id was >= `x.len()`. Reporting it costs nothing over the +/// indexing this already did — see [`range_scan_rows`] for why the check must +/// not be left to Rust's own panicking one. pub fn polygon_select( x: &[f64], y: &[f64], @@ -4814,13 +4818,13 @@ pub fn polygon_select( poly_x: &[f64], poly_y: &[f64], out: &mut [u32], -) -> usize { +) -> Option { assert_eq!(x.len(), y.len()); assert_eq!(poly_x.len(), poly_y.len()); assert!(out.len() >= rows.len()); let n = poly_x.len(); if n < 3 { - return 0; + return Some(0); } // (x_i, y_i, y_j, x_j - x_i, y_j - y_i) per edge, hoisted so the inner // loop is a bounds-check-free walk over a couple of KB. @@ -4839,7 +4843,10 @@ pub fn polygon_select( let index = PolygonSlabs::build(poly_y, &edges); let mut write = 0; for &r in rows { - let (xv, yv) = (x[r as usize], y[r as usize]); + let i = r as usize; + let (Some(&xv), Some(&yv)) = (x.get(i), y.get(i)) else { + return None; + }; let mut inside = false; // An edge can only flip parity for a point inside its own y-extent, // so the slab index hands back a superset of the edges that can @@ -4857,7 +4864,7 @@ pub fn polygon_select( write += 1; } } - write + Some(write) } /// Ceiling on `PolygonSlabs` CSR entries (4 MB of u32 at the worst case), so @@ -5085,7 +5092,7 @@ pub fn range_indices_rows( lo_y: f64, hi_y: f64, out: &mut [u32], -) -> usize { +) -> Option { assert_eq!(x.len(), y.len()); assert!(out.len() >= rows.len()); let n = rows.len(); @@ -5094,7 +5101,7 @@ pub fn range_indices_rows( return range_scan_rows(x, y, rows, lo_x, hi_x, lo_y, hi_y, out); } let chunk = n.div_ceil(threads); - let counts: Vec = std::thread::scope(|s| { + let counts: Vec> = std::thread::scope(|s| { let handles: Vec<_> = rows .chunks(chunk) .zip(out[..n].chunks_mut(chunk)) @@ -5107,6 +5114,10 @@ pub fn range_indices_rows( .map(|hd| hd.join().expect("range_indices_rows worker panicked")) .collect() }); + // An out-of-range id in any segment fails the whole call, before the + // gather-down reads a count that was never produced. + let counts: Option> = counts.into_iter().collect(); + let counts = counts?; // Same gather-down as `range_indices_impl`: each worker packed its hits at // the front of its own out-segment, so slide the later segments back. let mut write = counts[0]; @@ -5115,9 +5126,20 @@ pub fn range_indices_rows( out.copy_within(start..start + c, write); write += c; } - write + Some(write) } +/// One segment of the row-restricted scan; `None` if a row id was >= `x.len()`. +/// +/// The check is `get` rather than `x[i]` deliberately. Indexing bounds-checks +/// too, but reports by panicking, and `ffi_guard` only converts a panic into +/// the C ABI's error sentinel where panics unwind — the PyEmscripten wheel is +/// built `-C panic=abort` (`.github/workflows/release.yml`) precisely so they +/// cannot, and there an out-of-range id aborts the Pyodide instance instead of +/// returning. `xy.kernels` is public API, so row ids are caller data. Same two +/// bounds checks either way, so answering instead of aborting is free: a +/// separate validating pass over `rows` was not — one serial sweep of 5M u32 +/// cost more than this entire parallel scan (1.0 ms -> 2.4 ms). #[allow(clippy::too_many_arguments)] fn range_scan_rows( x: &[f64], @@ -5128,18 +5150,19 @@ fn range_scan_rows( lo_y: f64, hi_y: f64, out: &mut [u32], -) -> usize { +) -> Option { let mut n = 0usize; for &row in rows { let i = row as usize; - let xv = x[i]; - let yv = y[i]; + let (Some(&xv), Some(&yv)) = (x.get(i), y.get(i)) else { + return None; + }; if xv >= lo_x && xv <= hi_x && yv >= lo_y && yv <= hi_y { out[n] = row; n += 1; } } - n + Some(n) } /// Per-point log-normalized local density for a subset. This fuses the @@ -5390,7 +5413,7 @@ mod tests { fn polygon_case(poly_x: &[f64], poly_y: &[f64], x: &[f64], y: &[f64]) { let rows: Vec = (0..x.len() as u32).collect(); let mut out = vec![0u32; rows.len()]; - let n = polygon_select(x, y, &rows, poly_x, poly_y, &mut out); + let n = polygon_select(x, y, &rows, poly_x, poly_y, &mut out).expect("rows in bounds"); assert_eq!( &out[..n], &polygon_reference(x, y, &rows, poly_x, poly_y)[..], @@ -5459,12 +5482,44 @@ mod tests { let mut out = vec![0u32; rows.len()]; assert_eq!( polygon_select(&x, &y, &rows, &[0.0, 9.0], &[0.0, 9.0], &mut out), - 0 + Some(0) ); assert_eq!( polygon_select(&x, &y, &[], &[0.0, 9.0, 5.0], &[0.0, 9.0, 5.0], &mut out), - 0 + Some(0) ); + + // An id past the end of the column is reported, not indexed. The + // kernels answer this themselves so it holds where panics abort + // rather than unwind (the PyEmscripten wheel's `-C panic=abort`). + assert_eq!( + polygon_select( + &x, + &y, + &[x.len() as u32], + &[0.0, 9.0, 5.0], + &[0.0, 9.0, 5.0], + &mut out + ), + None + ); + assert_eq!( + range_indices_rows(&x, &y, &[x.len() as u32], 0.0, 9.0, 0.0, 9.0, &mut out), + None + ); + // The ceiling is `< len`, not `<= len`: the last id is in range, so it + // is answered rather than refused — here as "not selected", because + // that row's y is infinite. Row 0 of the lattice is genuinely inside. + let last = [x.len() as u32 - 1]; + assert_eq!( + range_indices_rows(&x, &y, &last, -1e9, 1e9, -1e9, 1e9, &mut out), + Some(0) + ); + assert_eq!( + range_indices_rows(&x, &y, &[0], -1e9, 1e9, -1e9, 1e9, &mut out), + Some(1) + ); + assert_eq!(out[0], 0); } /// `polygon_select` is a public export with no polygon-size ceiling of its diff --git a/src/lib.rs b/src/lib.rs index 9de2d61d..ef90df96 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2875,8 +2875,9 @@ pub unsafe extern "C" fn xy_range_indices( /// Canonical row ids from `rows` that fall inside the rectangular window — /// the row-restricted twin of `xy_range_indices`, shaped like /// `xy_polygon_select`. Returns the count written; `out` must hold `n_rows` -/// u32s. Row ids must be < `len`; an out-of-range id is caught by the bounds -/// check inside and returned as the error sentinel. +/// u32s. Row ids must be < `len`; an out-of-range id returns the error +/// sentinel on every target, including panic-abort ones where the kernel's own +/// indexing panic could not (see `kernels::range_scan_rows`). /// /// # Safety /// `x`/`y` must point to `len` readable f64s, `rows` to `n_rows` readable @@ -2913,13 +2914,15 @@ pub unsafe extern "C" fn xy_range_indices_rows( let rows = std::slice::from_raw_parts(rows, n_rows); let out = std::slice::from_raw_parts_mut(out, n_rows); ffi_guard(usize::MAX, || { - kernels::range_indices_rows(x, y, rows, lo_x, hi_x, lo_y, hi_y, out) + kernels::range_indices_rows(x, y, rows, lo_x, hi_x, lo_y, hi_y, out).unwrap_or(usize::MAX) }) } /// Canonical row ids from `rows` that fall inside the lasso polygon, by /// even-odd ray casting. Returns the count written; `out` must hold -/// `n_rows` u32s. A polygon of fewer than 3 vertices selects nothing. +/// `n_rows` u32s. A polygon of fewer than 3 vertices selects nothing. Row ids +/// must be < `len`; an out-of-range id returns the error sentinel on every +/// target (see `kernels::range_scan_rows`). /// /// # Safety /// `x`/`y` must point to `len` readable f64s, `rows` to `n_rows` readable @@ -2964,10 +2967,8 @@ pub unsafe extern "C" fn xy_polygon_select( let poly_x = std::slice::from_raw_parts(poly_x, n_poly); let poly_y = std::slice::from_raw_parts(poly_y, n_poly); let out = std::slice::from_raw_parts_mut(out, n_rows); - // An out-of-range row id panics the bounds check inside; ffi_guard turns - // that into the error sentinel rather than unwinding across the ABI. ffi_guard(usize::MAX, || { - kernels::polygon_select(x, y, rows, poly_x, poly_y, out) + kernels::polygon_select(x, y, rows, poly_x, poly_y, out).unwrap_or(usize::MAX) }) } diff --git a/tests/test_range_indices_rows.py b/tests/test_range_indices_rows.py index 5c45f887..040dbca7 100644 --- a/tests/test_range_indices_rows.py +++ b/tests/test_range_indices_rows.py @@ -101,11 +101,26 @@ def test_invalid_window_raises(bad: tuple[float, float, float, float]) -> None: def test_out_of_range_row_is_an_error_not_a_read() -> None: - """A row id past the column length must come back as the error sentinel.""" + """A row id past the column length must come back as the error sentinel. + + The entry point validates the ids rather than relying on the scan's own + indexing panic: `ffi_guard` turns a panic into the sentinel only where + panics unwind, and the PyEmscripten wheel is built `-C panic=abort` (see + `.github/workflows/release.yml`) precisely so they cannot. There the panic + aborted the Pyodide instance — and `xy.kernels` is public API, so the row + array is caller data. + """ x = np.zeros(4) rows = np.array([0, 9], dtype=np.uint32) with pytest.raises(ValueError): kernels.range_indices_rows(x, x, rows, -1.0, 1.0, -1.0, 1.0) + # The boundary id itself is valid; only past-the-end is an error. + np.testing.assert_array_equal( + kernels.range_indices_rows(x, x, np.array([3], dtype=np.uint32), -1.0, 1.0, -1.0, 1.0), + np.array([3], dtype=np.uint32), + ) + with pytest.raises(ValueError): + kernels.range_indices_rows(x, x, np.array([4], dtype=np.uint32), -1.0, 1.0, -1.0, 1.0) def test_select_range_agrees_across_the_pruning_branches() -> None: diff --git a/tests/test_select_polygon.py b/tests/test_select_polygon.py index 08ea099d..351db595 100644 --- a/tests/test_select_polygon.py +++ b/tests/test_select_polygon.py @@ -142,5 +142,9 @@ def test_polygon_vertex_and_row_validation(): with pytest.raises(ValueError): kernels.polygon_select(x, y, rows, square_x, square_y[:3]) # A row id past the end of the column is a caller bug, not a silent drop. + # The entry point checks the ids itself (`rows_in_bounds`) rather than + # leaving it to the kernel's indexing panic, so this holds on panic-abort + # targets too — under the PyEmscripten wheel's `-C panic=abort` the panic + # aborted the interpreter instead of raising. with pytest.raises(ValueError): kernels.polygon_select(x, y, np.array([99], dtype=np.uint32), square_x, square_y) From ace996be94848e0fc85283e40e060918ece7458b Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Sun, 26 Jul 2026 16:10:07 -0700 Subject: [PATCH 3/5] Take ABI 45: 43 shipped with #319 and #327 claims 44 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit claiming 43 for `xy_range_indices_rows`, and open PR #327 already claims 44. Both edits touch the same two lines with the same value, so such a merge is clean and nothing complains — which is the problem. One number would name two different ABIs, and `scripts/abi_smoke.py` compares the built library against the wrapper's constant, so it cannot catch a collision where both sides agree: a cdylib built from either tree would load against the other's wrapper and fail at dlsym instead of at the version check that exists to prevent exactly that. 45 is unclaimed across every open PR. It also fails loudly rather than silently if #327 lands second, because its 43 -> 44 edit then conflicts against a main at 45 instead of merging into agreement. --- python/xy/_native.py | 2 +- src/lib.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/xy/_native.py b/python/xy/_native.py index 509ccea5..30d7e0c6 100644 --- a/python/xy/_native.py +++ b/python/xy/_native.py @@ -24,7 +24,7 @@ from .config import MAX_CONTOUR_WORK, MAX_SCREEN_DIM -ABI_VERSION = 44 +ABI_VERSION = 45 # Rust reports invalid arguments (and, via the ffi_guard panic shield, any # internal panic) by returning `usize::MAX` from size-returning entry points. diff --git a/src/lib.rs b/src/lib.rs index ef90df96..109fafda 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -83,7 +83,7 @@ unsafe fn borrowed_byte_spans<'a>( /// ABI version — bumped on any signature change. The Python wrapper checks this /// at load time and refuses a mismatched library loudly (§33 comm-versioning /// rule, applied to the in-process boundary). -pub const ABI_VERSION: u32 = 44; +pub const ABI_VERSION: u32 = 45; const FACTORIZE_CAPACITY_EXCEEDED: usize = usize::MAX - 1; #[no_mangle] From 6c5e775617de94c497dff92226520f65232ae95d Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Sun, 26 Jul 2026 16:34:01 -0700 Subject: [PATCH 4/5] Range-check row ids before the u32 cast, not after MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The out-of-range guarantee the previous commit made stops at the ctypes boundary: the wrapper reached the kernel through `ascontiguousarray(rows, dtype=np.uint32)`, an unchecked C cast. An id of `2**32 + 3` arrives as row 3 — in range, indistinguishable from a real id, so the bounds check passes and the call returns a row the caller never asked for. `kernels.range_indices_rows(x, y, [2**32 + 3], ...)` answered `[3]`; a negative id only errored by accident, because `-1` wraps to 4294967295 and that happened to exceed the column length. Both row-restricted kernels now widen and range-check before casting. The internal caller passes uint32 already (`_expand_zone_chunks`), which takes the pass-through branch, so the hot path is untouched — 5M candidates still 1.00 ms. --- python/xy/_native.py | 31 +++++++++++++++++++++++++------ tests/test_range_indices_rows.py | 30 ++++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 6 deletions(-) diff --git a/python/xy/_native.py b/python/xy/_native.py index 30d7e0c6..3fdf3d29 100644 --- a/python/xy/_native.py +++ b/python/xy/_native.py @@ -34,6 +34,8 @@ # would be sliced as data. _USIZE_MAX = ctypes.c_size_t(-1).value _FACTORIZE_CAPACITY_EXCEEDED = _USIZE_MAX - 1 +# Canonical row ids ship as u32 (the index ceiling the kernels enforce). +_U32_MAX = 2**32 - 1 def _lib_filename() -> str: @@ -830,6 +832,27 @@ def _as_f64(arr: npt.NDArray[np.float64], label: str = "data") -> npt.NDArray[np return out +def _as_row_ids(rows: npt.NDArray[np.uint32], label: str = "rows") -> npt.NDArray[np.uint32]: + """Canonical row ids as contiguous u32, rejecting anything that would wrap. + + `ascontiguousarray(..., dtype=np.uint32)` is an unchecked C cast: an id of + `2**32 + 3` becomes 3 and a negative id becomes a huge one. The kernels + bounds-check the ids they are *given*, so a wrapped id is indistinguishable + from a real one there — the call would answer with a row the caller never + asked for instead of returning the error sentinel. Range-check before the + cast so the out-of-range contract holds for the value the caller passed. + """ + out = np.ascontiguousarray(rows) + if out.ndim != 1: + raise ValueError(f"{label} must be 1-D, got shape {out.shape}") + if out.dtype == np.uint32: + return out + widened = out.astype(np.int64, copy=False) + if widened.size and (widened.min() < 0 or widened.max() > _U32_MAX): + raise ValueError(f"{label} must be canonical row ids in [0, 2**32)") + return np.ascontiguousarray(widened, dtype=np.uint32) + + def _ptr_f64(arr: npt.NDArray[np.float64]) -> int: # Raw address int for a c_void_p parameter: ~2x cheaper per call than # `ctypes.data_as(...)`, which allocates a fresh pointer object. The @@ -2658,9 +2681,7 @@ def range_indices_rows( y = _as_f64(y, "y") if len(x) != len(y): raise ValueError("x and y must have equal length") - rows = np.ascontiguousarray(rows, dtype=np.uint32) - if rows.ndim != 1: - raise ValueError("rows must be 1-D") + rows = _as_row_ids(rows) out = np.empty(len(rows), dtype=np.uint32) if len(rows) == 0: return out @@ -2700,9 +2721,7 @@ def polygon_select( poly_y = _as_f64(poly_y, "polygon y") if len(poly_x) != len(poly_y): raise ValueError("polygon x and y must have equal length") - rows = np.ascontiguousarray(rows, dtype=np.uint32) - if rows.ndim != 1: - raise ValueError("rows must be 1-D") + rows = _as_row_ids(rows) out = np.empty(len(rows), dtype=np.uint32) if len(rows) == 0: return out diff --git a/tests/test_range_indices_rows.py b/tests/test_range_indices_rows.py index 040dbca7..1a752533 100644 --- a/tests/test_range_indices_rows.py +++ b/tests/test_range_indices_rows.py @@ -123,6 +123,36 @@ def test_out_of_range_row_is_an_error_not_a_read() -> None: kernels.range_indices_rows(x, x, np.array([4], dtype=np.uint32), -1.0, 1.0, -1.0, 1.0) +@pytest.mark.parametrize("bad", [2**32, 2**32 + 3, 2**40, -1, -(2**31)]) +def test_a_row_id_that_would_wrap_is_rejected_not_cast(bad: int) -> None: + """The u32 conversion is an unchecked C cast, so `2**32 + 3` would arrive as + row 3 — in range, and therefore answered rather than refused. The caller + would get a row it never asked for, which is worse than the error the + kernel's bounds check exists to produce.""" + x = np.arange(10.0) + y = np.zeros(10) + rows = np.array([bad], dtype=np.int64) + square_x = np.array([-1e9, 1e9, 1e9, -1e9]) + square_y = np.array([-1e9, -1e9, 1e9, 1e9]) + with pytest.raises(ValueError, match=r"canonical row ids in \[0, 2\*\*32\)"): + kernels.range_indices_rows(x, y, rows, -1e9, 1e9, -1e9, 1e9) + with pytest.raises(ValueError, match=r"canonical row ids in \[0, 2\*\*32\)"): + kernels.polygon_select(x, y, rows, square_x, square_y) + + +def test_in_range_ids_survive_every_integer_dtype() -> None: + """The range check must not start rejecting ordinary callers: any integer + dtype holding valid ids still answers, and a uint32 array is passed through + untouched.""" + x = np.arange(10.0) + y = np.zeros(10) + want = np.array([2, 5], dtype=np.uint32) + for dtype in (np.uint32, np.int64, np.int32, np.uint64, np.int16, np.uint8): + got = kernels.range_indices_rows(x, y, np.array([2, 5], dtype=dtype), -1e9, 1e9, -1e9, 1e9) + np.testing.assert_array_equal(got, want) + assert got.dtype == np.uint32 + + def test_select_range_agrees_across_the_pruning_branches() -> None: """Zone-pruned and whole-column selection must return the same rows. From c24f1b2f0273400410736362216bd105da2f15c1 Mon Sep 17 00:00:00 2001 From: Alek Petuskey Date: Sun, 26 Jul 2026 16:39:14 -0700 Subject: [PATCH 5/5] Refuse float row ids instead of casting them The range check widened through `astype(np.int64)`, which still let floats through silently and, worse, decided them differently per architecture: `[3.9]` selected row 3, and `[nan]` selected row 0 on arm64 because the NaN cast saturates, where x86_64 traps it to INT64_MIN and the range check would have caught it. A guard whose answer depends on the wheel is not a guard. The cast also emitted a RuntimeWarning from inside the wrapper. Integer dtypes only now, with an empty array exempted because `np.asarray([])` is float64 and selecting no rows is not a dtype error. --- python/xy/_native.py | 13 ++++++++++++- tests/test_range_indices_rows.py | 25 +++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/python/xy/_native.py b/python/xy/_native.py index 3fdf3d29..0c91c6d0 100644 --- a/python/xy/_native.py +++ b/python/xy/_native.py @@ -841,14 +841,25 @@ def _as_row_ids(rows: npt.NDArray[np.uint32], label: str = "rows") -> npt.NDArra from a real one there — the call would answer with a row the caller never asked for instead of returning the error sentinel. Range-check before the cast so the out-of-range contract holds for the value the caller passed. + + Integer dtypes only, and deliberately so: a float id has no row, and + deciding one by cast is both silent and platform-dependent. `3.9` would + become row 3 and `nan` would become row 0 on arm64, where the NaN cast + saturates, while the same `nan` traps as `INT64_MIN` on x86_64 — the guard + itself would disagree across the wheels we ship. """ out = np.ascontiguousarray(rows) if out.ndim != 1: raise ValueError(f"{label} must be 1-D, got shape {out.shape}") if out.dtype == np.uint32: return out + if out.size == 0: + # `np.asarray([])` is float64; selecting no rows is not a dtype error. + return np.empty(0, dtype=np.uint32) + if not np.issubdtype(out.dtype, np.integer): + raise ValueError(f"{label} must be an integer array of row ids, got dtype {out.dtype}") widened = out.astype(np.int64, copy=False) - if widened.size and (widened.min() < 0 or widened.max() > _U32_MAX): + if widened.min() < 0 or widened.max() > _U32_MAX: raise ValueError(f"{label} must be canonical row ids in [0, 2**32)") return np.ascontiguousarray(widened, dtype=np.uint32) diff --git a/tests/test_range_indices_rows.py b/tests/test_range_indices_rows.py index 1a752533..525823cd 100644 --- a/tests/test_range_indices_rows.py +++ b/tests/test_range_indices_rows.py @@ -140,6 +140,31 @@ def test_a_row_id_that_would_wrap_is_rejected_not_cast(bad: int) -> None: kernels.polygon_select(x, y, rows, square_x, square_y) +@pytest.mark.parametrize("bad", [[3.9], [np.nan], [np.inf], [-0.5], [3.0]]) +def test_a_float_row_id_is_refused_rather_than_rounded(bad: list[float]) -> None: + """A float id has no row, and choosing one by cast is platform-dependent: + `nan` saturates to 0 on arm64 but traps to INT64_MIN on x86_64, so a cast + would make the guard disagree across the wheels we ship. `3.9` silently + selecting row 3 is the same class of wrong answer as a wrapped id.""" + x = np.arange(8.0) + y = np.arange(8.0) + rows = np.asarray(bad) + assert rows.dtype == np.float64 + with pytest.raises(ValueError, match="must be an integer array of row ids"): + kernels.range_indices_rows(x, y, rows, 0.0, 8.0, 0.0, 8.0) + with pytest.raises(ValueError, match="must be an integer array of row ids"): + kernels.polygon_select( + x, y, rows, np.array([-1e9, 1e9, 1e9, -1e9]), np.array([-1e9, -1e9, 1e9, 1e9]) + ) + + +def test_an_empty_row_list_is_not_a_dtype_error() -> None: + """`np.asarray([])` is float64, but selecting nothing is legitimate.""" + x = np.arange(8.0) + got = kernels.range_indices_rows(x, x, np.asarray([]), 0.0, 8.0, 0.0, 8.0) + assert got.size == 0 and got.dtype == np.uint32 + + def test_in_range_ids_survive_every_integer_dtype() -> None: """The range check must not start rejecting ordinary callers: any integer dtype holding valid ids still answers, and a uint32 array is passed through