Minor sums memory

src/vsparse/_base.py

@@ -227,10 +227,11 @@ def _major_sums(self) -> np.ndarray:
         return np.bincount(group_of_major, weights=weighted, minlength=self.n_major)
 
     def _minor_sums(self) -> np.ndarray:
-        """Per-minor-index totals -- a scatter-add over every nonzero."""
-        group_sizes = np.diff(self.value_ptr)
-        expanded = np.repeat(self.values.astype(np.float64), group_sizes)
-        return np.bincount(self.indices, weights=expanded, minlength=self.n_minor)
+        """A parallel scatter-add over every nonzero to get per-minor-index totals."""
+        return _ops.minor_sums(
+            self.values, self.value_ptr, self.indices, self.n_minor,
+            _ops.accumulator_threads(self.n_minor),
+        )
 
     def sum(self, axis: int | None = None) -> np.ndarray | float:
         """Sum of (structural) values along ``axis`` (0=rows, 1=columns), or overall if ``None``."""
@@ -252,8 +253,11 @@ def _major_nnz(self) -> np.ndarray:
         )
 
     def _minor_nnz(self) -> np.ndarray:
-        """Per-minor-index stored-element counts -- a scatter-add over every nonzero."""
-        return np.bincount(self.indices, minlength=self.n_minor).astype(np.int64)
+        """A parallel scatter over every nonzero to get per-minor-index counts."""
+        return _ops.minor_counts(
+            self.value_ptr, self.indices, self.n_minor,
+            _ops.accumulator_threads(self.n_minor),
+        )
 
     def getnnz(self, axis: int | None = None) -> np.ndarray | int:
         """Count of stored elements along ``axis``, or overall if ``None``."""
@@ -298,11 +302,11 @@ def _major_reduce(self, ufunc: np.ufunc, initial: Any) -> np.ndarray:
 
     def _minor_reduce(self, ufunc: np.ufunc, initial: Any) -> np.ndarray:
         """Per-minor-index max/min, accounting for implicit zeros in sparse slices."""
-        group_sizes = np.diff(self.value_ptr)
-        expanded = np.repeat(self.values, group_sizes)
-        out = np.full(self.n_minor, initial, dtype=self.values.dtype)
-        ufunc.at(out, self.indices, expanded)
-        counts = np.bincount(self.indices, minlength=self.n_minor)
+        out, counts = _ops.minor_extrema(
+            self.values, self.value_ptr, self.indices, self.n_minor, initial, ufunc is np.maximum
+        )
+        # A minor index missing from some major slice has a structural zero
+        # competing in the reduction.
         not_dense = counts < self.n_major
         out[not_dense] = ufunc(out[not_dense], 0)
         return out

src/vsparse/_ops.py

@@ -6,14 +6,31 @@
 import numpy as np
 
 __all__ = [
+    "accumulator_threads",
     "major_matmat",
     "major_matvec",
+    "minor_counts",
+    "minor_extrema",
     "minor_matmat",
     "minor_matvec",
     "minor_select_counts",
     "minor_select_fill",
+    "minor_sums",
 ]
 
+# Thread-local accumulators for the minor-axis scatters below cost
+# ``nthreads * n_minor * bytes_per_element``. The thread count is capped to
+# keep that block under this budget.
+_ACCUMULATOR_BUDGET_BYTES = 64 << 20  # 64 MiB
+
+
+def accumulator_threads(n_minor: int, bytes_per_element: int = 8) -> int:
+    """Threads to run a minor-axis scatter with, capped by accumulator size."""
+    if n_minor <= 0:
+        return 1
+    affordable = max(1, _ACCUMULATOR_BUDGET_BYTES // (n_minor * max(1, bytes_per_element)))
+    return int(min(numba.get_num_threads(), affordable))
+
 
 @numba.njit(cache=True)
 def _major_matvec(major_ptr, values, value_ptr, indices, x, n_major, n_minor):
@@ -82,6 +99,76 @@ def _minor_matmat(major_ptr, values, value_ptr, indices, b, n_major):
     return y
 
 
+@numba.njit(cache=True, parallel=True)
+def minor_sums(values, value_ptr, indices, n_minor, nthreads):
+    n_groups = values.shape[0]
+    chunk = (n_groups + nthreads - 1) // nthreads
+    partial = np.zeros((nthreads, n_minor), dtype=np.float64)
+    for t in numba.prange(nthreads):  # ty: ignore[not-iterable]
+        start = t * chunk
+        end = min(n_groups, start + chunk)
+        local = partial[t]
+        for u in range(start, end):
+            v = np.float64(values[u])
+            for k in range(value_ptr[u], value_ptr[u + 1]):
+                local[indices[k]] += v
+    return partial.sum(axis=0)
+
+
+@numba.njit(cache=True, parallel=True)
+def _minor_extrema_kernel(values, value_ptr, indices, n_minor, nthreads, initial, is_max):
+    # The count comes along for free and tells the caller which minor indices
+    # have an implicit zero to fold in.
+    n_groups = values.shape[0]
+    chunk = (n_groups + nthreads - 1) // nthreads
+    part_val = np.full((nthreads, n_minor), initial, dtype=values.dtype)
+    part_cnt = np.zeros((nthreads, n_minor), dtype=np.int64)
+    for t in numba.prange(nthreads):  # ty: ignore[not-iterable]
+        start = t * chunk
+        end = min(n_groups, start + chunk)
+        local_val = part_val[t]
+        local_cnt = part_cnt[t]
+        for u in range(start, end):
+            v = values[u]
+            for k in range(value_ptr[u], value_ptr[u + 1]):
+                idx = indices[k]
+                local_cnt[idx] += 1
+                if is_max:
+                    if v > local_val[idx]:
+                        local_val[idx] = v
+                else:
+                    if v < local_val[idx]:
+                        local_val[idx] = v
+    return part_val, part_cnt
+
+
+def minor_extrema(values, value_ptr, indices, n_minor, initial, is_max):
+    """Extremum over the stored values, and the stored count, per minor index."""
+    nthreads = accumulator_threads(n_minor, values.dtype.itemsize + 8)
+    part_val, part_cnt = _minor_extrema_kernel(
+        values, value_ptr, indices, n_minor, nthreads, initial, is_max
+    )
+    extrema = part_val.max(axis=0) if is_max else part_val.min(axis=0)
+    return extrema, part_cnt.sum(axis=0)
+
+
+@numba.njit(cache=True, parallel=True)
+def minor_counts(value_ptr, indices, n_minor, nthreads):
+    n_groups = value_ptr.shape[0] - 1
+    chunk = (n_groups + nthreads - 1) // nthreads
+    partial = np.zeros((nthreads, n_minor), dtype=np.int64)
+    for t in numba.prange(nthreads):  # ty: ignore[not-iterable]
+        start = t * chunk
+        end = min(n_groups, start + chunk)
+        local = partial[t]
+        for u in range(start, end):
+            for k in range(value_ptr[u], value_ptr[u + 1]):
+                local[indices[k]] += 1
+    return partial.sum(axis=0)
+
+
+
+
 def major_matvec(major_ptr, values, value_ptr, indices, x, n_major, n_minor):
     return _major_matvec(major_ptr, values, value_ptr, indices, np.asarray(x), n_major, n_minor)
 

src/vsparse/_rapid_load.py

@@ -62,6 +62,7 @@
 from vsparse import _ivcsc
 from vsparse._anndata_class import VCSCAnnData
 from vsparse._indexutils import smallest_index_dtype
+from vsparse._ops import accumulator_threads
 
 if TYPE_CHECKING:
     from os import PathLike
@@ -469,11 +470,11 @@ def load_and_normalize(
         data = _build_data(values, value_ptr, indices.shape[0])
         row_indptr = value_ptr[major_ptr]
 
-        gene_totals_raw = _weighted_bincount(indices, data, n_genes, numba.get_num_threads())
+        gene_totals_raw = _weighted_bincount(indices, data, n_genes, accumulator_threads(n_genes))
         gene_mask = gene_totals_raw > (gene_threshold * n_cells)
         if min_cells is not None:
             gene_detection_counts = _gene_detection_counts(
-                indices, data, n_genes, numba.get_num_threads()
+                indices, data, n_genes, accumulator_threads(n_genes)
             )
             gene_mask &= gene_detection_counts >= min_cells
         metadata_cell_mask = cell_mask
@@ -500,11 +501,11 @@ def load_and_normalize(
         )
         del packed
 
-        gene_totals_raw = _weighted_bincount(indices, data, n_genes, numba.get_num_threads())
+        gene_totals_raw = _weighted_bincount(indices, data, n_genes, accumulator_threads(n_genes))
         gene_mask = gene_totals_raw > (gene_threshold * selected_rows.shape[0])
         if min_cells is not None:
             gene_detection_counts = _gene_detection_counts(
-                indices, data, n_genes, numba.get_num_threads()
+                indices, data, n_genes, accumulator_threads(n_genes)
             )
             gene_mask &= gene_detection_counts >= min_cells
 

tests/test_reduction_memory.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+import tracemalloc
+
+import numba
+import numpy as np
+import pytest
+import scipy.sparse as sp
+
+from vsparse import VCSCArray, VCSRArray
+from vsparse._ops import _ACCUMULATOR_BUDGET_BYTES, accumulator_threads
+
+
+@pytest.fixture(params=[VCSCArray, VCSRArray])
+def vcls(request):
+    return request.param
+
+
+def _scipy_for(vcls, dense):
+    return sp.csc_array(dense) if vcls is VCSCArray else sp.csr_array(dense)
+
+
+def _minor_axis(vcls):
+    return 1 if vcls is VCSCArray else 0
+
+
+@pytest.mark.parametrize("axis", [None, 0, 1])
+def test_sum_matches_dense(dense, vcls, axis):
+    """Totals along each axis, and overall."""
+    v = vcls.from_scipy(_scipy_for(vcls, dense))
+    np.testing.assert_allclose(v.sum(axis=axis), dense.sum(axis=axis))
+
+
+@pytest.mark.parametrize("axis", [0, 1])
+@pytest.mark.parametrize("kind", ["max", "min"])
+def test_extrema_match_dense(dense, vcls, axis, kind):
+    """Extrema have to fold in the implicit zeros the layout never stores."""
+    v = vcls.from_scipy(_scipy_for(vcls, dense))
+    np.testing.assert_allclose(getattr(v, kind)(axis=axis), getattr(dense, kind)(axis=axis))
+
+
+@pytest.mark.parametrize("axis", [0, 1])
+def test_getnnz_matches_dense(dense, vcls, axis):
+    v = vcls.from_scipy(_scipy_for(vcls, dense))
+    np.testing.assert_array_equal(v.getnnz(axis=axis), (dense != 0).sum(axis=axis))
+
+
+def test_reductions_on_an_all_zero_array(vcls):
+    """Every entry is an implicit zero, so the stored values are empty."""
+    v = vcls.from_scipy(_scipy_for(vcls, np.zeros((4, 3))))
+    np.testing.assert_allclose(v.max(axis=0), np.zeros(3))
+    np.testing.assert_allclose(v.min(axis=0), np.zeros(3))
+    np.testing.assert_allclose(v.sum(axis=0), np.zeros(3))
+    np.testing.assert_array_equal(v.getnnz(axis=0), np.zeros(3, dtype=np.int64))
+
+
+def test_integer_values_use_integer_sentinels(vcls):
+    """A float sentinel would make an integer max come back wrong or upcast."""
+    dense = np.array([[7, 0, 2], [0, 3, 9]], dtype=np.int32)
+    v = vcls.from_scipy(_scipy_for(vcls, dense.astype(np.float64)))
+    np.testing.assert_array_equal(v.max(axis=0), dense.max(axis=0))
+    np.testing.assert_array_equal(v.min(axis=0), dense.min(axis=0))
+
+
+def test_sums_accumulate_in_float64(vcls):
+    """Accumulating in the stored dtype would overflow or lose precision."""
+    v = vcls.from_scipy(_scipy_for(vcls, np.array([[1.0, 0.0, 2.0], [3.0, 4.0, 0.0]])))
+    assert v._minor_sums().dtype == np.float64
+
+
+@pytest.mark.parametrize("n_minor", [0, 1_000, 10**8])
+@pytest.mark.parametrize("bytes_per_element", [8, 16])
+def test_accumulator_block_stays_within_budget(n_minor, bytes_per_element):
+    """Thread-local accumulators must not become the new unbounded allocation."""
+    nthreads = accumulator_threads(n_minor, bytes_per_element)
+    assert 1 <= nthreads <= numba.get_num_threads()
+    if nthreads > 1:
+        assert nthreads * n_minor * bytes_per_element <= _ACCUMULATOR_BUDGET_BYTES
+
+
+@pytest.mark.parametrize(
+    ("label", "call"),
+    [
+        ("sum", lambda v: v.sum(axis=0)),
+        ("max", lambda v: v.max(axis=0)),
+        ("getnnz", lambda v: v.getnnz(axis=0)),
+    ],
+)
+def test_minor_axis_reductions_allocate_nothing_nnz_sized(label, call):
+    """An n_minor-sized result must not cost nnz-sized scratch."""
+    rng = np.random.default_rng(0)
+    n_rows, n_cols = 2_000, 500
+    dense = rng.integers(1, 5, size=(n_rows, n_cols)).astype(np.float64)
+    v = VCSRArray.from_scipy(sp.csr_array(dense))
+
+    call(v)  # warm up the JIT before measuring
+
+    tracemalloc.start()
+    try:
+        before = tracemalloc.get_traced_memory()[0]
+        tracemalloc.reset_peak()
+        out = call(v)
+        peak = tracemalloc.get_traced_memory()[1]
+    finally:
+        tracemalloc.stop()
+
+    assert peak - before < 1 << 20, f"{label} allocated {(peak - before) / 1e6:.1f} MB"
+    assert out.shape == (n_cols,)