Fix/dtype narrowing

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,12 @@ 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.
+        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 +131,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)
 

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,21 @@ 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`` picks the stored ``indices`` dtype; left at ``None`` the
+    input's dtype is preserved.
     """
     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 +197,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

src/vsparse/_indexutils.py

@@ -6,7 +6,14 @@
 
 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``."""
+    return np.dtype(np.int32) if n <= _INT32_MAX else np.dtype(np.int64)
 
 
 def is_full_slice(key: Any) -> bool:

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,12 @@ def write_ivcs_elem(
     """
     g = f.require_group(k)
     g.attrs["shape"] = v.shape
-    g.attrs["indices_dtype"] = np.dtype(v.indices.dtype).name
+    # 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
+    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)

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:])
@@ -256,18 +257,22 @@ 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
+
+    # ``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)
 
-    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)
 

tests/test_index_dtypes.py

@@ -0,0 +1,175 @@
+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):
+    out = mat.copy()
+    out.indices = out.indices.astype(np.int64)
+    out.indptr = out.indptr.astype(np.int64)
+    return out
+
+
+@pytest.mark.parametrize(
+    ("n", "expected"), [(INT32_MAX, np.int32), (INT32_MAX + 1, np.int64)]
+)
+def test_dtype_switches_at_the_int32_boundary(n, expected):
+    assert smallest_index_dtype(n) == np.dtype(expected)
+
+
+def test_from_scipy_narrows_int64_indices(dense, vcls):
+    """An int64-indexed input is stored as int32 when the minor axis fits."""
+    mat = _with_int64_indices(_scipy_for(vcls, dense))
+    v = vcls.from_scipy(mat)
+
+    assert v.indices.dtype == np.int32
+    assert v.indices.nbytes == 4 * v.nnz
+    np.testing.assert_allclose(v.toarray(), dense)
+
+
+def test_minor_axis_beyond_int32_keeps_int64(vcls):
+    """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 = (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)),
+        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):
+    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_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(
+            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),
+        )
+
+
+@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
+
+    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)
+
+    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)
+
+
+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
+
+    va = VCSCAnnData.from_anndata(ad.AnnData(X=sp.csr_array(dense)), format="csr")
+    assert isinstance(va.X, VCSRArray)
+    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)
+
+
+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,
+    )
+    return sp.csr_array(dense), np.array([True, False, True]), np.array([True, False, True, False])
+
+
+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
+    )
+
+    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_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
+
+    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
+    assert out_indices.dtype == np.int32
+    np.testing.assert_array_equal(out_indices, [0, 1, 0, 1])