diff --git a/python/xy/_native.py b/python/xy/_native.py index 9f514213..0c91c6d0 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. @@ -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: @@ -558,6 +560,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 @@ -817,6 +832,38 @@ 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. + + 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.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 @@ -2623,6 +2670,49 @@ 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 = _as_row_ids(rows) + 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], @@ -2642,9 +2732,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/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..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 @@ -5062,6 +5069,102 @@ 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], +) -> Option { + 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() + }); + // 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]; + for (t, &c) in counts.iter().enumerate().skip(1) { + let start = t * chunk; + out.copy_within(start..start + c, write); + write += c; + } + 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], + y: &[f64], + rows: &[u32], + lo_x: f64, + hi_x: f64, + lo_y: f64, + hi_y: f64, + out: &mut [u32], +) -> Option { + let mut n = 0usize; + for &row in rows { + let i = row as usize; + 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; + } + } + Some(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. @@ -5310,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)[..], @@ -5379,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 d6ecc734..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] @@ -2872,9 +2872,57 @@ 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 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 +/// 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).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 @@ -2919,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 new file mode 100644 index 00000000..525823cd --- /dev/null +++ b/tests/test_range_indices_rows.py @@ -0,0 +1,199 @@ +"""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. + + 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) + + +@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) + + +@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 + 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. + + `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(): 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)