diff --git a/src/vsparse/_base.py b/src/vsparse/_base.py index 9a3ed29..ae81e7c 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,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) diff --git a/src/vsparse/_construct.py b/src/vsparse/_construct.py index 9465ace..958bdd6 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,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 diff --git a/src/vsparse/_indexutils.py b/src/vsparse/_indexutils.py index a1edfb2..5dd7e84 100644 --- a/src/vsparse/_indexutils.py +++ b/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: diff --git a/src/vsparse/_io.py b/src/vsparse/_io.py index b5b77bd..002f8c8 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,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) diff --git a/src/vsparse/_rapid_load.py b/src/vsparse/_rapid_load.py index c5b5523..e09eaa8 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:]) @@ -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) diff --git a/tests/test_index_dtypes.py b/tests/test_index_dtypes.py new file mode 100644 index 0000000..0d3a410 --- /dev/null +++ b/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])