Scan selection candidates in place (ABI 42 -> 43)

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

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
 

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",

src/kernels.rs

@@ -4807,20 +4807,24 @@ 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],
     rows: &[u32],
     poly_x: &[f64],
     poly_y: &[f64],
     out: &mut [u32],
-) -> usize {
+) -> Option<usize> {
     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<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<Option<usize>> = 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<Vec<usize>> = 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<usize> {
+    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<u32> = (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