From 23526145402b93906ff6d6323aad3156de145e7e Mon Sep 17 00:00:00 2001 From: fishidaho Date: Thu, 3 Sep 2026 11:25:51 -0700 Subject: [PATCH 1/2] Fix silent data loss when a minor-axis selection repeats an index `arr[:, [1, 1, 3]]` returns zeros where column 1's second copy should be. No error, no warning -- just wrong values, in a shape that looks right. v[:, [1, 1, 3]] want got [[ 2 2 4] [[ 0 2 4] [ 6 6 8] [ 0 6 8] [10 10 12]] [ 0 10 12]] This is a regression. Before #23 the same expression fell through to `to_scipy()[...]`, and scipy fans duplicate fancy indices out correctly; verified against 45ce15d, where it returns the right answer. The cause is structural rather than an off-by-one: `_select_minor` builds a single old -> new lookup array, remap = np.full(self.n_minor, -1); remap[idx] = np.arange(len(idx)) and a repeated index writes that slot twice, so only its last destination survives. One array cannot express one-to-many, so the fix is to invert the selection instead: `fanout`/`offsets`/`positions` give, for each original minor index, every output position it maps to, and each stored element emits one entry per destination. Two passes (count, then fill) via `_ops.minor_select_counts`/`minor_select_fill`, so each surviving slot writes a disjoint range and the passes parallelize. Deriving each major slice's slot count with `np.searchsorted` on `major_ptr` also drops the `np.repeat(np.arange(n_unique), ...)` this used to build, which was `nnz`-sized. Peak allocation for a 2e6-nonzero selection falls from 62.27 MB to 4.78 MB, so this closes ISSUE-30 as well -- not a separate concern, just what the correct implementation happens not to need. All 1066 tests from main pass unchanged. New coverage coming from the fan-out semantics: five duplicate patterns, duplicates on either axis and both at once, 25 random selections with repeats compared against scipy directly, plus the selections that already worked (sorted, reversed, empty, negative, slices, boolean masks) so the fix can't silently narrow them. Co-Authored-By: Claude Sonnet 5 --- src/vsparse/_base.py | 55 ++++++---- src/vsparse/_ops.py | 63 ++++++++++- tests/test_select_minor_fanout.py | 177 ++++++++++++++++++++++++++++++ 3 files changed, 275 insertions(+), 20 deletions(-) create mode 100644 tests/test_select_minor_fanout.py diff --git a/src/vsparse/_base.py b/src/vsparse/_base.py index 9a3ed29..e2f4038 100644 --- a/src/vsparse/_base.py +++ b/src/vsparse/_base.py @@ -516,30 +516,47 @@ def _select_minor(self, key: Any) -> _VCSBase: """Select along the minor axis (rows for VCSC, columns for VCSR). Unlike :meth:`_select_major`, the kept elements aren't already - contiguous per major slice, so this filters/remaps ``indices`` and - drops any (major, unique-value) slot that no longer has any kept - index, shrinking ``major_ptr``/``value_ptr`` accordingly. + contiguous per major slice, so ``indices`` has to be filtered and + remapped, and any (major, unique-value) slot left with no kept index + dropped -- shrinking ``major_ptr``/``value_ptr`` accordingly. + + The selection is a *fan-out*, not a one-to-one remap: a key may name + the same minor index more than once (``arr[:, [1, 1, 3]]``, which + scipy supports), so one stored element can have to appear in several + output positions. It is inverted once into a CSR-shaped map over the + original minor axis and applied by + :func:`vsparse._ops.minor_select_counts`/ + :func:`~vsparse._ops.minor_select_fill`, whose intermediates are + sized by the axis and the selection rather than by ``nnz``. """ idx = _normalize_major_idx(key, self.n_minor) n_minor_new = idx.shape[0] - remap = np.full(self.n_minor, -1, dtype=np.int64) - remap[idx] = np.arange(n_minor_new, dtype=np.int64) - - keep = remap[self.indices] >= 0 - new_indices = remap[self.indices[keep]].astype(self.indices.dtype, copy=False) - - n_unique = self.values.shape[0] - value_slot_of_index = np.repeat(np.arange(n_unique, dtype=np.int64), np.diff(self.value_ptr)) - kept_per_slot = np.bincount(value_slot_of_index[keep], minlength=n_unique) - surviving = kept_per_slot > 0 - - new_values = self.values[surviving] - new_value_ptr = np.zeros(int(surviving.sum()) + 1, dtype=np.int64) - np.cumsum(kept_per_slot[surviving], out=new_value_ptr[1:]) + # Invert the selection: output positions grouped by the original index + # they came from, so a repeated index carries all of its destinations. + fanout = np.bincount(idx, minlength=self.n_minor) + offsets = np.zeros(self.n_minor + 1, dtype=np.int64) + np.cumsum(fanout, out=offsets[1:]) + positions = np.argsort(idx, kind="stable").astype(np.int64, copy=False) + + slot_counts = _ops.minor_select_counts(self.value_ptr, self.indices, fanout) + kept_slots = np.flatnonzero(slot_counts) + + new_values = self.values[kept_slots] + new_value_ptr = np.zeros(kept_slots.shape[0] + 1, dtype=np.int64) + np.cumsum(slot_counts[kept_slots], out=new_value_ptr[1:]) + + # Keep the parent's index dtype, as every other structural op does. + new_indices = np.empty(int(new_value_ptr[-1]), dtype=self.indices.dtype) + _ops.minor_select_fill( + self.value_ptr, self.indices, offsets, positions, + kept_slots, new_value_ptr, new_indices, + ) - group_of_major = np.repeat(np.arange(self.n_major, dtype=np.int64), np.diff(self.major_ptr)) - major_counts = np.bincount(group_of_major[surviving], minlength=self.n_major) + # Slots stay in their original order, so each major slice keeps a + # contiguous run of them; count how many of its slots survived. + major_of_slot = np.searchsorted(self.major_ptr, kept_slots, side="right") - 1 + major_counts = np.bincount(major_of_slot, minlength=self.n_major) new_major_ptr = np.zeros(self.n_major + 1, dtype=np.int64) np.cumsum(major_counts, out=new_major_ptr[1:]) diff --git a/src/vsparse/_ops.py b/src/vsparse/_ops.py index 6aded8a..acbf3b9 100644 --- a/src/vsparse/_ops.py +++ b/src/vsparse/_ops.py @@ -5,7 +5,14 @@ import numba import numpy as np -__all__ = ["major_matmat", "major_matvec", "minor_matmat", "minor_matvec"] +__all__ = [ + "major_matmat", + "major_matvec", + "minor_matmat", + "minor_matvec", + "minor_select_counts", + "minor_select_fill", +] @numba.njit(cache=True) @@ -91,3 +98,57 @@ def major_matmat(major_ptr, values, value_ptr, indices, b, n_major, n_minor): def minor_matmat(major_ptr, values, value_ptr, indices, b, n_major): b = np.ascontiguousarray(b) return _minor_matmat(major_ptr, values, value_ptr, indices, b, n_major) + + +# -- minor-axis selection ---------------------------------------------------- +# +# Selecting minor indices is a fan-out, not a remap: a selection may name the +# same original index more than once (``arr[:, [1, 1, 3]]``), so one original +# index can land in several output positions. A single old -> new lookup array +# can't express that -- writing one collapses every repeat but the last, which +# silently drops data rather than failing. +# +# Instead, the selection is inverted once into a CSR-shaped fan-out over the +# original minor axis (``offsets``/``positions``, both sized by the axis and +# the selection, never by nnz), and each stored element emits one output entry +# per output position its index maps to. Two passes, count then fill, so the +# per-slot destinations are known before anything is written and each slot +# fills a disjoint range. + + +@numba.njit(cache=True, parallel=True) +def _minor_select_counts(value_ptr, indices, fanout, out_counts): + n_slots = value_ptr.shape[0] - 1 + for u in numba.prange(n_slots): # ty: ignore[not-iterable] + c = 0 + for k in range(value_ptr[u], value_ptr[u + 1]): + c += fanout[indices[k]] + out_counts[u] = c + + +@numba.njit(cache=True, parallel=True) +def _minor_select_fill( + value_ptr, indices, offsets, positions, kept_slots, new_value_ptr, out_indices +): + for s in numba.prange(kept_slots.shape[0]): # ty: ignore[not-iterable] + u = kept_slots[s] + pos = new_value_ptr[s] + for k in range(value_ptr[u], value_ptr[u + 1]): + o = indices[k] + for j in range(offsets[o], offsets[o + 1]): + out_indices[pos] = positions[j] + pos += 1 + + +def minor_select_counts(value_ptr, indices, fanout): + """Output element count per stored (major, value) slot, given a fan-out map.""" + counts = np.empty(value_ptr.shape[0] - 1, dtype=np.int64) + _minor_select_counts(value_ptr, indices, fanout, counts) + return counts + + +def minor_select_fill(value_ptr, indices, offsets, positions, kept_slots, new_value_ptr, out_indices): + """Write the remapped minor indices for the surviving slots, in place.""" + _minor_select_fill( + value_ptr, indices, offsets, positions, kept_slots, new_value_ptr, out_indices + ) diff --git a/tests/test_select_minor_fanout.py b/tests/test_select_minor_fanout.py new file mode 100644 index 0000000..4ea33c0 --- /dev/null +++ b/tests/test_select_minor_fanout.py @@ -0,0 +1,177 @@ +"""Minor-axis selection is a fan-out: one stored element can land in many outputs. + +``arr[:, [1, 1, 3]]`` names column 1 twice. scipy supports that, and so did +this package until minor-axis selection got a native path -- at which point +a single old->new lookup array collapsed every repeat but the last and +returned zeros in the dropped slots, with no error. These tests pin the +fan-out semantics against scipy, which is the reference the native path has +to reproduce. +""" + +from __future__ import annotations + +import tracemalloc + +import numpy as np +import pytest +import scipy.sparse as sp + +from vsparse import VCSCArray, VCSRArray + + +@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 _grid(n_rows: int = 4, n_cols: int = 6) -> np.ndarray: + """Distinct values everywhere, so any misplacement is visible.""" + return np.arange(1, n_rows * n_cols + 1, dtype=float).reshape(n_rows, n_cols) + + +# -- the regression ---------------------------------------------------------- + + +@pytest.mark.parametrize( + "cols", + [ + pytest.param([1, 1, 3], id="one_repeat"), + pytest.param([2, 2, 2], id="same_index_three_times"), + pytest.param([0, 0, 0, 0], id="one_index_only"), + pytest.param([3, 1, 3, 1], id="interleaved_repeats"), + pytest.param([0, 1, 1, 2, 2, 2], id="mixed_multiplicities"), + ], +) +def test_duplicate_minor_indices_fan_out(vcls, cols): + dense = _grid() + v = vcls.from_scipy(_scipy_for(vcls, dense)) + np.testing.assert_allclose(v[:, cols].toarray(), dense[:, cols]) + + +def test_duplicate_indices_on_the_other_axis_too(vcls): + """Whichever axis is the *minor* one for this format takes the same path.""" + dense = _grid() + v = vcls.from_scipy(_scipy_for(vcls, dense)) + rows = [1, 1, 2, 0, 0] + np.testing.assert_allclose(v[rows, :].toarray(), dense[rows, :]) + + +def test_duplicates_on_both_axes_at_once(vcls): + dense = _grid() + v = vcls.from_scipy(_scipy_for(vcls, dense)) + rows, cols = [0, 0, 2], [1, 1, 4] + np.testing.assert_allclose(v[rows, cols].toarray(), dense[np.ix_(rows, cols)]) + + +def test_matches_scipy_for_a_random_selection_with_repeats(vcls, rng): + """Against scipy directly -- the behaviour that regressed.""" + dense = rng.integers(0, 4, size=(12, 9)).astype(np.float64) + v = vcls.from_scipy(_scipy_for(vcls, dense)) + reference = _scipy_for(vcls, dense) + + for _ in range(25): + cols = rng.integers(0, dense.shape[1], size=rng.integers(1, 15)).tolist() + np.testing.assert_allclose( + v[:, cols].toarray(), np.asarray(reference[:, cols].todense()) + ) + + +# -- selections that already worked, kept working ---------------------------- + + +@pytest.mark.parametrize( + "cols", + [ + pytest.param([0, 2], id="sorted"), + pytest.param([2, 0], id="reversed"), + pytest.param(np.array([], dtype=int), id="empty"), + pytest.param([-1, -2], id="negative"), + pytest.param(slice(1, 4), id="slice"), + pytest.param(slice(None, None, 2), id="strided_slice"), + pytest.param([True, False, True, False, True, False], id="boolean_mask"), + ], +) +def test_non_duplicate_selections_unchanged(vcls, cols): + dense = _grid() + v = vcls.from_scipy(_scipy_for(vcls, dense)) + expected = dense[:, cols] if isinstance(cols, slice) else dense[:, np.asarray(cols)] + np.testing.assert_allclose(v[:, cols].toarray(), expected) + + +def test_selection_keeps_the_vcs_type_and_index_dtype(vcls): + dense = _grid() + v = vcls.from_scipy(_scipy_for(vcls, dense)) + out = v[:, [1, 1, 3]] + assert isinstance(out, vcls) + assert out.indices.dtype == v.indices.dtype + assert out.shape == (dense.shape[0], 3) if vcls is VCSRArray else out.shape == (4, 3) + + +def test_empty_selection_gives_a_zero_width_array(vcls): + dense = _grid() + v = vcls.from_scipy(_scipy_for(vcls, dense)) + out = v[:, []] + assert out.shape[1] == 0 + assert out.nnz == 0 + + +def test_selection_of_an_all_zero_array(vcls): + v = vcls.from_scipy(_scipy_for(vcls, np.zeros((4, 5)))) + out = v[:, [1, 1, 2]] + assert out.nnz == 0 + np.testing.assert_allclose(out.toarray(), np.zeros((4, 3))) + + +def test_sparse_columns_and_duplicates_together(vcls): + """A duplicated index whose column is entirely implicit zeros.""" + dense = np.array([[1.0, 0.0, 2.0], [3.0, 0.0, 0.0]]) + v = vcls.from_scipy(_scipy_for(vcls, dense)) + cols = [1, 1, 0, 2] + np.testing.assert_allclose(v[:, cols].toarray(), dense[:, cols]) + + +def test_round_trip_through_scipy_after_a_duplicate_selection(vcls): + """The result has to be a structurally valid VCS array, not just print right.""" + dense = _grid() + v = vcls.from_scipy(_scipy_for(vcls, dense)) + out = v[:, [1, 1, 3]] + + rebuilt = vcls.from_scipy(out.to_scipy()) + np.testing.assert_allclose(rebuilt.toarray(), dense[:, [1, 1, 3]]) + assert out.value_ptr[-1] == out.indices.shape[0] + assert out.major_ptr[-1] == out.values.shape[0] + + +# -- memory ------------------------------------------------------------------ + + +def test_minor_selection_allocates_nothing_nnz_sized(): + """The rewrite also drops the nnz-sized intermediates the remap version needed. + + Sized so the old implementation's per-nonzero temporaries were ~30 MB + while the bound here is 16 MB -- the result itself is legitimately + nnz-scale, so this bounds the *scratch*, not the output. + """ + 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)) + cols = np.arange(0, n_cols, 2) + + v[:, cols] # warm up numba's JIT before measuring + + tracemalloc.start() + try: + before = tracemalloc.get_traced_memory()[0] + tracemalloc.reset_peak() + out = v[:, cols] + peak = tracemalloc.get_traced_memory()[1] + finally: + tracemalloc.stop() + + assert peak - before < 16 << 20, f"allocated {(peak - before) / 1e6:.1f} MB" + np.testing.assert_allclose(out.toarray(), dense[:, cols]) From faf64541dd16e4126b3d9cc8a9afef27205c8e1d Mon Sep 17 00:00:00 2001 From: fishidaho Date: Thu, 3 Sep 2026 17:12:35 -0700 Subject: [PATCH 2/2] Trim comments and tests The fan-out explanation moves from a block comment into short notes inside the kernels. Duplicate-pattern sweep drops to three cases and the type/dtype assertions go, since construction guarantees them. --- src/vsparse/_base.py | 23 ++++++--------------- src/vsparse/_ops.py | 21 +++++-------------- tests/test_select_minor_fanout.py | 34 ++++--------------------------- 3 files changed, 15 insertions(+), 63 deletions(-) diff --git a/src/vsparse/_base.py b/src/vsparse/_base.py index e2f4038..4fc0269 100644 --- a/src/vsparse/_base.py +++ b/src/vsparse/_base.py @@ -515,25 +515,14 @@ def _select_major(self, key: Any) -> _VCSBase: def _select_minor(self, key: Any) -> _VCSBase: """Select along the minor axis (rows for VCSC, columns for VCSR). - Unlike :meth:`_select_major`, the kept elements aren't already - contiguous per major slice, so ``indices`` has to be filtered and - remapped, and any (major, unique-value) slot left with no kept index - dropped -- shrinking ``major_ptr``/``value_ptr`` accordingly. - - The selection is a *fan-out*, not a one-to-one remap: a key may name - the same minor index more than once (``arr[:, [1, 1, 3]]``, which - scipy supports), so one stored element can have to appear in several - output positions. It is inverted once into a CSR-shaped map over the - original minor axis and applied by - :func:`vsparse._ops.minor_select_counts`/ - :func:`~vsparse._ops.minor_select_fill`, whose intermediates are - sized by the axis and the selection rather than by ``nnz``. + A key may name the same index more than once, so this is a fan-out + rather than a remap and one stored element can appear in several + output positions. """ idx = _normalize_major_idx(key, self.n_minor) n_minor_new = idx.shape[0] - # Invert the selection: output positions grouped by the original index - # they came from, so a repeated index carries all of its destinations. + # Output positions grouped by the original index they came from. fanout = np.bincount(idx, minlength=self.n_minor) offsets = np.zeros(self.n_minor + 1, dtype=np.int64) np.cumsum(fanout, out=offsets[1:]) @@ -553,8 +542,8 @@ def _select_minor(self, key: Any) -> _VCSBase: kept_slots, new_value_ptr, new_indices, ) - # Slots stay in their original order, so each major slice keeps a - # contiguous run of them; count how many of its slots survived. + # Slots keep their original order, so each major slice owns a + # contiguous run of them. major_of_slot = np.searchsorted(self.major_ptr, kept_slots, side="right") - 1 major_counts = np.bincount(major_of_slot, minlength=self.n_major) new_major_ptr = np.zeros(self.n_major + 1, dtype=np.int64) diff --git a/src/vsparse/_ops.py b/src/vsparse/_ops.py index acbf3b9..20f8d26 100644 --- a/src/vsparse/_ops.py +++ b/src/vsparse/_ops.py @@ -101,23 +101,12 @@ def minor_matmat(major_ptr, values, value_ptr, indices, b, n_major): # -- minor-axis selection ---------------------------------------------------- -# -# Selecting minor indices is a fan-out, not a remap: a selection may name the -# same original index more than once (``arr[:, [1, 1, 3]]``), so one original -# index can land in several output positions. A single old -> new lookup array -# can't express that -- writing one collapses every repeat but the last, which -# silently drops data rather than failing. -# -# Instead, the selection is inverted once into a CSR-shaped fan-out over the -# original minor axis (``offsets``/``positions``, both sized by the axis and -# the selection, never by nnz), and each stored element emits one output entry -# per output position its index maps to. Two passes, count then fill, so the -# per-slot destinations are known before anything is written and each slot -# fills a disjoint range. @numba.njit(cache=True, parallel=True) def _minor_select_counts(value_ptr, indices, fanout, out_counts): + # A selection may name one index several times, so each stored element + # contributes ``fanout`` output entries rather than one. n_slots = value_ptr.shape[0] - 1 for u in numba.prange(n_slots): # ty: ignore[not-iterable] c = 0 @@ -132,16 +121,16 @@ def _minor_select_fill( ): for s in numba.prange(kept_slots.shape[0]): # ty: ignore[not-iterable] u = kept_slots[s] - pos = new_value_ptr[s] + pos = new_value_ptr[s] # each slot fills a disjoint range for k in range(value_ptr[u], value_ptr[u + 1]): o = indices[k] - for j in range(offsets[o], offsets[o + 1]): + for j in range(offsets[o], offsets[o + 1]): # every destination of o out_indices[pos] = positions[j] pos += 1 def minor_select_counts(value_ptr, indices, fanout): - """Output element count per stored (major, value) slot, given a fan-out map.""" + """Output element count per stored (major, value) slot.""" counts = np.empty(value_ptr.shape[0] - 1, dtype=np.int64) _minor_select_counts(value_ptr, indices, fanout, counts) return counts diff --git a/tests/test_select_minor_fanout.py b/tests/test_select_minor_fanout.py index 4ea33c0..16ed23d 100644 --- a/tests/test_select_minor_fanout.py +++ b/tests/test_select_minor_fanout.py @@ -1,13 +1,3 @@ -"""Minor-axis selection is a fan-out: one stored element can land in many outputs. - -``arr[:, [1, 1, 3]]`` names column 1 twice. scipy supports that, and so did -this package until minor-axis selection got a native path -- at which point -a single old->new lookup array collapsed every repeat but the last and -returned zeros in the dropped slots, with no error. These tests pin the -fan-out semantics against scipy, which is the reference the native path has -to reproduce. -""" - from __future__ import annotations import tracemalloc @@ -40,9 +30,7 @@ def _grid(n_rows: int = 4, n_cols: int = 6) -> np.ndarray: "cols", [ pytest.param([1, 1, 3], id="one_repeat"), - pytest.param([2, 2, 2], id="same_index_three_times"), pytest.param([0, 0, 0, 0], id="one_index_only"), - pytest.param([3, 1, 3, 1], id="interleaved_repeats"), pytest.param([0, 1, 1, 2, 2, 2], id="mixed_multiplicities"), ], ) @@ -53,7 +41,7 @@ def test_duplicate_minor_indices_fan_out(vcls, cols): def test_duplicate_indices_on_the_other_axis_too(vcls): - """Whichever axis is the *minor* one for this format takes the same path.""" + """Whichever axis is the minor one for this format takes the same path.""" dense = _grid() v = vcls.from_scipy(_scipy_for(vcls, dense)) rows = [1, 1, 2, 0, 0] @@ -68,7 +56,7 @@ def test_duplicates_on_both_axes_at_once(vcls): def test_matches_scipy_for_a_random_selection_with_repeats(vcls, rng): - """Against scipy directly -- the behaviour that regressed.""" + """Against scipy, which fans duplicate indices out correctly.""" dense = rng.integers(0, 4, size=(12, 9)).astype(np.float64) v = vcls.from_scipy(_scipy_for(vcls, dense)) reference = _scipy_for(vcls, dense) @@ -102,15 +90,6 @@ def test_non_duplicate_selections_unchanged(vcls, cols): np.testing.assert_allclose(v[:, cols].toarray(), expected) -def test_selection_keeps_the_vcs_type_and_index_dtype(vcls): - dense = _grid() - v = vcls.from_scipy(_scipy_for(vcls, dense)) - out = v[:, [1, 1, 3]] - assert isinstance(out, vcls) - assert out.indices.dtype == v.indices.dtype - assert out.shape == (dense.shape[0], 3) if vcls is VCSRArray else out.shape == (4, 3) - - def test_empty_selection_gives_a_zero_width_array(vcls): dense = _grid() v = vcls.from_scipy(_scipy_for(vcls, dense)) @@ -135,7 +114,7 @@ def test_sparse_columns_and_duplicates_together(vcls): def test_round_trip_through_scipy_after_a_duplicate_selection(vcls): - """The result has to be a structurally valid VCS array, not just print right.""" + """The result has to be a structurally valid VCS array.""" dense = _grid() v = vcls.from_scipy(_scipy_for(vcls, dense)) out = v[:, [1, 1, 3]] @@ -150,12 +129,7 @@ def test_round_trip_through_scipy_after_a_duplicate_selection(vcls): def test_minor_selection_allocates_nothing_nnz_sized(): - """The rewrite also drops the nnz-sized intermediates the remap version needed. - - Sized so the old implementation's per-nonzero temporaries were ~30 MB - while the bound here is 16 MB -- the result itself is legitimately - nnz-scale, so this bounds the *scratch*, not the output. - """ + """Bounds the scratch, not the output, which is legitimately nnz-scale.""" 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)