Read the packed stream in slabs, so a whole-file pass isn't bounded by the file

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/baselines.json

@@ -12,7 +12,9 @@
     "cpu_ratio_1t_raw": 1.5,
     "cpu_ratio_1t_scanpy": 1.5,
     "cuda_time_ratio_kernel_over_csr": 2.0,
-    "cuda_device_bytes_ratio_vs_csr": 1.1
+    "cuda_device_bytes_ratio_vs_csr": 1.1,
+    "slab_peak_ratio_over_full": 1.1,
+    "slab_time_ratio_over_full": 1.5
   },
   "cases": {
     "layout_bytes_per_nonzero": {
@@ -133,6 +135,10 @@
     "cuda_parafac2_rmatmul_misaligned_vs_csr": {
       "cuda_time_ratio_kernel_over_csr": 0.4632,
       "cuda_device_bytes_ratio_vs_csr": 0.6792
+    },
+    "slab_stream_vs_full_decode": {
+      "slab_time_ratio_over_full": 1.1932,
+      "slab_peak_ratio_over_full": 0.7358
     }
   }
 }

benchmarks/cases.py

@@ -352,6 +352,137 @@ 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
+
+
+@slow
+def slab_stream_vs_full_decode() -> dict[str, float]:
+    """One pass over ``X`` in slabs, against decoding all of it, on a zstd IVCSR file.
+
+    Both read every packed byte front to back; the question is what each holds
+    while doing it. A full decode keeps the whole packed stream *and* the whole
+    decoded index array; the slab reader keeps one window and one block. The
+    throughput ratio is the thing to watch -- a windowed read re-decodes at most
+    one zstd chunk per window boundary, and if that overhead were large this
+    would be the wrong trade.
+
+    Memory is measured with ``tracemalloc`` rather than RSS: ``ru_maxrss`` is a
+    high-water mark that never falls, so on a fixture small enough to build
+    in-process it only ever reports the fixture. Timing and memory are measured
+    in separate passes, since tracing allocations distorts the timing of
+    allocation-heavy code.
+    """
+    import tempfile
+    import tracemalloc
+    from pathlib import Path as _Path
+
+    import anndata as ad
+    import h5py
+    import pandas as pd
+
+    from vsparse import VCSCAnnData, VCSRArray
+    from vsparse._stream import iter_blocks, open_packed, rows_per_window
+
+    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)),
+    )
+    nnz = float(mat.nnz)
+    del mat
+    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:
+            packed_bytes = float(f["X"]["packed_indices"].shape[0])
+            px = open_packed(f["X"])
+            chunk = rows_per_window(px)
+
+            def stream() -> None:
+                for _lo, _hi, block in iter_blocks(px, chunk_size=chunk):
+                    block.indices[0:1]  # touch it, so nothing is optimized away
+
+            full = best_time(lambda: ad.io.read_elem(f["X"]), repeat=3)
+            slab = best_time(stream, repeat=3)
+
+            tracemalloc.start()
+            ad.io.read_elem(f["X"])
+            full_peak = tracemalloc.get_traced_memory()[1]
+            tracemalloc.reset_peak()
+            stream()
+            slab_peak = tracemalloc.get_traced_memory()[1]
+            tracemalloc.stop()
+
+    out["full_decode_seconds"] = full
+    out["slab_stream_seconds"] = slab
+    out["full_decode_gb_per_second"] = packed_bytes / 1e9 / full
+    out["slab_stream_gb_per_second"] = packed_bytes / 1e9 / slab
+    out["slab_time_ratio_over_full"] = slab / full
+    out["full_decode_peak_gb"] = full_peak / 1e9
+    out["slab_stream_peak_gb"] = slab_peak / 1e9
+    out["slab_peak_ratio_over_full"] = slab_peak / full_peak
+    out["bytes_per_nonzero_read"] = packed_bytes / nnz
+    out["rows_per_window"] = float(chunk)
+    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/__init__.py

@@ -13,6 +13,7 @@
 from vsparse._norm_common import RECIPES, Recipe
 from vsparse._rapid_load import load_packed
 from vsparse._stats import VCSStats
+from vsparse._stream import iter_packed
 from vsparse._vcs_norm import VCSCArrayNormalized, VCSRArrayNormalized
 
 __all__ = [
@@ -27,6 +28,7 @@
     "VCSStats",
     "cuda_is_available",
     "from_anndata",
+    "iter_packed",
     "load_packed",
     "to_layer",
 ]

src/vsparse/_anndata_class.py

@@ -103,6 +103,39 @@ def _map_subset_2d(mapping: Mapping[Any, Any], oidx: Any, vidx: Any) -> dict[Any
     return {k: _subset_2d(v, oidx, vidx) for k, v in mapping.items() if k is not None}
 
 
+def subset_fields(fields: Mapping[str, Any], oidx: Any, vidx: Any) -> dict[str, Any]:
+    """Window a set of AnnData fields to positions ``oidx``/``vidx``.
+
+    Behind :meth:`VCSCAnnData._subset_fields`, and used directly by
+    :func:`vsparse.iter_packed`, which windows fields it read straight from a
+    file rather than off an existing object. Fields absent from ``fields`` are
+    absent from the result, so a file without e.g. ``layers`` is fine.
+
+    ``obs["vsparse_a"]``/``varm["vsparse_b"/"c"/"s"]`` get windowed along with
+    ``obs``/``varm``, same as any other per-cell/per-gene column -- each kept
+    cell/gene keeps its own precomputed factor. But population-level aspects of
+    those statistics (e.g. the median depth or per-gene mean baked into them)
+    were derived from the *pre-subset* population, so they no longer reflect
+    this narrower one -- mark them stale rather than silently pretending they
+    were computed fresh. Call ``.normalized(recalculate=True)`` to recompute for
+    this subset, or ``recalculate=False`` to keep using these (stale) values.
+    """
+    out = dict(fields)
+    for key, idx in (("obs", oidx), ("var", vidx)):
+        if key in out:
+            out[key] = cast(pd.DataFrame, out[key]).iloc[idx].copy()
+    for key, idx in (("obsm", oidx), ("varm", vidx)):
+        if key in out:
+            out[key] = _map_subset_1d(out[key], idx)
+    for key, (i, j) in (("obsp", (oidx, oidx)), ("varp", (vidx, vidx)), ("layers", (oidx, vidx))):
+        if key in out:
+            out[key] = _map_subset_2d(out[key], i, j)
+    uns = out.get("uns") or {}
+    if _VSPARSE_UNS_KEY in uns:
+        out["uns"] = {**uns, _VSPARSE_UNS_KEY: {**uns[_VSPARSE_UNS_KEY], "stale": True}}
+    return out
+
+
 def _map_copy(mapping: Mapping[Any, Any]) -> dict[Any, Any]:
     return {k: _copy_value(v) for k, v in mapping.items() if k is not None}
 
@@ -212,36 +245,20 @@ 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)
-
-        obs = cast(pd.DataFrame, self.obs).iloc[oidx].copy()
-        var = cast(pd.DataFrame, self.var).iloc[vidx].copy()
-
-        # obs["vsparse_a"]/varm["vsparse_b"/"c"/"s"] get windowed along with
-        # obs/varm below, same as any other per-cell/per-gene column -- each
-        # kept cell/gene keeps its own precomputed factor. But population-level
-        # aspects of those statistics (e.g. the median depth or per-gene mean
-        # baked into them) were derived from the *pre-subset* population, so
-        # they no longer reflect this narrower one -- mark them stale rather
-        # than silently pretending they were computed fresh. Call
-        # .normalized(recalculate=True) to recompute for this subset, or
-        # recalculate=False to keep using these (stale) values.
-        uns = self.uns
-        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),
+            **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.
+        """
+        return subset_fields({k: getattr(self, k) for k in _FIELD_KEYS}, oidx, vidx)
+
     def copy(self) -> VCSCAnnData:  # ty: ignore[invalid-method-override]
         """A deep copy, preserving the VCSC/VCSR-backed ``X``/``raw_X``.