Check VCS index bounds in one parallel pass

src/vsparse/_base.py

@@ -66,7 +66,7 @@ def __init__(
             raise ValueError("major_ptr[-1] must equal len(values)")
         if value_ptr[-1] != indices.shape[0]:
             raise ValueError("value_ptr[-1] must equal len(indices)")
-        if indices.shape[0] and (indices.min() < 0 or indices.max() >= n_minor):
+        if not _ops.indices_in_bounds(indices, n_minor):
             raise ValueError("indices out of bounds for the given shape")
 
         # Narrow after the bounds check above, so an out-of-range index is

src/vsparse/_ops.py

@@ -9,6 +9,7 @@
 
 __all__ = [
     "accumulator_threads",
+    "indices_in_bounds",
     "major_matmat",
     "major_matvec",
     "minor_counts",
@@ -34,6 +35,34 @@ def accumulator_threads(n_minor: int, bytes_per_element: int = 8) -> int:
     return int(min(numba.get_num_threads(), affordable))
 
 
+# Below this many indices, numpy's min/max beat the parallel kernel's launch cost.
+_PARALLEL_BOUNDS_MIN = 1 << 20
+
+
+@numba.njit(cache=True, parallel=True)
+def _count_out_of_bounds(indices: np.ndarray, n: int) -> int:
+    bad = 0
+    for k in numba.prange(indices.shape[0]):  # ty: ignore[not-iterable]
+        v = indices[k]
+        if v < 0 or v >= n:
+            bad += 1
+    return bad
+
+
+def indices_in_bounds(indices: np.ndarray, n: int) -> bool:
+    """True if every entry of ``indices`` lies in ``[0, n)``.
+
+    Large signed-integer arrays are checked in one parallel pass; numpy's
+    ``min()``/``max()`` take two single-threaded passes, which costs over a
+    second on a few billion indices.
+    """
+    if indices.shape[0] == 0:
+        return True
+    if indices.shape[0] >= _PARALLEL_BOUNDS_MIN and indices.dtype.kind == "i":
+        return _count_out_of_bounds(indices, n) == 0
+    return bool(indices.min() >= 0 and indices.max() < n)
+
+
 #: Tighter than `_ACCUMULATOR_BUDGET_BYTES`, because this scatter also pays to
 #: reduce its partials afterwards: a budget that admits enough threads to make
 #: that reduction dominate is slower than not parallelizing at all.

tests/test_index_dtypes.py

@@ -4,7 +4,7 @@
 import pytest
 import scipy.sparse as sp
 
-from vsparse import VCSCAnnData, VCSCArray, VCSRArray
+from vsparse import VCSCAnnData, VCSCArray, VCSRArray, _ops
 from vsparse._indexutils import smallest_index_dtype
 
 INT32_MAX = np.iinfo(np.int32).max
@@ -81,6 +81,38 @@ def test_out_of_bounds_index_raises_rather_than_truncating(vcls):
         )
 
 
+@pytest.mark.parametrize("dtype", [np.int32, np.int64])
+@pytest.mark.parametrize("bad", [None, -1, 5000, 2**31 - 1])
+def test_parallel_bounds_check(dtype, bad):
+    """Large index arrays take the one-pass parallel check; it agrees with min/max."""
+    n = 5000
+    if bad is not None and bad > np.iinfo(dtype).max:
+        pytest.skip("value does not fit the dtype")
+    idx = np.random.default_rng(0).integers(0, n, _ops._PARALLEL_BOUNDS_MIN + 3).astype(dtype)
+    idx[-1] = n - 1
+    if bad is not None:
+        idx[len(idx) // 2] = bad
+    assert _ops.indices_in_bounds(idx, n) == (bad is None)
+    assert _ops.indices_in_bounds(idx, n) == bool(idx.min() >= 0 and idx.max() < n)
+
+
+def test_large_out_of_bounds_index_raises(vcls):
+    """The constructor rejects an out-of-range index found by the parallel check."""
+    nnz = _ops._PARALLEL_BOUNDS_MIN + 1
+    n_minor = 10
+    shape = (n_minor, 1) if vcls is VCSCArray else (1, n_minor)
+    indices = np.zeros(nnz, dtype=np.int64)
+    indices[123] = n_minor
+    with pytest.raises(ValueError, match="out of bounds"):
+        vcls(
+            shape,
+            major_ptr=np.array([0, 1], dtype=np.int64),
+            values=np.array([1.0]),
+            value_ptr=np.array([0, nnz], dtype=np.int64),
+            indices=indices,
+        )
+
+
 @pytest.mark.parametrize(
     "op",
     [