diff --git a/README.md b/README.md index 13a01cf..4776e58 100644 --- a/README.md +++ b/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. diff --git a/benchmarks/baselines.json b/benchmarks/baselines.json index 38f2500..9a6ecdb 100644 --- a/benchmarks/baselines.json +++ b/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 } } } diff --git a/benchmarks/cases.py b/benchmarks/cases.py index b5bb151..cc19572 100644 --- a/benchmarks/cases.py +++ b/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 diff --git a/docs/usage.md b/docs/usage.md index 0f82ea7..5de94cf 100644 --- a/docs/usage.md +++ b/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`: diff --git a/src/vsparse/__init__.py b/src/vsparse/__init__.py index b3f2b49..e9ef159 100644 --- a/src/vsparse/__init__.py +++ b/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", ] diff --git a/src/vsparse/_anndata_class.py b/src/vsparse/_anndata_class.py index 69402e7..dbf7248 100644 --- a/src/vsparse/_anndata_class.py +++ b/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``. diff --git a/src/vsparse/_compression.py b/src/vsparse/_compression.py index 51d5037..fa4c504 100644 --- a/src/vsparse/_compression.py +++ b/src/vsparse/_compression.py @@ -11,10 +11,13 @@ HDF5 runs a dataset's filter pipeline one chunk at a time on a single thread, which makes decompression, not I/O, the bottleneck on read. -:func:`read_array` sidesteps that for zstd-filtered 1-D datasets: it looks up +:func:`read_range` sidesteps that for zstd-filtered 1-D datasets: it looks up every chunk's file offset in one pass (``chunk_iter``), reads the compressed chunks itself, and decodes them on a thread pool (``numcodecs.Zstd`` releases -the GIL). Anything else -- older Blosc2-compressed files, uncompressed or +the GIL). :func:`read_array` is that over the whole dataset; a caller streaming +a large array in windows reads one byte range at a time instead, and passes the +:func:`chunk_table` back in so the ``chunk_iter`` pass isn't repeated per +window. Anything else -- older Blosc2-compressed files, uncompressed or contiguous datasets, in-memory/non-POSIX drivers -- falls back to an ordinary h5py read, so those files keep loading exactly as before. @@ -51,7 +54,17 @@ if TYPE_CHECKING: from collections.abc import Iterator -__all__ = ["h5_dataset_kwargs", "numeric_only_compression", "read_array", "zarr_dataset_kwargs"] +#: Per-chunk ``(file offset, stored size, filter mask)``, as :func:`chunk_table` returns it. +_ChunkTable = tuple[np.ndarray, np.ndarray, np.ndarray] + +__all__ = [ + "chunk_table", + "h5_dataset_kwargs", + "numeric_only_compression", + "read_array", + "read_range", + "zarr_dataset_kwargs", +] _H5_COMPRESSION_KEYS = ("compression", "compression_opts") _ZARR_V2_COMPRESSION_KEYS = ("compressor",) @@ -228,13 +241,19 @@ def _is_parallel_zstd(ds: Any) -> bool: return plist.get_nfilters() == 1 and plist.get_filter(0)[0] == _H5_ZSTD_FILTER_ID -def _chunk_table(ds: Any) -> tuple[np.ndarray, np.ndarray, np.ndarray] | None: +def chunk_table(ds: Any) -> tuple[np.ndarray, np.ndarray, np.ndarray] | None: """Per-chunk ``(file offset, stored size, filter mask)``, in order. - ``None`` if any chunk was never written (HDF5 reads those as the fill - value, which only the ordinary read path reproduces) or this h5py/HDF5 - build lacks ``chunk_iter``. + ``None`` if ``ds`` isn't eligible for the parallel path, if any chunk was + never written (HDF5 reads those as the fill value, which only the ordinary + read path reproduces), or if this h5py/HDF5 build lacks ``chunk_iter``. + + Building it walks every chunk through a Python callback, so a caller making + repeated :func:`read_range` calls against one dataset should build it once + and pass it back in rather than pay that per window. """ + if not _is_parallel_zstd(ds): + return None chunk_len = ds.chunks[0] n_chunks = -(-ds.shape[0] // chunk_len) offsets = np.full(n_chunks, -1, dtype=np.int64) @@ -253,30 +272,35 @@ def record(info: Any) -> None: def _read_chunks_parallel( - ds: Any, offsets: np.ndarray, sizes: np.ndarray, masks: np.ndarray + ds: Any, offsets: np.ndarray, sizes: np.ndarray, masks: np.ndarray, lo: int, hi: int ) -> np.ndarray: + """Chunks ``[lo, hi)`` of ``ds``, decoded on a thread pool, as one flat array. + + The result holds whole chunks, so it runs past the dataset's last element + when the final chunk is a partial one -- HDF5 stores that chunk padded to + full size, and trimming is the caller's business. + """ import numcodecs - out = np.empty(ds.shape, dtype=ds.dtype) + chunk_len = ds.chunks[0] + out = np.empty((hi - lo) * chunk_len, dtype=ds.dtype) flat = out.view(np.uint8) - chunk_bytes = ds.chunks[0] * ds.dtype.itemsize + chunk_bytes = chunk_len * ds.dtype.itemsize codec = numcodecs.Zstd() fd = os.open(ds.file.filename, os.O_RDONLY) def load(i: int) -> None: raw = os.pread(fd, int(sizes[i]), int(offsets[i])) - dst = flat[i * chunk_bytes : (i + 1) * chunk_bytes] + dst = flat[(i - lo) * chunk_bytes : (i - lo + 1) * chunk_bytes] if masks[i] & 1: # filter-mask bit 0: HDF5 skipped zstd, stored the chunk as-is - dst[:] = np.frombuffer(raw, dtype=np.uint8, count=dst.size) - elif dst.size == chunk_bytes: + dst[: len(raw)] = np.frombuffer(raw, dtype=np.uint8) + else: codec.decode(raw, out=dst) - else: # the last chunk decodes to a full chunk; only its leading part is data - dst[:] = np.frombuffer(codec.decode(raw), dtype=np.uint8, count=dst.size) # One contiguous run of chunks per thread, rather than one task per # chunk, keeps the per-chunk (GIL-holding) Python overhead minimal. - n_workers = max(1, min(numba.get_num_threads(), offsets.shape[0])) - bounds = np.linspace(0, offsets.shape[0], n_workers + 1).astype(np.int64) + n_workers = max(1, min(numba.get_num_threads(), hi - lo)) + bounds = np.linspace(lo, hi, n_workers + 1).astype(np.int64) def load_run(t: int) -> None: for i in range(bounds[t], bounds[t + 1]): @@ -290,6 +314,40 @@ def load_run(t: int) -> None: return out +def read_range(ds: Any, start: int, stop: int, *, table: _ChunkTable | None = None) -> np.ndarray: + """``ds[start:stop]``, decoding only the zstd chunks that cover it, in parallel. + + Parameters + ---------- + ds + An h5py dataset or zarr array. + start, stop + Half-open element range; ``stop`` is clipped to ``len(ds)``. + table + A :func:`chunk_table` for ``ds``, to avoid rebuilding it per call. + ``None`` builds one, which also decides whether the parallel path + applies at all -- a Blosc2-compressed, contiguous, or non-h5py dataset + falls back to an ordinary read of the same range. + + ``chunk_table`` requires the file to be open read-only (see + :func:`_is_parallel_zstd`), so a range read taken while the file is open + ``"r+"`` silently drops to HDF5's single-threaded filter pipeline. + """ + stop = min(int(stop), ds.shape[0]) + start = int(start) + if table is None: + table = chunk_table(ds) + if table is None: + import hdf5plugin # noqa: F401 -- registers the zstd/Blosc2 HDF5 filters + + return np.asarray(ds[start:stop]) + chunk_len = ds.chunks[0] + lo = start // chunk_len + hi = max(lo, -(-stop // chunk_len)) + buf = _read_chunks_parallel(ds, *table, lo, hi) + return buf[start - lo * chunk_len : stop - lo * chunk_len] + + def read_array(ds: Any) -> np.ndarray: """Read a whole dataset into memory, decoding zstd chunks in parallel where possible. @@ -300,9 +358,4 @@ def read_array(ds: Any) -> np.ndarray: files written by earlier versions, is read the ordinary way. See the module docstring. """ - table = _chunk_table(ds) if _is_parallel_zstd(ds) else None - if table is None: - import hdf5plugin # noqa: F401 -- registers the zstd/Blosc2 HDF5 filters - - return np.asarray(ds[...]) - return _read_chunks_parallel(ds, *table) + return read_range(ds, 0, ds.shape[0]) diff --git a/src/vsparse/_io.py b/src/vsparse/_io.py index a8db421..3ed3ca1 100644 --- a/src/vsparse/_io.py +++ b/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 diff --git a/src/vsparse/_ivcsc.py b/src/vsparse/_ivcsc.py index a7b3a2e..edd51a3 100644 --- a/src/vsparse/_ivcsc.py +++ b/src/vsparse/_ivcsc.py @@ -42,7 +42,9 @@ import numba import numpy as np -__all__ = ["pack_indices", "unpack_indices"] +from vsparse._indexutils import smallest_index_dtype + +__all__ = ["pack_indices", "trim_to_boundary", "unpack_indices", "unpack_selected"] # Below this many packed bytes, decode serially -- the fixed overhead of the # counting pass plus a second (parallel) pass isn't worth it. @@ -50,6 +52,8 @@ # Lower bound on how small a per-thread chunk is allowed to get, so we don't # over-split modestly-sized arrays into slivers. _MIN_CHUNK_BYTES = 1 << 16 # 64 KiB +# Bytes :func:`_locate_ranks` counts at a time when skipping ahead. +_LOCATE_BLOCK = 256 @numba.njit(cache=True) @@ -140,14 +144,14 @@ def _realign_forward(buf: np.ndarray, guess: np.int64) -> np.int64: @numba.njit(cache=True) def _count_varints_range(buf: np.ndarray, start: np.int64, end: np.int64) -> np.int64: - """Count complete varints in buf[start:end); start/end must be true boundaries.""" - pos = start + """Count varints ending in buf[start:end) -- i.e. bytes with the continuation bit clear. + + Between true boundaries that's the number of complete varints. Counted + branch-free, one byte at a time, so the loop vectorizes. + """ count = np.int64(0) - while pos < end: - while buf[pos] & 0x80: - pos += 1 - pos += 1 - count += 1 + for pos in range(start, end): + count += 1 - (buf[pos] >> 7) return count @@ -189,20 +193,33 @@ def _decode_chunks( out[kk] = prev -def _group_chunk_boundaries( - value_ptr: np.ndarray, buf: np.ndarray, n_chunks: int -) -> tuple[np.ndarray, np.ndarray]: - """Split ``buf`` into ``n_chunks`` pieces landing on exact group boundaries. +@numba.njit(cache=True) +def trim_to_boundary(buf: np.ndarray) -> np.int64: + """Length of the longest prefix of ``buf`` that ends on a varint boundary. + + Everything below takes a buffer whose last varint is complete, which is free + when the buffer is the whole stream and false for a window read out of the + middle of one. :func:`_count_varints_range` walks its inner + continuation-bit loop past ``end``, which on a truncated final varint reads + out of bounds -- `njit` without ``boundscheck`` will not stop it. So a + window that didn't come from the writer intact is trimmed here first. + """ + n = buf.shape[0] + while n > 0 and (buf[n - 1] & 0x80) != 0: + n -= 1 + return n - Returns ``(chunk_group, chunk_byte)``, each of length ``n_chunks + 1``: - chunk ``t`` owns groups ``[chunk_group[t], chunk_group[t + 1])``, whose - packed bytes start at ``chunk_byte[t]``. Used by :func:`_unpack_parallel` - to parallelize decoding safely. + +def _varint_chunks(buf: np.ndarray, n_chunks: int) -> tuple[np.ndarray, np.ndarray]: + """Split ``buf`` into ``n_chunks`` spans starting on varint boundaries. + + ``buf`` must end on a varint boundary -- see :func:`trim_to_boundary`. + + Returns ``(real_starts, cum)``, each of length ``n_chunks + 1``: span ``t`` + covers bytes ``[real_starts[t], real_starts[t + 1])``, and its first varint + is varint number ``cum[t]`` of the whole stream. """ n_bytes = buf.shape[0] - n_groups = value_ptr.shape[0] - 1 - nnz = int(value_ptr[-1]) - guesses = np.linspace(0, n_bytes, n_chunks + 1).astype(np.int64) real_starts = np.empty(n_chunks + 1, dtype=np.int64) real_starts[0] = 0 @@ -215,7 +232,23 @@ def _group_chunk_boundaries( cum = np.empty(n_chunks + 1, dtype=np.int64) cum[0] = 0 cum[1:] = np.cumsum(counts) - assert cum[n_chunks] == nnz + return real_starts, cum + + +def _group_chunk_boundaries( + value_ptr: np.ndarray, buf: np.ndarray, n_chunks: int +) -> tuple[np.ndarray, np.ndarray]: + """Split ``buf`` into ``n_chunks`` pieces landing on exact group boundaries. + + Returns ``(chunk_group, chunk_byte)``, each of length ``n_chunks + 1``: + chunk ``t`` owns groups ``[chunk_group[t], chunk_group[t + 1])``, whose + packed bytes start at ``chunk_byte[t]``. Used by :func:`_unpack_parallel` + to parallelize decoding safely. + """ + n_bytes = buf.shape[0] + n_groups = value_ptr.shape[0] - 1 + real_starts, cum = _varint_chunks(buf, n_chunks) + assert cum[n_chunks] == int(value_ptr[-1]) chunk_group = np.empty(n_chunks + 1, dtype=np.int64) chunk_byte = np.empty(n_chunks + 1, dtype=np.int64) @@ -264,3 +297,178 @@ def unpack_indices(value_ptr: np.ndarray, packed: np.ndarray, dtype: np.dtype) - else: _unpack_parallel(value_ptr, packed, out, n_chunks) return out + + +# -- selective decode ----------------------------------------------------------- +# +# Decoding only some major slices (rows of an IVCSR array) still needs every +# byte of the stream in memory, since nothing records where a slice's bytes +# start. But it only has to *count* varints up to each selected slice, not +# decode them: :func:`_varint_chunks` splits the stream on varint boundaries +# with a parallel count, each span then walks forward to the selected slices +# that start inside it (:func:`_locate_ranks`), and only those slices are +# decoded (:func:`_count_kept`/:func:`_fill_selected`), in parallel. A minor-axis +# selection is applied during that same decode by renumbering each index +# through a lookup table, dropping the ones mapped to ``-1``. + + +@numba.njit(cache=True, parallel=True) +def _locate_ranks( + buf: np.ndarray, + real_starts: np.ndarray, + cum: np.ndarray, + targets: np.ndarray, + span_ptr: np.ndarray, + out: np.ndarray, +) -> None: + """Byte offset of varint number ``targets[i]``, for sorted ``targets``. + + Span ``t`` (from :func:`_varint_chunks`) owns targets + ``[span_ptr[t], span_ptr[t + 1])`` and walks forward to each in turn, so + the whole pass touches each byte at most once. + """ + for t in numba.prange(real_starts.shape[0] - 1): # ty: ignore[not-iterable] + pos = real_starts[t] + end = real_starts[t + 1] + rank = cum[t] + for i in range(span_ptr[t], span_ptr[t + 1]): + need = targets[i] - rank + # Skip whole blocks the target isn't in, counting with the + # vectorized kernel. That can stop partway through a varint; + # _advance_n_varints then finishes it as the first of `need`. + while need > 0 and pos + _LOCATE_BLOCK <= end: + ended = _count_varints_range(buf, pos, pos + _LOCATE_BLOCK) + if ended >= need: + break + need -= ended + pos += _LOCATE_BLOCK + pos = _advance_n_varints(buf, pos, need) + rank = targets[i] + out[i] = pos + + +@numba.njit(cache=True, parallel=True) +def _selected_groups(major_ptr, value_ptr, sel, group_offset, groups, sizes): + """Every stored group of the selected slices, in order, and how many indices each holds.""" + for i in numba.prange(sel.shape[0]): # ty: ignore[not-iterable] + out_g = group_offset[i] + for g in range(major_ptr[sel[i]], major_ptr[sel[i] + 1]): + groups[out_g] = g + sizes[out_g] = value_ptr[g + 1] - value_ptr[g] + out_g += 1 + + +@numba.njit(cache=True, parallel=True) +def _count_kept(major_ptr, value_ptr, buf, sel, starts, group_offset, remap, out_counts): + """How many indices of each selected group ``remap`` keeps (maps to ``>= 0``).""" + for i in numba.prange(sel.shape[0]): # ty: ignore[not-iterable] + pos = starts[i] + out_g = group_offset[i] + for g in range(major_ptr[sel[i]], major_ptr[sel[i] + 1]): + prev = np.int64(-1) + kept = 0 + for _ in range(value_ptr[g], value_ptr[g + 1]): + result, pos = _decode_varint(buf, pos) + prev = prev + 1 + np.int64(result) + if remap[prev] >= 0: + kept += 1 + out_counts[out_g] = kept + out_g += 1 + + +@numba.njit(cache=True, parallel=True) +def _fill_selected(major_ptr, value_ptr, buf, sel, starts, group_offset, remap, out_ptr, out): + """Decode the selected groups, writing each kept index's ``remap`` from ``out_ptr``.""" + for i in numba.prange(sel.shape[0]): # ty: ignore[not-iterable] + pos = starts[i] + out_g = group_offset[i] + for g in range(major_ptr[sel[i]], major_ptr[sel[i] + 1]): + prev = np.int64(-1) + k = out_ptr[out_g] + for _ in range(value_ptr[g], value_ptr[g + 1]): + result, pos = _decode_varint(buf, pos) + prev = prev + 1 + np.int64(result) + new = remap[prev] + if new >= 0: + out[k] = new + k += 1 + out_g += 1 + + +def _slice_starts( + major_ptr: np.ndarray, value_ptr: np.ndarray, buf: np.ndarray, sel: np.ndarray +) -> np.ndarray: + """Byte offset in ``buf`` where each selected major slice's indices start.""" + targets = value_ptr[major_ptr[sel]] + n_chunks = _num_chunks(buf.shape[0]) + real_starts, cum = _varint_chunks(buf, n_chunks) + # A target at the very end of the stream (an empty last slice) belongs to + # the last span, which walks it up to the final byte. + span_of = np.clip(np.searchsorted(cum, targets, side="right") - 1, 0, n_chunks - 1) + span_ptr = np.searchsorted(span_of, np.arange(n_chunks + 1)).astype(np.int64) + starts = np.empty(sel.shape[0], dtype=np.int64) + _locate_ranks(buf, real_starts, cum, targets, span_ptr, starts) + return starts + + +def unpack_selected( + major_ptr: np.ndarray, + value_ptr: np.ndarray, + packed: np.ndarray, + sel: np.ndarray, + remap: np.ndarray | None, + n_minor: int, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Decode only major slices ``sel``, and only the minor indices ``remap`` keeps. + + Parameters + ---------- + major_ptr, value_ptr, packed + The stored IVCS arrays. + sel + Sorted, unique major indices to decode. + remap + ``None`` to keep every minor index, or an ``int64`` array giving each + old minor index its new one, ``-1`` to drop it. + n_minor + Length of the minor axis afterwards. + + Returns + ------- + ``(major_ptr, groups, value_ptr, indices)`` for the selection, where + ``groups`` are the stored groups it keeps (so ``values[groups]`` are its + values). Groups left empty by ``remap`` are dropped, as indexing an + in-memory array would. + """ + major_ptr = np.ascontiguousarray(major_ptr, dtype=np.int64) + value_ptr = np.ascontiguousarray(value_ptr, dtype=np.int64) + buf = np.ascontiguousarray(packed, dtype=np.uint8) + sel = np.ascontiguousarray(sel, dtype=np.int64) + + # Every group of the selected slices, in order: group_offset is the + # selection's major_ptr before any group is emptied by `remap`. + group_offset = np.zeros(sel.shape[0] + 1, dtype=np.int64) + np.cumsum(major_ptr[sel + 1] - major_ptr[sel], out=group_offset[1:]) + groups = np.empty(int(group_offset[-1]), dtype=np.int64) + counts = np.empty(groups.shape[0], dtype=np.int64) + _selected_groups(major_ptr, value_ptr, sel, group_offset, groups, counts) + + starts = _slice_starts(major_ptr, value_ptr, buf, sel) + filtering = remap is not None + if remap is None: + remap = np.arange(n_minor, dtype=np.int64) + else: + _count_kept(major_ptr, value_ptr, buf, sel, starts, group_offset, remap, counts) + + out_ptr = np.zeros(groups.shape[0] + 1, dtype=np.int64) + np.cumsum(counts, out=out_ptr[1:]) + indices = np.empty(int(out_ptr[-1]), dtype=smallest_index_dtype(n_minor)) + _fill_selected(major_ptr, value_ptr, buf, sel, starts, group_offset, remap, out_ptr, indices) + if not filtering: + return group_offset, groups, out_ptr, indices + + nonempty = counts > 0 + nonempty_before = np.zeros(groups.shape[0] + 1, dtype=np.int64) + np.cumsum(nonempty, out=nonempty_before[1:]) + new_value_ptr = np.append(out_ptr[:-1][nonempty], out_ptr[-1]) + return nonempty_before[group_offset], groups[nonempty], new_value_ptr, indices diff --git a/src/vsparse/_rapid_load.py b/src/vsparse/_rapid_load.py index a1b4768..e2da7b7 100644 --- a/src/vsparse/_rapid_load.py +++ b/src/vsparse/_rapid_load.py @@ -2,18 +2,22 @@ :func:`load_packed` reads the packed ``X`` and decodes it immediately into an ordinary :class:`~vsparse.VCSCArray`/:class:`~vsparse.VCSRArray` inside a -:class:`~vsparse.VCSCAnnData`. Filtering and normalization are separate steps -(subsetting, then the ``.normalized()`` views), not part of loading. +:class:`~vsparse.VCSCAnnData`, optionally keeping only some cells and genes +as it decodes (``obs_idx``/``var_idx``). Normalization is a separate step (the +``.normalized()`` views), not part of loading. """ from __future__ import annotations from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import anndata as ad +import numpy as np +from vsparse import _io from vsparse._anndata_class import VCSCAnnData +from vsparse._indexutils import normalize_major_idx if TYPE_CHECKING: from os import PathLike @@ -23,12 +27,46 @@ _FIELD_KEYS = ("obs", "var", "obsm", "varm", "obsp", "varp", "layers", "uns") -def load_packed(path: str | PathLike[str], *, x_key: str = "X") -> VCSCAnnData: +def _positions(key: Any, n: int, name: str) -> np.ndarray | None: + """``key`` (a boolean mask, integer positions, or a slice) as int positions on an axis of ``n``.""" + if key is None: + return None + if isinstance(key, slice): + return np.arange(*key.indices(n), dtype=np.int64) + arr = np.asarray(key) + if arr.size == 0: + return np.empty(0, dtype=np.int64) + _check_selection(arr, n, name) + return normalize_major_idx(arr, n).astype(np.int64, copy=False) + + +def _check_selection(arr: np.ndarray, n: int, name: str) -> None: + if arr.ndim != 1 or not (arr.dtype == bool or np.issubdtype(arr.dtype, np.integer)): + raise TypeError( + f"{name} must be a 1-D boolean mask, an array of integer positions, " + f"or a slice; got {arr.dtype} with shape {arr.shape}" + ) + if arr.dtype == bool and arr.shape[0] != n: + raise IndexError(f"{name} is a boolean mask of length {arr.shape[0]}, expected {n}") + + +def load_packed( + path: str | PathLike[str], + *, + x_key: str = "X", + obs_idx: Any = None, + var_idx: Any = None, +) -> VCSCAnnData: """Load an IVCSR/IVCSC-backed ``.h5ad`` file, decoding ``X`` immediately. ``X`` comes back as an ordinary VCSCArray/VCSRArray with plain - ``indices``, decoded from the packed bytes read from disk. Subset and - normalize afterwards with the usual indexing and ``.normalized()`` views. + ``indices``, decoded from the packed bytes read from disk. Normalize + afterwards with the ``.normalized()`` views. + + ``load_packed(path, obs_idx=o, var_idx=v)`` equals ``load_packed(path)[o, v]``, + but decodes only the selected cells (for IVCSR; genes for IVCSC) and drops + unselected genes (cells) during that same decode, so the full matrix is + never built. Every packed byte is still read from disk. Parameters ---------- @@ -38,12 +76,33 @@ def load_packed(path: str | PathLike[str], *, x_key: str = "X") -> VCSCAnnData: :meth:`~vsparse.VCSCAnnData.write_h5ad`). x_key Top-level h5ad group holding the IVCSR/IVCSC array (``"X"`` by default). + obs_idx, var_idx + Cells/genes to keep: a boolean mask over the axis, an array of integer + positions (negative ones count from the end; unsorted or repeated + positions give that order), or a slice. ``None`` (the default) keeps + all. A mask built from metadata, e.g. ``obs["tissue"] == "blood"``, + works as is; to select by label, convert first with + ``obs.index.get_indexer(labels)``. """ import h5py import hdf5plugin # noqa: F401 -- registers the zstd/Blosc2 HDF5 filters with h5py.File(Path(path), "r") as f: - X = ad.io.read_elem(f[x_key]) kwargs = {k: ad.io.read_elem(f[k]) for k in _FIELD_KEYS if k in f} + if obs_idx is None and var_idx is None: + X = ad.io.read_elem(f[x_key]) + return VCSCAnnData(X=X, **kwargs) # ty: ignore[invalid-argument-type] + return _load_subset(f[x_key], kwargs, obs_idx, var_idx) + - return VCSCAnnData(X=X, **kwargs) # ty: ignore[invalid-argument-type] +def _load_subset(x_elem: Any, kwargs: dict[str, Any], obs_idx: Any, var_idx: Any) -> VCSCAnnData: + """:func:`load_packed` with a selection: ``X`` subset as it's decoded, the rest after reading.""" + n_obs, n_vars = _io._read_shape(x_elem) + rows = _positions(obs_idx, n_obs, "obs_idx") + cols = _positions(var_idx, n_vars, "var_idx") + X = _io.read_vcs_subset(x_elem, rows, cols) + meta = VCSCAnnData(**kwargs) + fields = meta._subset_fields( + slice(None) if rows is None else rows, slice(None) if cols is None else cols + ) + return VCSCAnnData(X=X, **fields) diff --git a/src/vsparse/_stream.py b/src/vsparse/_stream.py new file mode 100644 index 0000000..329245b --- /dev/null +++ b/src/vsparse/_stream.py @@ -0,0 +1,297 @@ +"""Read a packed IVCSR ``X`` a slab of bytes at a time, instead of all of it. + +:func:`vsparse.load_packed`'s ``obs_idx``/``var_idx`` decode only the slices you +ask for, but they still hold the whole packed stream while doing it (see +:func:`vsparse._ivcsc.unpack_selected`). That is the remaining half of +https://github.com/meyer-lab/vsparse/issues/67: 33 GB resident for the HEK293T +Xaira screen, 22 GB for Parse-10M, before the output is counted. + +This module reads the stream in windows instead. What makes that possible is +the layout (see :mod:`vsparse._ivcsc`): + +- ``major_ptr[r]:major_ptr[r + 1]`` are slice ``r``'s value groups, and groups + are written in slice order, so every slice owns one contiguous byte run; +- every value group restarts its delta at ``-1``, so a run of whole slices + decodes in isolation given the byte it starts at; +- so a window cut **at a slice boundary** needs only ``(next slice, next byte)`` + carried across it -- two integers, no decoder state. + +It is a sequential read, not a random one: ``value_ptr`` counts elements rather +than bytes, so nothing in the file says where slice ``r`` begins and the stream +is still walked front to back. Compression is not in the way -- +:func:`vsparse._compression.read_range` fetches whole zstd chunks through the +same ``chunk_table``/``pread``/``numcodecs`` path a full read uses, and never +seeks into a compressed chunk. The only overhead is re-decoding at most one +boundary chunk per window. + +``major_ptr``, ``value_ptr`` and ``values`` are read whole: ~12 bytes per value +group against 4 per nonzero, which on that screen is 4.1 GB rather than the +116 GB a decoded index array needs. That is this module's floor, and it scales +with value groups, not with ``nnz``. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import numpy as np + +from vsparse import _compression, _ivcsc +from vsparse._anndata_class import _FIELD_KEYS, VCSCAnnData, subset_fields +from vsparse._base import VCSCArray, VCSRArray + +if TYPE_CHECKING: + from collections.abc import Iterator + from os import PathLike + +__all__ = ["iter_blocks", "iter_packed", "open_packed", "rows_per_window"] + +IVCS_ENCODINGS = ("ivcsr", "ivcsc") + +#: Elements a window aims to cover. The packed bytes are ~1.1x this and a +#: decoded block is 4-8x it, so this sets the working set of any read the caller +#: didn't size: ~18 MB of bytes, ~67 MB decoded. Larger windows buy little -- +#: the per-window cost is one parallel byte scan and a few Python calls -- and +#: pay for it in residency. +_WINDOW_ELEMS = 1 << 24 +#: Bytes added on top of a sized read, so rounding never costs a whole retry. +_READ_SLACK = 1 << 12 +#: Headroom on a sized read, for a window whose varints run wider than the +#: file's average. +_SIZE_SLACK = 1.05 + + +@dataclass(frozen=True, slots=True) +class PackedX: + """An IVCSR/IVCSC group with everything but ``packed_indices`` in memory.""" + + #: ``(n_rows, n_cols)``, as stored. + shape: tuple[int, int] + #: ``"csr"`` or ``"csc"`` -- which axis the major one is. + format: str + #: Width the stored minor indices decode to. + indices_dtype: np.dtype + major_ptr: np.ndarray + values: np.ndarray + value_ptr: np.ndarray + #: ``value_ptr[major_ptr]``: the exact per-slice nnz indptr, free in + #: ``O(n_major)``, which sizes every read and output before a byte is read. + elem_ptr: np.ndarray + #: The undecoded ``packed_indices`` dataset. + packed: Any + #: Its :func:`vsparse._compression.chunk_table`, built once for all windows. + table: Any + #: Packed bytes per stored nonzero over the whole file, which sizes a read: + #: 1.0 if every delta fits one varint, up to :func:`_max_varint_bytes`. It is + #: known for free and a window sized from it lands first try. + bytes_per_elem: float + + @property + def n_major(self) -> int: + return self.major_ptr.shape[0] - 1 + + @property + def n_minor(self) -> int: + return self.shape[1] if self.format == "csr" else self.shape[0] + + def major_shape(self, n_major: int) -> tuple[int, int]: + """The shape of a block holding ``n_major`` major slices.""" + return (self.n_minor, n_major) if self.format == "csc" else (n_major, self.n_minor) + + +def open_packed(g: Any) -> PackedX: + """Read everything but ``packed_indices`` out of an IVCSR/IVCSC group.""" + encoding = g.attrs.get("encoding-type") + if encoding not in IVCS_ENCODINGS: + raise ValueError(f"{g.name} is stored {encoding!r}, not one of {IVCS_ENCODINGS}") + major_ptr = _compression.read_array(g["major_ptr"]) + value_ptr = _compression.read_array(g["value_ptr"]) + packed = g["packed_indices"] + return PackedX( + shape=(int(g.attrs["shape"][0]), int(g.attrs["shape"][1])), + format="csr" if encoding == "ivcsr" else "csc", + indices_dtype=np.dtype(g.attrs["indices_dtype"]), + major_ptr=major_ptr, + values=_compression.read_array(g["values"]), + value_ptr=value_ptr, + elem_ptr=value_ptr[major_ptr], + packed=packed, + table=_compression.chunk_table(packed), + bytes_per_elem=packed.shape[0] / max(1, int(value_ptr[-1])), + ) + + +def _max_varint_bytes(n_minor: int) -> int: + """Widest varint the encoder can emit, since every delta is ``< n_minor``.""" + return -(-max(1, int(n_minor - 1).bit_length()) // 7) + + +class _Cursor: + """Forward-only cursor over ``packed_indices``, carrying only ``(slice, byte)``.""" + + __slots__ = ("_px", "byte", "row") + + def __init__(self, px: PackedX) -> None: + self._px = px + self.row = 0 + self.byte = 0 + + def read_rows(self, row_hi: int) -> np.ndarray: + """Exactly the bytes of major slices ``[self.row, row_hi)``, and advance past them.""" + px = self._px + n_elem = int(px.elem_ptr[row_hi] - px.elem_ptr[self.row]) + buf, real_starts, cum = self._read_at_least(n_elem) + + # Where slice row_hi starts. The window's own span table already says + # how many varints precede each span, so this is a bounded walk inside + # one span (~a window/n_spans of bytes), not another pass over the + # window -- the counting pass is the expensive part and we have it. + span = int(np.searchsorted(cum, n_elem, side="left")) + if int(cum[span]) == n_elem: + end = int(real_starts[span]) + else: + end = int( + _ivcsc._advance_n_varints(buf, real_starts[span - 1], n_elem - int(cum[span - 1])) + ) + self.row = row_hi + self.byte += end + return buf[:end] + + def _read_at_least(self, n_elem: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """A window holding at least ``n_elem`` whole varints, with its span table. + + A short read costs a re-read *and* a re-count of the window, which is the + expensive half, so the first guess matters: it comes from the file's own + bytes-per-nonzero rather than from assuming single-byte varints. A retry + then rescales by what this window actually measured, which converges in + one step instead of creeping up by the shortfall. + """ + px = self._px + hard = max(int(n_elem * _max_varint_bytes(px.n_minor)), 1) + want = min(hard, max(int(n_elem * px.bytes_per_elem * _SIZE_SLACK) + _READ_SLACK, 1)) + while True: + buf = self._read(want) + real_starts, cum = _ivcsc._varint_chunks(buf, _ivcsc._num_chunks(buf.shape[0])) + have = int(cum[-1]) + if have >= n_elem: + return buf, real_starts, cum + if want >= hard or self.byte + want >= px.packed.shape[0]: + raise ValueError( + f"packed_indices holds {have} of the {n_elem} elements major slices " + f"starting at {self.row} should: the stream is truncated or corrupt" + ) + scaled = int(want * _SIZE_SLACK * n_elem / max(have, 1)) + _READ_SLACK + want = min(hard, max(scaled, want + 1)) + del buf # don't hold the short read while taking a longer one + + def _read(self, n_bytes: int) -> np.ndarray: + px = self._px + raw = _compression.read_range(px.packed, self.byte, self.byte + n_bytes, table=px.table) + # A window out of the middle of the stream usually ends mid-varint, and + # everything downstream of here assumes a true boundary. + return raw[: int(_ivcsc.trim_to_boundary(raw))] + + +def rows_per_window(px: PackedX) -> int: + """A major-slice count whose decoded block is roughly one window of elements.""" + per_row = max(1, int(px.elem_ptr[-1]) // max(1, px.n_major)) + return max(1, _WINDOW_ELEMS // per_row) + + +def iter_blocks(px: PackedX, *, chunk_size: int) -> Iterator[tuple[int, int, Any]]: + """``(row_lo, row_hi, array)`` for consecutive blocks of ``chunk_size`` major slices. + + The array layer of :func:`iter_packed`, without the AnnData fields -- what a + whole-file reduction that only needs the data wants, since reading ``obs`` + whole can cost more than the data does. + + A block is a contiguous run of slices, so its pointer arrays are slices of + the parent's and only the indices are decoded; nothing here needs + :func:`~vsparse._ivcsc.unpack_selected`'s selective machinery. + """ + cursor = _Cursor(px) + array_type = VCSRArray if px.format == "csr" else VCSCArray + for row_lo in range(0, px.n_major, chunk_size): + row_hi = min(row_lo + chunk_size, px.n_major) + buf = cursor.read_rows(row_hi) + g0, g1 = int(px.major_ptr[row_lo]), int(px.major_ptr[row_hi]) + value_local = px.value_ptr[g0 : g1 + 1] - int(px.elem_ptr[row_lo]) + yield ( + row_lo, + row_hi, + array_type( + px.major_shape(row_hi - row_lo), + px.major_ptr[row_lo : row_hi + 1] - g0, + px.values[g0:g1], + value_local, + _ivcsc.unpack_indices(value_local, buf, px.indices_dtype), + ), + ) + + +def iter_packed( + path: str | PathLike[str], *, chunk_size: int = 50_000, x_key: str = "X" +) -> Iterator[VCSCAnnData]: + """Yield consecutive row blocks of an IVCSR-backed ``.h5ad``, one block resident at a time. + + Block ``k`` is exactly ``load_packed(path)[lo:hi]``: ``obs`` and the other + per-cell fields are windowed to the block, ``var`` and the per-gene fields + come through whole, and stored normalization statistics are marked stale + (see :func:`vsparse._anndata_class.subset_fields`). + + This is what a whole-file reduction wants and ``obs_idx=`` cannot express. + Per-group means over every cell, say, are one pass here; as selective reads + they are one pass *per group*, each running the whole stream again, since + nothing in the file indexes a slice's bytes. + + Peak memory is about ``chunk_size * nnz_per_row * indices.itemsize`` for the + block plus the packed window, on top of the pointer arrays this module always + holds (see the module docstring). Blocks deliberately carry no statistics + sidecar: ``row_sums``/``row_nnz`` would window correctly but + ``col_sums``/``col_nnz``/``col_depth_sums`` are whole-matrix quantities, and + a block's are not the file's. + + Normalizing block by block is exact for a recipe whose scales are per-cell -- + ``cp10k_log1p`` is ``log1p(x * 1e4 / row_sums[i])``, so a block's own + statistics give the whole file's answer. It is **not** exact for ``parafac2`` + or ``pearson``, whose ``median(row_sums)`` target and ``col_depth_sums`` gene + scale are whole-population quantities; take those from the file's sidecar + (:func:`vsparse._stats.read_stats`) instead of per block. + + Parameters + ---------- + path + An ``.h5ad`` whose ``X`` (or the group named ``x_key``) was stored IVCSR. + chunk_size + Rows per block. The last block is whatever is left. + x_key + Top-level group holding the IVCSR array. + + Raises + ------ + NotImplementedError + If the array is stored IVCSC. There the major axis is genes, so a row + block is a *minor*-axis selection, which the packed layout can't serve + without decoding all of it; rewrite the file with + ``VCSCAnnData.from_anndata(..., format="csr")``. + """ + import anndata as ad + import h5py + import hdf5plugin # noqa: F401 -- registers the zstd/Blosc2 HDF5 filters + + if chunk_size <= 0: + raise ValueError(f"chunk_size must be positive, got {chunk_size}") + + with h5py.File(Path(path), "r") as f: + px = open_packed(f[x_key]) + if px.format != "csr": + raise NotImplementedError( + f"{x_key} is stored IVCSC (gene-major), so a row block is a minor-axis " + "selection and can't be read without decoding the whole array; rewrite it " + 'with VCSCAnnData.from_anndata(..., format="csr") to stream by cell' + ) + fields = {k: ad.io.read_elem(f[k]) for k in _FIELD_KEYS if k in f} + for row_lo, row_hi, block in iter_blocks(px, chunk_size=chunk_size): + yield VCSCAnnData(X=block, **subset_fields(fields, slice(row_lo, row_hi), slice(None))) diff --git a/tests/conftest.py b/tests/conftest.py index 7e0cce0..407db18 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -38,6 +38,35 @@ def csr(dense) -> sp.csr_array: return sp.csr_array(dense) +@pytest.fixture +def ivcs_file(tmp_path): + """Write a dense matrix to a packed ``.h5ad`` and hand back the path. + + ``write(dense, format="csr", obs=..., var=...)``. The streaming readers are + tested against ``load_packed`` on the same file, so the fixture exists to + keep the two sides from drifting apart in how the file was written. + """ + import anndata as ad + + from vsparse import VCSCAnnData + + count = 0 + + def write(dense, *, format="csr", obs=None, var=None, name=None): + nonlocal count + count += 1 + adata = ad.AnnData(X=sp.csr_array(dense)) + if obs is not None: + adata.obs = obs.set_index(adata.obs.index) + if var is not None: + adata.var = var.set_index(adata.var.index) + path = tmp_path / (name or f"data{count}.{format}.h5ad") + VCSCAnnData.from_anndata(adata, format=format).write_h5ad(path) + return path + + return write + + #: Threads the memory-limited tests run their kernels with. Their accumulators #: are ``nthreads * n_minor * width * 8`` bytes, so an unpinned count would make #: the same ceiling mean something different on every machine. diff --git a/tests/test_iter_packed.py b/tests/test_iter_packed.py new file mode 100644 index 0000000..3653233 --- /dev/null +++ b/tests/test_iter_packed.py @@ -0,0 +1,119 @@ +"""``iter_packed`` and the cursor underneath it (``vsparse._stream``). + +Block ``k`` has to equal ``load_packed(path)[lo:hi]``, so that is what most of +these assert. The rest force the cursor's window machinery to run many times +over a tiny fixture -- a window that ends mid-row, a read too small for the +elements it was sized for -- since on real files those paths run once in a +million rows and would otherwise never be exercised. +""" + +from __future__ import annotations + +import h5py +import numpy as np +import pandas as pd +import pytest + +from vsparse import _stream, iter_packed, load_packed + + +@pytest.fixture +def dense25(rng): + dense = (rng.random((25, 40)) < 0.35) * rng.integers(1, 30, (25, 40)) + dense[7] = 0 # an empty row owns no bytes at all + dense[11] = rng.integers(1, 30, 40) # a full one + return dense.astype(np.float32) + + +@pytest.fixture +def path(ivcs_file, dense25): + obs = pd.DataFrame({"grp": [f"c{i % 4}" for i in range(25)], "depth": dense25.sum(1)}) + return ivcs_file(dense25, obs=obs) + + +@pytest.mark.parametrize("chunk_size", [1, 2, 3, 7, 24, 25, 30]) +def test_blocks_reassemble_the_matrix(path, dense25, chunk_size): + blocks = list(iter_packed(path, chunk_size=chunk_size)) + assert len(blocks) == -(-25 // chunk_size) + np.testing.assert_array_equal(np.vstack([b.X.toarray() for b in blocks]), dense25) + + +@pytest.mark.parametrize("chunk_size", [1, 4, 25]) +def test_each_block_equals_the_same_in_memory_slice(path, chunk_size): + full = load_packed(path) + for i, block in enumerate(iter_packed(path, chunk_size=chunk_size)): + want = full[i * chunk_size : (i + 1) * chunk_size] + np.testing.assert_array_equal(block.X.toarray(), want.X.toarray()) + pd.testing.assert_frame_equal(block.obs, want.obs) + pd.testing.assert_frame_equal(block.var, want.var) + + +def test_blocks_carry_no_statistics_sidecar(path): + """``col_sums`` and friends are whole-matrix quantities; a block's are not the file's.""" + block = next(iter(iter_packed(path, chunk_size=5))) + assert block.X._stats is None + + +def test_a_wide_matrix_needs_more_bytes_than_elements(ivcs_file, rng): + """Column gaps over 127 take a second varint byte, so the sized read comes up short.""" + dense = ((rng.random((30, 5000)) < 0.01) * rng.integers(1, 9, (30, 5000))).astype(np.float32) + assert dense.sum(axis=0).max() > 0 + p = ivcs_file(dense) + with h5py.File(p, "r") as f: + px = _stream.open_packed(f["X"]) + assert px.packed.shape[0] > int(px.elem_ptr[-1]) # multi-byte varints really are present + np.testing.assert_array_equal( + np.vstack([b.X.toarray() for b in iter_packed(p, chunk_size=4)]), dense + ) + + +@pytest.mark.parametrize("chunk_size", [1, 4, 25]) +def test_the_cursor_accounts_for_every_byte(path, dense25, chunk_size): + """Blocks partition the stream exactly -- the cursor ends on its last byte. + + This is what catches an off-by-one in locating where the next slice starts: + a window one varint short or long still decodes its own block correctly, but + the byte total at the end will not land on ``len(packed_indices)``. + """ + with h5py.File(path, "r") as f: + px = _stream.open_packed(f["X"]) + cursor = _stream._Cursor(px) + for row_lo in range(0, px.n_major, chunk_size): + cursor.read_rows(min(row_lo + chunk_size, px.n_major)) + assert cursor.row == px.n_major + assert cursor.byte == px.packed.shape[0] + + for row_lo, row_hi, block in _stream.iter_blocks(px, chunk_size=chunk_size): + np.testing.assert_array_equal(block.toarray(), dense25[row_lo:row_hi]) + + +def test_rows_per_window_covers_the_matrix_in_a_few_blocks(path): + """Sized by elements, so a dense file gets fewer rows per window than a sparse one.""" + with h5py.File(path, "r") as f: + px = _stream.open_packed(f["X"]) + per_window = _stream.rows_per_window(px) + assert per_window >= 1 + # This fixture is far smaller than one window, so one block covers it. + assert per_window >= px.n_major + + +def test_an_all_zero_matrix_streams(ivcs_file): + """Every row is empty, so no row owns any bytes and the stream is empty.""" + p = ivcs_file(np.zeros((6, 4), dtype=np.float32)) + blocks = list(iter_packed(p, chunk_size=4)) + np.testing.assert_array_equal(np.vstack([b.X.toarray() for b in blocks]), np.zeros((6, 4))) + + +def test_a_file_with_no_rows_yields_nothing(ivcs_file): + assert list(iter_packed(ivcs_file(np.zeros((0, 4), dtype=np.float32)))) == [] + + +def test_chunk_size_must_be_positive(path): + with pytest.raises(ValueError, match="chunk_size must be positive"): + next(iter(iter_packed(path, chunk_size=0))) + + +def test_ivcsc_files_cannot_be_streamed_by_row(ivcs_file, dense25): + p = ivcs_file(dense25, format="csc") + with pytest.raises(NotImplementedError, match="stored IVCSC"): + next(iter(iter_packed(p))) diff --git a/tests/test_load_subset.py b/tests/test_load_subset.py new file mode 100644 index 0000000..0e4e9eb --- /dev/null +++ b/tests/test_load_subset.py @@ -0,0 +1,158 @@ +"""``load_packed(obs_idx=..., var_idx=...)``: subsetting while unpacking IVCSR/IVCSC.""" + +from __future__ import annotations + +import warnings +from typing import Any + +import numpy as np +import pandas as pd +import pytest +import scipy.sparse as sp + +from vsparse import VCSCAnnData, VCSCArray, VCSRArray, _ivcsc, load_packed + +N_OBS, N_VARS = 120, 3000 # wide, so many index deltas need multi-byte varints + + +def _counts(rng: np.random.Generator) -> sp.csr_array: + mat = sp.random_array((N_OBS, N_VARS), density=0.02, format="lil", random_state=0) + mat[:7] = 0 # empty rows, at both ends and in the middle + mat[-5:] = 0 + mat[50:53] = 0 + mat[:, 100:300] = 0 # and empty columns + out = sp.csr_array(mat) + out.data = rng.integers(1, 40, size=out.data.shape[0]).astype(np.float64) + return out + + +@pytest.fixture(params=[VCSRArray, VCSCArray], ids=["ivcsr", "ivcsc"]) +def packed_path(request, tmp_path): + rng = np.random.default_rng(0) + adata = VCSCAnnData( + X=request.param.from_scipy(_counts(rng)), + obs=pd.DataFrame( + {"group": rng.integers(0, 4, N_OBS)}, index=[f"c{i}" for i in range(N_OBS)] + ), + var=pd.DataFrame(index=[f"g{i}" for i in range(N_VARS)]), + obsm={"pca": rng.normal(size=(N_OBS, 3))}, + varm={"loadings": rng.normal(size=(N_VARS, 2))}, + obsp={"dist": sp.random_array((N_OBS, N_OBS), density=0.1, format="csr")}, + uns={"note": "kept"}, + ) + path = tmp_path / "packed.h5ad" + adata.write_h5ad(path) + return path + + +@pytest.fixture(params=[False, True], ids=["serial", "many_chunks"]) +def chunking(request, monkeypatch): + """Exercise both a single-span scan and many spans (normally only past 1 MiB).""" + if request.param: + monkeypatch.setattr(_ivcsc, "_num_chunks", lambda n_bytes: max(1, min(13, n_bytes))) + + +def _dense(mat: Any) -> np.ndarray: + return mat.toarray() + + +def _assert_same_x(left, right) -> None: + assert left.shape == right.shape + assert type(left) is type(right) + for name in ("major_ptr", "values", "value_ptr", "indices"): + np.testing.assert_array_equal(getattr(left, name), getattr(right, name), err_msg=name) + + +def _assert_matches_indexing(path, obs_idx, var_idx) -> None: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", UserWarning) # repeated positions -> duplicate names + subset = load_packed(path, obs_idx=obs_idx, var_idx=var_idx) + expected = load_packed(path)[ + slice(None) if obs_idx is None else obs_idx, + slice(None) if var_idx is None else var_idx, + ] + _assert_same_x(subset.X, expected.X) + pd.testing.assert_frame_equal(subset.obs, expected.obs) + pd.testing.assert_frame_equal(subset.var, expected.var) + np.testing.assert_array_equal(subset.obsm["pca"], expected.obsm["pca"]) + np.testing.assert_array_equal(subset.varm["loadings"], expected.varm["loadings"]) + np.testing.assert_array_equal(_dense(subset.obsp["dist"]), _dense(expected.obsp["dist"])) + assert subset.uns["note"] == "kept" + + +_RNG = np.random.default_rng(1) +SELECTIONS = { + "obs_mask": (_RNG.random(N_OBS) < 0.2, None), + "var_mask": (None, _RNG.random(N_VARS) < 0.3), + "both_masks": (_RNG.random(N_OBS) < 0.5, _RNG.random(N_VARS) < 0.1), + "sorted_positions": (np.array([0, 6, 7, 40, N_OBS - 1]), np.arange(0, N_VARS, 7)), + "unsorted_repeated_negative": (np.array([9, 3, 3, -1, 60]), np.array([2999, 5, 5, -3, 0])), + "slices": (slice(5, 100, 3), slice(None, None, -2)), + "empty_rows_only": (np.array([0, 1, 50, 51, N_OBS - 1]), None), + "nothing": (np.array([], dtype=np.int64), np.zeros(N_VARS, dtype=bool)), + "everything": (np.ones(N_OBS, dtype=bool), np.arange(N_VARS)), + "list_of_ints": ([4, 8, 15, 16, 23, 42], None), +} + + +@pytest.mark.parametrize("selection", SELECTIONS.values(), ids=SELECTIONS.keys()) +def test_subset_load_matches_indexing_a_full_load(packed_path, chunking, selection): + _assert_matches_indexing(packed_path, *selection) + + +def test_subset_recomputes_rather_than_inherits_the_stats_sidecar(packed_path): + mask = np.zeros(N_OBS, dtype=bool) + mask[::3] = True + subset: Any = load_packed(packed_path, obs_idx=mask).X + full: Any = load_packed(packed_path).X + assert subset._stats is None + expected = full[mask] + np.testing.assert_array_equal(subset.stats.row_sums, expected.stats.row_sums) + np.testing.assert_array_equal(subset.stats.col_depth_sums, expected.stats.col_depth_sums) + + +def test_metadata_masks_work_directly(packed_path): + full = load_packed(packed_path) + mask = full.obs["group"] == 2 # a pandas Series + _assert_matches_indexing(packed_path, mask, None) + + +@pytest.mark.parametrize( + ("kwargs", "error"), + [ + ({"obs_idx": np.array(["c1", "c2"])}, TypeError), + ({"var_idx": np.array([0.5, 1.0])}, TypeError), + ({"obs_idx": np.ones((2, 2), dtype=bool)}, TypeError), + ({"obs_idx": np.ones(N_OBS - 1, dtype=bool)}, IndexError), + ({"var_idx": np.array([N_VARS])}, IndexError), + ({"obs_idx": np.array([-N_OBS - 1])}, IndexError), + ], +) +def test_invalid_selections_are_rejected(packed_path, kwargs, error): + with pytest.raises(error): + load_packed(packed_path, **kwargs) + + +@pytest.mark.parametrize("seed", range(5)) +def test_unpack_selected_matches_a_full_decode(chunking, seed): + rng = np.random.default_rng(seed) + arr = VCSRArray.from_scipy(_counts(rng)) + packed = _ivcsc.pack_indices(arr.value_ptr, arr.indices) + full = VCSRArray( + arr.shape, + arr.major_ptr, + arr.values, + arr.value_ptr, + _ivcsc.unpack_indices(arr.value_ptr, packed, arr.indices.dtype), + ) + sel = np.flatnonzero(rng.random(N_OBS) < 0.3) + cols = np.flatnonzero(rng.random(N_VARS) < 0.4) + remap = np.full(N_VARS, -1, dtype=np.int64) + remap[cols] = np.arange(cols.shape[0]) + + major_ptr, groups, value_ptr, indices = _ivcsc.unpack_selected( + arr.major_ptr, arr.value_ptr, packed, sel, remap, cols.shape[0] + ) + values = arr.values[groups] # noqa: PD011 -- a VCSRArray, not pandas + got = VCSRArray((sel.shape[0], cols.shape[0]), major_ptr, values, value_ptr, indices) + _assert_same_x(got, full[sel][:, cols]) diff --git a/tests/test_parallel_read.py b/tests/test_parallel_read.py index 9c73ba0..326a483 100644 --- a/tests/test_parallel_read.py +++ b/tests/test_parallel_read.py @@ -1,4 +1,8 @@ -"""Tests for the parallel zstd chunk reader and its fallbacks (``_compression.read_array``).""" +"""Tests for the parallel zstd chunk reader, its byte-range form, and their fallbacks. + +See ``vsparse._compression``: :func:`read_range` is the primitive, +:func:`read_array` is it over a whole dataset. +""" from __future__ import annotations @@ -47,6 +51,60 @@ def test_parallel_read_matches(tmp_path, spy, dtype, n): assert spy == ["/a"] +@pytest.mark.parametrize( + ("start", "stop"), + [ + (0, 4321), # the whole dataset + (0, 1), # one element, first chunk + (999, 1001), # across a chunk boundary + (1000, 2000), # exactly one whole chunk + (1500, 4321), # partial first chunk through the partial last one + (4320, 4321), # last element + (2000, 2000), # empty, mid-stream + (0, 0), # empty, at the start + ], +) +def test_read_range_matches_slicing(tmp_path, spy, start, stop): + """A byte range decodes only the chunks that cover it, and matches h5py's slice.""" + data = np.arange(4321, dtype=np.int64) + path = tmp_path / "r.h5" + with h5py.File(path, "w") as f: + f.create_dataset("a", data=data, chunks=(1000,), **ZSTD) + with h5py.File(path, "r") as f: + np.testing.assert_array_equal( + _compression.read_range(f["a"], start, stop), data[start:stop] + ) + assert spy == ["/a"] + + +def test_read_range_reuses_a_cached_chunk_table(tmp_path, spy): + """A caller streaming in windows builds the chunk table once and passes it back.""" + data = np.arange(5000, dtype=np.int32) + path = tmp_path / "cached.h5" + with h5py.File(path, "w") as f: + f.create_dataset("a", data=data, chunks=(1000,), **ZSTD) + with h5py.File(path, "r") as f: + table = _compression.chunk_table(f["a"]) + assert table is not None + for lo in range(0, 5000, 700): + np.testing.assert_array_equal( + _compression.read_range(f["a"], lo, lo + 700, table=table), + data[lo : lo + 700], + ) + assert spy == ["/a"] * 8 + + +def test_chunk_table_is_none_for_other_layouts(tmp_path): + """``chunk_table`` is the eligibility test as well as the table.""" + data = np.arange(1234, dtype=np.int64) + path = tmp_path / "plain.h5" + with h5py.File(path, "w") as f: + f.create_dataset("a", data=data) + with h5py.File(path, "r") as f: + assert _compression.chunk_table(f["a"]) is None + np.testing.assert_array_equal(_compression.read_range(f["a"], 10, 20), data[10:20]) + + def test_userblock_offsets(tmp_path, spy): """Chunk offsets from chunk_iter already include any userblock.""" data = np.arange(50_000, dtype=np.int64)