From f189233612ae0114d176ad3d495f85bcbf886656 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Wed, 14 Oct 2026 07:59:32 -0700 Subject: [PATCH 1/3] Replace Blosc2 with zstd-1 and decode chunks in parallel on read Blosc2+LZ4 compresses the IVCSR packed_indices stream by ~1.00x on real files (the delta+varint bytes have no repeats for an LZ codec to find), while zstd's entropy coding gets ~1.15-1.5x. zstd at level 1 is now the default for every numeric array written to .h5ad and zarr. HDF5 runs a dataset's filter pipeline one chunk at a time on one thread, so reads are decode-bound. read_array() looks up every chunk's file offset in one chunk_iter pass, preads the compressed chunks, and decodes them with numcodecs.Zstd on a thread pool sized by numba's thread count. VCS/IVCS arrays read through the anndata registry and the rapid loader both use it. Anything else -- Blosc2 files from earlier versions, uncompressed or contiguous datasets, unwritten chunks, files open for writing or on non-sec2 drivers -- falls back to an ordinary h5py read. read_h5ad now imports hdf5plugin itself: in a fresh process it previously failed on compressed obs/var columns, which are read before X. Claude-Session: https://claude.ai/code/session_01QBeTZy6BKKGKGJ7fwvkWgM --- README.md | 2 +- docs/usage.md | 11 +- pyproject.toml | 1 + src/vsparse/_anndata_class.py | 16 ++- src/vsparse/_compression.py | 146 +++++++++++++++++++++++---- src/vsparse/_io.py | 11 +- src/vsparse/_rapid_load.py | 13 +-- tests/test_anndata_class.py | 4 +- tests/test_compression.py | 16 +-- tests/test_parallel_read.py | 185 ++++++++++++++++++++++++++++++++++ uv.lock | 2 + 11 files changed, 357 insertions(+), 50 deletions(-) create mode 100644 tests/test_parallel_read.py diff --git a/README.md b/README.md index 323866d..eedb428 100644 --- a/README.md +++ b/README.md @@ -108,7 +108,7 @@ va = vsparse.VCSCAnnData.from_anndata(adata) # Compresses X and raw.X va.X # VCSCArray va.raw_X # VCSCArray (separate from anndata's .raw) -# Persist to HDF5 (.h5ad) or Zarr with default Blosc2+LZ4 compression +# Persist to HDF5 (.h5ad) or Zarr with default zstd compression # (stored as IVCSC/IVCSR, see below) va.write_h5ad("compressed.h5ad") # Read back with VCSCAnnData.read_h5ad va2 = vsparse.VCSCAnnData.read_h5ad("compressed.h5ad") diff --git a/docs/usage.md b/docs/usage.md index 4fdfb88..7f5bd0a 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -84,10 +84,13 @@ allowlist of types, so a plain `AnnData` cannot hold a `VCSCArray` in `X`. ### On-disk compression and the IVCSC/IVCSR storage format -`write_h5ad`/`write_zarr` compress every array with Blosc2+LZ4 by default -(`h5py`/`hdf5plugin` for `.h5ad`, `numcodecs`/zarr's native Blosc codec for -zarr). Pass `dataset_kwargs={}` to write uncompressed, or your own -`dataset_kwargs` to use a different codec. +`write_h5ad`/`write_zarr` compress every numeric array with zstd at level 1 +by default (`h5py`/`hdf5plugin` for `.h5ad`, `numcodecs`/zarr's native zstd +codec for zarr). On read, `.h5ad` arrays are decompressed chunk-by-chunk in +parallel (`numba.get_num_threads()` threads) rather than through HDF5's +single-threaded filter pipeline. Files written with the earlier Blosc2+LZ4 +default still read, through the ordinary h5py path. Pass `dataset_kwargs={}` +to write uncompressed, or your own `dataset_kwargs` to use a different codec. `X`/`raw_X` are always stored as **IVCSC/IVCSR** -- the VCSC/VCSR layout, but with the minor-axis `indices` array delta+varint byte-packed diff --git a/pyproject.toml b/pyproject.toml index 1010d52..b89cba6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ dependencies = [ "numba>=0.60", "anndata>=0.13", "hdf5plugin>=4.0", + "numcodecs>=0.13", ] [project.urls] diff --git a/src/vsparse/_anndata_class.py b/src/vsparse/_anndata_class.py index 98b6137..69402e7 100644 --- a/src/vsparse/_anndata_class.py +++ b/src/vsparse/_anndata_class.py @@ -470,9 +470,10 @@ def write_h5ad( # ty: ignore[invalid-method-override] with fewer categories than rows are converted. dataset_kwargs Passed to ``h5py.Group.create_dataset`` for every array written. - Defaults to Blosc2+LZ4 compression; pass ``{}`` to store - uncompressed. Either way, compression is only ever applied to - numeric arrays -- see :func:`vsparse._compression.numeric_only_compression`. + Defaults to zstd (level 1) compression, which :meth:`read_h5ad` + decodes in parallel; pass ``{}`` to store uncompressed. Either + way, compression is only ever applied to numeric arrays -- see + :func:`vsparse._compression.numeric_only_compression`. """ import h5py @@ -490,8 +491,13 @@ def write_h5ad( # ty: ignore[invalid-method-override] @classmethod def read_h5ad(cls, filename: str | PathLike[str]) -> VCSCAnnData: - """Read a file written by :meth:`write_h5ad`.""" + """Read a file written by :meth:`write_h5ad`. + + Reads files compressed with the current zstd default and with the + Blosc2 default of earlier versions alike. + """ import h5py + import hdf5plugin # noqa: F401 -- registers the zstd/Blosc2 HDF5 filters with h5py.File(filename, "r") as f: return cls._read_group(f) @@ -508,7 +514,7 @@ def write_zarr( See :meth:`write_h5ad` for the IVCSC/IVCSR storage format and ``convert_strings_to_categoricals``/``dataset_kwargs`` (including the numeric-only compression behavior); - the default compression here is Blosc+LZ4 via ``numcodecs``. + the default compression here is zstd (level 1) via ``numcodecs``. """ import zarr diff --git a/src/vsparse/_compression.py b/src/vsparse/_compression.py index 22c116d..bb53379 100644 --- a/src/vsparse/_compression.py +++ b/src/vsparse/_compression.py @@ -1,16 +1,30 @@ -"""Default on-disk compression for arrays written via anndata's IO registry. - -Blosc2+LZ4 gives a good default speed/ratio tradeoff and, unlike gzip, is -fast enough to leave on by default. It requires the ``hdf5plugin`` filter to -be registered for HDF5; zarr already ships blosc support via ``numcodecs``. - -Blosc is only ever applied to *numeric* arrays. anndata applies whatever -``dataset_kwargs`` it's given uniformly to every array it writes, including -the small variable-length string arrays nested inside dataframes and -categoricals (category labels, ``obs``/``var`` string columns, the +"""On-disk compression: zstd-1 on write, parallel chunk decoding on read. + +zstd at level 1 is the default for every numeric array written. On the +IVCSR ``packed_indices`` stream -- delta+varint bytes that are mostly +single-byte values with a skewed distribution -- byte-oriented LZ codecs +find nothing to match (Blosc2+LZ4, the previous default, stored real files +at 1.00x), while zstd's entropy coding still gets ~1.15-1.5x. Higher zstd +levels don't help, since there are no long-range repeats to find. It +requires the ``hdf5plugin`` filter to be registered for HDF5; zarr ships zstd +support via ``numcodecs``. + +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 +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 +contiguous datasets, in-memory/non-POSIX drivers -- falls back to an ordinary +h5py read, so those files keep loading exactly as before. + +Compression is only ever applied to *numeric* arrays. anndata applies +whatever ``dataset_kwargs`` it's given uniformly to every array it writes, +including the small variable-length string arrays nested inside dataframes +and categoricals (category labels, ``obs``/``var`` string columns, the ``_index``). At least some HDF5 filter-plugin builds (seen in this environment: h5py 3.16 / HDF5 2.0.0 / hdf5plugin's Blosc2) segfault -(``SIGFPE``) when the Blosc2 filter is applied to a variable-length-string +(``SIGFPE``) when a plugin filter is applied to a variable-length-string dataset. :func:`numeric_only_compression` patches the relevant ``create_dataset``/``create_array`` calls for the duration of a write so string/object arrays always land uncompressed, regardless of what @@ -19,33 +33,40 @@ Strings are left uncompressed rather than given a different codec: gzip and lzf are safe on variable-length strings here, but the size problem is the per-row string itself, not its compression. Encoding low-cardinality -columns as categoricals turns the per-row data numeric, which the existing -Blosc2 path then compresses -- see +columns as categoricals turns the per-row data numeric, which the numeric +path then compresses -- see :meth:`vsparse.VCSCAnnData.write_h5ad`'s ``convert_strings_to_categoricals``. """ from __future__ import annotations import contextlib +import os +from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any, cast +import numba import numpy as np if TYPE_CHECKING: from collections.abc import Iterator -__all__ = ["h5_dataset_kwargs", "numeric_only_compression", "zarr_dataset_kwargs"] +__all__ = ["h5_dataset_kwargs", "numeric_only_compression", "read_array", "zarr_dataset_kwargs"] _H5_COMPRESSION_KEYS = ("compression", "compression_opts") _ZARR_V2_COMPRESSION_KEYS = ("compressor",) _ZARR_V3_COMPRESSION_KEYS = ("compressors", "compressor") +_ZSTD_LEVEL = 1 +# Registered HDF5 filter id for zstd (https://github.com/HDFGroup/hdf5_plugins). +_H5_ZSTD_FILTER_ID = 32015 + def h5_dataset_kwargs() -> dict[str, Any]: - """``h5py.Group.create_dataset`` kwargs for Blosc2+LZ4 compression.""" + """``h5py.Group.create_dataset`` kwargs for zstd-1 compression.""" import hdf5plugin - return dict(hdf5plugin.Blosc2(cname="lz4")) + return dict(hdf5plugin.Zstd(clevel=_ZSTD_LEVEL)) def _is_zarr_v2() -> bool: @@ -58,15 +79,15 @@ def _is_zarr_v2() -> bool: def zarr_dataset_kwargs() -> dict[str, Any]: - """``dataset_kwargs`` for Blosc+LZ4 compression, understood by anndata's zarr writers.""" + """``dataset_kwargs`` for zstd-1 compression, understood by anndata's zarr writers.""" if _is_zarr_v2(): import numcodecs - return {"compressor": numcodecs.Blosc(cname="lz4")} + return {"compressor": numcodecs.Zstd(level=_ZSTD_LEVEL)} import zarr.codecs - return {"compressor": zarr.codecs.BloscCodec(cname="lz4")} + return {"compressor": zarr.codecs.ZstdCodec(level=_ZSTD_LEVEL)} def _is_string_like(dtype: Any, data: Any) -> bool: @@ -166,3 +187,90 @@ def numeric_only_compression(store_kind: str) -> contextlib.AbstractContextManag if store_kind == "zarr": return _numeric_only_zarr() raise ValueError(f"store_kind must be 'h5' or 'zarr', got {store_kind!r}") + + +# -- parallel zstd read --------------------------------------------------------- + + +def _is_parallel_zstd(ds: Any) -> bool: + """True if ``ds`` is a chunked, zstd-only, 1-D dataset in a plain read-only HDF5 file.""" + import h5py + + if not isinstance(ds, h5py.Dataset) or ds.ndim != 1 or ds.chunks is None: + return False + # The chunks are read with os.pread on the file itself, which needs a + # POSIX file on disk that nothing is still writing to. + if not hasattr(os, "pread") or ds.file.mode != "r" or ds.file.driver != "sec2": + return False + plist = ds.id.get_create_plist() + 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: + """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``. + """ + chunk_len = ds.chunks[0] + n_chunks = -(-ds.shape[0] // chunk_len) + offsets = np.full(n_chunks, -1, dtype=np.int64) + sizes = np.zeros(n_chunks, dtype=np.int64) + masks = np.zeros(n_chunks, dtype=np.int64) + + def record(info: Any) -> None: + i = info.chunk_offset[0] // chunk_len + offsets[i], sizes[i], masks[i] = info.byte_offset, info.size, info.filter_mask + + try: + ds.id.chunk_iter(record) + except (AttributeError, NotImplementedError): + return None + return None if (offsets < 0).any() else (offsets, sizes, masks) + + +def _read_chunks_parallel( + ds: Any, offsets: np.ndarray, sizes: np.ndarray, masks: np.ndarray +) -> np.ndarray: + import numcodecs + + out = np.empty(ds.shape, dtype=ds.dtype) + flat = out.view(np.uint8) + chunk_bytes = ds.chunks[0] * 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])) + # Filter-mask bit 0 set means HDF5 skipped zstd and stored the chunk as-is. + data = raw if masks[i] & 1 else codec.decode(raw) + dst = flat[i * chunk_bytes : (i + 1) * chunk_bytes] + # The last chunk decodes to a full chunk; only its leading part is data. + dst[:] = np.frombuffer(data, dtype=np.uint8, count=dst.size) + + try: + n_workers = max(1, min(numba.get_num_threads(), offsets.shape[0])) + with ThreadPoolExecutor(max_workers=n_workers) as pool: + list(pool.map(load, range(offsets.shape[0]))) + finally: + os.close(fd) + return out + + +def read_array(ds: Any) -> np.ndarray: + """Read a whole dataset into memory, decoding zstd chunks in parallel where possible. + + ``ds`` is an h5py dataset or zarr array. A zstd-compressed, chunked, 1-D + h5py dataset (what :func:`h5_dataset_kwargs` writes for every VCS array) + is decoded chunk-by-chunk on a thread pool sized by + ``numba.get_num_threads()``; everything else, including Blosc2-compressed + 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) diff --git a/src/vsparse/_io.py b/src/vsparse/_io.py index 9d1fccf..e8f5fc9 100644 --- a/src/vsparse/_io.py +++ b/src/vsparse/_io.py @@ -25,6 +25,7 @@ from vsparse import _ivcsc from vsparse._base import VCSCArray, VCSRArray, _VCSBase +from vsparse._compression import read_array from vsparse._indexutils import smallest_index_dtype if TYPE_CHECKING: @@ -60,7 +61,7 @@ def _make_read_vcs(cls: type[_VCSBase]): def _read(elem: GroupStorageType, *, _reader: Reader) -> _VCSBase: shape_vals = [int(s) for s in np.asarray(elem.attrs["shape"]).tolist()] shape = (shape_vals[0], shape_vals[1]) - arrays = {name: cast(np.ndarray, _reader.read_elem(elem[name])) for name in _ARRAY_KEYS} + arrays = {name: read_array(elem[name]) for name in _ARRAY_KEYS} return cls( shape, arrays["major_ptr"], @@ -113,10 +114,10 @@ def _make_read_ivcs(cls: type[_VCSBase]): def _read(elem: GroupStorageType, *, _reader: Reader) -> _VCSBase: shape_vals = [int(s) for s in np.asarray(elem.attrs["shape"]).tolist()] shape = (shape_vals[0], shape_vals[1]) - major_ptr = cast(np.ndarray, _reader.read_elem(elem["major_ptr"])) - values = cast(np.ndarray, _reader.read_elem(elem["values"])) - value_ptr = cast(np.ndarray, _reader.read_elem(elem["value_ptr"])) - packed = cast(np.ndarray, _reader.read_elem(elem["packed_indices"])) + major_ptr = read_array(elem["major_ptr"]) + values = read_array(elem["values"]) + value_ptr = read_array(elem["value_ptr"]) + packed = read_array(elem["packed_indices"]) dtype = np.dtype(cast(str, elem.attrs["indices_dtype"])) indices = _ivcsc.unpack_indices(value_ptr, packed, dtype) return cls(shape, major_ptr, values, value_ptr, indices) diff --git a/src/vsparse/_rapid_load.py b/src/vsparse/_rapid_load.py index aa0455d..28ea23a 100644 --- a/src/vsparse/_rapid_load.py +++ b/src/vsparse/_rapid_load.py @@ -61,6 +61,7 @@ from vsparse import _ivcsc from vsparse._anndata_class import VCSCAnnData +from vsparse._compression import read_array from vsparse._indexutils import smallest_index_dtype from vsparse._ops import accumulator_threads @@ -451,10 +452,10 @@ def _read_ivcsr_group( g = f[x_key] shape = (int(g.attrs["shape"][0]), int(g.attrs["shape"][1])) indices_dtype = np.dtype(g.attrs["indices_dtype"]) - major_ptr = g["major_ptr"][...] - values = g["values"][...] - value_ptr = g["value_ptr"][...] - packed = g["packed_indices"][...] + major_ptr = read_array(g["major_ptr"]) + values = read_array(g["values"]) + value_ptr = read_array(g["value_ptr"]) + packed = read_array(g["packed_indices"]) kwargs = {k: ad.io.read_elem(f[k]) for k in _FIELD_KEYS if k in f} if "raw" in f: @@ -527,7 +528,7 @@ def load_and_normalize( ) """ import h5py - import hdf5plugin # noqa: F401 -- registers the Blosc2 HDF5 filter + import hdf5plugin # noqa: F401 -- registers the zstd/Blosc2 HDF5 filters _validate_min_cells(min_cells) @@ -600,7 +601,7 @@ def load_packed(path: str | PathLike[str], *, x_key: str = "X") -> VCSCAnnData: Top-level h5ad group holding the IVCSR/IVCSC array (``"X"`` by default). """ import h5py - import hdf5plugin # noqa: F401 -- registers the Blosc2 HDF5 filter + 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]) diff --git a/tests/test_anndata_class.py b/tests/test_anndata_class.py index 85b7bb3..73d4367 100644 --- a/tests/test_anndata_class.py +++ b/tests/test_anndata_class.py @@ -145,9 +145,9 @@ def test_write_h5ad_default_compression(base_adata, tmp_path): import h5py with h5py.File(path, "r") as f: - # Blosc2 is registered as a dynamically-loaded HDF5 filter, so h5py + # zstd is registered as a dynamically-loaded HDF5 filter, so h5py # doesn't recognize it via `compression` -- check the raw filter id. - assert any(int(fid) >= 32000 for fid in f["X"]["values"]._filters) + assert "32015" in f["X"]["values"]._filters def test_write_h5ad_dataset_kwargs_override(base_adata, tmp_path): diff --git a/tests/test_compression.py b/tests/test_compression.py index 44394ba..f91f932 100644 --- a/tests/test_compression.py +++ b/tests/test_compression.py @@ -10,16 +10,16 @@ def test_h5_dataset_kwargs_options(): - """Verify h5_dataset_kwargs correctly builds filter parameters.""" - blosc_kwargs = _compression.h5_dataset_kwargs() - assert "compression" in blosc_kwargs - assert "compression_opts" in blosc_kwargs + """Verify h5_dataset_kwargs selects the zstd filter at level 1.""" + kwargs = _compression.h5_dataset_kwargs() + assert kwargs["compression"] == _compression._H5_ZSTD_FILTER_ID + assert tuple(kwargs["compression_opts"]) == (1,) def test_zarr_dataset_kwargs_options(): """Verify zarr_dataset_kwargs correctly builds compressor dict.""" - blosc_kwargs = _compression.zarr_dataset_kwargs() - assert "compressor" in blosc_kwargs or "codecs" in blosc_kwargs + kwargs = _compression.zarr_dataset_kwargs() + assert "zstd" in type(kwargs["compressor"]).__name__.lower() def test_numeric_only_compression_invalid_backend(): @@ -55,7 +55,7 @@ def test_numeric_only_compression_skips_strings(tmp_path): f.create_dataset("strings", data=np.array(["a", "b", "c"], dtype=object), **kwargs) with h5py.File(path, "r") as f: - # Numeric dataset has blosc2 filter - assert any(int(fid) >= 32000 for fid in f["numbers"]._filters) + # Numeric dataset has the zstd filter + assert str(_compression._H5_ZSTD_FILTER_ID) in f["numbers"]._filters # String dataset has no compression filters assert dict(f["strings"]._filters) == {} diff --git a/tests/test_parallel_read.py b/tests/test_parallel_read.py new file mode 100644 index 0000000..9c73ba0 --- /dev/null +++ b/tests/test_parallel_read.py @@ -0,0 +1,185 @@ +"""Tests for the parallel zstd chunk reader and its fallbacks (``_compression.read_array``).""" + +from __future__ import annotations + +import subprocess +import sys + +import h5py +import hdf5plugin +import numcodecs +import numpy as np +import pandas as pd +import pytest +import scipy.sparse as sp + +from vsparse import VCSRArray, VCSCAnnData, _compression + +ZSTD = _compression.h5_dataset_kwargs() + + +@pytest.fixture +def spy(monkeypatch): + """Record whether read_array took the parallel path.""" + calls = [] + real = _compression._read_chunks_parallel + + def wrapped(*args): + calls.append(args[0].name) + return real(*args) + + monkeypatch.setattr(_compression, "_read_chunks_parallel", wrapped) + return calls + + +@pytest.mark.parametrize("dtype", [np.uint8, np.uint16, np.int32, np.int64, np.float32]) +@pytest.mark.parametrize("n", [0, 1, 999, 1000, 4321]) +def test_parallel_read_matches(tmp_path, spy, dtype, n): + """Chunked zstd datasets decode in parallel, including a partial final chunk.""" + data = (np.random.default_rng(0).integers(0, 200, n)).astype(dtype) + path = tmp_path / "a.h5" + with h5py.File(path, "w") as f: + f.create_dataset("a", data=data, chunks=(1000,) if n else None, maxshape=(None,), **ZSTD) + with h5py.File(path, "r") as f: + out = _compression.read_array(f["a"]) + assert out.dtype == data.dtype + np.testing.assert_array_equal(out, data) + assert spy == ["/a"] + + +def test_userblock_offsets(tmp_path, spy): + """Chunk offsets from chunk_iter already include any userblock.""" + data = np.arange(50_000, dtype=np.int64) + path = tmp_path / "ub.h5" + with h5py.File(path, "w", userblock_size=512) as f: + f.create_dataset("a", data=data, chunks=(4096,), **ZSTD) + with h5py.File(path, "r") as f: + np.testing.assert_array_equal(_compression.read_array(f["a"]), data) + assert spy == ["/a"] + + +def test_chunk_stored_without_filter(tmp_path, spy): + """A chunk HDF5 stored raw (filter-mask bit set) is copied, not decoded.""" + data = np.arange(3000, dtype=np.int32) + path = tmp_path / "mask.h5" + with h5py.File(path, "w") as f: + ds = f.create_dataset("a", shape=(3000,), dtype=np.int32, chunks=(1000,), **ZSTD) + ds[:] = data + ds.id.write_direct_chunk((1000,), data[1000:2000].tobytes(), filter_mask=1) + with h5py.File(path, "r") as f: + np.testing.assert_array_equal(_compression.read_array(f["a"]), data) + assert spy == ["/a"] + + +def test_unwritten_chunk_falls_back(tmp_path, spy): + """A never-written chunk reads as the fill value via the ordinary path.""" + path = tmp_path / "sparse.h5" + with h5py.File(path, "w") as f: + ds = f.create_dataset( + "a", shape=(3000,), dtype=np.int64, chunks=(1000,), fillvalue=7, **ZSTD + ) + ds[:1000] = 1 + with h5py.File(path, "r") as f: + out = _compression.read_array(f["a"]) + np.testing.assert_array_equal(out, np.r_[np.ones(1000), np.full(2000, 7)]) + assert spy == [] + + +@pytest.mark.parametrize( + "kwargs", + [ + {}, # contiguous, uncompressed + {"chunks": (100,)}, # chunked, uncompressed + dict(hdf5plugin.Blosc2(cname="lz4")), # the previous default + {"chunks": (100,), "shuffle": True, **ZSTD}, # zstd plus another filter + ], + ids=["contiguous", "chunked", "blosc2", "zstd+shuffle"], +) +def test_other_layouts_fall_back(tmp_path, spy, kwargs): + data = np.arange(1234, dtype=np.int64) + path = tmp_path / "other.h5" + with h5py.File(path, "w") as f: + f.create_dataset("a", data=data, **kwargs) + with h5py.File(path, "r") as f: + np.testing.assert_array_equal(_compression.read_array(f["a"]), data) + assert spy == [] + + +def test_writable_or_in_memory_file_falls_back(tmp_path, spy): + data = np.arange(5000, dtype=np.int64) + path = tmp_path / "rw.h5" + with h5py.File(path, "w") as f: + f.create_dataset("a", data=data, chunks=(1000,), **ZSTD) + np.testing.assert_array_equal(_compression.read_array(f["a"]), data) + with h5py.File(path, "r", driver="core") as f: + np.testing.assert_array_equal(_compression.read_array(f["a"]), data) + assert spy == [] + + +def test_zarr_array_reads(tmp_path): + zarr = pytest.importorskip("zarr") + data = np.arange(5000, dtype=np.int64) + arr = zarr.create_array(tmp_path / "z.zarr", shape=data.shape, dtype=data.dtype, chunks=(1000,)) + arr[:] = data + np.testing.assert_array_equal(_compression.read_array(arr), data) + + +def _vcsr_anndata(n_obs=400, n_vars=60) -> VCSCAnnData: + rng = np.random.default_rng(1) + m = sp.random(n_obs, n_vars, density=0.2, format="csr", random_state=rng) + m.data = rng.integers(1, 20, m.nnz).astype(np.float32) + return VCSCAnnData( + X=VCSRArray.from_scipy(m), + obs=pd.DataFrame( + {"depth": np.asarray(m.sum(axis=1)).ravel()}, index=[f"c{i}" for i in range(n_obs)] + ), + var=pd.DataFrame(index=[f"g{j}" for j in range(n_vars)]), + ) + + +def test_write_h5ad_reads_in_parallel(tmp_path, spy): + """Every IVCSR array written by write_h5ad takes the parallel read path.""" + va = _vcsr_anndata() + path = tmp_path / "x.h5ad" + va.write_h5ad(path) + back = VCSCAnnData.read_h5ad(path) + assert back.X is not None and va.X is not None + np.testing.assert_array_equal(back.X.toarray(), va.X.toarray()) + assert {"/X/major_ptr", "/X/values", "/X/value_ptr", "/X/packed_indices"} <= set(spy) + + +def test_blosc2_file_still_reads(tmp_path, spy): + """Files written with the earlier Blosc2+LZ4 default keep reading.""" + va = _vcsr_anndata() + path = tmp_path / "legacy.h5ad" + va.write_h5ad(path, dataset_kwargs=dict(hdf5plugin.Blosc2(cname="lz4"))) + with h5py.File(path, "r") as f: + assert "32026" in f["X"]["packed_indices"]._filters + back = VCSCAnnData.read_h5ad(path) + assert back.X is not None and va.X is not None + np.testing.assert_array_equal(back.X.toarray(), va.X.toarray()) + np.testing.assert_array_equal(back.obs["depth"], va.obs["depth"]) + assert spy == [] + + +def test_read_h5ad_in_fresh_process(tmp_path): + """read_h5ad registers the HDF5 filters itself; compressed obs columns read before X does.""" + path = tmp_path / "fresh.h5ad" + _vcsr_anndata().write_h5ad(path) + code = f"import vsparse; print(vsparse.VCSCAnnData.read_h5ad({str(path)!r}).shape)" + result = subprocess.run( + [sys.executable, "-c", code], capture_output=True, text=True, check=False + ) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "(400, 60)" + + +def test_numcodecs_decodes_hdf5_zstd_chunks(tmp_path): + """The parallel path relies on the HDF5 zstd filter writing plain zstd frames.""" + data = np.arange(1000, dtype=np.int64) + path = tmp_path / "frame.h5" + with h5py.File(path, "w") as f: + f.create_dataset("a", data=data, chunks=(1000,), **ZSTD) + with h5py.File(path, "r") as f: + _mask, raw = f["a"].id.read_direct_chunk((0,)) + np.testing.assert_array_equal(np.frombuffer(numcodecs.Zstd().decode(raw), np.int64), data) diff --git a/uv.lock b/uv.lock index edf4973..c6c2705 100644 --- a/uv.lock +++ b/uv.lock @@ -1827,6 +1827,7 @@ dependencies = [ { name = "anndata" }, { name = "hdf5plugin" }, { name = "numba" }, + { name = "numcodecs" }, { name = "numpy" }, { name = "scipy" }, ] @@ -1864,6 +1865,7 @@ requires-dist = [ { name = "hdf5plugin", specifier = ">=4.0" }, { name = "myst-parser", marker = "extra == 'docs'", specifier = ">=3.0" }, { name = "numba", specifier = ">=0.60" }, + { name = "numcodecs", specifier = ">=0.13" }, { name = "numpy", specifier = ">=2.0" }, { name = "scipy", specifier = ">=1.12" }, { name = "sphinx", marker = "extra == 'docs'", specifier = ">=7.0" }, From cca3f93e29353848a8f9dfb4ef379c4a8577d6a2 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Wed, 14 Oct 2026 08:11:52 -0700 Subject: [PATCH 2/3] Restore main's _rapid_load.py after the merge with #66 Merging main into this branch kept this branch's _rapid_load.py, which re-added load_and_normalize, _read_ivcsr_group and _validate_min_cells after #66 removed them -- without the numba kernels and imports they call, so ruff flagged 19 undefined names and calling it would raise NameError. Take main's version. load_packed already reads X through the anndata registry, which uses the parallel reader, so the only change this branch needs here is the hdf5plugin comment. Claude-Session: https://claude.ai/code/session_01QBeTZy6BKKGKGJ7fwvkWgM --- src/vsparse/_rapid_load.py | 153 ------------------------------------- 1 file changed, 153 deletions(-) diff --git a/src/vsparse/_rapid_load.py b/src/vsparse/_rapid_load.py index d2f5ec9..a1b4768 100644 --- a/src/vsparse/_rapid_load.py +++ b/src/vsparse/_rapid_load.py @@ -14,9 +14,6 @@ import anndata as ad from vsparse._anndata_class import VCSCAnnData -from vsparse._compression import read_array -from vsparse._indexutils import smallest_index_dtype -from vsparse._ops import accumulator_threads if TYPE_CHECKING: from os import PathLike @@ -26,156 +23,6 @@ _FIELD_KEYS = ("obs", "var", "obsm", "varm", "obsp", "varp", "layers", "uns") -def _validate_min_cells(min_cells: int | None) -> None: - if min_cells is None: - return - if isinstance(min_cells, bool) or not isinstance(min_cells, int): - raise TypeError("min_cells must be an integer or None") - if min_cells < 0: - raise ValueError("min_cells must be non-negative") - - -def _read_ivcsr_group( - f: Any, x_key: str -) -> tuple[ - tuple[int, int], np.dtype, np.ndarray, np.ndarray, np.ndarray, np.ndarray, dict[str, Any] -]: - """Read the packed IVCSR array at ``f[x_key]`` plus the surrounding obs/var/etc. fields.""" - g = f[x_key] - shape = (int(g.attrs["shape"][0]), int(g.attrs["shape"][1])) - indices_dtype = np.dtype(g.attrs["indices_dtype"]) - major_ptr = read_array(g["major_ptr"]) - values = read_array(g["values"]) - value_ptr = read_array(g["value_ptr"]) - packed = read_array(g["packed_indices"]) - - kwargs = {k: ad.io.read_elem(f[k]) for k in _FIELD_KEYS if k in f} - if "raw" in f: - kwargs["raw"] = ad.io.read_elem(f["raw"]) - return shape, indices_dtype, major_ptr, values, value_ptr, packed, kwargs - - -def load_and_normalize( - path: str | PathLike[str], - *, - min_cell_counts: float = 10.0, - gene_threshold: float = 0.0, - min_cells: int | None = None, - obs_filter: Callable[[pd.DataFrame], object] | None = None, - x_key: str = "X", -) -> ad.AnnData: - """Load, filter, and depth-normalize a VCSR/IVCSR-backed ``.h5ad`` file. - - Reproduces ``parafac2.normalize.prepare_dataset``: cells with total - counts <= ``min_cell_counts`` and genes with total counts <= - ``gene_threshold * n_cells`` are dropped. When ``min_cells`` is given, - genes expressed in fewer than ``min_cells`` cells are also dropped. Gene - filters are measured on the raw counts after any ``obs_filter``. - The remaining matrix is row-normalized to the median per-cell depth, then - column-normalized by gene sum, then transformed as ``log10(1000x + 1)``. - Surrounding metadata (``obs``, ``var``, ``obsm``, etc.) is sliced to - match the retained cells and genes. - - Parameters - ---------- - path - Path to an ``.h5ad`` file whose ``X`` (or ``layers[x_key]``) was - stored as IVCSC/IVCSR (see - :meth:`~vsparse.VCSCAnnData.write_h5ad`). - min_cell_counts - Cells with total raw counts <= this are dropped. - gene_threshold - Minimum threshold fraction for gene inclusion, as in - ``parafac2.normalize.prepare_dataset``: genes with total raw counts - <= ``gene_threshold * n_cells`` are dropped. - min_cells - Optional gene filter. Genes expressed in fewer than this - many cells are dropped. Expression is defined as a raw count > 0. - obs_filter - Optional callable receiving ``obs`` and returning a one-dimensional - boolean mask. When provided, rows are subset before cell filtering, - gene filtering, and normalization, so gene totals and - ``gene_threshold`` are computed using only the selected cells. The - packed IVCSR stream is still read in full, but indices and values are - materialized only for selected rows. - x_key - Top-level h5ad group holding the IVCSR array (``"X"`` by default). - - Returns - ------- - ad.AnnData - Filtered, depth-normalized AnnData object with ``X`` as a CSR array - and sliced metadata. - - Examples - -------- - Select cells using multiple ``obs`` columns and multiple accepted values:: - - load_and_normalize( - path, - obs_filter=lambda obs: ( - obs["condition"].isin(["control", "vehicle"]) - & (obs["timepoint"] == "T3") - ), - ) - """ - import h5py - import hdf5plugin # noqa: F401 -- registers the zstd/Blosc2 HDF5 filters - - _validate_min_cells(min_cells) - - with h5py.File(Path(path), "r") as f: - shape, indices_dtype, major_ptr, values, value_ptr, packed, kwargs = _read_ivcsr_group( - f, x_key - ) - - n_cells, n_genes = shape - - # Cell mask: no decode needed. - cell_totals = _cell_totals(major_ptr, values, value_ptr) - - if obs_filter is None: - cell_mask, row_indptr, indices, data, metadata_cell_mask = _rows_without_obs_filter( - major_ptr, values, value_ptr, packed, indices_dtype, cell_totals, min_cell_counts - ) - gene_denom = n_cells - else: - cell_mask, row_indptr, indices, data, metadata_cell_mask, gene_denom = ( - _rows_with_obs_filter( - major_ptr, - values, - value_ptr, - packed, - indices_dtype, - cell_totals, - kwargs.get("obs"), - obs_filter, - min_cell_counts, - n_cells, - ) - ) - del packed - - gene_mask = _compute_gene_mask(indices, data, n_genes, gene_denom, gene_threshold, min_cells) - - new_indptr, out_indices, out_data, kept_rows, n_kept_genes = _filter_and_compact( - row_indptr, indices, data, cell_mask, gene_mask - ) - del indices, data, row_indptr - - normalized = _normalize_and_transform(new_indptr, out_indices, out_data, n_kept_genes) - X = csr_array((normalized, out_indices, new_indptr), shape=(kept_rows.shape[0], n_kept_genes)) - - if "obs" in kwargs and "var" in kwargs: - adata = ad.AnnData(**kwargs) - else: - adata = ad.AnnData(shape=shape, **kwargs) - adata = adata[metadata_cell_mask, gene_mask].copy() - adata.X = X - - return adata - - def load_packed(path: str | PathLike[str], *, x_key: str = "X") -> VCSCAnnData: """Load an IVCSR/IVCSC-backed ``.h5ad`` file, decoding ``X`` immediately. From ab37ec833b5265a328edb2b088623af0fb57d583 Mon Sep 17 00:00:00 2001 From: Aaron Meyer Date: Wed, 14 Oct 2026 08:13:52 -0700 Subject: [PATCH 3/3] Write large arrays in 1 MiB chunks and decode one chunk run per thread h5py's automatic chunking gave the VCS arrays ~64-330 KiB chunks, and the parallel reader pays a fixed, GIL-holding Python cost per chunk, so on lupus (1.26M cells) reading packed_indices took 0.33 s instead of the ~0.08 s the zstd decode itself needs. numeric_only_compression now gives compressed 1-D HDF5 arrays larger than one chunk ~1 MiB chunks (unless the caller picked chunks), and read_array hands each thread one contiguous run of chunks and decodes full chunks straight into the output. On lupus: packed_indices 1.21 s (Blosc2) -> 0.095 s, value_ptr 0.10 s -> 0.031 s, load_packed 3.74 s -> 2.76 s, at the same zstd ratio. Claude-Session: https://claude.ai/code/session_01QBeTZy6BKKGKGJ7fwvkWgM --- src/vsparse/_compression.py | 50 ++++++++++++++++++++++++++++++------- tests/test_compression.py | 24 ++++++++++++++++++ 2 files changed, 65 insertions(+), 9 deletions(-) diff --git a/src/vsparse/_compression.py b/src/vsparse/_compression.py index bb53379..51d5037 100644 --- a/src/vsparse/_compression.py +++ b/src/vsparse/_compression.py @@ -60,6 +60,8 @@ _ZSTD_LEVEL = 1 # Registered HDF5 filter id for zstd (https://github.com/HDFGroup/hdf5_plugins). _H5_ZSTD_FILTER_ID = 32015 +# Chunk size for large compressed 1-D arrays -- see _chunk_large_1d. +_H5_CHUNK_BYTES = 1 << 20 def h5_dataset_kwargs() -> dict[str, Any]: @@ -115,6 +117,23 @@ def _is_string_like(dtype: Any, data: Any) -> bool: return False +def _chunk_large_1d(kwargs: dict[str, Any], data: Any) -> None: + """Give a compressed 1-D array larger than one chunk ~1 MiB chunks, unless chunks were chosen. + + h5py's automatic chunking picks ~64-330 KiB chunks for the VCS arrays, + and :func:`read_array` pays a fixed, GIL-bound cost per chunk -- so + fewer, larger chunks read several times faster, at the same zstd ratio. + """ + if "compression" not in kwargs or kwargs.get("chunks") is not None: + return + shape, dtype = getattr(data, "shape", None), getattr(data, "dtype", None) + if shape is None or dtype is None or len(shape) != 1: + return + per_chunk = max(1, _H5_CHUNK_BYTES // np.dtype(dtype).itemsize) + if shape[0] > per_chunk: + kwargs["chunks"] = (per_chunk,) + + @contextlib.contextmanager def _numeric_only_h5() -> Iterator[None]: import h5py @@ -126,6 +145,8 @@ def patched(self: h5py.Group, name: str, *args: Any, **kwargs: Any) -> h5py.Data if _is_string_like(kwargs.get("dtype"), data): for key in _H5_COMPRESSION_KEYS: kwargs.pop(key, None) + else: + _chunk_large_1d(kwargs, data) return original(self, name, *args, **kwargs) h5py.Group.create_dataset = patched @@ -176,11 +197,12 @@ def patched(self: zarr.Group, name: str, **kwargs: Any) -> Any: def numeric_only_compression(store_kind: str) -> contextlib.AbstractContextManager[None]: - """Scoped patch making Blosc compression apply to numeric arrays only. + """Scoped patch making compression apply to numeric arrays only. ``store_kind`` is ``"h5"`` or ``"zarr"``. See the module docstring for why - this is needed. Use as ``with numeric_only_compression("h5"): ...`` around - the write. + this is needed. For HDF5 it also gives large compressed 1-D arrays ~1 MiB + chunks (see :func:`_chunk_large_1d`). Use as + ``with numeric_only_compression("h5"): ...`` around the write. """ if store_kind == "h5": return _numeric_only_h5() @@ -243,16 +265,26 @@ def _read_chunks_parallel( def load(i: int) -> None: raw = os.pread(fd, int(sizes[i]), int(offsets[i])) - # Filter-mask bit 0 set means HDF5 skipped zstd and stored the chunk as-is. - data = raw if masks[i] & 1 else codec.decode(raw) dst = flat[i * chunk_bytes : (i + 1) * chunk_bytes] - # The last chunk decodes to a full chunk; only its leading part is data. - dst[:] = np.frombuffer(data, dtype=np.uint8, count=dst.size) + 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: + 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) + + def load_run(t: int) -> None: + for i in range(bounds[t], bounds[t + 1]): + load(i) try: - n_workers = max(1, min(numba.get_num_threads(), offsets.shape[0])) with ThreadPoolExecutor(max_workers=n_workers) as pool: - list(pool.map(load, range(offsets.shape[0]))) + list(pool.map(load_run, range(n_workers))) finally: os.close(fd) return out diff --git a/tests/test_compression.py b/tests/test_compression.py index f91f932..438c1cc 100644 --- a/tests/test_compression.py +++ b/tests/test_compression.py @@ -59,3 +59,27 @@ def test_numeric_only_compression_skips_strings(tmp_path): assert str(_compression._H5_ZSTD_FILTER_ID) in f["numbers"]._filters # String dataset has no compression filters assert dict(f["strings"]._filters) == {} + + +def _written_chunks(tmp_path, data, **extra): + path = tmp_path / "chunks.h5" + with _compression.numeric_only_compression("h5"), h5py.File(path, "w") as f: + f.create_dataset("a", data=data, **_compression.h5_dataset_kwargs(), **extra) + with h5py.File(path, "r") as f: + return f["a"].chunks + + +def test_large_1d_arrays_get_1mib_chunks(tmp_path): + """Compressed 1-D arrays larger than one chunk get ~1 MiB chunks for the parallel reader.""" + for dtype in (np.uint8, np.int64): + per_chunk = (1 << 20) // np.dtype(dtype).itemsize + assert _written_chunks(tmp_path, np.zeros(3 * per_chunk + 7, dtype=dtype)) == (per_chunk,) + + +def test_small_explicit_and_2d_chunks_untouched(tmp_path): + """Arrays that fit in one chunk, explicit chunks, and 2-D arrays keep h5py's choice.""" + small = np.zeros(1000, dtype=np.int64) + big = np.zeros(1 << 20, dtype=np.int64) + assert _written_chunks(tmp_path, small) != ((1 << 20) // 8,) + assert _written_chunks(tmp_path, big, chunks=(4096,)) == (4096,) + assert len(_written_chunks(tmp_path, np.zeros((1 << 18, 4)))) == 2