Subset cells and genes while unpacking IVCSR/IVCSC

README.md

@@ -143,6 +143,16 @@ va = vsparse.load_packed("archived.h5ad")
 norm = va.X.normalized("parafac2")  # lazy view; norm @ B works without densifying
 ```
 
+`load_packed` can also subset as it decodes: `obs_idx`/`var_idx` take a boolean mask, an array of integer positions, or a slice, and give the same result as indexing a full load, without ever building the full matrix. With IVCSR, only the selected cells are decoded and unselected genes are dropped in the same pass (IVCSC is the same with the axes swapped). Every packed byte is still read, so reading the stream sets the floor: on a 1.26M-cell file, a 1% cell subset decodes `X` in 0.21 s against 0.50 s for all of it.
+
+```python
+import anndata, h5py
+
+with h5py.File("archived.h5ad") as f:  # read only obs to build the mask
+    obs = anndata.io.read_elem(f["obs"])
+va_blood = vsparse.load_packed("archived.h5ad", obs_idx=obs["tissue"] == "blood")
+```
+
 ### Statistics sidecar
 
 Every `VCSCArray`/`VCSRArray` carries per-row and per-column summary statistics as `arr.stats` (a `vsparse.VCSStats`): `row_sums`/`col_sums` (total counts), `row_nnz`/`col_nnz` (n-genes/n-cells by counts), and `col_depth_sums` (per-gene sums of depth-normalized counts). They are computed on first access and written next to `X`/`raw_X` by `write_h5ad`/`write_zarr` whether or not they were ever computed or a normalized view was ever made, so a file read back has them without touching the data. Every recipe's per-cell and per-gene scales (`a`/`b`) are derived from them, so only centering/scaling recipes still need a pass over the data. Files without a usable sidecar (older versions, other tools) load as before, and the statistics are recomputed when first needed.

benchmarks/cases.py

@@ -352,6 +352,62 @@ def large_layout_and_matmul() -> dict[str, float]:
     }
 
 
+@slow
+def subset_load_vs_full_load() -> dict[str, float]:
+    """Reading ``X`` subset as it's unpacked, against reading all of it, from a zstd IVCSR file.
+
+    Random selections of 0.1/1/10% of cells, a "clustered" 1% (one contiguous
+    run, as selecting a single sample tends to be), and 1% of cells with 10%
+    of genes. Each ratio is that read's time over a full read of ``X``. Every
+    packed byte is still read either way, so reading the stream is the floor.
+    Whole ``load_packed`` calls also read ``obs``/``var``, which is recorded
+    separately since it costs the same either way (and, with a string index of
+    this many cells, more than ``X`` does).
+    """
+    import tempfile
+    from pathlib import Path
+
+    import anndata as ad
+    import h5py
+    import pandas as pd
+
+    from vsparse import VCSCAnnData, VCSRArray, load_packed
+    from vsparse._io import read_vcs_subset
+
+    n_obs, n_vars = 400_000, 3_000
+    mat = integer_counts_csr(n_obs, n_vars, density=0.02)
+    adata = VCSCAnnData(
+        X=VCSRArray.from_scipy(mat),
+        obs=pd.DataFrame(index=np.arange(n_obs).astype(str)),
+        var=pd.DataFrame(index=np.arange(n_vars).astype(str)),
+    )
+    del mat
+    rng = np.random.default_rng(0)
+    start = n_obs // 3
+    selections = {
+        f"random_{frac:g}": (np.flatnonzero(rng.random(n_obs) < frac), None)
+        for frac in (0.001, 0.01, 0.1)
+    }
+    selections["clustered_0.01"] = (np.arange(start, start + n_obs // 100), None)
+    selections["cells_0.01_genes_0.1"] = (
+        np.flatnonzero(rng.random(n_obs) < 0.01),
+        np.flatnonzero(rng.random(n_vars) < 0.1),
+    )
+    out: dict[str, float] = {}
+    with tempfile.TemporaryDirectory() as tmp:
+        path = Path(tmp) / "bench.ivcsr.h5ad"
+        adata.write_h5ad(path)
+        del adata
+        with h5py.File(path, "r") as f:
+            full = best_time(lambda: ad.io.read_elem(f["X"]), repeat=5)
+            out["full_x_seconds"] = full
+            for name, (rows, cols) in selections.items():
+                t = best_time(lambda rows=rows, cols=cols: read_vcs_subset(f["X"], rows, cols), 5)
+                out[f"time_ratio_{name}"] = t / full
+        out["full_load_packed_seconds"] = best_time(lambda: load_packed(path), repeat=3)
+    return out
+
+
 # -- CUDA kernels vs materializing to a sparse matrix ------------------------
 #
 # The question `vsparse._cuda` exists to answer: is walking the

docs/usage.md

@@ -127,6 +127,31 @@ va = vsparse.load_packed("archived.ivcsr.h5ad")
 norm = va.X.normalized("parafac2")
 ```
 
+### Subsetting while loading
+
+`load_packed(path, obs_idx=..., var_idx=...)` returns exactly
+`load_packed(path)[obs_idx, var_idx]`, but subsets `X` as it decodes rather than
+after. Each selection is a boolean mask over the axis, an array of integer positions
+(negative positions count from the end; unsorted or repeated positions come back in
+that order), or a slice; select by label by converting first, e.g. with
+`obs.index.get_indexer(labels)`.
+
+For an IVCSR file only the selected cells' indices are decoded, and a gene selection
+drops and renumbers indices in that same pass (IVCSC: the same, with cells and genes
+swapped). The full matrix is never allocated. Every packed byte is still read from
+disk, though: with no stored per-row byte offsets, finding where each selected cell
+starts takes a parallel scan of the stream. The subset doesn't carry the file's
+statistics sidecar (see below), which describes the whole matrix; its own `stats` are
+computed on first use.
+
+```python
+import anndata, h5py
+
+with h5py.File("archived.ivcsr.h5ad") as f:  # read only obs to build the mask
+    mask = anndata.io.read_elem(f["obs"])["condition"] == "control"
+controls = vsparse.load_packed("archived.ivcsr.h5ad", obs_idx=mask, var_idx=hvg_positions)
+```
+
 ### Statistics sidecar
 
 Each `VCSCArray`/`VCSRArray` carries a {class}`vsparse.VCSStats` as `arr.stats`:

src/vsparse/_anndata_class.py

@@ -212,7 +212,18 @@ def __getitem__(self, index: Any) -> VCSCAnnData:  # ty: ignore[invalid-method-o
         oidx, vidx = self._normalize_indices(index)
         oidx = _as_slice_index(oidx, self.n_obs)
         vidx = _as_slice_index(vidx, self.n_vars)
+        return type(self)(
+            X=_subset_2d(self._vcs_X, oidx, vidx),
+            raw_X=_subset_2d(self._vcs_raw_X, oidx, vidx),
+            **self._subset_fields(oidx, vidx),
+        )
+
+    def _subset_fields(self, oidx: Any, vidx: Any) -> dict[str, Any]:
+        """Every field but ``X``/``raw_X``, subset to positions ``oidx``/``vidx``.
 
+        Shared by :meth:`__getitem__` and :func:`~vsparse.load_packed`'s
+        ``obs_idx``/``var_idx``, which decodes ``X`` already subset.
+        """
         obs = cast(pd.DataFrame, self.obs).iloc[oidx].copy()
         var = cast(pd.DataFrame, self.var).iloc[vidx].copy()
 
@@ -229,18 +240,16 @@ def __getitem__(self, index: Any) -> VCSCAnnData:  # ty: ignore[invalid-method-o
         if _VSPARSE_UNS_KEY in uns:
             uns = {**uns, _VSPARSE_UNS_KEY: {**uns[_VSPARSE_UNS_KEY], "stale": True}}
 
-        return type(self)(
-            X=_subset_2d(self._vcs_X, oidx, vidx),
-            raw_X=_subset_2d(self._vcs_raw_X, oidx, vidx),
-            obs=obs,
-            var=var,
-            uns=uns,
-            obsm=_map_subset_1d(self.obsm, oidx),
-            varm=_map_subset_1d(self.varm, vidx),
-            obsp=_map_subset_2d(self.obsp, oidx, oidx),
-            varp=_map_subset_2d(self.varp, vidx, vidx),
-            layers=_map_subset_2d(self.layers, oidx, vidx),
-        )
+        return {
+            "obs": obs,
+            "var": var,
+            "uns": uns,
+            "obsm": _map_subset_1d(self.obsm, oidx),
+            "varm": _map_subset_1d(self.varm, vidx),
+            "obsp": _map_subset_2d(self.obsp, oidx, oidx),
+            "varp": _map_subset_2d(self.varp, vidx, vidx),
+            "layers": _map_subset_2d(self.layers, oidx, vidx),
+        }
 
     def copy(self) -> VCSCAnnData:  # ty: ignore[invalid-method-override]
         """A deep copy, preserving the VCSC/VCSR-backed ``X``/``raw_X``.

src/vsparse/_io.py

@@ -41,7 +41,7 @@
 
     GroupStorageType = h5py.Group | zarr.Group
 
-__all__ = ["write_ivcs_elem"]
+__all__ = ["read_vcs_subset", "write_ivcs_elem"]
 
 _ARRAY_KEYS = ("major_ptr", "values", "value_ptr", "indices")
 _SPEC_VERSION = "0.1.0"
@@ -145,6 +145,85 @@ def _read(elem: GroupStorageType, *, _reader: Reader) -> _VCSBase:
     return _read
 
 
+def _sorted_unique(idx: np.ndarray) -> tuple[np.ndarray, np.ndarray | None]:
+    """``idx`` as sorted unique positions, plus how to recover ``idx`` from them (``None``: no need)."""
+    if idx.shape[0] < 2 or bool(np.all(idx[1:] > idx[:-1])):
+        return idx, None
+    uniq, inverse = np.unique(idx, return_inverse=True)
+    return uniq, inverse
+
+
+def read_vcs_subset(
+    elem: GroupStorageType, rows: np.ndarray | None, cols: np.ndarray | None
+) -> _VCSBase:
+    """Read the stored VCS array ``elem``, keeping only ``rows``/``cols`` (int positions, or ``None`` for all).
+
+    Equal to reading the whole array and indexing it with ``[rows, cols]``. An
+    IVCSC/IVCSR array is subset as it's decoded (see
+    :func:`vsparse._ivcsc.unpack_selected`): every packed byte is still read,
+    but only the selected major slices are decoded, and a minor-axis selection
+    is applied in the same pass, so the full matrix is never built. A
+    selection that isn't sorted and unique is decoded as one that is, then
+    reordered in memory.
+
+    A subset doesn't carry the stored statistics sidecar, which describes the
+    whole array; its own :attr:`~vsparse.VCSCArray.stats` are computed on first use.
+    """
+    spec = cast(str, elem.attrs.get("encoding-type", ""))
+    if spec in _IVCS_SPEC_NAMES.values():
+        return _read_ivcs_subset(elem, VCSCArray if spec == "ivcsc" else VCSRArray, rows, cols)
+    arr = cast(_VCSBase, ad.io.read_elem(elem))
+    return arr[slice(None) if rows is None else rows, slice(None) if cols is None else cols]
+
+
+def _major_selection(key: np.ndarray | None, n_major: int) -> tuple[np.ndarray, np.ndarray | None]:
+    """The sorted unique major slices to decode, and how to reorder them into ``key``."""
+    if key is None:
+        return np.arange(n_major, dtype=np.int64), None
+    return _sorted_unique(key)
+
+
+def _minor_remap(
+    key: np.ndarray | None, n_minor: int
+) -> tuple[np.ndarray | None, np.ndarray | None, int]:
+    """Each minor index's new one (``-1``: dropped), how to reorder into ``key``, and the new length."""
+    if key is None:
+        return None, None, n_minor
+    kept, inverse = _sorted_unique(key)
+    remap = np.full(n_minor, -1, dtype=np.int64)
+    remap[kept] = np.arange(kept.shape[0])
+    return remap, inverse, kept.shape[0]
+
+
+def _read_ivcs_subset(
+    elem: GroupStorageType,
+    cls: type[_VCSBase],
+    rows: np.ndarray | None,
+    cols: np.ndarray | None,
+) -> _VCSBase:
+    n_major, n_minor = cls._swap(_read_shape(elem))
+    major_key, minor_key = (cols, rows) if cls is VCSCArray else (rows, cols)
+    sel, major_inverse = _major_selection(major_key, n_major)
+    remap, minor_inverse, n_minor_new = _minor_remap(minor_key, n_minor)
+
+    values = read_array(elem["values"])
+    major_ptr, groups, value_ptr, indices = _ivcsc.unpack_selected(
+        read_array(elem["major_ptr"]),
+        read_array(elem["value_ptr"]),
+        read_array(elem["packed_indices"]),
+        sel,
+        remap,
+        n_minor_new,
+    )
+    n_rows, n_cols = cls._swap((sel.shape[0], n_minor_new))
+    arr = cls((n_rows, n_cols), major_ptr, values[groups], value_ptr, indices)
+    if major_inverse is not None:
+        arr = arr._select_major(major_inverse)
+    if minor_inverse is not None:
+        arr = arr._select_minor(minor_inverse)
+    return arr
+
+
 def _register() -> None:
     from anndata._io.specs.registry import _REGISTRY, IOSpec