From c4c8c90e449b65d446fd3d7df59121f694265e33 Mon Sep 17 00:00:00 2001
From: fishidaho <etshurt@gmail.com>
Date: Wed, 2 Sep 2026 16:48:53 -0700
Subject: [PATCH 1/3] Size index arrays by the axis they address, not the array
 beside them

`indices` is the only nnz-sized array in the VCS layout, and nothing ever
chose its dtype: `_construct.compress` copied whatever dtype the input
scipy array carried, and `write_ivcs_elem` recorded the same. scipy hands
out int64 indices for any array with enough nonzeros, so a 33k-gene minor
axis was routinely stored -- in memory, on disk, and in every kernel that
walks it -- at 8 bytes per nonzero instead of 4.

`_rapid_load._filter_and_compact` had the same conflation in a sharper
form: one `idx_dtype`, keyed off nnz, applied to both `new_indptr` (indexed
by nonzero count, genuinely needs int64 at scale) and `out_indices` (gene
indices, bounded by the gene count). Crossing INT32_MAX nonzeros silently
doubled the largest allocation in the function for no reason.

Adds `_indexutils.smallest_index_dtype(n)` as the single place that rule
lives, and applies it at each point an index array is sized:

- `_construct.compress` takes `n_minor` and narrows up front, so the wide
  buffer is never allocated rather than allocated and then shrunk.
- `_VCSBase.__init__` narrows as the construction choke point -- after the
  bounds check, so an out-of-range index is still rejected rather than
  truncated, and a no-op (no copy) when the dtype is already right.
- `write_ivcs_elem` re-derives the dtype a reader will rebuild `indices`
  as, instead of trusting the array it was handed.
- `_filter_and_compact` keys its pointer and column-index dtypes off nnz
  and the kept-gene count separately.
- `transpose_major` now uses the shared helper for the rule it already
  applied inline.

Narrowing never widens: indices already stored in something smaller than
int32 are left alone.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 src/vsparse/_base.py       |  16 ++-
 src/vsparse/_construct.py  |  20 +++-
 src/vsparse/_indexutils.py |  17 ++-
 src/vsparse/_io.py         |  12 +-
 src/vsparse/_rapid_load.py |  23 ++--
 tests/test_index_dtypes.py | 221 +++++++++++++++++++++++++++++++++++++
 6 files changed, 295 insertions(+), 14 deletions(-)
 create mode 100644 tests/test_index_dtypes.py

diff --git a/src/vsparse/_base.py b/src/vsparse/_base.py
index eb89f05..12e63bf 100644
--- a/src/vsparse/_base.py
+++ b/src/vsparse/_base.py
@@ -20,6 +20,7 @@
 from vsparse import _construct, _ops
 from vsparse._indexutils import is_full_slice as _is_full_slice
 from vsparse._indexutils import normalize_major_idx as _normalize_major_idx
+from vsparse._indexutils import smallest_index_dtype as _smallest_index_dtype
 
 __all__ = ["VCSCArray", "VCSRArray"]
 
@@ -67,6 +68,17 @@ def __init__(
         if indices.shape[0] and (indices.min() < 0 or indices.max() >= 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
+        # rejected rather than silently truncated by the cast. ``indices`` is
+        # the only nnz-sized array in the layout, so storing it wider than
+        # ``n_minor`` requires is the largest avoidable cost here; this is a
+        # no-op (no copy) whenever it already has the right dtype, which is
+        # the normal case now that :func:`vsparse._construct.compress` picks
+        # the dtype up front.
+        idx_dtype = _smallest_index_dtype(n_minor)
+        if idx_dtype.itemsize < indices.dtype.itemsize:
+            indices = indices.astype(idx_dtype, copy=False)
+
         self.shape = (int(shape[0]), int(shape[1]))
         self.major_ptr = major_ptr
         self.values = values
@@ -124,9 +136,9 @@ def copy(self):
     def from_scipy(cls, mat) -> _VCSBase:
         """Build from any scipy sparse array/matrix (converted internally)."""
         mat = mat.tocsc() if cls._format == "csc" else mat.tocsr()
-        n_major, _n_minor = cls._swap(mat.shape)
+        n_major, n_minor = cls._swap(mat.shape)
         major_ptr, values, value_ptr, indices = _construct.compress(
-            mat.indptr, mat.indices, mat.data, n_major
+            mat.indptr, mat.indices, mat.data, n_major, n_minor
         )
         return cls(mat.shape, major_ptr, values, value_ptr, indices)
 
diff --git a/src/vsparse/_construct.py b/src/vsparse/_construct.py
index 9465ace..9545e98 100644
--- a/src/vsparse/_construct.py
+++ b/src/vsparse/_construct.py
@@ -27,6 +27,8 @@
 import numba
 import numpy as np
 
+from vsparse._indexutils import smallest_index_dtype
+
 __all__ = ["compress", "decompress", "transpose_major"]
 
 
@@ -108,14 +110,26 @@ def compress(
     minor_indices: np.ndarray,
     data: np.ndarray,
     n_major: int,
+    n_minor: int | None = None,
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
     """Build the VCS layout from a standard compressed-sparse layout.
 
     Parameters mirror scipy's ``indptr``/``indices``/``data`` for either a
     CSC (major axis = columns) or CSR (major axis = rows) matrix.
+
+    ``n_minor`` (the length of the axis ``minor_indices`` points into) picks
+    the stored ``indices`` dtype: scipy hands out int64 indices for any array
+    with many nonzeros, but the values themselves only have to address
+    ``n_minor``, so a 33k-gene axis is stored as int32 no matter how large
+    the input's own index dtype was. This is an ``nnz``-sized array, so the
+    difference is the single largest term in an array's memory footprint.
+    Left at ``None``, the input's dtype is preserved (no narrowing).
     """
     major_ptr = np.ascontiguousarray(major_ptr, dtype=np.int64)
-    minor_indices = np.ascontiguousarray(minor_indices)
+    idx_dtype = None if n_minor is None else smallest_index_dtype(n_minor)
+    if idx_dtype is not None and idx_dtype.itemsize > minor_indices.dtype.itemsize:
+        idx_dtype = None  # never widen a caller's already-narrower indices
+    minor_indices = np.ascontiguousarray(minor_indices, dtype=idx_dtype)
     data = np.ascontiguousarray(data)
     return _compress(major_ptr, minor_indices, data, n_major)
 
@@ -188,7 +202,7 @@ def transpose_major(
     value_ptr_out = np.concatenate([group_starts, [nnz]]).astype(np.int64)
     major_ptr_out = np.searchsorted(sorted_major[group_starts], np.arange(n_minor + 1)).astype(np.int64)
 
-    minor_dtype = np.int32 if n_major <= np.iinfo(np.int32).max else np.int64
-    indices_out = sorted_minor.astype(minor_dtype, copy=False)
+    # The output's minor axis is the input's major axis, so that's the bound.
+    indices_out = sorted_minor.astype(smallest_index_dtype(n_major), copy=False)
 
     return major_ptr_out, values_out, value_ptr_out, indices_out
diff --git a/src/vsparse/_indexutils.py b/src/vsparse/_indexutils.py
index a1edfb2..d6d9ef5 100644
--- a/src/vsparse/_indexutils.py
+++ b/src/vsparse/_indexutils.py
@@ -6,7 +6,22 @@
 
 import numpy as np
 
-__all__ = ["is_full_slice", "normalize_major_idx"]
+__all__ = ["is_full_slice", "normalize_major_idx", "smallest_index_dtype"]
+
+_INT32_MAX = np.iinfo(np.int32).max
+
+
+def smallest_index_dtype(n: int) -> np.dtype:
+    """Narrowest signed integer dtype that can address an axis of length ``n``.
+
+    Index arrays are sized by the axis they point *into*, not by the array
+    they live next to: minor-axis ``indices`` are bounded by ``n_minor``
+    (typically a gene count, comfortably int32) even when the array holds
+    more than ``INT32_MAX`` nonzeros. Keying each index array off its own
+    bound is what keeps a large-nnz array from paying int64 for indices that
+    never need it.
+    """
+    return np.dtype(np.int32) if n <= _INT32_MAX else np.dtype(np.int64)
 
 
 def is_full_slice(key: Any) -> bool:
diff --git a/src/vsparse/_io.py b/src/vsparse/_io.py
index b5b77bd..55c9266 100644
--- a/src/vsparse/_io.py
+++ b/src/vsparse/_io.py
@@ -25,6 +25,7 @@
 
 from vsparse import _ivcsc
 from vsparse._base import VCSCArray, VCSRArray, _VCSBase
+from vsparse._indexutils import smallest_index_dtype
 
 if TYPE_CHECKING:
     from collections.abc import Mapping
@@ -96,7 +97,16 @@ def write_ivcs_elem(
     """
     g = f.require_group(k)
     g.attrs["shape"] = v.shape
-    g.attrs["indices_dtype"] = np.dtype(v.indices.dtype).name
+    # ``indices`` isn't stored directly here (it's delta+varint packed), so
+    # this attribute is purely the dtype a reader rebuilds it as. Record the
+    # narrowest dtype that can address the minor axis rather than whatever
+    # the in-memory array happens to carry: an array built by some other
+    # route can still be holding int64 indices for a small minor axis, and
+    # there's no reason to make every future read pay for that.
+    in_memory = v.indices.dtype
+    narrowest = smallest_index_dtype(v.n_minor)
+    stored_dtype = narrowest if narrowest.itemsize < in_memory.itemsize else in_memory
+    g.attrs["indices_dtype"] = stored_dtype.name
     for name in ("major_ptr", "values", "value_ptr"):
         ad.io.write_elem(g, name, getattr(v, name), dataset_kwargs=dataset_kwargs)
     packed = _ivcsc.pack_indices(v.value_ptr, v.indices)
diff --git a/src/vsparse/_rapid_load.py b/src/vsparse/_rapid_load.py
index 35d8c35..127eb34 100644
--- a/src/vsparse/_rapid_load.py
+++ b/src/vsparse/_rapid_load.py
@@ -61,6 +61,7 @@
 
 from vsparse import _ivcsc
 from vsparse._anndata_class import VCSCAnnData
+from vsparse._indexutils import smallest_index_dtype
 
 if TYPE_CHECKING:
     from os import PathLike
@@ -159,7 +160,7 @@ def _build_selected_rows(
     full_indptr = value_ptr[major_ptr]
     row_nnz = np.diff(full_indptr)[row_mask]
     nnz = int(row_nnz.sum())
-    ptr_dtype = np.int64 if nnz > np.iinfo(np.int32).max else np.int32
+    ptr_dtype = smallest_index_dtype(nnz)
 
     indptr = np.zeros(row_nnz.shape[0] + 1, dtype=ptr_dtype)
     np.cumsum(row_nnz, out=indptr[1:])
@@ -238,18 +239,26 @@ def _filter_and_compact(
     gene_mask: np.ndarray,
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, int]:
     kept_rows = np.nonzero(cell_mask)[0]
-    gene_remap = (np.cumsum(gene_mask) - 1).astype(np.int32)
-    gene_remap[~gene_mask] = -1
     n_kept_genes = int(gene_mask.sum())
+    gene_remap = (np.cumsum(gene_mask) - 1).astype(smallest_index_dtype(n_kept_genes))
+    gene_remap[~gene_mask] = -1
+
+    # Two index arrays, two different bounds. ``new_indptr`` is indexed by
+    # nonzero count and genuinely needs int64 once nnz passes INT32_MAX;
+    # ``out_indices`` holds *gene* indices, bounded by ``n_kept_genes``, and
+    # is the nnz-sized one. Sizing both off nnz (as this used to) silently
+    # doubles the largest allocation in the function the moment a big enough
+    # dataset pushes the pointer array over the int32 line.
+    ptr_dtype = smallest_index_dtype(int(indices.shape[0]))
+    col_dtype = smallest_index_dtype(n_kept_genes)
 
-    idx_dtype = np.int64 if indices.shape[0] > np.iinfo(np.int32).max else np.int32
-    counts = np.empty(kept_rows.shape[0], dtype=idx_dtype)
+    counts = np.empty(kept_rows.shape[0], dtype=ptr_dtype)
     _count_kept(row_indptr, indices, gene_mask, kept_rows, counts)
-    new_indptr = np.zeros(kept_rows.shape[0] + 1, dtype=idx_dtype)
+    new_indptr = np.zeros(kept_rows.shape[0] + 1, dtype=ptr_dtype)
     np.cumsum(counts, out=new_indptr[1:])
 
     nnz_filtered = int(new_indptr[-1])
-    out_indices = np.empty(nnz_filtered, dtype=idx_dtype)
+    out_indices = np.empty(nnz_filtered, dtype=col_dtype)
     out_data = np.empty(nnz_filtered, dtype=np.float32)
     _fill_kept(row_indptr, indices, data, gene_remap, kept_rows, new_indptr, out_indices, out_data)
 
diff --git a/tests/test_index_dtypes.py b/tests/test_index_dtypes.py
new file mode 100644
index 0000000..9338a2c
--- /dev/null
+++ b/tests/test_index_dtypes.py
@@ -0,0 +1,221 @@
+"""Index arrays are sized by the axis they address, not by the array beside them.
+
+Covers both halves of that rule: ``indices`` narrowed at construction/write
+time (so a small gene axis never costs int64), and ``_filter_and_compact``
+choosing its pointer and column-index dtypes from separate bounds.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import scipy.sparse as sp
+
+from vsparse import VCSCAnnData, VCSCArray, VCSRArray
+from vsparse._indexutils import smallest_index_dtype
+from vsparse._rapid_load import _filter_and_compact
+
+INT32_MAX = np.iinfo(np.int32).max
+
+
+@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 _with_int64_indices(mat):
+    """The same matrix, forced to carry int64 ``indices``/``indptr``."""
+    out = mat.copy()
+    out.indices = out.indices.astype(np.int64)
+    out.indptr = out.indptr.astype(np.int64)
+    return out
+
+
+# -- the rule itself ---------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    ("n", "expected"),
+    [
+        (0, np.int32),
+        (1, np.int32),
+        (INT32_MAX - 1, np.int32),
+        (INT32_MAX, np.int32),
+        (INT32_MAX + 1, np.int64),
+        (2**40, np.int64),
+    ],
+)
+def test_smallest_index_dtype_boundary(n, expected):
+    assert smallest_index_dtype(n) == np.dtype(expected)
+
+
+# -- construction ------------------------------------------------------------
+
+
+def test_from_scipy_narrows_int64_indices(dense, vcls):
+    """A minor axis that fits int32 is stored as int32, whatever the input carried."""
+    mat = _with_int64_indices(_scipy_for(vcls, dense))
+    assert mat.indices.dtype == np.int64
+
+    v = vcls.from_scipy(mat)
+    assert v.indices.dtype == np.int32
+    np.testing.assert_allclose(v.toarray(), dense)
+
+
+def test_narrowing_halves_the_nnz_sized_array(vcls, rng):
+    """``indices`` is the only nnz-sized array, so this is the whole point."""
+    dense = rng.integers(0, 4, size=(60, 40)).astype(np.float64)
+    mat = _with_int64_indices(_scipy_for(vcls, dense))
+
+    v = vcls.from_scipy(mat)
+    assert v.nnz > 0
+    assert v.indices.nbytes == 4 * v.nnz
+    assert v.indices.nbytes < mat.indices.nbytes
+
+
+def test_minor_axis_beyond_int32_keeps_int64(vcls):
+    """The bound is the axis length, so a genuinely huge axis still gets int64."""
+    n_huge = INT32_MAX + 10
+    # Two populated major slices against an enormous minor axis: shape is
+    # large, nnz is 4, so this stays a tiny allocation.
+    minor_idx = np.array([0, INT32_MAX + 5, 1, INT32_MAX + 9], dtype=np.int64)
+    shape = (n_huge, 2) if vcls is VCSCArray else (2, n_huge)
+    mat_cls = sp.csc_array if vcls is VCSCArray else sp.csr_array
+    mat = mat_cls(
+        (np.array([1.0, 2.0, 3.0, 4.0]), minor_idx, np.array([0, 2, 4], dtype=np.int64)),
+        shape=shape,
+    )
+
+    v = vcls.from_scipy(mat)
+    assert v.indices.dtype == np.int64
+    np.testing.assert_array_equal(np.sort(v.indices), np.sort(minor_idx))
+
+
+def test_construction_never_widens_narrower_indices(vcls):
+    """A caller who already stored something narrower than int32 keeps it."""
+    shape = (4, 3) if vcls is VCSCArray else (3, 4)
+    v = vcls(
+        shape,
+        major_ptr=np.array([0, 1, 2, 3], dtype=np.int64),
+        values=np.array([1.0, 2.0, 3.0]),
+        value_ptr=np.array([0, 1, 2, 3], dtype=np.int64),
+        indices=np.array([0, 1, 2], dtype=np.int16),
+    )
+    assert v.indices.dtype == np.int16
+
+
+def test_out_of_bounds_index_still_raises_rather_than_truncating(vcls):
+    """Narrowing happens after validation, so a bad index is rejected, not wrapped."""
+    shape = (4, 1) if vcls is VCSCArray else (1, 4)
+    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, 1], dtype=np.int64),
+            indices=np.array([2**32 + 1], dtype=np.int64),
+        )
+
+
+def test_transpose_major_narrows_indices(vcls, dense):
+    v = vcls.from_scipy(_scipy_for(vcls, dense))
+    dual = v._transpose_major()
+    assert dual.indices.dtype == np.int32
+    np.testing.assert_allclose(dual.toarray(), dense)
+
+
+# -- write / read round trip -------------------------------------------------
+
+
+@pytest.mark.parametrize("fmt", ["vcsc", "ivcsc"])
+def test_roundtrip_preserves_values_and_stores_narrow_indices(tmp_path, dense, fmt):
+    import anndata as ad
+
+    adata = ad.AnnData(X=sp.csr_array(dense))
+    va = VCSCAnnData.from_anndata(adata, format="csr")
+    path = tmp_path / f"data.{fmt}.h5ad"
+    va.write_h5ad(path, format=fmt)
+
+    back = VCSCAnnData.read_h5ad(path)
+    assert isinstance(back.X, VCSRArray)
+    assert back.X.indices.dtype == np.int32
+    np.testing.assert_allclose(back.X.toarray(), dense)
+    assert back.X.shape == dense.shape
+
+
+def test_packed_write_records_narrow_dtype_for_a_wide_in_memory_array(tmp_path, dense):
+    """The write path re-derives the dtype rather than trusting the array it's handed."""
+    import anndata as ad
+
+    adata = ad.AnnData(X=sp.csr_array(dense))
+    va = VCSCAnnData.from_anndata(adata, format="csr")
+    assert isinstance(va.X, VCSRArray)
+    # Simulate an array built by some other route that kept int64 indices.
+    va.X.indices = va.X.indices.astype(np.int64)
+
+    path = tmp_path / "wide.h5ad"
+    va.write_h5ad(path, format="ivcsc")
+
+    back = VCSCAnnData.read_h5ad(path)
+    assert isinstance(back.X, VCSRArray)
+    assert back.X.indices.dtype == np.int32
+    np.testing.assert_allclose(back.X.toarray(), dense)
+
+
+# -- _filter_and_compact: two bounds, two dtypes -----------------------------
+
+
+def _small_filter_inputs():
+    dense = np.array(
+        [[1.0, 0.0, 2.0, 0.0], [0.0, 3.0, 0.0, 4.0], [5.0, 0.0, 6.0, 0.0]],
+        dtype=np.float32,
+    )
+    X = sp.csr_array(dense)
+    cell_mask = np.array([True, False, True])
+    gene_mask = np.array([True, False, True, False])
+    return X, cell_mask, gene_mask
+
+
+def test_filter_and_compact_uses_int32_for_both_when_both_fit():
+    X, cell_mask, gene_mask = _small_filter_inputs()
+    new_indptr, out_indices, out_data, kept_rows, n_kept = _filter_and_compact(
+        X.indptr, X.indices, X.data, cell_mask, gene_mask
+    )
+
+    assert new_indptr.dtype == np.int32
+    assert out_indices.dtype == np.int32
+    assert n_kept == 2
+    np.testing.assert_array_equal(kept_rows, [0, 2])
+    np.testing.assert_allclose(out_data, [1.0, 2.0, 5.0, 6.0])
+    np.testing.assert_array_equal(out_indices, [0, 1, 0, 1])
+
+
+def test_filter_and_compact_gene_indices_stay_int32_when_pointers_need_int64(monkeypatch):
+    """The regression: a big-nnz dataset must not drag the gene indices up with it.
+
+    Allocating a genuinely >INT32_MAX-nonzero matrix isn't testable, so the
+    nnz-keyed half of the decision is forced instead -- exactly the situation
+    a full-scale dataset produces, where the old shared dtype doubled the
+    nnz-sized ``out_indices`` for no reason.
+    """
+    import vsparse._rapid_load as rapid_load
+
+    X, cell_mask, gene_mask = _small_filter_inputs()
+    nnz_in = int(X.indices.shape[0])
+    real = rapid_load.smallest_index_dtype
+
+    def forced(n: int) -> np.dtype:
+        return np.dtype(np.int64) if n == nnz_in else real(n)
+
+    monkeypatch.setattr(rapid_load, "smallest_index_dtype", forced)
+    new_indptr, out_indices, _, _, _ = _filter_and_compact(
+        X.indptr, X.indices, X.data, cell_mask, gene_mask
+    )
+
+    assert new_indptr.dtype == np.int64  # keyed off nnz: correctly widened
+    assert out_indices.dtype == np.int32  # keyed off the gene axis: unaffected
+    np.testing.assert_array_equal(out_indices, [0, 1, 0, 1])

From 1c119fea23e457f42d82654d4d57fd826779bbad Mon Sep 17 00:00:00 2001
From: fishidaho <etshurt@gmail.com>
Date: Thu, 3 Sep 2026 11:10:24 -0700
Subject: [PATCH 2/3] Extend dtype-narrowing coverage to main's new
 construction paths

#22/#23 added astype, native minor-axis selection, and elementwise
add/sub/multiply (which rebuild through scipy). Every one of them
constructs its result with `type(self)(...)`, so all of them inherit the
narrowing in `_VCSBase.__init__` for free -- no per-method change was
needed, which is the payoff for putting the rule at the construction choke
point rather than in `from_scipy`.

Verified rather than assumed: an int64-indexed input stays int32 through
astype, `v[:, cols]`, both-axes selection, scalar mul/div/neg, add/sub
against another VCS array, copy, log1p, `_transpose_major` and `T`. The
scipy round-trip in the elementwise path is the one that could plausibly
have handed back int64, so it gets its own case.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 tests/test_index_dtypes.py | 71 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 71 insertions(+)

diff --git a/tests/test_index_dtypes.py b/tests/test_index_dtypes.py
index 9338a2c..1e2c5dd 100644
--- a/tests/test_index_dtypes.py
+++ b/tests/test_index_dtypes.py
@@ -219,3 +219,74 @@ def forced(n: int) -> np.dtype:
     assert new_indptr.dtype == np.int64  # keyed off nnz: correctly widened
     assert out_indices.dtype == np.int32  # keyed off the gene axis: unaffected
     np.testing.assert_array_equal(out_indices, [0, 1, 0, 1])
+
+
+# -- the invariant holds across every construction path ----------------------
+#
+# Narrowing lives in `_VCSBase.__init__` rather than in each method, so any
+# operation that builds a new array through `type(self)(...)` inherits it.
+# These pin that for the operations added since -- astype, minor-axis
+# selection, and the elementwise ops that round-trip through scipy -- so a
+# future method can't quietly reintroduce int64 indices for a small axis.
+
+
+def _int64_indexed(vcls, dense):
+    """A VCS array built from an input carrying int64 indices."""
+    return vcls.from_scipy(_with_int64_indices(_scipy_for(vcls, dense)))
+
+
+def test_astype_keeps_narrow_indices(vcls, dense):
+    v = _int64_indexed(vcls, dense)
+    out = v.astype(np.float32)
+    assert out.indices.dtype == np.int32
+    assert out.dtype == np.float32
+    np.testing.assert_allclose(out.toarray(), dense.astype(np.float32))
+
+
+def test_minor_axis_selection_keeps_narrow_indices(vcls, dense):
+    v = _int64_indexed(vcls, dense)
+    n_cols = dense.shape[1]
+    cols = np.arange(0, n_cols, 2)
+
+    out = v[:, cols]
+    assert out.indices.dtype == np.int32
+    np.testing.assert_allclose(out.toarray(), dense[:, cols])
+
+
+def test_both_axes_selection_keeps_narrow_indices(vcls, dense):
+    if dense.shape[0] < 2 or dense.shape[1] < 2:
+        pytest.skip("shape too small")
+    v = _int64_indexed(vcls, dense)
+    rows, cols = np.arange(0, dense.shape[0], 2), np.arange(0, dense.shape[1], 2)
+
+    out = v[rows, :][:, cols]
+    assert out.indices.dtype == np.int32
+    np.testing.assert_allclose(out.toarray(), dense[np.ix_(rows, cols)])
+
+
+@pytest.mark.parametrize(
+    "op",
+    [
+        pytest.param(lambda v, d: v * 2.0, id="scalar_mul"),
+        pytest.param(lambda v, d: v / 2.0, id="scalar_div"),
+        pytest.param(lambda v, d: -v, id="neg"),
+        pytest.param(lambda v, d: v + v, id="add_vcs"),
+        pytest.param(lambda v, d: v - v, id="sub_vcs"),
+        pytest.param(lambda v, d: v.copy(), id="copy"),
+        pytest.param(lambda v, d: v.log1p(), id="log1p"),
+        pytest.param(lambda v, d: v._transpose_major(), id="transpose_major"),
+        pytest.param(lambda v, d: v.T, id="T"),
+    ],
+)
+def test_derived_arrays_keep_narrow_indices(vcls, dense, op):
+    """Every op returning a VCS array goes through __init__, so all of them narrow."""
+    result = op(_int64_indexed(vcls, dense), dense)
+    assert result.indices.dtype == np.int32
+
+
+def test_scipy_roundtrip_ops_narrow_a_wide_result(vcls, dense):
+    """The elementwise ops rebuild via from_scipy, where scipy may hand back int64."""
+    v = _int64_indexed(vcls, dense)
+    summed = v + v
+    assert summed.indices.dtype == np.int32
+    np.testing.assert_allclose(summed.toarray(), dense * 2)

From 0ce7e6b9a71f3ee0e0fac958684527562a232874 Mon Sep 17 00:00:00 2001
From: fishidaho <etshurt@gmail.com>
Date: Thu, 3 Sep 2026 17:04:40 -0700
Subject: [PATCH 3/3] Trim comments and tests

Comments now state what the code does rather than the reasoning behind it.
Index-dtype tests drop the cases the code trivially guarantees, keeping the
boundary, the truncation risks, the roundtrips, and the paths that build a
new array.
---
 src/vsparse/_base.py       |   9 +-
 src/vsparse/_construct.py  |   9 +-
 src/vsparse/_indexutils.py |  10 +-
 src/vsparse/_io.py         |   8 +-
 src/vsparse/_rapid_load.py |   8 +-
 tests/test_index_dtypes.py | 193 ++++++++-----------------------------
 6 files changed, 47 insertions(+), 190 deletions(-)

diff --git a/src/vsparse/_base.py b/src/vsparse/_base.py
index c6a39ba..ae81e7c 100644
--- a/src/vsparse/_base.py
+++ b/src/vsparse/_base.py
@@ -68,13 +68,8 @@ def __init__(
         if indices.shape[0] and (indices.min() < 0 or indices.max() >= 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
-        # rejected rather than silently truncated by the cast. ``indices`` is
-        # the only nnz-sized array in the layout, so storing it wider than
-        # ``n_minor`` requires is the largest avoidable cost here; this is a
-        # no-op (no copy) whenever it already has the right dtype, which is
-        # the normal case now that :func:`vsparse._construct.compress` picks
-        # the dtype up front.
+        # Narrow after the bounds check above, so an out-of-range index is
+        # rejected rather than silently truncated by the cast.
         idx_dtype = _smallest_index_dtype(n_minor)
         if idx_dtype.itemsize < indices.dtype.itemsize:
             indices = indices.astype(idx_dtype, copy=False)
diff --git a/src/vsparse/_construct.py b/src/vsparse/_construct.py
index 9545e98..958bdd6 100644
--- a/src/vsparse/_construct.py
+++ b/src/vsparse/_construct.py
@@ -117,13 +117,8 @@ def compress(
     Parameters mirror scipy's ``indptr``/``indices``/``data`` for either a
     CSC (major axis = columns) or CSR (major axis = rows) matrix.
 
-    ``n_minor`` (the length of the axis ``minor_indices`` points into) picks
-    the stored ``indices`` dtype: scipy hands out int64 indices for any array
-    with many nonzeros, but the values themselves only have to address
-    ``n_minor``, so a 33k-gene axis is stored as int32 no matter how large
-    the input's own index dtype was. This is an ``nnz``-sized array, so the
-    difference is the single largest term in an array's memory footprint.
-    Left at ``None``, the input's dtype is preserved (no narrowing).
+    ``n_minor`` picks the stored ``indices`` dtype; left at ``None`` the
+    input's dtype is preserved.
     """
     major_ptr = np.ascontiguousarray(major_ptr, dtype=np.int64)
     idx_dtype = None if n_minor is None else smallest_index_dtype(n_minor)
diff --git a/src/vsparse/_indexutils.py b/src/vsparse/_indexutils.py
index d6d9ef5..5dd7e84 100644
--- a/src/vsparse/_indexutils.py
+++ b/src/vsparse/_indexutils.py
@@ -12,15 +12,7 @@
 
 
 def smallest_index_dtype(n: int) -> np.dtype:
-    """Narrowest signed integer dtype that can address an axis of length ``n``.
-
-    Index arrays are sized by the axis they point *into*, not by the array
-    they live next to: minor-axis ``indices`` are bounded by ``n_minor``
-    (typically a gene count, comfortably int32) even when the array holds
-    more than ``INT32_MAX`` nonzeros. Keying each index array off its own
-    bound is what keeps a large-nnz array from paying int64 for indices that
-    never need it.
-    """
+    """Narrowest signed integer dtype that can address an axis of length ``n``."""
     return np.dtype(np.int32) if n <= _INT32_MAX else np.dtype(np.int64)
 
 
diff --git a/src/vsparse/_io.py b/src/vsparse/_io.py
index 55c9266..002f8c8 100644
--- a/src/vsparse/_io.py
+++ b/src/vsparse/_io.py
@@ -97,12 +97,8 @@ def write_ivcs_elem(
     """
     g = f.require_group(k)
     g.attrs["shape"] = v.shape
-    # ``indices`` isn't stored directly here (it's delta+varint packed), so
-    # this attribute is purely the dtype a reader rebuilds it as. Record the
-    # narrowest dtype that can address the minor axis rather than whatever
-    # the in-memory array happens to carry: an array built by some other
-    # route can still be holding int64 indices for a small minor axis, and
-    # there's no reason to make every future read pay for that.
+    # Record the narrowest dtype that can address the minor axis rather than
+    # whatever the in-memory array happens to carry.
     in_memory = v.indices.dtype
     narrowest = smallest_index_dtype(v.n_minor)
     stored_dtype = narrowest if narrowest.itemsize < in_memory.itemsize else in_memory
diff --git a/src/vsparse/_rapid_load.py b/src/vsparse/_rapid_load.py
index debf496..e09eaa8 100644
--- a/src/vsparse/_rapid_load.py
+++ b/src/vsparse/_rapid_load.py
@@ -261,12 +261,8 @@ def _filter_and_compact(
     gene_remap = (np.cumsum(gene_mask) - 1).astype(smallest_index_dtype(n_kept_genes))
     gene_remap[~gene_mask] = -1
 
-    # Two index arrays, two different bounds. ``new_indptr`` is indexed by
-    # nonzero count and genuinely needs int64 once nnz passes INT32_MAX;
-    # ``out_indices`` holds *gene* indices, bounded by ``n_kept_genes``, and
-    # is the nnz-sized one. Sizing both off nnz (as this used to) silently
-    # doubles the largest allocation in the function the moment a big enough
-    # dataset pushes the pointer array over the int32 line.
+    # ``new_indptr`` is indexed by nonzero count and needs int64 once nnz
+    # passes INT32_MAX; ``out_indices`` holds gene indices and never does.
     ptr_dtype = smallest_index_dtype(int(indices.shape[0]))
     col_dtype = smallest_index_dtype(n_kept_genes)
 
diff --git a/tests/test_index_dtypes.py b/tests/test_index_dtypes.py
index 1e2c5dd..0d3a410 100644
--- a/tests/test_index_dtypes.py
+++ b/tests/test_index_dtypes.py
@@ -1,10 +1,3 @@
-"""Index arrays are sized by the axis they address, not by the array beside them.
-
-Covers both halves of that rule: ``indices`` narrowed at construction/write
-time (so a small gene axis never costs int64), and ``_filter_and_compact``
-choosing its pointer and column-index dtypes from separate bounds.
-"""
-
 from __future__ import annotations
 
 import numpy as np
@@ -28,62 +21,33 @@ def _scipy_for(vcls, dense):
 
 
 def _with_int64_indices(mat):
-    """The same matrix, forced to carry int64 ``indices``/``indptr``."""
     out = mat.copy()
     out.indices = out.indices.astype(np.int64)
     out.indptr = out.indptr.astype(np.int64)
     return out
 
 
-# -- the rule itself ---------------------------------------------------------
-
-
 @pytest.mark.parametrize(
-    ("n", "expected"),
-    [
-        (0, np.int32),
-        (1, np.int32),
-        (INT32_MAX - 1, np.int32),
-        (INT32_MAX, np.int32),
-        (INT32_MAX + 1, np.int64),
-        (2**40, np.int64),
-    ],
+    ("n", "expected"), [(INT32_MAX, np.int32), (INT32_MAX + 1, np.int64)]
 )
-def test_smallest_index_dtype_boundary(n, expected):
+def test_dtype_switches_at_the_int32_boundary(n, expected):
     assert smallest_index_dtype(n) == np.dtype(expected)
 
 
-# -- construction ------------------------------------------------------------
-
-
 def test_from_scipy_narrows_int64_indices(dense, vcls):
-    """A minor axis that fits int32 is stored as int32, whatever the input carried."""
+    """An int64-indexed input is stored as int32 when the minor axis fits."""
     mat = _with_int64_indices(_scipy_for(vcls, dense))
-    assert mat.indices.dtype == np.int64
-
     v = vcls.from_scipy(mat)
-    assert v.indices.dtype == np.int32
-    np.testing.assert_allclose(v.toarray(), dense)
-
-
-def test_narrowing_halves_the_nnz_sized_array(vcls, rng):
-    """``indices`` is the only nnz-sized array, so this is the whole point."""
-    dense = rng.integers(0, 4, size=(60, 40)).astype(np.float64)
-    mat = _with_int64_indices(_scipy_for(vcls, dense))
 
-    v = vcls.from_scipy(mat)
-    assert v.nnz > 0
+    assert v.indices.dtype == np.int32
     assert v.indices.nbytes == 4 * v.nnz
-    assert v.indices.nbytes < mat.indices.nbytes
+    np.testing.assert_allclose(v.toarray(), dense)
 
 
 def test_minor_axis_beyond_int32_keeps_int64(vcls):
-    """The bound is the axis length, so a genuinely huge axis still gets int64."""
-    n_huge = INT32_MAX + 10
-    # Two populated major slices against an enormous minor axis: shape is
-    # large, nnz is 4, so this stays a tiny allocation.
+    """Narrowing an axis that genuinely needs int64 would truncate the indices."""
     minor_idx = np.array([0, INT32_MAX + 5, 1, INT32_MAX + 9], dtype=np.int64)
-    shape = (n_huge, 2) if vcls is VCSCArray else (2, n_huge)
+    shape = (INT32_MAX + 10, 2) if vcls is VCSCArray else (2, INT32_MAX + 10)
     mat_cls = sp.csc_array if vcls is VCSCArray else sp.csr_array
     mat = mat_cls(
         (np.array([1.0, 2.0, 3.0, 4.0]), minor_idx, np.array([0, 2, 4], dtype=np.int64)),
@@ -96,7 +60,6 @@ def test_minor_axis_beyond_int32_keeps_int64(vcls):
 
 
 def test_construction_never_widens_narrower_indices(vcls):
-    """A caller who already stored something narrower than int32 keeps it."""
     shape = (4, 3) if vcls is VCSCArray else (3, 4)
     v = vcls(
         shape,
@@ -108,8 +71,8 @@ def test_construction_never_widens_narrower_indices(vcls):
     assert v.indices.dtype == np.int16
 
 
-def test_out_of_bounds_index_still_raises_rather_than_truncating(vcls):
-    """Narrowing happens after validation, so a bad index is rejected, not wrapped."""
+def test_out_of_bounds_index_raises_rather_than_truncating(vcls):
+    """Narrowing happens after validation, so a bad index cannot wrap silently."""
     shape = (4, 1) if vcls is VCSCArray else (1, 4)
     with pytest.raises(ValueError, match="out of bounds"):
         vcls(
@@ -121,22 +84,27 @@ def test_out_of_bounds_index_still_raises_rather_than_truncating(vcls):
         )
 
 
-def test_transpose_major_narrows_indices(vcls, dense):
-    v = vcls.from_scipy(_scipy_for(vcls, dense))
-    dual = v._transpose_major()
-    assert dual.indices.dtype == np.int32
-    np.testing.assert_allclose(dual.toarray(), dense)
-
-
-# -- write / read round trip -------------------------------------------------
+@pytest.mark.parametrize(
+    "op",
+    [
+        pytest.param(lambda v, d: v.astype(np.float32), id="astype"),
+        pytest.param(lambda v, d: v[:, ::2], id="select_minor"),
+        pytest.param(lambda v, d: v + v, id="add"),
+        pytest.param(lambda v, d: v * 2.0, id="scalar_mul"),
+        pytest.param(lambda v, d: v._transpose_major(), id="transpose_major"),
+    ],
+)
+def test_derived_arrays_stay_narrow(vcls, dense, op):
+    """Narrowing lives in __init__, so every op building a new array inherits it."""
+    v = vcls.from_scipy(_with_int64_indices(_scipy_for(vcls, dense)))
+    assert op(v, dense).indices.dtype == np.int32
 
 
 @pytest.mark.parametrize("fmt", ["vcsc", "ivcsc"])
 def test_roundtrip_preserves_values_and_stores_narrow_indices(tmp_path, dense, fmt):
     import anndata as ad
 
-    adata = ad.AnnData(X=sp.csr_array(dense))
-    va = VCSCAnnData.from_anndata(adata, format="csr")
+    va = VCSCAnnData.from_anndata(ad.AnnData(X=sp.csr_array(dense)), format="csr")
     path = tmp_path / f"data.{fmt}.h5ad"
     va.write_h5ad(path, format=fmt)
 
@@ -144,17 +112,14 @@ def test_roundtrip_preserves_values_and_stores_narrow_indices(tmp_path, dense, f
     assert isinstance(back.X, VCSRArray)
     assert back.X.indices.dtype == np.int32
     np.testing.assert_allclose(back.X.toarray(), dense)
-    assert back.X.shape == dense.shape
 
 
-def test_packed_write_records_narrow_dtype_for_a_wide_in_memory_array(tmp_path, dense):
-    """The write path re-derives the dtype rather than trusting the array it's handed."""
+def test_packed_write_narrows_a_wide_in_memory_array(tmp_path, dense):
+    """The write path re-derives the dtype rather than trusting the array it is handed."""
     import anndata as ad
 
-    adata = ad.AnnData(X=sp.csr_array(dense))
-    va = VCSCAnnData.from_anndata(adata, format="csr")
+    va = VCSCAnnData.from_anndata(ad.AnnData(X=sp.csr_array(dense)), format="csr")
     assert isinstance(va.X, VCSRArray)
-    # Simulate an array built by some other route that kept int64 indices.
     va.X.indices = va.X.indices.astype(np.int64)
 
     path = tmp_path / "wide.h5ad"
@@ -166,21 +131,15 @@ def test_packed_write_records_narrow_dtype_for_a_wide_in_memory_array(tmp_path,
     np.testing.assert_allclose(back.X.toarray(), dense)
 
 
-# -- _filter_and_compact: two bounds, two dtypes -----------------------------
-
-
 def _small_filter_inputs():
     dense = np.array(
         [[1.0, 0.0, 2.0, 0.0], [0.0, 3.0, 0.0, 4.0], [5.0, 0.0, 6.0, 0.0]],
         dtype=np.float32,
     )
-    X = sp.csr_array(dense)
-    cell_mask = np.array([True, False, True])
-    gene_mask = np.array([True, False, True, False])
-    return X, cell_mask, gene_mask
+    return sp.csr_array(dense), np.array([True, False, True]), np.array([True, False, True, False])
 
 
-def test_filter_and_compact_uses_int32_for_both_when_both_fit():
+def test_filter_and_compact_uses_int32_when_both_bounds_fit():
     X, cell_mask, gene_mask = _small_filter_inputs()
     new_indptr, out_indices, out_data, kept_rows, n_kept = _filter_and_compact(
         X.indptr, X.indices, X.data, cell_mask, gene_mask
@@ -194,99 +153,23 @@ def test_filter_and_compact_uses_int32_for_both_when_both_fit():
     np.testing.assert_array_equal(out_indices, [0, 1, 0, 1])
 
 
-def test_filter_and_compact_gene_indices_stay_int32_when_pointers_need_int64(monkeypatch):
-    """The regression: a big-nnz dataset must not drag the gene indices up with it.
-
-    Allocating a genuinely >INT32_MAX-nonzero matrix isn't testable, so the
-    nnz-keyed half of the decision is forced instead -- exactly the situation
-    a full-scale dataset produces, where the old shared dtype doubled the
-    nnz-sized ``out_indices`` for no reason.
-    """
+def test_gene_indices_stay_int32_when_pointers_need_int64(monkeypatch):
+    """Forced, since a >INT32_MAX-nonzero matrix cannot be allocated in a test."""
     import vsparse._rapid_load as rapid_load
 
     X, cell_mask, gene_mask = _small_filter_inputs()
     nnz_in = int(X.indices.shape[0])
     real = rapid_load.smallest_index_dtype
 
-    def forced(n: int) -> np.dtype:
-        return np.dtype(np.int64) if n == nnz_in else real(n)
-
-    monkeypatch.setattr(rapid_load, "smallest_index_dtype", forced)
+    monkeypatch.setattr(
+        rapid_load,
+        "smallest_index_dtype",
+        lambda n: np.dtype(np.int64) if n == nnz_in else real(n),
+    )
     new_indptr, out_indices, _, _, _ = _filter_and_compact(
         X.indptr, X.indices, X.data, cell_mask, gene_mask
     )
 
-    assert new_indptr.dtype == np.int64  # keyed off nnz: correctly widened
-    assert out_indices.dtype == np.int32  # keyed off the gene axis: unaffected
+    assert new_indptr.dtype == np.int64
+    assert out_indices.dtype == np.int32
     np.testing.assert_array_equal(out_indices, [0, 1, 0, 1])
-
-
-# -- the invariant holds across every construction path ----------------------
-#
-# Narrowing lives in `_VCSBase.__init__` rather than in each method, so any
-# operation that builds a new array through `type(self)(...)` inherits it.
-# These pin that for the operations added since -- astype, minor-axis
-# selection, and the elementwise ops that round-trip through scipy -- so a
-# future method can't quietly reintroduce int64 indices for a small axis.
-
-
-def _int64_indexed(vcls, dense):
-    """A VCS array built from an input carrying int64 indices."""
-    return vcls.from_scipy(_with_int64_indices(_scipy_for(vcls, dense)))
-
-
-def test_astype_keeps_narrow_indices(vcls, dense):
-    v = _int64_indexed(vcls, dense)
-    out = v.astype(np.float32)
-    assert out.indices.dtype == np.int32
-    assert out.dtype == np.float32
-    np.testing.assert_allclose(out.toarray(), dense.astype(np.float32))
-
-
-def test_minor_axis_selection_keeps_narrow_indices(vcls, dense):
-    v = _int64_indexed(vcls, dense)
-    n_cols = dense.shape[1]
-    cols = np.arange(0, n_cols, 2)
-
-    out = v[:, cols]
-    assert out.indices.dtype == np.int32
-    np.testing.assert_allclose(out.toarray(), dense[:, cols])
-
-
-def test_both_axes_selection_keeps_narrow_indices(vcls, dense):
-    if dense.shape[0] < 2 or dense.shape[1] < 2:
-        pytest.skip("shape too small")
-    v = _int64_indexed(vcls, dense)
-    rows, cols = np.arange(0, dense.shape[0], 2), np.arange(0, dense.shape[1], 2)
-
-    out = v[rows, :][:, cols]
-    assert out.indices.dtype == np.int32
-    np.testing.assert_allclose(out.toarray(), dense[np.ix_(rows, cols)])
-
-
-@pytest.mark.parametrize(
-    "op",
-    [
-        pytest.param(lambda v, d: v * 2.0, id="scalar_mul"),
-        pytest.param(lambda v, d: v / 2.0, id="scalar_div"),
-        pytest.param(lambda v, d: -v, id="neg"),
-        pytest.param(lambda v, d: v + v, id="add_vcs"),
-        pytest.param(lambda v, d: v - v, id="sub_vcs"),
-        pytest.param(lambda v, d: v.copy(), id="copy"),
-        pytest.param(lambda v, d: v.log1p(), id="log1p"),
-        pytest.param(lambda v, d: v._transpose_major(), id="transpose_major"),
-        pytest.param(lambda v, d: v.T, id="T"),
-    ],
-)
-def test_derived_arrays_keep_narrow_indices(vcls, dense, op):
-    """Every op returning a VCS array goes through __init__, so all of them narrow."""
-    result = op(_int64_indexed(vcls, dense), dense)
-    assert result.indices.dtype == np.int32
-
-
-def test_scipy_roundtrip_ops_narrow_a_wide_result(vcls, dense):
-    """The elementwise ops rebuild via from_scipy, where scipy may hand back int64."""
-    v = _int64_indexed(vcls, dense)
-    summed = v + v
-    assert summed.indices.dtype == np.int32
-    np.testing.assert_allclose(summed.toarray(), dense * 2)