diff --git a/src/vsparse/_base.py b/src/vsparse/_base.py index c6cde21..6a116e1 100644 --- a/src/vsparse/_base.py +++ b/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 diff --git a/src/vsparse/_ops.py b/src/vsparse/_ops.py index 6a686c2..3827803 100644 --- a/src/vsparse/_ops.py +++ b/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. diff --git a/tests/test_index_dtypes.py b/tests/test_index_dtypes.py index d4b546c..a062ace 100644 --- a/tests/test_index_dtypes.py +++ b/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", [