Select minor duplicate indices

src/vsparse/_base.py

@@ -515,31 +515,37 @@ 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 this filters/remaps ``indices`` and
-        drops any (major, unique-value) slot that no longer has any kept
-        index, shrinking ``major_ptr``/``value_ptr`` accordingly.
+        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]
 
-        remap = np.full(self.n_minor, -1, dtype=np.int64)
-        remap[idx] = np.arange(n_minor_new, dtype=np.int64)
+        # 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:])
+        positions = np.argsort(idx, kind="stable").astype(np.int64, copy=False)
 
-        keep = remap[self.indices] >= 0
-        new_indices = remap[self.indices[keep]].astype(self.indices.dtype, copy=False)
+        slot_counts = _ops.minor_select_counts(self.value_ptr, self.indices, fanout)
+        kept_slots = np.flatnonzero(slot_counts)
 
-        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[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:])
 
-        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:])
+        # 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 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)
         np.cumsum(major_counts, out=new_major_ptr[1:])
 

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,46 @@ 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 ----------------------------------------------------
+
+
+@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
+        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]  # 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]):  # 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."""
+    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
+    )

tests/test_select_minor_fanout.py

@@ -0,0 +1,151 @@
+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([0, 0, 0, 0], id="one_index_only"),
+        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, 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)
+
+    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_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."""
+    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():
+    """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)
+    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])