From 274790913f7ea0f93ad5d372db7aeafccbd7107b Mon Sep 17 00:00:00 2001 From: Aljes Date: Fri, 30 Oct 2026 18:18:50 +0000 Subject: [PATCH 1/5] B6a: memory harness, then recast streaming The perf suite counts bytes (returned, decoded in whole chunks, re-encoded in whole chunks or shards), measures resident memory in a subprocess, builds every layout in h5ad, Zarr v2 and anndata's sharded Zarr v3, and grows n_var as well as n_obs (F091, F146). Recast within it: - check_cast reads --chunk rows of a dense matrix (F029) - CSC densify and dense-to-CSC go through a CSR scratch (F030) - the transpose steps by entries and streams huge majors (F077) - sparsify's 1-D chunks are elements, not a row height (F078) - the index range check always runs, and on dense sources (F013) - float16 sparse output is refused (F058) - a lossy cast drops the zeros it makes (F121) - nested matrix paths are converted (F043) h5py 3.15 leaks on every dataset call; the floor is 3.16.0. Co-Authored-By: Claude Opus 5.5 --- docs/reference/commands.md | 17 +- docs/reference/testing.md | 89 +++++- pyproject.toml | 4 +- src/adata/commands/recast.py | 83 +++--- src/adata/core/recast.py | 419 ++++++++++++++++++++------- tests/perf_counters.py | 451 +++++++++++++++++++++++++++-- tests/perf_fixtures.py | 147 ++++++++++ tests/test_performance.py | 353 +++++++++++++++++++++++ tests/test_recast_streaming.py | 498 +++++++++++++++++++++++++++++++++ uv.lock | 62 ++-- 10 files changed, 1928 insertions(+), 195 deletions(-) create mode 100644 tests/perf_fixtures.py create mode 100644 tests/test_recast_streaming.py diff --git a/docs/reference/commands.md b/docs/reference/commands.md index 6ed3a31..2e14d16 100644 --- a/docs/reference/commands.md +++ b/docs/reference/commands.md @@ -220,13 +220,26 @@ Two things are refused before anything is written. A cast that would not round-trip -- checked by casting every value and casting it back, not by comparing dtypes -- and a densification that would inflate the store beyond four times its size. `--force` overrides either, and the message says how -many values would change or how large the result would be. +many values would change or how large the result would be. A cast that loses +values drops the entries it turns to zero, as scipy would, rather than storing +them as explicit zeros. + +A sparse matrix of **float16 is refused**, whatever `--force` says: scipy has +no float16 sparse matrix, so neither anndata nor scanpy could read the result. +Use `--dtype float32`, or `--layout dense`. The index dtype is **preserved** unless `--indices-dtype` asks otherwise, so narrowing the values does not silently widen the indices and leave the file bigger than it started. `indices` and `indptr` keep their own widths, which can differ: a narrow matrix with more than 2^31 nonzeros needs int64 offsets -over int32 coordinates, and both are range-checked before writing. +over int32 coordinates, and both are range-checked before writing. The check +always runs: `--force` does not waive it, since an offset that overflows is a +corrupt matrix, not a lossy value. A dense matrix made sparse has its +coordinates checked before writing and its nonzero count as it is counted. + +`--chunk` counts rows of a dense matrix, in every layout: densifying a CSC +matrix or writing a dense one as CSC goes through a CSR in scratch, so no block +spans every row. Sparse matrices are read about a million entries at a time. The output path may not name the input; use `--inplace`. Either way the new store is written beside its destination and moved into place only once the diff --git a/docs/reference/testing.md b/docs/reference/testing.md index 0393c76..d68e89b 100644 --- a/docs/reference/testing.md +++ b/docs/reference/testing.md @@ -43,6 +43,7 @@ is at its floor and imports. | `test_commands_phase2.py`, `test_commands_coverage.py`, `test_cli.py` | Command surfaces and error paths | | `test_convert.py` | `convert` between every backend pair, read back by anndata; its refusals, staging, and the forwarding of recast's old flags | | `test_recast.py` | `recast` against scipy: casts, transposes, density changes and every refusal | +| `test_recast_streaming.py` | `recast` on HDF5, Zarr v2 and Zarr v3 in blocks of a few rows and entries: every layout to every other, every dtype in every layout, matrices at any depth, index ranges with and without `--force`; against scipy and read back by anndata | | `test_tenx.py` | 10x HDF5 import and export against Cell Ranger-shaped fixtures (v2, v3, multimodal, probe-based): mapping, filters, round trips, refusals | | `test_mtx.py` | 10x / STARsolo MTX folders against hand-written fixtures (Cell Ranger 2 and 3, STARsolo, velocyto, GEO prefixes): mapping, entry orders, refusals that leave nothing, round trips | | `test_cellbender.py` | CellBender 0.2 and 0.3+ output, full and filtered: each latent at its barcode, `--min-cell-probability`, `--raw` alignment, refusals | @@ -60,7 +61,7 @@ is at its floor and imports. | `test_table.py` | The streaming CSV reader against pandas: column types, quoting, row subsets, blocks, refusals | | `test_mtx_codec.py`, `test_csr_build.py` | The streaming Matrix Market reader and writer against scipy, and the bucket sort they and `recast` share | | `test_external_readers.py` | 10x conversions, HDF5 and MTX, Parse runs and looms, checked by scanpy (Parse's own recipe for Parse), anndata with loompy (the pySCENIC tutorial's recipe for SCope looms), anndata's `read_csv`, `read_text` and `read_hdf`, a pandas 1.5 pycisTopic pickle read back there, mudata (a file it writes read as it reads it; ours read by it), and PyTables in their own uv environments (`integration`) | -| `test_performance.py` | Complexity guards -- what an operation costs, not how long it takes (see below) | +| `test_performance.py` | Complexity guards -- what an operation costs, not how long it takes (see below); its fixtures are in `perf_fixtures.py` | | `test_benchmark_harness.py` | That the benchmark's measurement is trustworthy | | `test_subset.py`, `test_export.py`, `test_import.py`, `test_info_read.py`, `test_zarr.py`, `test_query.py` | Per-feature unit tests | @@ -109,15 +110,16 @@ be rather than only that it did not grow. ### What is counted -`tests/perf_counters.py` patches four seams, all verified against the pinned -h5py 3.15.1 and zarr 3.1.5: +`tests/perf_counters.py` patches these seams, all verified against the pinned +h5py 3.16.0 and zarr 3.1.5: | Hook | Sees | |---|---| | `h5py.Dataset.__getitem__` | HDF5 reads | | `h5py.Dataset.__setitem__` | HDF5 writes | | `h5py.Group.create_dataset` | HDF5 writes made at creation | -| `zarr.Array.__getitem__` | Zarr reads, v2 and v3 alike | +| `zarr.Array.__getitem__`, `.oindex[...]` | Zarr reads, v2 and v3 alike | +| `zarr.Array.__setitem__`, `zarr.Group.create_array` | Zarr writes (bytes only) | | `zarr.storage.LocalStore.get` | chunk and metadata fetches | | `LocalStore.set` / `.delete` | write storms | @@ -138,9 +140,41 @@ Elements are counted, not just calls: the `--merge` bug was n calls each reading n elements, and counting calls alone would miss a vectorised variant that reads n elements n times in one call. +**Bytes, three ways.** Elements alone hid a 70x read amplification and a +shard rewritten on every block, so `IOCounts` also counts bytes: + +| Field | Counts | Catches | +|---|---|---| +| `read_bytes`, `largest_read`, `largest_by_name` | bytes the reads returned; the largest one read | a block that `--chunk` does not bound, with no allocator involved | +| `decoded_bytes` | bytes of every whole chunk the reads touched | reads that decode a chunk to use a row of it | +| `written_bytes`, `largest_write` | bytes handed to writes | | +| `encoded_bytes` | bytes of every whole chunk the writes touched, and every whole shard of a sharded Zarr array | writes smaller than the unit, each rewriting all of it | +| `store_get_bytes`, `store_set_bytes` | bytes Zarr's store fetched and stored | | + +Decoded and encoded bytes are a model worked out from the selection, the chunk +grid and the codecs, not a measurement of the library, which is what keeps them +deterministic. It follows HDF5's chunk cache (per dataset, 1 MiB unless the +file says otherwise, least recently used), so appending a block at a time to a +growing compressed dataset costs each chunk once. A chunk with no filter costs +only the bytes asked for: HDF5 reads and writes it in place. Zarr caches +nothing, and writing part of a shard rewrites the shard. +`test_the_byte_counters_see_whole_chunks_decoded_and_rewritten` checks the +model against numbers worked out by hand. + +**Resident memory.** `count_allocations` uses tracemalloc, which sees numpy's +buffers but not what HDF5, blosc or zarr's codecs allocate in C. +`peak_rss(argv)` runs `adata ` in a fresh interpreter and returns how far +the command raised the operating system's high-water mark (`ru_maxrss`, read +once the CLI is imported and again when the command ends). A subprocess, +because the high-water mark only rises. It has a few megabytes of noise, so +guard it with generous ratios against inputs of tens of megabytes, as +`test_a_dense_cast_adds_a_block_of_resident_memory_not_the_matrix` does; +`test_peak_rss_sees_a_command_that_holds_the_matrix` checks that it sees +`--in-memory` hold the matrix. + **Known bypasses.** `Dataset.read_direct`, `np.asarray(dataset)` (it goes via -`__array__`), `dataset.asstr()[...]` and `.fields()` reach the file without -passing any hook. None are used today. If you add one, the counters will +`__array__`), `dataset.asstr()[...]`, `.fields()` and Zarr's `vindex` reach +the file without passing any hook. None are used today. If you add one, the counters will quietly report less work than happened -- which is why every guard also asserts a lower bound, and why there are two canary tests with known absolute counts. If a canary fails, fix the hooks before trusting anything else here. @@ -184,7 +218,8 @@ accept a 64x increase in a command that is supposed to read nothing. |---|---|---| | `assert_grows_linearly` | no worse than linear in one axis | most commands | | `assert_grows_slower_than_input` | grows at least N times slower than the input | streaming at a fixed chunk | -| `assert_independent_of` | does not grow at all | inspection, and grouping work per row | +| `assert_memory_is_a_fraction_of_the_input` | adds at most a share (a quarter) of what the input adds, in bytes | memory, along either axis | +| `assert_independent_of` | does not grow at all | inspection, grouping work per row, the largest read | **Inspection is free, and that is an exact number.** `view` and `ls` reach only `.shape`, `.dtype` and `.attrs`: `axis_len` goes through `element_len`, which reads a @@ -211,6 +246,34 @@ guards conceal. Streamed `export sparse` is separately asserted to stay under a quarter of what `--in-memory` costs, both measured in the same run so the factor holds on any machine. +### Memory along every axis and layout + +The guards above grow n_obs of an eight-column CSR store, the one shape every +path handles well. A CSC densify holding n_obs x `--chunk` values, a dense cast +check reading the whole matrix and a transpose stepping by the average major +all passed them. `tests/perf_fixtures.py` builds stores the way users have +them -- every layout, written by anndata as h5ad, Zarr v2 and anndata 0.13's +default sharded Zarr v3 (inner chunks of a few hundred bytes in shards of the +whole array) -- and `test_peak_memory_is_a_small_share_of_the_input` runs +`recast` (each dtype and layout change), `subset` and `concat` over each +layout, growing n_obs and then n_var. + +A growth ratio cannot be the claim along n_var: widen a matrix and a block of +`--chunk` rows widens with it, as fast as the input, so a ratio cannot tell it +from the whole matrix. Peak memory added per byte of input added can: a block +of 32 rows of a matrix thousands of rows tall adds a few hundredths, and a path +holding the matrix adds a whole byte and more for each copy. + +The known gaps are marked `xfail`, each with its owner: sparse `subset` +(F032, strict, so the fix has to flip it) and `concat` of a sparse matrix grown +in n_obs (not strict: concat keeps obs names by design, and B3b's row pointers +change the figure). + +The recast guards beside it pin each block in bytes, on every format: the cast +check reads `--chunk` rows of a dense matrix, CSC densify and sparsify read +blocks set by the chunk, and a skewed transpose never reads more than about two +buckets of entries at once. + ### Where counting is not enough Two real defects in `concat.py` were invisible to all of the above, and both @@ -277,9 +340,17 @@ should scale with, scale only that, and hold everything else fixed. | `convert` to 10x MTX | n_obs | **flat** read calls | | `convert` CellBender (plain, `--min-cell-probability`, `--raw` full and filtered) | n_barcodes | linear; **flat** write calls | | `recast --dtype`, `--layout csc`, `--layout dense` | nnz, n_obs | linear; the streaming transpose's **peak memory** is bounded | +| `recast` (every dtype and layout change), `subset`, `concat` | n_obs, n_var; CSR, CSC, dense | **peak memory** a small share of the input | +| `recast --dtype` (dense), CSC `--layout dense`, dense `--layout csc`, a skewed transpose | n_obs; h5ad, Zarr v2, sharded Zarr v3 | the **largest read** set by the chunk | + +The `slow` marker is on the guards that build a 65,536-row store or run `peak_rss`; +they still gate merges, and `-m "not slow"` skips them locally. `perf` is on the +tracing guards. -The `slow` marker is on the three guards that build a 65,536-row store; they still gate -merges, and `-m "not slow"` skips them locally. `perf` is on the tracing guards. +h5py 3.15 kept 100 to 160 bytes of every dataset read, write and `.shape` for good, so +any long streaming loop grew with its block count; +`test_h5py_calls_keep_no_memory_once_they_return` keeps the floor at 3.16.0, which +does not. ### Reading a failure diff --git a/pyproject.toml b/pyproject.toml index b08310d..3c19719 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,9 @@ classifiers = [ # and can break in a minor (0.26 stopped installing click), but a cap at the # next minor would turn away each new typer until we release, so it is 1.0. dependencies = [ - "h5py>=3.15.1,<4", + # 3.15 leaks 100 to 160 bytes of memory on every dataset read, write and + # `.shape`, so streaming used memory in proportion to its blocks. + "h5py>=3.16.0,<4", # Zarr v2 compressors (storage/__init__.py). "numcodecs>=0.14", "numpy>=2.3.5,<3", diff --git a/src/adata/commands/recast.py b/src/adata/commands/recast.py index e1c1188..1cbc5d1 100644 --- a/src/adata/commands/recast.py +++ b/src/adata/commands/recast.py @@ -22,6 +22,7 @@ _stored_bytes, convert_matrix, parse_dtype, + parse_index_dtype, plan_conversion, ) from adata.elements import spec @@ -86,11 +87,7 @@ def recast_store( raise ValueError("--chunk must be at least 1.") data_dtype = parse_dtype(dtype) if dtype else None - index_dtype = ( - parse_dtype(indices_dtype, allowed=("int32", "int64")) - if indices_dtype - else None - ) + index_dtype = parse_index_dtype(indices_dtype) if indices_dtype else None # Refuses -o naming the input: opening the destination "w" clears it # while the source is still being read from it. HDF5 refuses; Zarr does @@ -122,6 +119,7 @@ def recast_store( index_dtype=index_dtype, layout=layout, chunk=chunk, + chunk_rows=chunk_rows, force=force, console=console, ) @@ -130,7 +128,9 @@ def recast_store( check_free_space( dst_path, - lambda: disk_size(file) + sum(_extra(p) for p in plans.values()), + lambda: disk_size(file) + sum( + _extra(p) + _scratch(p, in_memory) for p in plans.values() + ), check=space_check, ) @@ -182,6 +182,27 @@ def _extra(plan: Plan) -> float: return max(0.0, projected - current) * _on_disk(src.obj) / current +def _scratch(plan: Plan, in_memory: bool) -> float: + """Bytes of scratch a streaming transpose writes beside the output. + + Every entry passes through a bucket as an int64 major, an int64 minor + and its value; a CSC matrix being densified is also sorted into a CSR + first. A dense matrix sparsified to CSC goes through the same, but its + nnz is not known until it is read, so it is not counted. + """ + src = plan.source + if in_memory or not src.sparse: + return 0.0 + transposed = plan.layout in spec.SPARSE_TYPES and plan.layout != src.kind + densified_csc = plan.layout == "dense" and src.kind == spec.CSC_MATRIX + if not (transposed or densified_csc): + return 0.0 + entry = 16 + src.dtype.itemsize + if densified_csc: + entry += 8 + src.dtype.itemsize + return float(src.nnz * entry) + + def _on_disk(obj: Any) -> int: """Bytes a matrix's datasets take on disk, compressed.""" targets = [obj] if is_dataset(obj) else [ @@ -201,27 +222,32 @@ def _on_disk(obj: Any) -> int: def _write(src: Any, dst: Any, plans: dict, **options: Any) -> None: """Copy the store across, converting the targeted entries as they pass.""" - by_parent: dict = {} - for path in plans: - parent, _, leaf = path.rpartition("/") - by_parent.setdefault(parent, {})[leaf] = plans[path] - - for key in src.keys(): - if key in by_parent.get("", {}): - convert_matrix( - src[key], dst, key, plan=by_parent[""][key], **options - ) - elif key in by_parent: - # Any parent, not just layers and raw. Restricting it to those - # two meant an explicitly named `obsm/X_pca` was copied - # unconverted and the command still reported success. - _write_group(src[key], dst, key, by_parent, **options) + _write_members(src, dst, "", plans, **options) + + +def _write_members( + group: Any, out: Any, prefix: str, plans: dict, **options: Any +) -> None: + """Copy `group`'s members into `out`, descending towards planned paths. + + At any depth, not just one level down. Only matrices directly below the + root, or one group down, were converted; one at + `uns/neighbors/connectivities` was copied unconverted while the command + reported converting it (F043). Restricting it further, to `layers` and + `raw`, had already done the same to an explicitly named `obsm/X_pca`. + """ + for key in group.keys(): + path = f"{prefix}{key}" + if path in plans: + convert_matrix(group[key], out, key, plan=plans[path], **options) + elif any(planned.startswith(f"{path}/") for planned in plans): + _write_group(group[key], out, key, path, plans, **options) else: - copy_tree(src[key], dst, key) + copy_tree(group[key], out, key) def _write_group( - group: Any, dst: Any, name: str, by_parent: dict, **options: Any + group: Any, dst: Any, name: str, path: str, plans: dict, **options: Any ) -> None: """Recreate one container, converting the members that were named.""" from adata.storage import copy_attrs, is_zarr_group @@ -233,12 +259,5 @@ def _write_group( target_backend="zarr" if is_zarr_group(dst) else "hdf5", ) if not spec.encoding_type(out): - spec.set_encoding(out, spec.RAW if name == "raw" else spec.DICT) - - - wanted = by_parent.get(name, {}) - for key in group.keys(): - if key in wanted: - convert_matrix(group[key], out, key, plan=wanted[key], **options) - else: - copy_tree(group[key], out, key) + spec.set_encoding(out, spec.RAW if path == "raw" else spec.DICT) + _write_members(group, out, f"{path}/", plans, **options) diff --git a/src/adata/core/recast.py b/src/adata/core/recast.py index 7c227fa..c874225 100644 --- a/src/adata/core/recast.py +++ b/src/adata/core/recast.py @@ -24,7 +24,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, Iterator, Optional, Tuple import numpy as np from rich.console import Console @@ -32,12 +32,13 @@ from adata.elements import spec from adata.elements.write import set_shape_attr from adata.core.csr_build import ( + MIN_BUCKET_ENTRIES, BucketSorter, bucket_bounds, count_keys, scratch_group, ) -from adata.core.subset import _growable_like, _sparse_dataset +from adata.core.subset import _sparse_dataset from adata.storage import ( copy_attrs, create_dataset, @@ -74,6 +75,9 @@ #: Values per read while streaming `data`/`indices`. DEFAULT_CHUNK = 1 << 20 +#: Elements per chunk of a `data` or `indices` built from a dense matrix, +#: whose own chunks are a grid of rows and columns, not of elements. +ELEMENT_CHUNK = 1 << 16 def parse_dtype(text: str, *, allowed: Tuple[str, ...] = DTYPES) -> np.dtype: @@ -86,6 +90,11 @@ def parse_dtype(text: str, *, allowed: Tuple[str, ...] = DTYPES) -> np.dtype: return np.dtype(name) +def parse_index_dtype(text: str) -> np.dtype: + """Resolve `--indices-dtype`, against `INDEX_DTYPES` as it stands.""" + return parse_dtype(text, allowed=INDEX_DTYPES) + + # --------------------------------------------------------------------------- # safety @@ -118,24 +127,43 @@ def describe(self, source: Any, target: Any) -> str: ) -def check_cast(dataset: Any, target: np.dtype, *, chunk: int = DEFAULT_CHUNK) -> CastReport: +def check_cast( + dataset: Any, + target: np.dtype, + *, + chunk: int = DEFAULT_CHUNK, + rows: Optional[int] = None, +) -> CastReport: """Would casting `dataset` to `target` lose anything? Streams the values, casts each block and casts it back. An exact round-trip is the only honest test: whether float64 counts survive float32 depends on the counts, not on the dtypes, and that is precisely the question issue #13 is asking. + + A block is `chunk` elements of a 1-D array (a sparse matrix's `data`), + or `rows` rows of a 2-D one -- `--chunk`'s rows, the block the cast then + writes. It was `chunk` rows of a dense matrix too, which is every row of + any real one: the whole matrix read at once, and two more copies of it + made, whatever `--chunk` said (F029). Without `rows`, as many rows as + hold `chunk` elements. """ report = CastReport() source = np.dtype(getattr(dataset, "dtype", "float64")) - n = int(dataset.shape[0]) if getattr(dataset, "shape", None) else 0 + shape = tuple(getattr(dataset, "shape", None) or ()) + n = int(shape[0]) if shape else 0 + step = int(chunk) + if len(shape) > 1: + width = int(np.prod(shape[1:], dtype=np.int64)) + step = int(rows) if rows is not None else chunk // max(1, width) + step = max(1, step) info = np.finfo(target) if target.kind == "f" else ( np.iinfo(target) if target.kind in "iu" else None ) - for start in range(0, n, chunk): - block = np.asarray(dataset[start : min(start + chunk, n)]) + for start in range(0, n, step): + block = np.asarray(dataset[start : min(start + step, n)]) report.total += block.size if block.size == 0: continue @@ -186,20 +214,36 @@ def check_index_dtype( data is needed. `indices` holds coordinates, bounded by the larger dimension; `indptr` holds offsets, bounded by nnz. They are checked separately because they can legitimately need different widths. + + A dense source has no nnz to read: pass `group=None` and only the + coordinates are checked here, and `sparsify` checks the offsets as it + counts them. """ - nnz = int(group["indices"].shape[0]) - for name, dtype, largest, what in ( - ("indices", index_dtype, max(int(shape[0]), int(shape[1])), "coordinates"), - ("indptr", pointer_dtype, nnz, "offsets"), - ): - limit = int(np.iinfo(dtype).max) - if largest > limit: - raise ValueError( - f"{dtype} cannot hold this matrix's {what}: {name} must reach " - f"{largest:,} for a {shape[0]:,} x {shape[1]:,} matrix with " - f"{nnz:,} nonzeros, and {dtype} tops out at {limit:,}. " - "Use int64." - ) + nnz = None if group is None else int(group["indices"].shape[0]) + check_index_range( + "indices", index_dtype, max(int(shape[0]), int(shape[1])), shape, nnz + ) + if nnz is not None: + check_index_range("indptr", pointer_dtype, nnz, shape, nnz) + + +def check_index_range( + name: str, dtype: np.dtype, largest: int, shape: Tuple[int, int], nnz: Optional[int] +) -> None: + """Refuse `dtype` for `indices` or `indptr` if `largest` is past its range.""" + limit = int(np.iinfo(dtype).max) + if largest <= limit: + return + what = "offsets" if name == "indptr" else "coordinates" + nonzeros = ( + f" with {nnz:,} nonzeros" if nnz is not None + else f" with at least {largest:,} nonzeros" if name == "indptr" else "" + ) + raise ValueError( + f"{dtype} cannot hold this matrix's {what}: {name} must reach " + f"{largest:,} for a {shape[0]:,} x {shape[1]:,} matrix{nonzeros}, " + f"and {dtype} tops out at {limit:,}. Use int64." + ) def _stored_bytes(obj: Any) -> int: @@ -427,27 +471,29 @@ def transpose_sparse_streaming( bucket_entries: Optional[int] = None, console: Optional[Console] = None, compression: Optional[str] = None, + keep_layout: bool = True, ) -> None: """Swap CSR<->CSC without loading the matrix. A transpose cannot be done in one pass: the first entry of the output may come from the last row of the input. Three passes instead, with - memory set by `bucket_entries` rather than by nnz: + memory set by `chunk` and `bucket_entries` rather than by nnz: 1. count nonzeros per output major, by streaming `indices` alone. That gives the output `indptr` by cumulative sum. - 2. stream the input again, splitting each block's entries into buckets by - which slice of the output they land in, and append each bucket to - scratch datasets. + 2. stream the input again, `chunk` entries at a time, splitting each + block's entries into buckets by which slice of the output they land + in, and append each bucket to scratch datasets. 3. read one bucket at a time, sort it, and append to the output in order. Costs about two extra passes over nnz, which is the price of not holding the matrix. `--in-memory` is there for when you would rather pay in RAM. + `keep_layout=False` writes plain arrays rather than the source's chunks + and codecs, for a transpose into scratch. """ target = ( spec.CSC_MATRIX if src.kind == spec.CSR_MATRIX else spec.CSR_MATRIX ) - n_major_in = src.shape[0] if src.kind == spec.CSR_MATRIX else src.shape[1] n_major_out = src.shape[1] if src.kind == spec.CSR_MATRIX else src.shape[0] indptr = np.asarray(src.obj["indptr"][...], dtype=np.int64) @@ -461,7 +507,7 @@ def transpose_sparse_streaming( count_keys(counts, block) out_indptr = np.concatenate(([0], np.cumsum(counts))) - bounds = bucket_bounds(out_indptr, chunk=chunk, bucket_entries=bucket_entries) + bounds = _transpose_bounds(out_indptr, chunk=chunk, bucket_entries=bucket_entries) n_buckets = len(bounds) - 1 if console is not None and n_buckets > 1: console.print( @@ -471,38 +517,33 @@ def transpose_sparse_streaming( # nnz is invariant under a transpose, so the output can be sized now and # written by slice rather than grown block by block. group = _new_sparse_group(src, dst_parent, name, target, src.shape) + templates = (source_data, source_indices, src.obj["indptr"]) if keep_layout else (None,) * 3 out_data = _sparse_dataset( - group, "data", data_dtype, nnz, source_data, compression=compression + group, "data", data_dtype, nnz, templates[0], compression=compression ) out_indices = _sparse_dataset( - group, "indices", index_dtype, nnz, source_indices, compression=compression + group, "indices", index_dtype, nnz, templates[1], compression=compression ) written = 0 with scratch_group(dst_parent, f"__{name}_transpose_scratch__") as scratch: sorter = BucketSorter(scratch, bounds, src.dtype) - # Pass 2: scatter the input into buckets, one input block at a time. - # Step over the major axis in blocks holding about `chunk` nonzeros, - # so the read size is set by the data rather than by how many rows - # happen to be empty. - average = max(1, nnz // max(1, n_major_in)) - major_step = max(1, chunk // average) - for lo in range(0, n_major_in, major_step): - hi = min(lo + major_step, n_major_in) - start, end = int(indptr[lo]), int(indptr[hi]) - if end <= start: - continue + # Pass 2: scatter the input into buckets, `chunk` entries at a time. + # It stepped by majors of average size, so on a skewed matrix a step + # covered several majors far above the average -- adjacent genes + # expressed in every cell -- and read many times the chunk at once, + # more as n_obs grew (F077). A major holding more than the chunk is + # split across blocks. + for start in range(0, nnz, chunk): + end = min(start + chunk, nnz) minor = np.asarray(source_indices[start:end], dtype=np.int64) values = np.asarray(source_data[start:end]) - major = np.repeat( - np.arange(lo, hi, dtype=np.int64), np.diff(indptr[lo : hi + 1]) - ) # The input's minor axis is the output's major one. - sorter.add(minor, major, values) + sorter.add(minor, _majors_of(indptr, start, end), values) - # Pass 3: each bucket in turn, sorted into output order. - for _, new_minor, values in sorter.runs(): + # Pass 3: each bucket in turn, in output order. + for new_minor, values in _bucket_runs(sorter, chunk): count = new_minor.size out_indices[written : written + count] = new_minor.astype( index_dtype, copy=False @@ -514,12 +555,78 @@ def transpose_sparse_streaming( cast_indptr = out_indptr.astype(pointer_dtype, copy=False) dataset = _sparse_dataset( - group, "indptr", pointer_dtype, cast_indptr.size, src.obj["indptr"], + group, "indptr", pointer_dtype, cast_indptr.size, templates[2], compression=compression, ) dataset[:] = cast_indptr +def _majors_of(indptr: np.ndarray, start: int, end: int) -> np.ndarray: + """The major each entry in `[start, end)` belongs to, by `indptr`.""" + first = int(np.searchsorted(indptr, start, side="right")) - 1 + last = int(np.searchsorted(indptr, end - 1, side="right")) - 1 + spans = np.diff(np.clip(indptr[first : last + 2], start, end)) + return np.repeat(np.arange(first, last + 1, dtype=np.int64), spans) + + +def _transpose_bounds( + out_indptr: np.ndarray, *, chunk: int, bucket_entries: Optional[int] +) -> np.ndarray: + """Buckets of about one read each, and a bucket of its own for any + output major holding more than that, so `_bucket_runs` can stream it. + + `bucket_bounds` balances by entries, but a bound can only fall between + majors, so a gene in every cell shared a bucket with its neighbours and + the bucket held n_obs entries or more. Concat's CSC buckets do the same + (`core/concat.py`). + """ + per_bucket = ( + max(1, int(bucket_entries)) if bucket_entries + else max(MIN_BUCKET_ENTRIES, int(chunk)) + ) + bounds = bucket_bounds(out_indptr, chunk=chunk, bucket_entries=bucket_entries) + big = np.flatnonzero(np.diff(out_indptr) > per_bucket) + return np.unique(np.concatenate((bounds, big, big + 1))) + + +def _bucket_runs(sorter: BucketSorter, chunk: int) -> Iterator[Tuple[np.ndarray, np.ndarray]]: + """Every bucket's `(new_minor, values)`, in output order. + + A bucket of several majors holds about a read's worth: it is read whole + and sorted. A bucket of one major can hold far more -- n_obs entries for + a gene in every cell -- and is streamed a chunk at a time instead. It + needs no sort: pass 2 adds entries in input order, so one output major's + entries arrive in the order of the input majors they came from, which is + their new minor order. + """ + for b, (majors, minors, values) in enumerate(sorter.buckets): + n = int(majors.shape[0]) + if n == 0: + continue + if int(sorter.bounds[b + 1] - sorter.bounds[b]) == 1 and n > chunk: + for lo in range(0, n, chunk): + yield ( + np.asarray(minors[lo : lo + chunk], dtype=np.int64), + np.asarray(values[lo : lo + chunk]), + ) + continue + major = np.asarray(majors[...], dtype=np.int64) + minor = np.asarray(minors[...], dtype=np.int64) + order = np.lexsort((minor, major)) + yield minor[order], np.asarray(values[...])[order] + + +def _transpose(src: Matrix, dst_parent: Any, name: str, *, in_memory: bool, + chunk: int, **kwargs: Any) -> None: + """`transpose_sparse_in_memory` or `_streaming`, as asked.""" + if in_memory: + kwargs.pop("keep_layout", None) + kwargs.pop("compression", None) + transpose_sparse_in_memory(src, dst_parent, name, **kwargs) + else: + transpose_sparse_streaming(src, dst_parent, name, chunk=chunk, **kwargs) + + def densify( src: Matrix, dst_parent: Any, @@ -527,6 +634,8 @@ def densify( *, data_dtype: np.dtype, chunk_rows: int = 1024, + chunk: int = DEFAULT_CHUNK, + in_memory: bool = False, ) -> None: """Write a sparse matrix out as a dense array, a block of rows at a time. @@ -534,13 +643,38 @@ def densify( matrix too large to densify in memory still converts -- it just produces a file that is honestly much larger, which `check_growth` warns about before any of it is written. + + A block is `chunk_rows` rows whatever the layout. A CSC matrix is first + transposed to CSR in scratch: its blocks were `chunk_rows` columns of + every row, n_obs x --chunk values, so memory followed n_obs and not + --chunk (F030). Writing narrow column slabs instead would rewrite every + chunk of the dense output once per slab. """ + if src.kind != spec.CSC_MATRIX: + _densify_rows(src, dst_parent, name, src.obj.attrs, + data_dtype=data_dtype, chunk_rows=chunk_rows) + return + with scratch_group(dst_parent, f"__{name}_densify_scratch__") as scratch: + _transpose( + src, scratch, "csr", in_memory=in_memory, chunk=chunk, + data_dtype=src.dtype, index_dtype=np.dtype("int64"), + pointer_dtype=np.dtype("int64"), keep_layout=False, + ) + _densify_rows(describe(scratch["csr"]), dst_parent, name, src.obj.attrs, + data_dtype=data_dtype, chunk_rows=chunk_rows) + + +def _densify_rows( + src: Matrix, dst_parent: Any, name: str, attrs: Any, *, + data_dtype: np.dtype, chunk_rows: int, +) -> None: + """Densify a CSR matrix `chunk_rows` rows at a time, with `attrs`.""" n_rows, n_cols = src.shape backend = "zarr" if is_zarr_group(dst_parent) else "hdf5" dst = create_dataset( dst_parent, name, shape=(n_rows, n_cols), dtype=data_dtype ) - copy_attrs(src.obj.attrs, dst.attrs, target_backend=backend) + copy_attrs(attrs, dst.attrs, target_backend=backend) spec.set_encoding(dst, spec.ARRAY) # `shape` belongs to the sparse encoding and would contradict the array's # own shape if it were carried over. @@ -549,15 +683,11 @@ def densify( indptr = np.asarray(src.obj["indptr"][...], dtype=np.int64) indices, values = src.obj["indices"], src.obj["data"] - csr = src.kind == spec.CSR_MATRIX - n_major = n_rows if csr else n_cols - for lo in range(0, n_major, chunk_rows): - hi = min(lo + chunk_rows, n_major) + for lo in range(0, n_rows, chunk_rows): + hi = min(lo + chunk_rows, n_rows) start, end = int(indptr[lo]), int(indptr[hi]) - block = np.zeros( - (hi - lo, n_cols) if csr else (n_rows, hi - lo), dtype=data_dtype - ) + block = np.zeros((hi - lo, n_cols), dtype=data_dtype) if end > start: minor = np.asarray(indices[start:end], dtype=np.int64) data = np.asarray(values[start:end]) @@ -568,14 +698,8 @@ def densify( # their sum, which is what scipy's own `toarray` produces. # Plain assignment keeps whichever came last, so a # non-canonical input silently changed value on densifying. - if csr: - np.add.at(block, (major, minor), data) - else: - np.add.at(block, (minor, major), data) - if csr: - dst[lo:hi, :] = block - else: - dst[:, lo:hi] = block + np.add.at(block, (major, minor), data) + dst[lo:hi, :] = block def sparsify( @@ -588,48 +712,103 @@ def sparsify( index_dtype: np.dtype, pointer_dtype: np.dtype = np.dtype("int64"), chunk_rows: int = 1024, + chunk: int = DEFAULT_CHUNK, + in_memory: bool = False, console: Optional[Console] = None, compression: Optional[str] = None, ) -> Tuple[int, float]: - """Write a dense array out as CSR or CSC, a block at a time. + """Write a dense array out as CSR or CSC, `chunk_rows` rows at a time. Returns `(nnz, density)` so the caller can say whether it was worth it -- a matrix that is half nonzero gets bigger, not smaller, and the user should hear that from us rather than from `du`. - """ - from adata.core.subset import _append, _growable + CSC goes through a CSR in scratch and the streaming transpose. It read + `chunk_rows` columns of every row at a time, n_obs x --chunk values + (F030), and each of those reads decoded every chunk of a row-chunked + source, once per block of columns. + """ n_rows, n_cols = src.shape - csr = enc == spec.CSR_MATRIX - n_major = n_rows if csr else n_cols + if enc == spec.CSC_MATRIX: + with scratch_group(dst_parent, f"__{name}_sparsify_scratch__") as scratch: + nnz = _sparsify_rows( + src, scratch, "csr", data_dtype=data_dtype, + index_dtype=np.dtype("int64"), pointer_dtype=np.dtype("int64"), + chunk_rows=chunk_rows, compression=compression, + ) + # The scratch's offsets are int64; the output's must hold nnz too. + check_index_range("indptr", pointer_dtype, nnz, src.shape, nnz) + _transpose( + describe(scratch["csr"]), dst_parent, name, in_memory=in_memory, + chunk=chunk, data_dtype=data_dtype, index_dtype=index_dtype, + pointer_dtype=pointer_dtype, compression=compression, + ) + else: + nnz = _sparsify_rows( + src, dst_parent, name, data_dtype=data_dtype, index_dtype=index_dtype, + pointer_dtype=pointer_dtype, chunk_rows=chunk_rows, + compression=compression, + ) + + density = nnz / max(1, n_rows * n_cols) + if console is not None and density > 0.5: + console.print( + f"[yellow]{name} is {density:.0%} nonzero; the sparse form is " + f"larger than the dense one below about 33%.[/]" + ) + return nnz, density + +def _sparsify_rows( + src: Matrix, + dst_parent: Any, + name: str, + *, + data_dtype: np.dtype, + index_dtype: np.dtype, + pointer_dtype: np.dtype, + chunk_rows: int, + compression: Optional[str], +) -> int: + """Write a dense matrix out as CSR, `chunk_rows` rows at a time; its nnz.""" + from adata.core.subset import _append + + n_rows, n_cols = src.shape group = dst_parent.create_group(name) backend = "zarr" if is_zarr_group(dst_parent) else "hdf5" copy_attrs(src.obj.attrs, group.attrs, target_backend=backend) - spec.set_encoding(group, enc) + spec.set_encoding(group, spec.CSR_MATRIX) set_shape_attr(group, src.shape) - out_data = _growable_like( + out_data = _growable_from( group, "data", data_dtype, src.obj, compression=compression ) - out_indices = _growable_like( + out_indices = _growable_from( group, "indices", index_dtype, src.obj, compression=compression ) - counts: List[int] = [] - - for lo in range(0, n_major, chunk_rows): - hi = min(lo + chunk_rows, n_major) - block = np.asarray(src.obj[lo:hi, :] if csr else src.obj[:, lo:hi]) - if not csr: - block = block.T # iterate majors as rows either way - nonzero_major, nonzero_minor = np.nonzero(block) - counts.extend(np.bincount(nonzero_major, minlength=hi - lo).tolist()) - _append(out_indices, nonzero_minor.astype(index_dtype, copy=False)) - _append(out_data, block[nonzero_major, nonzero_minor].astype( - data_dtype, copy=False - )) + # Per-row counts in one int64 array: a Python list of int objects took + # 36 to 40 bytes a row against 8. + counts = np.zeros(n_rows, dtype=np.int64) + pointer_limit = int(np.iinfo(pointer_dtype).max) + nnz = 0 - indptr = np.concatenate(([0], np.cumsum(counts, dtype=np.int64))) + for lo in range(0, n_rows, chunk_rows): + hi = min(lo + chunk_rows, n_rows) + # Cast before looking for nonzeros: a lossy cast turns 0.3 into 0, + # and finding the nonzeros first stored those as explicit zeros + # (F121), 616 entries where scipy has 194. + with np.errstate(invalid="ignore", over="ignore"): + block = np.asarray(src.obj[lo:hi, :]).astype(data_dtype, copy=False) + rows, cols = np.nonzero(block) + nnz += int(rows.size) + if nnz > pointer_limit: + # Offsets cast with a plain astype wrapped negative (F013). + check_index_range("indptr", pointer_dtype, nnz, src.shape, None) + counts[lo:hi] = np.bincount(rows, minlength=hi - lo) + _append(out_indices, cols.astype(index_dtype, copy=False)) + _append(out_data, block[rows, cols]) + + indptr = np.concatenate(([0], np.cumsum(counts))) kw = {} if compression is not None: kw = with_compression( @@ -639,16 +818,39 @@ def sparsify( create_dataset( group, "indptr", data=indptr.astype(pointer_dtype, copy=False), **kw ) + return nnz - nnz = int(indptr[-1]) - density = nnz / max(1, n_rows * n_cols) - if console is not None and density > 0.5: - console.print( - f"[yellow]{name} is {density:.0%} nonzero; the sparse form is " - f"larger than the dense one below about 33%.[/]" +def _growable_from( + group: Any, + name: str, + dtype: np.dtype, + template: Any, + *, + compression: Optional[str] = None, +) -> Any: + """A growable 1-D array with `template`'s codecs and never its chunks. + + `_growable_like` takes a template's first chunk dimension as the element + chunk, which for a dense matrix is a row height: a source chunked 16 + rows tall gave `data` 16-element chunks, 281,188 of them for 4.5M + nonzeros and a file three times the size (F078). Its fill value and + scale-offset filter belong to the template's dtype, not this one's. + """ + backend = "zarr" if is_zarr_group(group) else "hdf5" + kw = dataset_create_kwargs(template, target_backend=backend, dst_parent=group) + for key in ("chunks", "shards", "fillvalue", "scaleoffset"): + kw.pop(key, None) + if compression is not None: + kw = with_compression( + kw, compression, shape=(0,), dtype=dtype, backend=backend, + zarr_format=zarr_format_of(group), ) - return nnz, density + if backend == "zarr": + return group.create_array(name, shape=(0,), dtype=dtype, chunks=(ELEMENT_CHUNK,), **kw) + return group.create_dataset( + name, shape=(0,), maxshape=(None,), dtype=dtype, chunks=(ELEMENT_CHUNK,), **kw + ) def cast_dense( @@ -716,10 +918,15 @@ def plan_conversion( index_dtype: Optional[np.dtype] = None, layout: Optional[str] = None, chunk: int = DEFAULT_CHUNK, + chunk_rows: int = 1024, force: bool = False, console: Optional[Console] = None, ) -> Plan: - """Decide what to do, and refuse here if it should not be done.""" + """Decide what to do, and refuse here if it should not be done. + + `chunk_rows` is the block of a dense matrix the cast check reads, as the + conversion will. + """ src = describe(obj) target_layout = resolve_layout(src, layout) data_dtype = np.dtype(dtype) if dtype is not None else src.dtype @@ -736,10 +943,20 @@ def plan_conversion( else: idx_dtype = ptr_dtype = np.dtype("int64") + if target_layout in spec.SPARSE_TYPES and data_dtype == np.dtype("float16"): + # scipy has no float16 sparse matrix, so anndata, scanpy and this + # tool's own readers fail on one. It was written, with a message + # that every value round-tripped (F058). + raise ValueError( + f"{name}: a sparse matrix cannot hold float16 -- scipy, and so " + "anndata and scanpy, cannot read one. Use --dtype float32, or " + "--layout dense." + ) + report: Optional[CastReport] = None if dtype is not None and data_dtype != src.dtype: values = src.obj["data"] if src.sparse else src.obj - report = check_cast(values, data_dtype, chunk=chunk) + report = check_cast(values, data_dtype, chunk=chunk, rows=chunk_rows) message = report.describe(src.dtype, data_dtype) if not report.lossless and not force: raise ValueError(f"{name}: {message}. Pass --force to convert anyway.") @@ -749,8 +966,13 @@ def plan_conversion( # Always, not only when asked: an inferred dtype can be too narrow too, # and a silently overflowed offset is indistinguishable from corruption. - if src.sparse and not force: - check_index_dtype(src.obj, idx_dtype, ptr_dtype, src.shape) + # Nor does --force waive it: it accepts lossy values and a large store, + # and it let indptr wrap negative (F013). A dense source's coordinates + # are checked here and its offsets as `sparsify` counts them. + if target_layout in spec.SPARSE_TYPES: + check_index_dtype( + src.obj if src.sparse else None, idx_dtype, ptr_dtype, src.shape + ) if target_layout == "dense" and src.sparse: projected = src.shape[0] * src.shape[1] * data_dtype.itemsize @@ -785,7 +1007,7 @@ def convert_matrix( if plan is None: plan = plan_conversion( obj, name, dtype=dtype, index_dtype=index_dtype, layout=layout, - chunk=chunk, force=force, console=console, + chunk=chunk, chunk_rows=chunk_rows, force=force, console=console, ) src = plan.source target_layout = plan.layout @@ -796,7 +1018,10 @@ def convert_matrix( # --- then write ------------------------------------------------------- if target_layout == "dense": if src.sparse: - densify(src, dst_parent, name, data_dtype=data_dtype, chunk_rows=chunk_rows) + densify( + src, dst_parent, name, data_dtype=data_dtype, + chunk_rows=chunk_rows, chunk=chunk, in_memory=in_memory, + ) else: cast_dense(src, dst_parent, name, data_dtype=data_dtype, chunk_rows=chunk_rows) return @@ -811,6 +1036,8 @@ def convert_matrix( index_dtype=idx_dtype, pointer_dtype=ptr_dtype, chunk_rows=chunk_rows, + chunk=chunk, + in_memory=in_memory, console=console, ) return diff --git a/tests/perf_counters.py b/tests/perf_counters.py index cd7e257..1145f26 100644 --- a/tests/perf_counters.py +++ b/tests/perf_counters.py @@ -19,6 +19,31 @@ n elements; counting calls alone would also miss a vectorised-but-quadratic variant that reads n elements n times in one call. +Bytes come in three kinds, because each catches a different defect: + +* **returned** (`read_bytes`, `largest_read`): what the caller holds. The + largest single read is the block a streaming loop keeps in memory, so a + guard on it guards memory without needing an allocator hook at all. +* **decoded** (`decoded_bytes`): every chunk a read touches, whole. The + backends decompress whole chunks, so taking one row from each of a thousand + tall chunks costs a thousand chunks, however little the read returned. +* **encoded** (`encoded_bytes`): every chunk a write touches, whole, and every + shard for a sharded Zarr array, since writing part of a shard rewrites all + of it. A loop writing blocks smaller than that unit re-encodes each unit + once per block. + +Decoded and encoded bytes are a model worked out from the selection, the +chunk grid and the codecs, not a measurement of the library, which is what +keeps them deterministic. It follows HDF5's per-dataset chunk cache, and +counts an HDF5 chunk with no filter as read or written in place; see +`_ChunkCache`. + +Memory has two instruments. `count_allocations` (tracemalloc) is exact and +in-process, and sees numpy's buffers but not what HDF5, blosc or zarr's codecs +allocate in C. `peak_rss` runs a command in a fresh interpreter and reads the +operating system's high-water mark, which sees everything, with a few +megabytes of noise: guard it with generous ratios. + Known bypasses -------------- These reach the file without passing through any hook below: @@ -28,15 +53,19 @@ ``np.asarray`` adds nothing to the element count) * ``dataset.asstr()[...]``, which goes through ``AsStrWrapper.__getitem__`` * ``dataset.fields(...)`` +* ``zarr.Array.vindex[...]`` and the ``get_*_selection`` methods + +``zarr.Array.oindex[...]`` is counted: `core/subset.py` uses it. -None are used in `src/adata/` today. If you add one, the counters will silently -report less work than actually happened -- which is why every guard also +None of the others are used in `src/adata/` today. If you add one, the +counters will silently report less work than actually happened -- which is why every guard also asserts a *lower* bound, and why `test_performance.py` keeps a canary test with a known absolute count. """ from __future__ import annotations +import itertools import sys import tracemalloc from contextlib import contextmanager @@ -63,10 +92,27 @@ class IOCounts: store_get: int = 0 store_get_bytes: int = 0 store_set: int = 0 + store_set_bytes: int = 0 store_delete: int = 0 + #: Bytes the reads handed back, through either backend. + read_bytes: int = 0 + #: The most bytes one read handed back: the block a loop holds. + largest_read: int = 0 + #: Bytes of the whole chunks those reads touched (see the module notes). + decoded_bytes: int = 0 + #: Bytes handed to writes, including data given at creation. + written_bytes: int = 0 + #: The most bytes one write was handed. + largest_write: int = 0 + #: Bytes of the whole chunks, or shards, those writes touched. + encoded_bytes: int = 0 + #: Per-dataset element counts, for diagnosing which column blew up. by_name: Dict[str, int] = field(default_factory=dict) + #: The largest read per dataset, in bytes and in elements. + largest_by_name: Dict[str, int] = field(default_factory=dict) + largest_elements_by_name: Dict[str, int] = field(default_factory=dict) @property def elements(self) -> int: @@ -106,12 +152,25 @@ def reads(self) -> int: """ return self.elements + self.store_get + @property + def read_amplification(self) -> float: + """Bytes decoded per byte returned: 1.0 when reads follow the chunks.""" + return self.decoded_bytes / self.read_bytes if self.read_bytes else 0.0 + + @property + def write_amplification(self) -> float: + """Bytes re-encoded per byte written: 1.0 when writes fill whole units.""" + return self.encoded_bytes / self.written_bytes if self.written_bytes else 0.0 + def __str__(self) -> str: # pragma: no cover - diagnostic only return ( f"read={self.elements} (h5={self.h5_elements} " f"zarr={self.zarr_elements}) wrote={self.h5_written_elements} " f"calls={self.calls} store: get={self.store_get} " - f"set={self.store_set} delete={self.store_delete}" + f"set={self.store_set} delete={self.store_delete} " + f"bytes: read={self.read_bytes} (largest {self.largest_read}) " + f"decoded={self.decoded_bytes} written={self.written_bytes} " + f"(largest {self.largest_write}) encoded={self.encoded_bytes}" ) @@ -122,24 +181,192 @@ def _size(value: Any) -> int: return 1 +def _nbytes(value: Any) -> int: + try: + return int(np.asarray(value).nbytes) + except Exception: # pragma: no cover - exotic dtypes + return 0 + + +def _axis_chunks(key: Any, extent: int, size: int) -> int: + """How many chunks of `size` along an axis of `extent` a key touches.""" + if extent == 0 or size <= 0: + return 0 + if isinstance(key, (int, np.integer)): + return 1 + if isinstance(key, slice): + start, stop, step = key.indices(extent) + if step == 1: + return (stop - 1) // size - start // size + 1 if stop > start else 0 + picked = np.arange(start, stop, step) + else: + picked = np.asarray(key) + if picked.dtype == bool: + picked = np.flatnonzero(picked) + picked = picked.astype(np.int64, copy=False).ravel() + picked = np.where(picked < 0, picked + extent, picked) + return int(np.unique(picked // size).size) if picked.size else 0 + + +def _axis_keys(key: Any, ndim: int) -> tuple: + """One key per axis: Ellipsis expanded, missing trailing axes whole.""" + keys = key if isinstance(key, tuple) else (key,) + keys = tuple(k for k in keys if k is not None and not isinstance(k, str)) + if any(k is Ellipsis for k in keys): + at = next(i for i, k in enumerate(keys) if k is Ellipsis) + keys = keys[:at] + (slice(None),) * (ndim - len(keys) + 1) + keys[at + 1 :] + return keys + (slice(None),) * (ndim - len(keys)) + + +def chunks_touched(key: Any, shape: Any, chunks: Any) -> int: + """How many chunks of a grid a selection touches, axis by axis. + + Orthogonal selections only, which is what h5py and `oindex` take: the + count is the product of the chunks touched along each axis. + """ + shape = tuple(int(n) for n in shape) + total = 1 + for k, extent, size in zip(_axis_keys(key, len(shape)), shape, chunks): + total *= _axis_chunks(k, extent, int(size)) + return total + + +def _unit_bytes(array: Any, unit: Any) -> int: + count = 1 + for c in unit: + count *= int(c) + return count * int(np.dtype(array.dtype).itemsize) + + +def _chunk_grid(array: Any, *, write: bool) -> Any: + """The unit a read decodes, or a write re-encodes: None if contiguous.""" + if write: + shards = getattr(array, "shards", None) + if shards: + return shards + return getattr(array, "chunks", None) + + +def _axis_indices(key: Any, extent: int, size: int) -> np.ndarray: + """The chunk indices along one axis a key touches, sorted.""" + if isinstance(key, (int, np.integer)): + return np.array([(int(key) % extent) // size]) if extent else np.array([], int) + if isinstance(key, slice): + start, stop, step = key.indices(extent) + if step == 1: + if stop <= start: + return np.array([], dtype=np.int64) + return np.arange(start // size, (stop - 1) // size + 1) + picked = np.arange(start, stop, step) + else: + picked = np.asarray(key) + if picked.dtype == bool: + picked = np.flatnonzero(picked) + picked = picked.astype(np.int64, copy=False).ravel() + picked = np.where(picked < 0, picked + extent, picked) + return np.unique(picked // size) + + +#: Past this many chunks in one selection, the cache is not modelled: a read +#: that large evicts any cache a backend has anyway. +_MODELLED_CHUNKS = 1 << 16 + + +class _ChunkCache: + """Which chunks a backend still holds decoded, dataset by dataset. + + HDF5 keeps a chunk cache per dataset (1 MiB by default, none at all for a + chunk bigger than that), and serves a read or a write from it without + touching the file: appending a block at a time to a growing dataset + encodes each chunk once, when it leaves the cache, not once a block. + Modelled as least recently used, at the dataset's own cache size. A + chunk with no filter has nothing to decode, and HDF5 reads or writes + just the part asked for. Zarr caches nothing, so every chunk a selection + touches is decoded and every shard it writes is rewritten. + """ + + def __init__(self) -> None: + self.held: Dict[Any, Dict[Any, None]] = {} + + @staticmethod + def _unfiltered_hdf5(array: Any) -> bool: + """An HDF5 chunk with no filter is read and written in place, in part.""" + plist = getattr(getattr(array, "id", None), "get_create_plist", None) + if plist is None: + return False # Zarr + try: + return plist().get_nfilters() == 0 + except Exception: # pragma: no cover - an HDF5 without the call + return False + + @staticmethod + def _capacity(array: Any, unit_bytes: int) -> int: + plist = getattr(getattr(array, "id", None), "get_access_plist", None) + if plist is None or unit_bytes <= 0: + return 0 # Zarr + try: + cache_bytes = int(plist().get_chunk_cache()[1]) + except Exception: # pragma: no cover - an HDF5 without the call + return 0 + return cache_bytes // unit_bytes + + def cost(self, array: Any, key: Any, returned: int, *, write: bool) -> int: + """Bytes of the whole chunks (or shards) `key` makes the backend code.""" + try: + grid = _chunk_grid(array, write=write) + shape = tuple(int(n) for n in array.shape) + except Exception: # pragma: no cover - an object without a layout + return returned + if not grid or not shape or self._unfiltered_hdf5(array): + # Contiguous, a scalar, or chunks with no codec: what was asked + # for is what moved. + return returned + unit = _unit_bytes(array, grid) + touched = chunks_touched(key, shape, grid) + capacity = self._capacity(array, unit) + if capacity == 0 or touched > _MODELLED_CHUNKS: + return touched * unit + + axes = [ + _axis_indices(k, extent, int(size)) + for k, extent, size in zip(_axis_keys(key, len(shape)), shape, grid) + ] + held = self.held.setdefault((str(array.file.filename), array.name), {}) + missed = 0 + for coordinate in itertools.product(*(a.tolist() for a in axes)): + if coordinate in held: + del held[coordinate] # re-inserted below, as most recent + else: + missed += 1 + held[coordinate] = None + if len(held) > capacity: + del held[next(iter(held))] + return missed * unit + + @contextmanager def count_io() -> Iterator[IOCounts]: """Count backend reads and writes for the duration of the block. - Patches four seams, all verified against h5py 3.15.1 and zarr 3.1.5: + Patches these seams, verified against h5py 3.15.1 and zarr 3.1.5: ====================================== ================================= ``h5py.Dataset.__getitem__`` HDF5 reads ``h5py.Dataset.__setitem__`` HDF5 writes ``h5py.Group.create_dataset`` HDF5 writes made at creation ``zarr.Array.__getitem__`` Zarr reads, v2 and v3 alike + ``zarr.Array.oindex[...]`` Zarr orthogonal reads + ``zarr.Array.__setitem__`` Zarr writes (bytes only) + ``zarr.Group.create_array`` Zarr writes made at creation ``zarr.storage.LocalStore.get`` chunk and metadata fetches ``LocalStore.set`` / ``.delete`` write storms ====================================== ================================= - Both write seams are needed: `create_dataset(name, data=...)` writes its - payload during creation and never touches `__setitem__`, so hooking only - the latter leaves `import` and `create` measuring zero. + Both HDF5 write seams are needed: `create_dataset(name, data=...)` writes + its payload during creation and never touches `__setitem__`, so hooking + only the latter leaves `import` and `create` measuring zero. Zarr's array + writes add to the byte counts but not to `writes`, which counts them in + store keys (see `IOCounts.writes`). The store hooks are the ones that see a metadata write storm: attribute writes never touch `Array.__getitem__`, so an O(k^2) attribute rewrite is @@ -149,14 +376,19 @@ def count_io() -> Iterator[IOCounts]: """ import h5py import zarr + from zarr.core.indexing import OIndex from zarr.storage import LocalStore counts = IOCounts() + cache = _ChunkCache() h5_get = h5py.Dataset.__getitem__ h5_set = h5py.Dataset.__setitem__ h5_create = h5py.Group.create_dataset zarr_get = zarr.Array.__getitem__ + zarr_set = zarr.Array.__setitem__ + zarr_create = zarr.Group.create_array + oindex_get = OIndex.__getitem__ store_get = LocalStore.get store_set = LocalStore.set store_delete = LocalStore.delete @@ -164,27 +396,54 @@ def count_io() -> Iterator[IOCounts]: def _record(name: str, n: int) -> None: counts.by_name[name] = counts.by_name.get(name, 0) + n + def _read(array: Any, key: Any, result: Any, name: str) -> int: + n = _size(result) + returned = _nbytes(result) + counts.read_bytes += returned + counts.largest_read = max(counts.largest_read, returned) + counts.largest_by_name[name] = max(counts.largest_by_name.get(name, 0), returned) + counts.largest_elements_by_name[name] = max( + counts.largest_elements_by_name.get(name, 0), n + ) + counts.decoded_bytes += cache.cost(array, key, returned, write=False) + _record(name, n) + return n + + def _write(array: Any, key: Any, value: Any) -> None: + given = _nbytes(value) + counts.written_bytes += given + counts.largest_write = max(counts.largest_write, given) + counts.encoded_bytes += cache.cost(array, key, given, write=True) + def h5_wrapper(self: Any, key: Any) -> Any: result = h5_get(self, key) - n = _size(result) counts.h5_calls += 1 - counts.h5_elements += n - _record(getattr(self, "name", "?"), n) + counts.h5_elements += _read(self, key, result, getattr(self, "name", "?")) return result def zarr_wrapper(self: Any, key: Any) -> Any: result = zarr_get(self, key) - n = _size(result) counts.zarr_calls += 1 - counts.zarr_elements += n - _record(getattr(self, "path", "?") or "/", n) + counts.zarr_elements += _read(self, key, result, getattr(self, "path", "?") or "/") + return result + + def oindex_wrapper(self: Any, key: Any) -> Any: + result = oindex_get(self, key) + array = self.array + counts.zarr_calls += 1 + counts.zarr_elements += _read(array, key, result, getattr(array, "path", "?") or "/") return result def h5_set_wrapper(self: Any, key: Any, value: Any) -> Any: counts.h5_write_calls += 1 counts.h5_written_elements += _size(value) + _write(self, key, value) return h5_set(self, key, value) + def zarr_set_wrapper(self: Any, key: Any, value: Any) -> Any: + _write(self, key, value) + return zarr_set(self, key, value) + def h5_create_wrapper(self: Any, name: Any, *args: Any, **kwargs: Any) -> Any: counts.h5_write_calls += 1 data = kwargs.get("data") @@ -193,7 +452,16 @@ def h5_create_wrapper(self: Any, name: Any, *args: Any, **kwargs: Any) -> Any: data = args[2] if len(args) > 2 else None if data is not None: counts.h5_written_elements += _size(data) - return h5_create(self, name, *args, **kwargs) + created = h5_create(self, name, *args, **kwargs) + if data is not None: + _write(created, Ellipsis, data) + return created + + def zarr_create_wrapper(self: Any, name: Any, *args: Any, **kwargs: Any) -> Any: + created = zarr_create(self, name, *args, **kwargs) + if kwargs.get("data") is not None: + _write(created, Ellipsis, kwargs["data"]) + return created async def get_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: result = await store_get(self, *args, **kwargs) @@ -207,29 +475,36 @@ async def get_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: async def set_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: counts.store_set += 1 + value = args[1] if len(args) > 1 else kwargs.get("value") + try: + counts.store_set_bytes += len(value) + except TypeError: # pragma: no cover - prototype buffers + pass return await store_set(self, *args, **kwargs) async def delete_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: counts.store_delete += 1 return await store_delete(self, *args, **kwargs) - h5py.Dataset.__getitem__ = h5_wrapper - h5py.Dataset.__setitem__ = h5_set_wrapper - h5py.Group.create_dataset = h5_create_wrapper - zarr.Array.__getitem__ = zarr_wrapper - LocalStore.get = get_wrapper - LocalStore.set = set_wrapper - LocalStore.delete = delete_wrapper + patched = [ + (h5py.Dataset, "__getitem__", h5_get, h5_wrapper), + (h5py.Dataset, "__setitem__", h5_set, h5_set_wrapper), + (h5py.Group, "create_dataset", h5_create, h5_create_wrapper), + (zarr.Array, "__getitem__", zarr_get, zarr_wrapper), + (zarr.Array, "__setitem__", zarr_set, zarr_set_wrapper), + (zarr.Group, "create_array", zarr_create, zarr_create_wrapper), + (OIndex, "__getitem__", oindex_get, oindex_wrapper), + (LocalStore, "get", store_get, get_wrapper), + (LocalStore, "set", store_set, set_wrapper), + (LocalStore, "delete", store_delete, delete_wrapper), + ] + for owner, name, _, wrapper in patched: + setattr(owner, name, wrapper) try: yield counts finally: - h5py.Dataset.__getitem__ = h5_get - h5py.Dataset.__setitem__ = h5_set - h5py.Group.create_dataset = h5_create - zarr.Array.__getitem__ = zarr_get - LocalStore.get = store_get - LocalStore.set = store_set - LocalStore.delete = store_delete + for owner, name, original, _ in patched: + setattr(owner, name, original) @contextmanager @@ -258,6 +533,83 @@ def count_allocations() -> Iterator[List[int]]: tracemalloc.stop() +#: Run by `peak_rss` in a fresh interpreter. The high-water mark is read once +#: the CLI and everything it imports is loaded, and again after the command, +#: so the difference is what the command itself added. +_RSS_RUNNER = r""" +import json, resource, sys + +def high_water(): + value = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + return value if sys.platform == "darwin" else value * 1024 # Linux: KiB + +argv, report = json.loads(sys.argv[1]), sys.argv[2] +from adata.cli import app + +before = high_water() +try: + code = app(args=argv, prog_name="adata", standalone_mode=False) or 0 +except SystemExit as stop: + code = stop.code or 0 +with open(report, "w") as handle: + json.dump({"before": before, "after": high_water(), "code": code}, handle) +""" + + +@dataclass +class PeakRSS: + """Resident memory of one command, in bytes.""" + + #: The high-water mark once the CLI was imported, before the command ran. + before: int + #: The high-water mark when the command had finished. + after: int + #: What the command's stderr said, for a failure message. + stderr: str = "" + + @property + def added(self) -> int: + """How far the command raised the high-water mark.""" + return max(0, self.after - self.before) + + +def peak_rss(argv: List[str], *, cwd: Any = None) -> PeakRSS: + """Peak resident memory of `adata `, run in a fresh interpreter. + + The instrument that sees everything: HDF5's chunk caches, the buffers + blosc and zarr's codecs decode into, and anything else allocated in C, + none of which tracemalloc can see. A subprocess because the high-water + mark only rises: in this process it would hold whatever the test suite + reached before. Not exact -- the allocator keeps a few megabytes it has + freed -- so guard it with ratios that leave room, against inputs large + enough (tens of megabytes) that holding one would be unmistakable. + """ + import json + import subprocess + import tempfile + from pathlib import Path + + with tempfile.TemporaryDirectory() as scratch: + report = Path(scratch) / "rss.json" + done = subprocess.run( + [sys.executable, "-c", _RSS_RUNNER, json.dumps([str(a) for a in argv]), + str(report)], + cwd=cwd, capture_output=True, text=True, check=False, + ) + if done.returncode != 0 or not report.exists(): + raise RuntimeError( + f"adata {' '.join(map(str, argv))} could not be measured " + f"(exit {done.returncode}):\n{done.stderr[-2000:]}" + ) + measured = json.loads(report.read_text()) + if measured["code"]: + raise RuntimeError( + f"adata {' '.join(map(str, argv))} failed (exit {measured['code']}):\n" + f"{done.stdout[-2000:]}{done.stderr[-2000:]}" + ) + return PeakRSS(measured["before"], measured["after"], done.stderr) + + @contextmanager def count_lines(prefix: str) -> Iterator[List[int]]: """Count executed Python line events in files under `prefix`, as `[n]`. @@ -527,3 +879,46 @@ def assert_grows_slower_than_input( f"{budget:.0f}x. Series: {small}->{c_small}, {large}->{c_large}" ) + + +#: The most memory a streaming path may add per byte of input added. A path +#: holding a block of `--chunk` rows of a matrix many blocks tall adds a few +#: hundredths; one holding the matrix adds a whole byte, and more for every +#: copy it makes. A quarter leaves room for the per-row arrays a path keeps +#: by design (an indptr, a row's count) on a narrow fixture. +INPUT_SHARE = 0.25 + + +def assert_memory_is_a_fraction_of_the_input( + measure: Callable[[int], int], + input_bytes: Callable[[int], int], + *, + what: str, + axis: str, + sizes: tuple, + at_most: float = INPUT_SHARE, +) -> None: + """Assert peak memory grows by at most `at_most` of what the input grows by. + + `measure(n)` is peak memory at size `n` (allocations or RSS), and + `input_bytes(n)` the size of the input in memory at `n`. Increments, as + in `assert_grows_linearly`, so fixed costs cancel. + + This holds along either axis, which a growth ratio does not: widen a + matrix and a block of `--chunk` rows widens with it, so its memory grows + exactly as fast as the input does, and a ratio cannot tell it from the + whole matrix. As a share of the input it is still `--chunk` / n_obs, + against 1 or more for a path that holds the matrix. + """ + small, large = sizes[0], sizes[-1] + m_small, m_large = measure(small), measure(large) + i_small, i_large = input_bytes(small), input_bytes(large) + added = i_large - i_small + assert added > 0, f"{what}: the input did not grow between {small} and {large}" + share = (m_large - m_small) / added + assert share <= at_most, ( + f"{what} holds memory in proportion to the input, scaling {axis}: the " + f"input grew by {added / 1e6:.2f} MB ({small}->{large}) and peak memory by " + f"{(m_large - m_small) / 1e6:.2f} MB, {share:.2f} of it, against at most " + f"{at_most}. Peak {m_small / 1e6:.2f} MB -> {m_large / 1e6:.2f} MB." + ) diff --git a/tests/perf_fixtures.py b/tests/perf_fixtures.py new file mode 100644 index 0000000..f00f75e --- /dev/null +++ b/tests/perf_fixtures.py @@ -0,0 +1,147 @@ +"""Stores for the memory and I/O guards, built the way users' stores are. + +`test_performance.py`'s `_store` builds one shape: a CSR X eight columns wide, +grown in n_obs. A path whose memory scales with the width, with the CSC +layout, or with a dense X passes every guard built on it, which is how the +recast and subset paths that `--chunk` does not bound reached a release (F091, +F146). These builders scale every axis and layout instead, in each of the +formats the tool reads: + +* ``h5ad``: as anndata writes it -- a contiguous dense X, auto-chunked sparse + arrays; +* ``zarr2``: anndata's Zarr v2; +* ``zarr3``: anndata 0.13's default Zarr v3, sharded (`auto_shard_zarr_v3`): + inner chunks of a few hundred bytes inside shards of the whole array. Any + read or write that does not line up with those shows here and nowhere else. + +anndata writes every one, so a fixture is what a user would hand the tool, +not what the tool would write itself. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Optional, Tuple + +import numpy as np + +#: Matrix layouts a store can hold. +LAYOUTS = ("csr", "csc", "dense") + +#: Store formats, as anndata writes each. +FORMATS = ("h5ad", "zarr2", "zarr3") + + +def suffix(fmt: str) -> str: + return ".h5ad" if fmt == "h5ad" else ".zarr" + + +def counts_matrix( + n_obs: int, + n_var: int, + layout: str = "csr", + *, + density: float = 0.1, + dtype: str = "float32", + seed: int = 0, +): + """Whole-number counts, so any float cast round-trips and needs no --force.""" + from scipy import sparse + + rng = np.random.default_rng(seed) + matrix = sparse.random( + n_obs, n_var, density=density, format="csr", dtype="float64", random_state=rng + ) + matrix.data = np.ceil(matrix.data * 100) + matrix = matrix.astype(dtype) + if layout == "dense": + return matrix.toarray() + return matrix.tocsc() if layout == "csc" else matrix + + +def skewed_matrix(n_obs: int, n_var: int = 64, *, dense_genes: Tuple[int, int] = (8, 16), + layout: str = "csc", seed: int = 0): + """Counts where a run of adjacent genes is in every cell, the rest sparse. + + Ribosomal genes sort together and are expressed nearly everywhere, so a + real matrix has exactly this shape: a few columns hold most entries. + """ + from scipy import sparse + + matrix = counts_matrix(n_obs, n_var, "csr", density=0.02, seed=seed).tolil() + lo, hi = dense_genes + matrix[:, lo:hi] = np.arange(1, n_obs * (hi - lo) + 1, dtype="float32").reshape( + n_obs, hi - lo + ) % 7 + 1 + matrix = sparse.csr_matrix(matrix) + return matrix.tocsc() if layout == "csc" else matrix + + +def write_anndata(obj, path: Path, fmt: str) -> Path: + """Write `obj` as anndata would in `fmt`.""" + import anndata as ad + + if fmt == "h5ad": + obj.write_h5ad(path) + elif fmt == "zarr2": + with ad.settings.override(zarr_write_format=2): + obj.write_zarr(path) + elif fmt == "zarr3": + with ad.settings.override(zarr_write_format=3, auto_shard_zarr_v3=True): + obj.write_zarr(path) + else: # pragma: no cover - a typo in a test + raise ValueError(f"unknown format {fmt!r}") + return path + + +def matrix_store( + directory: Path, + *, + n_obs: int, + n_var: int, + layout: str = "csr", + fmt: str = "h5ad", + density: float = 0.1, + dtype: str = "float32", + matrix=None, + dense_chunks: Optional[Tuple[int, int]] = None, + tag: str = "m", +) -> Path: + """A store holding X of the given shape and layout, and a numeric obs column. + + `dense_chunks` rechunks a dense h5ad X, as a store written by another + tool might be (anndata writes it contiguous). + """ + import anndata as ad + import pandas as pd + + directory.mkdir(parents=True, exist_ok=True) + if matrix is None: + matrix = counts_matrix(n_obs, n_var, layout, density=density, dtype=dtype) + n_obs, n_var = matrix.shape + obj = ad.AnnData( + X=matrix, + obs=pd.DataFrame( + {"col": np.arange(n_obs, dtype="int32")}, + index=[f"c{i}" for i in range(n_obs)], + ), + var=pd.DataFrame(index=[f"g{i}" for i in range(n_var)]), + ) + path = directory / f"{tag}-{layout}-{n_obs}x{n_var}{suffix(fmt)}" + write_anndata(obj, path, fmt) + if dense_chunks is not None: + rechunk_dense(path, dense_chunks) + return path + + +def rechunk_dense(path: Path, chunks: Tuple[int, int], key: str = "X") -> None: + """Rewrite a dense h5ad matrix with `chunks`, keeping its attributes.""" + import h5py + + with h5py.File(path, "a") as handle: + values = handle[key][...] + attrs = dict(handle[key].attrs) + del handle[key] + dataset = handle.create_dataset(key, data=values, chunks=chunks) + for name, value in attrs.items(): + dataset.attrs[name] = value diff --git a/tests/test_performance.py b/tests/test_performance.py index 4871722..22c8d12 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -57,6 +57,7 @@ GROWTH_LIMIT, SIZES, assert_grows_linearly, + assert_memory_is_a_fraction_of_the_input, assert_grows_slower_than_input, assert_independent_of, count_allocations, @@ -195,6 +196,114 @@ def test_the_zarr_store_hooks_fire(tmp_path): assert io.store_get > 0, "LocalStore.get hook is not firing" +def test_the_byte_counters_see_whole_chunks_decoded_and_rewritten(tmp_path): + """F146: the counters summed the elements a read returned, never the + chunks decoded to serve it, nor the shards a write rewrote, so a 70x read + amplification and a shard rewritten once a block both passed. Known + reads and writes here, with the numbers worked out by hand.""" + import h5py + import zarr + + row, chunk = 40 * 8, 10 * 40 * 8 + with h5py.File(tmp_path / "c.h5", "w") as handle: + for name, extra in (("gz", {"compression": "gzip"}), ("raw", {})): + handle.create_dataset(name, data=np.zeros((100, 40)), chunks=(10, 40), **extra) + + with h5py.File(tmp_path / "c.h5", "r") as handle, count_io() as io: + handle["gz"][5:6] # one row: the whole compressed chunk is decoded + handle["gz"][6:7] # the same chunk again, from HDF5's chunk cache + handle["raw"][5:6] # no filter: HDF5 reads just the row + assert (io.read_bytes, io.largest_read) == (3 * row, row) + assert io.decoded_bytes == chunk + row + + with h5py.File(tmp_path / "c.h5", "r", rdcc_nbytes=0) as handle, count_io() as io: + handle["gz"][5:6] + handle["gz"][6:7] + assert io.decoded_bytes == 2 * chunk, "with no cache, each read decodes the chunk" + + array = zarr.create_array( + str(tmp_path / "s.zarr"), shape=(64, 64), chunks=(8, 8), shards=(32, 64), + dtype="float32", zarr_format=3, + ) + with count_io() as io: + array[0:4, :] = np.ones((4, 64), dtype="float32") + array.oindex[np.array([1, 9, 40]), :] + assert io.written_bytes == 4 * 64 * 4 + assert io.encoded_bytes == 32 * 64 * 4, "a write to part of a shard rewrites all of it" + assert io.read_bytes == 3 * 64 * 4 + assert io.decoded_bytes == 3 * 8 * (8 * 8 * 4), "three rows of eight inner chunks" + assert io.zarr_calls == 1, "oindex reads are counted" + + +def test_h5py_calls_keep_no_memory_once_they_return(tmp_path): + """h5py 3.15 kept 100 to 160 bytes of every dataset read, write and + `.shape` for good: 48 MB of resident memory after 300,000 writes. Every + streaming loop makes calls in proportion to its blocks, so memory grew + with the input, and in the transpose, whose bucket writes grow with the + square of the bucket count, faster still. 3.16.0, the floor now, keeps + nothing; this keeps the floor honest.""" + import gc + import tracemalloc + + import h5py + + calls = 20_000 + with h5py.File(tmp_path / "calls.h5", "w") as handle: + dataset = handle.create_dataset("x", shape=(100,), dtype="int64") + values = np.arange(3) + + def work() -> None: + dataset[10:13] = values + dataset[10:13] + dataset.shape + + work() + tracemalloc.start() + try: + gc.collect() + before = tracemalloc.get_traced_memory()[0] + for _ in range(calls): + work() + gc.collect() + kept = tracemalloc.get_traced_memory()[0] - before + finally: + tracemalloc.stop() + + assert kept <= calls, ( + f"h5py {h5py.__version__} kept {kept / calls:.0f} bytes a call " + f"({kept / 1e6:.1f} MB over {calls:,} calls of a read, a write and .shape)" + ) + + +@pytest.mark.slow +def test_peak_rss_sees_a_command_that_holds_the_matrix(tmp_path): + """F146: memory was measured only with tracemalloc, which cannot see what + HDF5, blosc or zarr's codecs allocate in C. `peak_rss` runs a command in + a fresh interpreter and reads the high-water mark. It must see the + in-memory transpose hold the matrix (that is what `--in-memory` is), + and the streaming one hold far less.""" + from tests.perf_counters import peak_rss + from tests.perf_fixtures import counts_matrix, matrix_store + + matrix = counts_matrix(40_000, 1000, "csr", density=0.1) + source = matrix_store(tmp_path, n_obs=40_000, n_var=1000, matrix=matrix) + size = matrix.data.nbytes + matrix.indices.nbytes + matrix.indptr.nbytes + + def added(*extra: str) -> int: + out = tmp_path / f"t{len(extra)}.h5ad" + return peak_rss(["recast", source, "X", "-o", out, "--layout", "csc", *extra]).added + + loaded, streamed = added("--in-memory"), added() + assert loaded >= 2 * size, ( + f"the in-memory transpose of a {size / 1e6:.0f} MB matrix added only " + f"{loaded / 1e6:.0f} MB of resident memory: peak_rss is not seeing it" + ) + assert streamed < loaded, ( + f"streaming added {streamed / 1e6:.0f} MB and the in-memory transpose " + f"{loaded / 1e6:.0f} MB, for a {size / 1e6:.0f} MB matrix" + ) + + # --------------------------------------------------------------------------- # concat: reads must not grow faster than linearly on any axis # @@ -1949,3 +2058,247 @@ def record(*args, _name=name, _original=original, **kwargs): called.clear() _recast(source, tmp_path / "asked.h5ad", layout="csc", in_memory=True) assert called == ["transpose_sparse_in_memory"], called + + +# --------------------------------------------------------------------------- +# memory set by --chunk, in every layout and along either axis +# +# The peak guards above grow n_obs of an eight-column CSR store, the one +# shape every path handles well (F091). A CSC densify holding n_obs x --chunk +# values, a dense cast check reading the whole matrix and a transpose +# stepping by the average major all passed them. These build every layout, +# grow either axis, and state the claim as directly as it can be stated: peak +# memory added per byte of input added. A path holding a block of --chunk +# rows adds a few hundredths of it; one holding the matrix adds all of it. + +#: Rows in a dense block, and entries in a sparse read, while the grid runs. +GRID_ROWS, GRID_ENTRIES = 32, 16_384 + +#: Sizes along the grown axis, and the other axis's fixed size. Small enough +#: that a transpose needs no more than a few dozen buckets (see +#: `test_a_transpose_buckets_hold_about_a_chunk_on_a_skewed_matrix`). +GRID_SIZES = {"n_obs": (1024, 8192), "n_var": (128, 1024)} +GRID_FIXED = {"n_obs": 1024, "n_var": 4096} +GRID_DENSITY = 0.05 + + +def _grid_shape(axis: str, n: int) -> tuple: + return (n, GRID_FIXED["n_obs"]) if axis == "n_obs" else (GRID_FIXED["n_var"], n) + + +def _in_memory_bytes(matrix) -> int: + if sparse.issparse(matrix): + return matrix.data.nbytes + matrix.indices.nbytes + matrix.indptr.nbytes + return int(matrix.nbytes) + + +#: What each command does to a store, at the grid's chunk sizes. +GRID_COMMANDS = { + "recast --dtype": lambda s, o: _recast( + s, o, dtype="float64", chunk_rows=GRID_ROWS, chunk=GRID_ENTRIES + ), + "recast --layout csr": lambda s, o: _recast( + s, o, layout="csr", force=True, chunk_rows=GRID_ROWS, chunk=GRID_ENTRIES + ), + "recast --layout csc": lambda s, o: _recast( + s, o, layout="csc", force=True, chunk_rows=GRID_ROWS, chunk=GRID_ENTRIES + ), + "recast --layout dense": lambda s, o: _recast( + s, o, layout="dense", force=True, chunk_rows=GRID_ROWS, chunk=GRID_ENTRIES + ), + "subset": lambda s, o: subset_h5ad( + s, o, None, None, console=QUIET, obs_query="col >= 0", chunk_rows=GRID_ROWS + ), + "concat": lambda s, o: concat_on_disk( + [s, s.with_name(f"b{s.name}")], o, QUIET, chunk_rows=GRID_ROWS, index_unique="-" + ), +} + +#: Known gaps, each owned by another batch, and whether the fix must flip it. +_F032 = "F032: sparse subset blocks are 4,096 majors, whatever --chunk (B6b, #97)" +_CONCAT = ( + "concat keeps every obs name by design, and its row pointers in Python " + "lists until B3b (#116, F132); together a third of this fixture's input" +) +GRID_GAPS = { + ("subset", "csr", "n_obs"): (_F032, True), + ("subset", "csr", "n_var"): (_F032, True), + ("subset", "csc", "n_obs"): (_F032, True), + ("subset", "csc", "n_var"): (_F032, True), + ("concat", "csr", "n_obs"): (_CONCAT, False), + ("concat", "csc", "n_obs"): (_CONCAT, False), +} + + +def _grid_cases(): + for command in GRID_COMMANDS: + for layout in ("csr", "csc", "dense"): + if command == f"recast --layout {layout}": + continue # a no-op, refused + for axis in GRID_SIZES: + gap = GRID_GAPS.get((command, layout, axis)) + marks = [pytest.mark.xfail(reason=gap[0], strict=gap[1])] if gap else [] + yield pytest.param(command, layout, axis, marks=marks, + id=f"{command.replace(' ', '')}-{layout}-{axis}") + + +@pytest.mark.parametrize("command,layout,axis", list(_grid_cases())) +def test_peak_memory_is_a_small_share_of_the_input(tmp_path, command, layout, axis): + """F091: every layout, both axes, against what the input itself costs.""" + import shutil + + from tests.perf_fixtures import counts_matrix, matrix_store + + run = GRID_COMMANDS[command] + matrices = {} + + def matrix(n: int): + if n not in matrices: + matrices[n] = counts_matrix(*_grid_shape(axis, n), layout, density=GRID_DENSITY) + return matrices[n] + + def measure(n: int) -> int: + n_obs, n_var = _grid_shape(axis, n) + source = matrix_store(tmp_path / f"g{n}", n_obs=n_obs, n_var=n_var, + layout=layout, matrix=matrix(n)) + if command == "concat": + shutil.copy(source, source.with_name(f"b{source.name}")) + with count_allocations() as peak: + run(source, tmp_path / f"g{n}" / "out.h5ad") + return peak[0] + + copies = 2 if command == "concat" else 1 + assert_memory_is_a_fraction_of_the_input( + measure, + lambda n: copies * _in_memory_bytes(matrix(n)), + what=f"{command} of a {layout} X", + axis=axis, + sizes=GRID_SIZES[axis], + ) + + +# --------------------------------------------------------------------------- +# recast's blocks, one finding at a time +# +# The grid above shows memory; these pin down the block each path reads, in +# bytes, which needs no allocator at all: the largest single read must be +# set by --chunk, not by the input. Every format, since anndata 0.13's +# sharded Zarr v3 lays a matrix out unlike anything else. + +STORE_FORMATS = ("h5ad", "zarr2", "zarr3") + + +def _largest_matrix_read(io, *, elements: bool = False) -> int: + """The largest read of X's values or coordinates, or of its scratch, in + bytes or elements. + + Not its indptr, which every path reads whole by design (8 bytes a row), + nor obs and var, which are copied. + """ + def of_x(name: str) -> bool: + name = name.strip("/") + return name.split("/")[0] in ("X", "csr") or name.startswith("__X_") + + sizes = io.largest_elements_by_name if elements else io.largest_by_name + return max( + (size for name, size in sizes.items() + if of_x(name) and not name.endswith("indptr")), + default=0, + ) + + +@pytest.mark.parametrize("fmt", STORE_FORMATS) +def test_a_dense_cast_check_reads_a_block_of_rows_not_the_matrix(tmp_path, fmt): + """F029: `recast --dtype` on a dense matrix read `dataset[0:1048576]` -- + every row and column -- to check the cast, then made two more copies, + whatever --chunk said: 1,500,000 elements in one read of a 3000 x 500 + matrix at --chunk 100. The largest read is --chunk rows now.""" + from tests.perf_fixtures import matrix_store + + def measure(n: int) -> int: + # Narrow, and 256-row blocks: on sharded Zarr each block still + # rewrites the whole shard (F036, B6b), which is slow, not wrong. + source = matrix_store(tmp_path, n_obs=n, n_var=16, layout="dense", fmt=fmt) + out = tmp_path / f"o{n}{source.suffix}" + with count_io() as io: + _recast(source, out, dtype="float64", chunk_rows=256) + assert io.read_bytes >= n * 16 * 4, "the matrix was read" + return _largest_matrix_read(io) + + assert_independent_of( + measure, what=f"the largest read of recast --dtype ({fmt})", axis="n_obs", + sizes=(1024, 16384), + ) + + +@pytest.mark.parametrize("fmt", STORE_FORMATS) +@pytest.mark.parametrize("layout,source", [("dense", "csc"), ("csc", "dense")]) +def test_csc_densify_and_sparsify_read_blocks_set_by_the_chunk(tmp_path, fmt, layout, source): + """F030: densifying a CSC matrix allocated n_rows x --chunk per block, + and sparsifying to CSC read --chunk columns of every row, so memory + followed n_obs: 8 GB a block at 2M cells and the default --chunk. Both go + through a CSR in scratch now, in --chunk rows and chunk-sized reads.""" + from tests.perf_fixtures import matrix_store + + def measure(n: int) -> int: + store = matrix_store(tmp_path, n_obs=n, n_var=64, layout=source, fmt=fmt) + with count_allocations() as peak, count_io() as io: + _recast(store, tmp_path / f"o{n}{store.suffix}", layout=layout, force=True, + chunk_rows=GRID_ROWS, chunk=4096) + assert _largest_matrix_read(io) <= 4096 * 8 + GRID_ROWS * 64 * 8, io + return peak[0] + + assert_grows_slower_than_input( + measure, what=f"recast {source} -> {layout} ({fmt})", axis="n_obs", + sizes=(512, 8192), at_least=4.0, + ) + + +@pytest.mark.parametrize("fmt", STORE_FORMATS) +@pytest.mark.parametrize("source", ["csc", "csr"]) +def test_a_transpose_reads_about_a_chunk_on_a_skewed_matrix(tmp_path, fmt, source): + """F077: pass 2 of the streaming transpose stepped by the average nnz per + major. Eight adjacent genes in every cell made one step many times the + chunk -- 7.7x at 200k cells, growing with them. It steps by entries + now, and a gene in every cell gets a bucket of its own, streamed.""" + from adata.core.csr_build import MIN_BUCKET_ENTRIES + from tests.perf_fixtures import matrix_store, skewed_matrix + + chunk = 4096 + # A read of the source holds a chunk of entries; a bucket of several + # majors, read whole, up to one bucket plus one major that fits in one. + bound = 2 * max(chunk, MIN_BUCKET_ENTRIES) + + for n in (1024, 16384): + store = matrix_store(tmp_path, n_obs=n, n_var=64, layout=source, fmt=fmt, + matrix=skewed_matrix(n, 64, layout=source), tag="skew") + target = "csr" if source == "csc" else "csc" + with count_io() as io: + _recast(store, tmp_path / f"o{n}{store.suffix}", layout=target, chunk=chunk) + assert io.read_bytes >= n * 8 * 4, "the dense genes were read" + largest = _largest_matrix_read(io, elements=True) + assert largest <= bound, ( + f"a skewed {source} transpose of {n:,} cells ({fmt}) read " + f"{largest:,} entries at once, against {bound:,} for a chunk of " + f"{chunk:,}: a run of genes in every cell is being read whole" + ) + + +@pytest.mark.slow +def test_a_dense_cast_adds_a_block_of_resident_memory_not_the_matrix(tmp_path): + """F029 as the operating system sees it, through `peak_rss`: at 64 MB of + float32 the cast check added several times the matrix in resident + memory (the block, its cast, the cast back and a mask), at --chunk 64.""" + from tests.perf_counters import peak_rss + from tests.perf_fixtures import matrix_store + + def measure(n: int) -> int: + store = matrix_store(tmp_path, n_obs=n, n_var=256, layout="dense", density=0.5) + out = tmp_path / f"o{n}.h5ad" + return peak_rss(["recast", store, "X", "-o", out, "--dtype", "float64", + "--chunk", "64"]).added + + assert_memory_is_a_fraction_of_the_input( + measure, lambda n: n * 256 * 4, what="recast --dtype of a dense X (RSS)", + axis="n_obs", sizes=(8192, 65536), + ) diff --git a/tests/test_recast_streaming.py b/tests/test_recast_streaming.py new file mode 100644 index 0000000..4b4173a --- /dev/null +++ b/tests/test_recast_streaming.py @@ -0,0 +1,498 @@ +"""`recast` on every backend, with blocks small enough to cross every boundary. + +`test_recast.py` checks recast on h5ad with its default blocks, where most +matrices fit in one. These run on h5ad, Zarr v2 and Zarr v3 (`new_store`), +at a `--chunk` of a few rows and a few entries, so every path splits the +matrix into many blocks, and check the result two ways: against scipy's own +conversion of the same matrix, and by reading it back with anndata. + +Memory is checked in `test_performance.py`; these are the correctness half of +the same fixes (issue #96). +""" + +from __future__ import annotations + +import re +from pathlib import Path +from typing import Any, Optional + +import numpy as np +import pytest +from rich.console import Console +from typer.testing import CliRunner + +from adata.cli import app +from adata.commands.recast import recast_store +from adata.core import recast as core +from adata.elements import spec +from adata.elements import write as ew + +ad = pytest.importorskip("anndata", reason="anndata reads every result back") +sparse = pytest.importorskip("scipy.sparse") + +runner = CliRunner() +QUIET = Console(quiet=True) +_ANSI = re.compile(r"\x1b\[[0-9;]*[a-zA-Z]") + + +def _out(result) -> str: + text = result.stdout + (result.stderr or "") + return " ".join(_ANSI.sub("", text).split()) + + +def _matrix(n_obs: int = 23, n_var: int = 17, *, density: float = 0.3, seed: int = 0, + integral: bool = True): + rng = np.random.default_rng(seed) + matrix = sparse.random(n_obs, n_var, density=density, format="csr", + dtype="float64", random_state=rng) + # Counts from 1, so no stored entry is an explicit zero. + matrix.data = np.ceil(matrix.data * 100) if integral else matrix.data + return matrix + + +def _as(matrix, layout: str): + if layout == "dense": + return matrix.toarray() + return matrix.tocsc() if layout == "csc" else sparse.csr_matrix(matrix) + + +def _write_matrix(parent: Any, name: str, value: Any) -> None: + if sparse.issparse(value): + ew.write_sparse( + parent, name, value.data, value.indices, value.indptr, value.shape, + enc_type=spec.CSC_MATRIX if value.format == "csc" else spec.CSR_MATRIX, + ) + else: + ew.write_dense(parent, name, value) + + +def _put(new_store, value: Any, at: str = "X", *, stem: str = "in") -> Path: + """A store on this test's backend holding `value` at `at`, obs and var to match.""" + path, opener = new_store(stem) + n_obs, n_var = value.shape + with opener("a") as root: + ew.write_dataframe_header(root, "obs", [f"c{i}" for i in range(n_obs)], []) + ew.write_dataframe_header(root, "var", [f"g{i}" for i in range(n_var)], []) + parent = root + *groups, leaf = at.split("/") + for name in groups: + parent = parent[name] if name in parent else ew.write_mapping(parent, name) + _write_matrix(parent, leaf, value) + return path + + +def _output(source: Path, stem: str = "out") -> Path: + return source.with_name(stem + source.suffix) + + +def _read(path: Path): + return ad.read_h5ad(path) if path.suffix == ".h5ad" else ad.read_zarr(path) + + +def _at(obj: Any, path: str) -> Any: + """`X`, `obsm/X_pca` or `uns/a/b` of an AnnData read back.""" + if path == "uns/neighbors/connectivities": + # anndata moves the graph older versions kept in uns to obsp on reading. + return obj.obsp["connectivities"] + first, *rest = path.split("/") + value = getattr(obj, first) + for part in rest: + value = value[part] + return value + + +def _recast(source: Path, *, out: Optional[Path] = None, entries=("X",), **kwargs) -> Path: + out = out or _output(source) + kwargs.setdefault("chunk", 7) + kwargs.setdefault("chunk_rows", 3) + recast_store(source, list(entries), out, QUIET, **kwargs) + return out + + +def _same(got, expected) -> None: + """Equal values, equal dtype, and the layout scipy would give.""" + if sparse.issparse(expected): + assert sparse.issparse(got), f"expected {expected.format}, got a dense array" + assert got.format == expected.format + assert got.has_sorted_indices, "indices out of order within a major" + assert got.nnz == expected.nnz, f"nnz {got.nnz}, scipy has {expected.nnz}" + got, expected = got.toarray(), expected.toarray() + else: + assert not sparse.issparse(got), f"expected a dense array, got {got.format}" + got = np.asarray(got) + assert got.dtype == expected.dtype, f"{got.dtype} against {expected.dtype}" + assert np.array_equal(got, expected) + + +# --------------------------------------------------------------------------- +# every conversion, in many blocks + + +@pytest.mark.parametrize("target", ["csr", "csc", "dense"]) +@pytest.mark.parametrize("source", ["csr", "csc", "dense"]) +def test_every_layout_conversion_matches_scipy_in_blocks(new_store, source, target): + """F030, F077: CSC densify and dense-to-CSC sparsify now go through a CSR + scratch and the streaming transpose, and the transpose steps by entries. + Whatever route a conversion takes, in blocks of three rows and seven + entries the result must be what scipy makes of the same matrix.""" + matrix = _matrix() + source_path = _put(new_store, _as(matrix, source)) + + out = _recast(source_path, layout=target, force=True) + + expected = _as(matrix, target) + if target != "dense": + expected.sort_indices() + _same(_read(out).X, expected) + + +@pytest.mark.parametrize("source", ["csr", "csc"]) +def test_a_transpose_splitting_and_streaming_majors_matches_scipy(new_store, source): + """F077: pass 2 now reads blocks of entries, splitting a major across + them, and pass 3 streams a major larger than the chunk from a bucket of + its own instead of sorting it whole. Empty majors at both ends, a full + row and a full column, in blocks of five entries and buckets of four.""" + from adata.storage import open_store + + dense = np.zeros((30, 12)) + dense[1:29, 4] = np.arange(1, 29) # a full column; rows 0 and 29 empty + dense[3, 1:11] = 2 # a full row; columns 0 and 11 empty + dense[7, 9] = 5 + matrix = _as(sparse.csr_matrix(dense), source) + path = _put(new_store, matrix) + out = _output(path) + + with open_store(path, "r") as src, open_store(out, "w") as dst: + core.transpose_sparse_streaming( + core.describe(src.root["X"]), dst.root, "X", + data_dtype=np.dtype("float64"), index_dtype=np.dtype("int64"), + chunk=5, bucket_entries=4, + ) + group = dst.root["X"] + got = (sparse.csc_matrix if source == "csr" else sparse.csr_matrix)( + (group["data"][...], group["indices"][...], group["indptr"][...]), + shape=dense.shape, + ) + + expected = _as(sparse.csr_matrix(dense), "csc" if source == "csr" else "csr") + _same(got, expected) + + +@pytest.mark.parametrize("layout", ["csr", "csc", "dense"]) +def test_a_cast_in_many_blocks_matches_scipy(new_store, layout): + """F029: the round-trip check now reads a dense matrix a block of rows at a + time. It must still see every value: one lossy value in the last block is + enough to refuse.""" + matrix = _matrix() + value = _as(matrix, layout) + source_path = _put(new_store, value) + + out = _recast(source_path, dtype="float32") + _same(_read(out).X, value.astype("float32")) + + lossy = _matrix(integral=False) + lossy_path = _put(new_store, _as(lossy, layout), stem="lossy") + with pytest.raises(ValueError, match="is lossy"): + _recast(lossy_path, dtype="float32") + + +def test_the_cast_check_reads_one_block_of_rows_at_a_time(new_store): + """F029: check_cast read `dataset[0:1048576]` of a dense matrix, every row + and column at once, whatever --chunk said, then held two more copies of + it. It reads --chunk rows now.""" + matrix = _matrix(n_obs=40, n_var=9).toarray() + source_path = _put(new_store, matrix) + reads = [] + + from adata.storage import open_store + + with open_store(source_path, "r") as store: + dataset = store.root["X"] + + class Watched: + dtype, shape, ndim = dataset.dtype, dataset.shape, dataset.ndim + + def __getitem__(self, key): + block = dataset[key] + reads.append(np.asarray(block).shape) + return block + + report = core.check_cast(Watched(), np.dtype("float32"), rows=4) + + assert report.lossless and report.total == matrix.size + assert max(shape[0] for shape in reads) == 4, reads + + +# --------------------------------------------------------------------------- +# index ranges +# +# A matrix needing int64 offsets has more than 2^31 nonzeros, which a test +# cannot build. int16 stands in for int32: the checks compare against the +# dtype's own range, so the logic is the same at 32,767 as at 2^31 - 1. + + +@pytest.fixture +def int16_indices(monkeypatch): + """Let `--indices-dtype` take int16, as a stand-in for int32.""" + monkeypatch.setattr(core, "INDEX_DTYPES", core.INDEX_DTYPES + ("int16",)) + + +@pytest.mark.parametrize("force", [False, True], ids=["plain", "forced"]) +def test_force_does_not_waive_the_index_range_check(new_store, int16_indices, force): + """F013: `--force` skipped check_index_dtype altogether. Given only to + accept a lossy float cast, it also let indptr be cast to a dtype that + cannot hold nnz: offsets wrapped negative, a corrupt matrix with exit 0. + The range check is not a judgement about values, so nothing waives it.""" + matrix = sparse.csr_matrix(np.ones((300, 200))) # 60,000 nonzeros + source_path = _put(new_store, matrix) + out = _output(source_path) + + with pytest.raises(ValueError, match="cannot hold this matrix's offsets"): + recast_store(source_path, ["X"], out, QUIET, indices_dtype="int16", + dtype="float32", force=force) + assert not out.exists() + + +def test_a_dense_source_has_its_coordinates_checked_before_writing(new_store, int16_indices): + """F013: dense to sparse never ran the range check, so coordinates past + the index dtype's range wrapped. They are known from the shape, so the + refusal comes before anything is written.""" + source_path = _put(new_store, np.ones((2, 40_000), dtype="float32")) + out = _output(source_path) + + with pytest.raises(ValueError, match="cannot hold this matrix's coordinates"): + recast_store(source_path, ["X"], out, QUIET, layout="csr", indices_dtype="int16") + assert not out.exists() + + +@pytest.mark.parametrize("layout", ["csr", "csc"]) +def test_a_dense_source_whose_nonzeros_overflow_the_offsets_is_refused( + new_store, int16_indices, layout +): + """F013: sparsify cast indptr with a plain astype, so once nnz passed the + dtype's range the offsets wrapped and the command still succeeded. nnz is + only known by reading, so this one is refused while writing; staging then + removes the half-written output.""" + source_path = _put(new_store, np.ones((300, 200), dtype="float32")) + out = _output(source_path) + + with pytest.raises(ValueError, match="cannot hold this matrix's offsets"): + recast_store(source_path, ["X"], out, QUIET, layout=layout, + indices_dtype="int16", chunk_rows=16) + assert not out.exists() + leftovers = [p.name for p in source_path.parent.iterdir() if "adata-" in p.name] + assert leftovers == [], leftovers + + +def test_offsets_that_fit_are_written_in_the_narrow_dtype(new_store, int16_indices): + """The other side of F013's check: a dense source whose nonzeros fit the + narrow dtype is converted, with both arrays in it.""" + dense = np.eye(100, dtype="float32") + source_path = _put(new_store, dense) + + out = _recast(source_path, layout="csc", indices_dtype="int16") + got = _read(out).X + _same(got, sparse.csc_matrix(dense)) + + from adata.storage import open_store + + with open_store(out, "r") as store: + assert store.root["X/indices"].dtype == np.dtype("int16") + assert store.root["X/indptr"].dtype == np.dtype("int16") + + +# --------------------------------------------------------------------------- +# float16 + + +@pytest.mark.parametrize("layout", ["csr", "csc"]) +@pytest.mark.parametrize("source", ["csr", "csc", "dense"]) +def test_float16_is_refused_for_a_sparse_output(new_store, source, layout): + """F058: `--dtype float16` to a sparse layout said every value round-trips + and wrote a matrix that scipy, and so anndata, cannot read: 'scipy.sparse + does not support dtype float16'. With --inplace it replaced a readable + source with an unreadable one. It is refused now, before anything is + written, with what to do instead.""" + source_path = _put(new_store, _as(_matrix(), source)) + out = _output(source_path) + args = ["recast", str(source_path), "X", "-o", str(out), "--dtype", "float16"] + if source != layout: + args += ["--layout", layout] + + result = runner.invoke(app, args) + + assert result.exit_code == 1 + text = _out(result) + assert "float16" in text and "float32" in text and "--layout dense" in text, text + assert not out.exists() + + +def test_float16_is_refused_for_every_sparse_matrix_all_selects(new_store): + """F058, through --all: one sparse layer is enough to refuse the lot.""" + path, opener = new_store() + matrix = _matrix(n_obs=3, n_var=2) + with opener("a") as root: + _write_matrix(root, "X", matrix.toarray()) + _write_matrix(root["layers"], "counts", matrix) + out = _output(path) + + result = runner.invoke(app, ["recast", str(path), "--all", "-o", str(out), + "--dtype", "float16"]) + + assert result.exit_code == 1 and "layers/counts" in _out(result) + assert not out.exists() + + +@pytest.mark.parametrize("layout", ["csr", "csc", "dense"]) +@pytest.mark.parametrize("dtype", [d for d in core.DTYPES if d != "float16"] + ["float16"]) +def test_anndata_reads_back_every_dtype_and_layout_recast_writes( + request, new_store, backend, dtype, layout +): + """F058's general rule: every dtype `--dtype` accepts, in every layout it + can be written in, gives a store anndata reads back with the values + scipy gives. float16 only dense, the one layout scipy can hold it in.""" + if (dtype, backend) == ("bool", "h5ad") and layout != "dense": + request.applymarker(pytest.mark.xfail( + reason="F071: the float source's fill value cannot become h5py's " + "bool (B5, #95)", strict=False, + )) + matrix = _matrix(n_obs=6, n_var=5) + matrix.data = np.minimum(matrix.data, 1) if dtype == "bool" else matrix.data + source_path = _put(new_store, matrix) + if dtype == "float16" and layout != "dense": + pytest.skip("refused, see test_float16_is_refused_for_a_sparse_output") + + out = _recast(source_path, dtype=dtype, layout=layout, force=True) + + if layout == "dense": + expected = matrix.toarray().astype(dtype) + else: + expected = _as(matrix.astype(dtype), layout) + _same(_read(out).X, expected) + + +# --------------------------------------------------------------------------- +# a lossy cast drops the zeros it makes + + +@pytest.mark.parametrize("layout", ["csr", "csc"]) +def test_a_lossy_cast_from_dense_keeps_no_entries_that_became_zero(new_store, layout): + """F121: sparsify found the nonzeros before casting, so 0.3 became an + explicit 0 entry: 616 stored entries where scipy has 194. The values read + the same, but the store is larger and nnz-based QC (n_genes_by_counts) is + wrong. The block is cast first now, and its zeros dropped.""" + rng = np.random.default_rng(0) + dense = np.where(rng.random((30, 20)) < 0.6, rng.random((30, 20)) * 3, 0.0) + source_path = _put(new_store, dense) + + out = _recast(source_path, layout=layout, dtype="int32", force=True) + + expected = _as(sparse.csr_matrix(dense.astype("int32")), layout) + got = _read(out).X + assert (got.data != 0).all(), "explicit zeros were written" + _same(got, expected) + + +# --------------------------------------------------------------------------- +# nested paths + + +@pytest.mark.parametrize( + "at", + ["uns/neighbors/connectivities", "uns/a/b/c", "obsm/X_pca"], +) +def test_a_matrix_at_any_depth_is_converted_not_copied(new_store, at): + """F043: only a matrix one level below the root was converted. One at + `uns/neighbors/connectivities`, where older anndata keeps the neighbour + graph, was copied unconverted while the command reported 'Converting' + and 'Wrote', and even warned about the cast it had not made.""" + from adata.storage import open_store + + matrix = _matrix(n_obs=12, n_var=12, density=0.2, integral=False) + path = _put(new_store, matrix, at=at) + parent = at.rpartition("/")[0] + with open_store(path, "a") as store: + # A sibling at the same level and an entry above it, both of which + # must come through untouched. + _write_matrix(store.root[parent], "sibling", matrix) + ew.write_scalar(store.root["uns"], "note", 3) + + out = _recast(path, entries=[at], dtype="float32", force=True) + + got = _read(out) + _same(_at(got, at), matrix.astype("float32")) + _same(_at(got, f"{parent}/sibling"), matrix) + assert got.uns["note"] == 3 + + +# --------------------------------------------------------------------------- +# output layout + + +@pytest.mark.parametrize("layout", ["csr", "csc"]) +def test_a_sparse_matrix_from_a_dense_one_has_element_sized_chunks(new_store, layout): + """F078: sparsify gave `data` and `indices` the dense source's chunk + height as their element chunk: a source chunked 16 rows tall gave 16- + element chunks, 281,188 of them for 4.5M nonzeros and a file three times + the size. A row height is not an element count.""" + dense = _as(_matrix(n_obs=400, n_var=50, density=0.3), "dense") + path, opener = new_store() + with opener("a") as root: + ew.write_dataframe_header(root, "obs", [f"c{i}" for i in range(400)], []) + ew.write_dataframe_header(root, "var", [f"g{i}" for i in range(50)], []) + ew.write_dense(root, "X", dense, chunks=(4, 50)) + + out = _recast(path, layout=layout, chunk_rows=64) + + from adata.storage import open_store + + with open_store(out, "r") as store: + for name in ("data", "indices"): + chunks = store.root[f"X/{name}"].chunks + assert chunks[0] >= min(4096, store.root[f"X/{name}"].shape[0]), ( + f"X/{name} chunked {chunks} for {store.root[f'X/{name}'].shape} " + "elements: the source's 4-row chunk height became an element count" + ) + _same(_read(out).X, _as(sparse.csr_matrix(dense), layout)) + + +# --------------------------------------------------------------------------- +# free space + + +@pytest.mark.parametrize( + "layout,source,scratch", + [("csc", "csr", True), ("dense", "csc", True), ("dense", "csr", False), ("csr", "csr", False)], +) +def test_the_free_space_estimate_counts_the_transpose_scratch( + tmp_path, monkeypatch, layout, source, scratch +): + """A streaming transpose writes every entry to scratch beside the output, + as an int64 major, an int64 minor and its value, and a CSC matrix being + densified goes through a CSR first. The estimate counted only the output, + so a disk with room for that ran out part way through the scratch.""" + import adata.commands.recast as command + + matrix = _matrix(n_obs=50, n_var=40) + value = _as(matrix, source) + path = tmp_path / "in.h5ad" + ad.AnnData(X=value).write_h5ad(path) + asked = [] + monkeypatch.setattr( + command, "check_free_space", + lambda output, needed, check=True: asked.append(needed()), + ) + + def estimate(in_memory: bool) -> float: + asked.clear() + out = tmp_path / f"o{len(asked)}{in_memory}{layout}.h5ad" + recast_store(path, ["X"], out, QUIET, layout=layout, force=True, + dtype="float32", in_memory=in_memory) + return asked[0] + + extra = estimate(False) - estimate(True) + if not scratch: + assert extra == 0 + return + per_entry = 16 + 8 + (8 + 8 if layout == "dense" else 0) # float64 source + assert extra == pytest.approx(matrix.nnz * per_entry), extra diff --git a/uv.lock b/uv.lock index 573b408..64ac41a 100644 --- a/uv.lock +++ b/uv.lock @@ -182,37 +182,45 @@ wheels = [ [[package]] name = "h5py" -version = "3.15.1" +version = "3.16.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/4d/6a/0d79de0b025aa85dc8864de8e97659c94cf3d23148394a954dc5ca52f8c8/h5py-3.15.1.tar.gz", hash = "sha256:c86e3ed45c4473564de55aa83b6fc9e5ead86578773dfbd93047380042e26b69", size = 426236, upload-time = "2025-10-16T10:35:27.404Z" } +sdist = { url = "https://files.pythonhosted.org/packages/db/33/acd0ce6863b6c0d7735007df01815403f5589a21ff8c2e1ee2587a38f548/h5py-3.16.0.tar.gz", hash = "sha256:a0dbaad796840ccaa67a4c144a0d0c8080073c34c76d5a6941d6818678ef2738", size = 446526, upload-time = "2026-03-06T13:49:08.07Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/62/b8/c0d9aa013ecfa8b7057946c080c0c07f6fa41e231d2e9bd306a2f8110bdc/h5py-3.15.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:316dd0f119734f324ca7ed10b5627a2de4ea42cc4dfbcedbee026aaa361c238c", size = 3399089, upload-time = "2025-10-16T10:34:12.135Z" }, - { url = "https://files.pythonhosted.org/packages/a4/5e/3c6f6e0430813c7aefe784d00c6711166f46225f5d229546eb53032c3707/h5py-3.15.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b51469890e58e85d5242e43aab29f5e9c7e526b951caab354f3ded4ac88e7b76", size = 2847803, upload-time = "2025-10-16T10:34:14.564Z" }, - { url = "https://files.pythonhosted.org/packages/00/69/ba36273b888a4a48d78f9268d2aee05787e4438557450a8442946ab8f3ec/h5py-3.15.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8a33bfd5dfcea037196f7778534b1ff7e36a7f40a89e648c8f2967292eb6898e", size = 4914884, upload-time = "2025-10-16T10:34:18.452Z" }, - { url = "https://files.pythonhosted.org/packages/3a/30/d1c94066343a98bb2cea40120873193a4fed68c4ad7f8935c11caf74c681/h5py-3.15.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:25c8843fec43b2cc368aa15afa1cdf83fc5e17b1c4e10cd3771ef6c39b72e5ce", size = 5109965, upload-time = "2025-10-16T10:34:21.853Z" }, - { url = "https://files.pythonhosted.org/packages/81/3d/d28172116eafc3bc9f5991b3cb3fd2c8a95f5984f50880adfdf991de9087/h5py-3.15.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:a308fd8681a864c04423c0324527237a0484e2611e3441f8089fd00ed56a8171", size = 4561870, upload-time = "2025-10-16T10:34:26.69Z" }, - { url = "https://files.pythonhosted.org/packages/a5/83/393a7226024238b0f51965a7156004eaae1fcf84aa4bfecf7e582676271b/h5py-3.15.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f4a016df3f4a8a14d573b496e4d1964deb380e26031fc85fb40e417e9131888a", size = 5037161, upload-time = "2025-10-16T10:34:30.383Z" }, - { url = "https://files.pythonhosted.org/packages/cf/51/329e7436bf87ca6b0fe06dd0a3795c34bebe4ed8d6c44450a20565d57832/h5py-3.15.1-cp312-cp312-win_amd64.whl", hash = "sha256:59b25cf02411bf12e14f803fef0b80886444c7fe21a5ad17c6a28d3f08098a1e", size = 2874165, upload-time = "2025-10-16T10:34:33.461Z" }, - { url = "https://files.pythonhosted.org/packages/09/a8/2d02b10a66747c54446e932171dd89b8b4126c0111b440e6bc05a7c852ec/h5py-3.15.1-cp312-cp312-win_arm64.whl", hash = "sha256:61d5a58a9851e01ee61c932bbbb1c98fe20aba0a5674776600fb9a361c0aa652", size = 2458214, upload-time = "2025-10-16T10:34:35.733Z" }, - { url = "https://files.pythonhosted.org/packages/88/b3/40207e0192415cbff7ea1d37b9f24b33f6d38a5a2f5d18a678de78f967ae/h5py-3.15.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c8440fd8bee9500c235ecb7aa1917a0389a2adb80c209fa1cc485bd70e0d94a5", size = 3376511, upload-time = "2025-10-16T10:34:38.596Z" }, - { url = "https://files.pythonhosted.org/packages/31/96/ba99a003c763998035b0de4c299598125df5fc6c9ccf834f152ddd60e0fb/h5py-3.15.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:ab2219dbc6fcdb6932f76b548e2b16f34a1f52b7666e998157a4dfc02e2c4123", size = 2826143, upload-time = "2025-10-16T10:34:41.342Z" }, - { url = "https://files.pythonhosted.org/packages/6a/c2/fc6375d07ea3962df7afad7d863fe4bde18bb88530678c20d4c90c18de1d/h5py-3.15.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d8cb02c3a96255149ed3ac811eeea25b655d959c6dd5ce702c9a95ff11859eb5", size = 4908316, upload-time = "2025-10-16T10:34:44.619Z" }, - { url = "https://files.pythonhosted.org/packages/d9/69/4402ea66272dacc10b298cca18ed73e1c0791ff2ae9ed218d3859f9698ac/h5py-3.15.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:121b2b7a4c1915d63737483b7bff14ef253020f617c2fb2811f67a4bed9ac5e8", size = 5103710, upload-time = "2025-10-16T10:34:48.639Z" }, - { url = "https://files.pythonhosted.org/packages/e0/f6/11f1e2432d57d71322c02a97a5567829a75f223a8c821764a0e71a65cde8/h5py-3.15.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:59b0d63b318bf3cc06687def2b45afd75926bbc006f7b8cd2b1a231299fc8599", size = 4556042, upload-time = "2025-10-16T10:34:51.841Z" }, - { url = "https://files.pythonhosted.org/packages/18/88/3eda3ef16bfe7a7dbc3d8d6836bbaa7986feb5ff091395e140dc13927bcc/h5py-3.15.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:e02fe77a03f652500d8bff288cbf3675f742fc0411f5a628fa37116507dc7cc0", size = 5030639, upload-time = "2025-10-16T10:34:55.257Z" }, - { url = "https://files.pythonhosted.org/packages/e5/ea/fbb258a98863f99befb10ed727152b4ae659f322e1d9c0576f8a62754e81/h5py-3.15.1-cp313-cp313-win_amd64.whl", hash = "sha256:dea78b092fd80a083563ed79a3171258d4a4d307492e7cf8b2313d464c82ba52", size = 2864363, upload-time = "2025-10-16T10:34:58.099Z" }, - { url = "https://files.pythonhosted.org/packages/5d/c9/35021cc9cd2b2915a7da3026e3d77a05bed1144a414ff840953b33937fb9/h5py-3.15.1-cp313-cp313-win_arm64.whl", hash = "sha256:c256254a8a81e2bddc0d376e23e2a6d2dc8a1e8a2261835ed8c1281a0744cd97", size = 2449570, upload-time = "2025-10-16T10:35:00.473Z" }, - { url = "https://files.pythonhosted.org/packages/a0/2c/926eba1514e4d2e47d0e9eb16c784e717d8b066398ccfca9b283917b1bfb/h5py-3.15.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:5f4fb0567eb8517c3ecd6b3c02c4f4e9da220c8932604960fd04e24ee1254763", size = 3380368, upload-time = "2025-10-16T10:35:03.117Z" }, - { url = "https://files.pythonhosted.org/packages/65/4b/d715ed454d3baa5f6ae1d30b7eca4c7a1c1084f6a2edead9e801a1541d62/h5py-3.15.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:954e480433e82d3872503104f9b285d369048c3a788b2b1a00e53d1c47c98dd2", size = 2833793, upload-time = "2025-10-16T10:35:05.623Z" }, - { url = "https://files.pythonhosted.org/packages/ef/d4/ef386c28e4579314610a8bffebbee3b69295b0237bc967340b7c653c6c10/h5py-3.15.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fd125c131889ebbef0849f4a0e29cf363b48aba42f228d08b4079913b576bb3a", size = 4903199, upload-time = "2025-10-16T10:35:08.972Z" }, - { url = "https://files.pythonhosted.org/packages/33/5d/65c619e195e0b5e54ea5a95c1bb600c8ff8715e0d09676e4cce56d89f492/h5py-3.15.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:28a20e1a4082a479b3d7db2169f3a5034af010b90842e75ebbf2e9e49eb4183e", size = 5097224, upload-time = "2025-10-16T10:35:12.808Z" }, - { url = "https://files.pythonhosted.org/packages/30/30/5273218400bf2da01609e1292f562c94b461fcb73c7a9e27fdadd43abc0a/h5py-3.15.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:fa8df5267f545b4946df8ca0d93d23382191018e4cda2deda4c2cedf9a010e13", size = 4551207, upload-time = "2025-10-16T10:35:16.24Z" }, - { url = "https://files.pythonhosted.org/packages/d3/39/a7ef948ddf4d1c556b0b2b9559534777bccc318543b3f5a1efdf6b556c9c/h5py-3.15.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:99d374a21f7321a4c6ab327c4ab23bd925ad69821aeb53a1e75dd809d19f67fa", size = 5025426, upload-time = "2025-10-16T10:35:19.831Z" }, - { url = "https://files.pythonhosted.org/packages/b6/d8/7368679b8df6925b8415f9dcc9ab1dab01ddc384d2b2c24aac9191bd9ceb/h5py-3.15.1-cp314-cp314-win_amd64.whl", hash = "sha256:9c73d1d7cdb97d5b17ae385153472ce118bed607e43be11e9a9deefaa54e0734", size = 2865704, upload-time = "2025-10-16T10:35:22.658Z" }, - { url = "https://files.pythonhosted.org/packages/d3/b7/4a806f85d62c20157e62e58e03b27513dc9c55499768530acc4f4c5ce4be/h5py-3.15.1-cp314-cp314-win_arm64.whl", hash = "sha256:a6d8c5a05a76aca9a494b4c53ce8a9c29023b7f64f625c6ce1841e92a362ccdf", size = 2465544, upload-time = "2025-10-16T10:35:25.695Z" }, + { url = "https://files.pythonhosted.org/packages/c8/c0/5d4119dba94093bbafede500d3defd2f5eab7897732998c04b54021e530b/h5py-3.16.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:c5313566f4643121a78503a473f0fb1e6dcc541d5115c44f05e037609c565c4d", size = 3685604, upload-time = "2026-03-06T13:48:04.198Z" }, + { url = "https://files.pythonhosted.org/packages/b0/42/c84efcc1d4caebafb1ecd8be4643f39c85c47a80fe254d92b8b43b1eadaf/h5py-3.16.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:42b012933a83e1a558c673176676a10ce2fd3759976a0fedee1e672d1e04fc9d", size = 3061940, upload-time = "2026-03-06T13:48:05.783Z" }, + { url = "https://files.pythonhosted.org/packages/89/84/06281c82d4d1686fde1ac6b0f307c50918f1c0151062445ab3b6fa5a921d/h5py-3.16.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:ff24039e2573297787c3063df64b60aab0591980ac898329a08b0320e0cf2527", size = 5198852, upload-time = "2026-03-06T13:48:07.482Z" }, + { url = "https://files.pythonhosted.org/packages/9e/e9/1a19e42cd43cc1365e127db6aae85e1c671da1d9a5d746f4d34a50edb577/h5py-3.16.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:dfc21898ff025f1e8e67e194965a95a8d4754f452f83454538f98f8a3fcb207e", size = 5405250, upload-time = "2026-03-06T13:48:09.628Z" }, + { url = "https://files.pythonhosted.org/packages/b7/8e/9790c1655eabeb85b92b1ecab7d7e62a2069e53baefd58c98f0909c7a948/h5py-3.16.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:698dd69291272642ffda44a0ecd6cd3bda5faf9621452d255f57ce91487b9794", size = 5190108, upload-time = "2026-03-06T13:48:11.26Z" }, + { url = "https://files.pythonhosted.org/packages/51/d7/ab693274f1bd7e8c5f9fdd6c7003a88d59bedeaf8752716a55f532924fbb/h5py-3.16.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:2b2c02b0a160faed5fb33f1ba8a264a37ee240b22e049ecc827345d0d9043074", size = 5419216, upload-time = "2026-03-06T13:48:13.322Z" }, + { url = "https://files.pythonhosted.org/packages/03/c1/0976b235cf29ead553e22f2fb6385a8252b533715e00d0ae52ed7b900582/h5py-3.16.0-cp312-cp312-win_amd64.whl", hash = "sha256:96b422019a1c8975c2d5dadcf61d4ba6f01c31f92bbde6e4649607885fe502d6", size = 3182868, upload-time = "2026-03-06T13:48:15.759Z" }, + { url = "https://files.pythonhosted.org/packages/14/d9/866b7e570b39070f92d47b0ff1800f0f8239b6f9e45f02363d7112336c1f/h5py-3.16.0-cp312-cp312-win_arm64.whl", hash = "sha256:39c2838fb1e8d97bcf1755e60ad1f3dd76a7b2a475928dc321672752678b96db", size = 2653286, upload-time = "2026-03-06T13:48:17.279Z" }, + { url = "https://files.pythonhosted.org/packages/0f/9e/6142ebfda0cb6e9349c091eae73c2e01a770b7659255248d637bec54a88b/h5py-3.16.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:370a845f432c2c9619db8eed334d1e610c6015796122b0e57aa46312c22617d9", size = 3671808, upload-time = "2026-03-06T13:48:19.737Z" }, + { url = "https://files.pythonhosted.org/packages/b0/65/5e088a45d0f43cd814bc5bec521c051d42005a472e804b1a36c48dada09b/h5py-3.16.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:42108e93326c50c2810025aade9eac9d6827524cdccc7d4b75a546e5ab308edb", size = 3045837, upload-time = "2026-03-06T13:48:21.854Z" }, + { url = "https://files.pythonhosted.org/packages/da/1e/6172269e18cc5a484e2913ced33339aad588e02ba407fafd00d369e22ef3/h5py-3.16.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:099f2525c9dcf28de366970a5fb34879aab20491589fa89ce2863a84218bb524", size = 5193860, upload-time = "2026-03-06T13:48:24.071Z" }, + { url = "https://files.pythonhosted.org/packages/bd/98/ef2b6fe2903e377cbe870c3b2800d62552f1e3dbe81ce49e1923c53d1c5c/h5py-3.16.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:9300ad32dea9dfc5171f94d5f6948e159ed93e4701280b0f508773b3f582f402", size = 5400417, upload-time = "2026-03-06T13:48:25.728Z" }, + { url = "https://files.pythonhosted.org/packages/bc/81/5b62d760039eed64348c98129d17061fdfc7839fc9c04eaaad6dee1004e4/h5py-3.16.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:171038f23bccddfc23f344cadabdfc9917ff554db6a0d417180d2747fe4c75a7", size = 5185214, upload-time = "2026-03-06T13:48:27.436Z" }, + { url = "https://files.pythonhosted.org/packages/28/c4/532123bcd9080e250696779c927f2cb906c8bf3447df98f5ceb8dcded539/h5py-3.16.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:7e420b539fb6023a259a1b14d4c9f6df8cf50d7268f48e161169987a57b737ff", size = 5414598, upload-time = "2026-03-06T13:48:29.49Z" }, + { url = "https://files.pythonhosted.org/packages/c3/d9/a27997f84341fc0dfcdd1fe4179b6ba6c32a7aa880fdb8c514d4dad6fba3/h5py-3.16.0-cp313-cp313-win_amd64.whl", hash = "sha256:18f2bbcd545e6991412253b98727374c356d67caa920e68dc79eab36bf5fedad", size = 3175509, upload-time = "2026-03-06T13:48:31.131Z" }, + { url = "https://files.pythonhosted.org/packages/a5/23/bb8647521d4fd770c30a76cfc6cb6a2f5495868904054e92f2394c5a78ff/h5py-3.16.0-cp313-cp313-win_arm64.whl", hash = "sha256:656f00e4d903199a1d58df06b711cf3ca632b874b4207b7dbec86185b5c8c7d4", size = 2647362, upload-time = "2026-03-06T13:48:33.411Z" }, + { url = "https://files.pythonhosted.org/packages/48/3c/7fcd9b4c9eed82e91fb15568992561019ae7a829d1f696b2c844355d95dd/h5py-3.16.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:9c9d307c0ef862d1cd5714f72ecfafe0a5d7529c44845afa8de9f46e5ba8bd65", size = 3678608, upload-time = "2026-03-06T13:48:35.183Z" }, + { url = "https://files.pythonhosted.org/packages/6a/b7/9366ed44ced9b7ef357ab48c94205280276db9d7f064aa3012a97227e966/h5py-3.16.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:8c1eff849cdd53cbc73c214c30ebdb6f1bb8b64790b4b4fc36acdb5e43570210", size = 3054773, upload-time = "2026-03-06T13:48:37.139Z" }, + { url = "https://files.pythonhosted.org/packages/58/a5/4964bc0e91e86340c2bbda83420225b2f770dcf1eb8a39464871ad769436/h5py-3.16.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:e2c04d129f180019e216ee5f9c40b78a418634091c8782e1f723a6ca3658b965", size = 5198886, upload-time = "2026-03-06T13:48:38.879Z" }, + { url = "https://files.pythonhosted.org/packages/f1/16/d905e7f53e661ce2c24686c38048d8e2b750ffc4350009d41c4e6c6c9826/h5py-3.16.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:e4360f15875a532bc7b98196c7592ed4fc92672a57c0a621355961cafb17a6dd", size = 5404883, upload-time = "2026-03-06T13:48:41.324Z" }, + { url = "https://files.pythonhosted.org/packages/4b/f2/58f34cb74af46d39f4cd18ea20909a8514960c5a3e5b92fd06a28161e0a8/h5py-3.16.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:3fae9197390c325e62e0a1aa977f2f62d994aa87aab182abbea85479b791197c", size = 5192039, upload-time = "2026-03-06T13:48:43.117Z" }, + { url = "https://files.pythonhosted.org/packages/ce/ca/934a39c24ce2e2db017268c08da0537c20fa0be7e1549be3e977313fc8f5/h5py-3.16.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:43259303989ac8adacc9986695b31e35dba6fd1e297ff9c6a04b7da5542139cc", size = 5421526, upload-time = "2026-03-06T13:48:44.838Z" }, + { url = "https://files.pythonhosted.org/packages/3e/14/615a450205e1b56d16c6783f5ccd116cde05550faad70ae077c955654a75/h5py-3.16.0-cp314-cp314-win_amd64.whl", hash = "sha256:fa48993a0b799737ba7fd21e2350fa0a60701e58180fae9f2de834bc39a147ab", size = 3183263, upload-time = "2026-03-06T13:48:47.117Z" }, + { url = "https://files.pythonhosted.org/packages/7b/48/a6faef5ed632cae0c65ac6b214a6614a0b510c3183532c521bdb0055e117/h5py-3.16.0-cp314-cp314-win_arm64.whl", hash = "sha256:1897a771a7f40d05c262fc8f37376ec37873218544b70216872876c627640f63", size = 2663450, upload-time = "2026-03-06T13:48:48.707Z" }, + { url = "https://files.pythonhosted.org/packages/5d/32/0c8bb8aedb62c772cf7c1d427c7d1951477e8c2835f872bc0a13d1f85f86/h5py-3.16.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:15922e485844f77c0b9d275396d435db3baa58292a9c2176a386e072e0cf2491", size = 3760693, upload-time = "2026-03-06T13:48:50.453Z" }, + { url = "https://files.pythonhosted.org/packages/1d/1f/fcc5977d32d6387c5c9a694afee716a5e20658ac08b3ff24fdec79fb05f2/h5py-3.16.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:df02dd29bd247f98674634dfe41f89fd7c16ba3d7de8695ec958f58404a4e618", size = 3181305, upload-time = "2026-03-06T13:48:52.221Z" }, + { url = "https://files.pythonhosted.org/packages/f5/a1/af87f64b9f986889884243643621ebbd4ac72472ba8ec8cec891ac8e2ca1/h5py-3.16.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:0f456f556e4e2cebeebd9d66adf8dc321770a42593494a0b6f0af54a7567b242", size = 5074061, upload-time = "2026-03-06T13:48:54.089Z" }, + { url = "https://files.pythonhosted.org/packages/cc/d0/146f5eaff3dc246a9c7f6e5e4f42bd45cc613bce16693bcd4d1f7c958bf5/h5py-3.16.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:3e6cb3387c756de6a9492d601553dffea3fe11b5f22b443aac708c69f3f55e16", size = 5279216, upload-time = "2026-03-06T13:48:56.75Z" }, + { url = "https://files.pythonhosted.org/packages/a1/9d/12a13424f1e604fc7df9497b73c0356fb78c2fb206abd7465ce47226e8fd/h5py-3.16.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:8389e13a1fd745ad2856873e8187fd10268b2d9677877bb667b41aebd771d8b7", size = 5070068, upload-time = "2026-03-06T13:48:59.169Z" }, + { url = "https://files.pythonhosted.org/packages/41/8c/bbe98f813722b4873818a8db3e15aa3e625b59278566905ac439725e8070/h5py-3.16.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:346df559a0f7dcb31cf8e44805319e2ab24b8957c45e7708ce503b2ec79ba725", size = 5300253, upload-time = "2026-03-06T13:49:02.033Z" }, + { url = "https://files.pythonhosted.org/packages/32/9e/87e6705b4d6890e7cecdf876e2a7d3e40654a2ae37482d79a6f1b87f7b92/h5py-3.16.0-cp314-cp314t-win_amd64.whl", hash = "sha256:4c6ab014ab704b4feaa719ae783b86522ed0bf1f82184704ed3c9e4e3228796e", size = 3381671, upload-time = "2026-03-06T13:49:04.351Z" }, + { url = "https://files.pythonhosted.org/packages/96/91/9fad90cfc5f9b2489c7c26ad897157bce82f0e9534a986a221b99760b23b/h5py-3.16.0-cp314-cp314t-win_arm64.whl", hash = "sha256:faca8fb4e4319c09d83337adc80b2ca7d5c5a343c2d6f1b6388f32cfecca13c1", size = 2740706, upload-time = "2026-03-06T13:49:06.347Z" }, ] [[package]] @@ -527,7 +535,7 @@ spatial = [ [package.metadata] requires-dist = [ { name = "anndata", marker = "extra == 'dev'", specifier = ">=0.13" }, - { name = "h5py", specifier = ">=3.15.1,<4" }, + { name = "h5py", specifier = ">=3.16.0,<4" }, { name = "numcodecs", specifier = ">=0.14" }, { name = "numpy", specifier = ">=2.3.5,<3" }, { name = "pandas", marker = "extra == 'dev'", specifier = ">=3.0" }, From db33a987815e5f1a3829a74405c1d7c0bb031a2b Mon Sep 17 00:00:00 2001 From: Aljes Date: Sun, 1 Nov 2026 23:40:01 +0000 Subject: [PATCH 2/5] Make the CSC densify/sparsify guard per row, and measured twice Its growth ratio missed main's n_obs x --chunk block on Zarr, where the store's fixed cost hid it, and failed once in a full run under coverage when a one-time growth landed in one measurement. Co-Authored-By: Claude Opus 5.5 --- tests/test_performance.py | 33 ++++++++++++++++++++++++--------- 1 file changed, 24 insertions(+), 9 deletions(-) diff --git a/tests/test_performance.py b/tests/test_performance.py index 22c8d12..12709bd 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -2240,17 +2240,32 @@ def test_csc_densify_and_sparsify_read_blocks_set_by_the_chunk(tmp_path, fmt, la through a CSR in scratch now, in --chunk rows and chunk-sized reads.""" from tests.perf_fixtures import matrix_store + rows = 256 # --chunk + def measure(n: int) -> int: store = matrix_store(tmp_path, n_obs=n, n_var=64, layout=source, fmt=fmt) - with count_allocations() as peak, count_io() as io: - _recast(store, tmp_path / f"o{n}{store.suffix}", layout=layout, force=True, - chunk_rows=GRID_ROWS, chunk=4096) - assert _largest_matrix_read(io) <= 4096 * 8 + GRID_ROWS * 64 * 8, io - return peak[0] - - assert_grows_slower_than_input( - measure, what=f"recast {source} -> {layout} ({fmt})", axis="n_obs", - sizes=(512, 8192), at_least=4.0, + peaks = [] + # Twice, keeping the smaller: once in a full run, a one-time growth + # of the process's own structures landed in one measurement (4 MB). + for attempt in range(2): + with count_allocations() as peak, count_io() as io: + _recast(store, tmp_path / f"o{n}-{attempt}{store.suffix}", layout=layout, + force=True, chunk_rows=rows, chunk=GRID_ENTRIES) + peaks.append(peak[0]) + # No read larger than a chunk of entries or a block of rows. + bound = GRID_ENTRIES * 8 + rows * 64 * 8 + assert io.read_bytes > 0 and _largest_matrix_read(io) <= bound, io + return min(peaks) + + # Per row added, so the fixed cost of the store cancels. A block of + # --chunk columns over every row holds `rows` float32 values a row; what + # is left is the per-row index arrays (an indptr, a count), a few int64. + small, large = 1024, 8192 + per_row = (measure(large) - measure(small)) / (large - small) + assert per_row <= rows * 4 / 4, ( + f"recast {source} -> {layout} ({fmt}) held {per_row:.0f} bytes more per " + f"row of input; a block of {rows} columns over every row is " + f"{rows * 4} a row, and per-row index arrays a few times 8" ) From 8ac6d13684aad5a46ca488f8045b433326f45bf5 Mon Sep 17 00:00:00 2001 From: Aljes Date: Mon, 2 Nov 2026 00:32:58 +0000 Subject: [PATCH 3/5] Review fixes: HDF5 2.0's cache, plain scratch, free space, steadier grid - open_store keeps HDF5's chunk cache at 1 MiB a dataset: HDF5 2.0 (h5py 3.16) made it 8 MiB, and split grew with its outputs - dense to CSC writes a plain CSR scratch (scratch has no chunk cache, so a gzip one was re-encoded per block) and gives the output the source's codecs; it checks free space again once nnz is known, and refuses an nnz overflow as soon as the count passes the range - cheap refusals before the cast check; coordinates checked against the output's minor axis; in-memory CSC densify's scratch estimated - chunk model: chunks go into HDF5's cache whole unless too big for it - grid: bigger chunk and density, each size measured twice; concat no longer marked; lower bounds on the largest-read guards - docs: F121 is about dense sources Co-Authored-By: Claude Opus 5.5 --- docs/reference/commands.md | 9 +- docs/reference/testing.md | 24 ++--- pyproject.toml | 5 +- src/adata/commands/recast.py | 28 +++--- src/adata/core/recast.py | 155 +++++++++++++++++++++++---------- src/adata/storage/__init__.py | 9 +- tests/perf_counters.py | 35 ++++---- tests/test_performance.py | 90 ++++++++++++++----- tests/test_recast_streaming.py | 105 ++++++++++++++++------ 9 files changed, 324 insertions(+), 136 deletions(-) diff --git a/docs/reference/commands.md b/docs/reference/commands.md index 2e14d16..81f31db 100644 --- a/docs/reference/commands.md +++ b/docs/reference/commands.md @@ -220,9 +220,10 @@ Two things are refused before anything is written. A cast that would not round-trip -- checked by casting every value and casting it back, not by comparing dtypes -- and a densification that would inflate the store beyond four times its size. `--force` overrides either, and the message says how -many values would change or how large the result would be. A cast that loses -values drops the entries it turns to zero, as scipy would, rather than storing -them as explicit zeros. +many values would change or how large the result would be. A dense matrix made +sparse by a cast that loses values drops the entries the cast turns to zero, +as scipy would, rather than storing them as explicit zeros; a sparse matrix +keeps the entries it has, as scipy's `astype` does. A sparse matrix of **float16 is refused**, whatever `--force` says: scipy has no float16 sparse matrix, so neither anndata nor scanpy could read the result. @@ -240,6 +241,8 @@ coordinates checked before writing and its nonzero count as it is counted. `--chunk` counts rows of a dense matrix, in every layout: densifying a CSC matrix or writing a dense one as CSC goes through a CSR in scratch, so no block spans every row. Sparse matrices are read about a million entries at a time. +The free-space check counts that scratch; a dense matrix written as CSC learns +its nonzero count only by reading it, so it checks again before its transpose. The output path may not name the input; use `--inplace`. Either way the new store is written beside its destination and moved into place only once the diff --git a/docs/reference/testing.md b/docs/reference/testing.md index d68e89b..d58033c 100644 --- a/docs/reference/testing.md +++ b/docs/reference/testing.md @@ -153,11 +153,13 @@ shard rewritten on every block, so `IOCounts` also counts bytes: Decoded and encoded bytes are a model worked out from the selection, the chunk grid and the codecs, not a measurement of the library, which is what keeps them -deterministic. It follows HDF5's chunk cache (per dataset, 1 MiB unless the -file says otherwise, least recently used), so appending a block at a time to a -growing compressed dataset costs each chunk once. A chunk with no filter costs -only the bytes asked for: HDF5 reads and writes it in place. Zarr caches -nothing, and writing part of a shard rewrites the shard. +deterministic. It follows HDF5's chunk cache (per dataset, least recently used, +at the dataset's own size: 1 MiB in the stores `open_store` opens, 8 MiB by +default since HDF5 2.0), so appending a block at a time to a growing dataset +costs each chunk once. A chunk is read into the cache whole, filtered or not; +only an unfiltered chunk too big for the cache is read and written in place, +just the bytes asked for. Zarr caches nothing, and writing part of a shard +rewrites the shard. `test_the_byte_counters_see_whole_chunks_decoded_and_rewritten` checks the model against numbers worked out by hand. @@ -264,10 +266,8 @@ from the whole matrix. Peak memory added per byte of input added can: a block of 32 rows of a matrix thousands of rows tall adds a few hundredths, and a path holding the matrix adds a whole byte and more for each copy. -The known gaps are marked `xfail`, each with its owner: sparse `subset` -(F032, strict, so the fix has to flip it) and `concat` of a sparse matrix grown -in n_obs (not strict: concat keeps obs names by design, and B3b's row pointers -change the figure). +The known gap is marked `xfail` with its owner: sparse `subset` (F032, B6b), +strict, so the fix has to flip it. The recast guards beside it pin each block in bytes, on every format: the cast check reads `--chunk` rows of a dense matrix, CSC densify and sparsify read @@ -347,10 +347,12 @@ The `slow` marker is on the guards that build a 65,536-row store or run `peak_rs they still gate merges, and `-m "not slow"` skips them locally. `perf` is on the tracing guards. -h5py 3.15 kept 100 to 160 bytes of every dataset read, write and `.shape` for good, so +h5py 3.15 kept 32 to 160 bytes of every dataset read, write and `.shape` for good, so any long streaming loop grew with its block count; `test_h5py_calls_keep_no_memory_once_they_return` keeps the floor at 3.16.0, which -does not. +does not. 3.16.0 brings HDF5 2.0, whose chunk cache is 8 MiB a dataset rather than 1; +`open_store` keeps 1 MiB, which +`test_stores_give_each_hdf5_dataset_a_one_mebibyte_chunk_cache` checks. ### Reading a failure diff --git a/pyproject.toml b/pyproject.toml index 3c19719..2d4899d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,8 +43,9 @@ classifiers = [ # and can break in a minor (0.26 stopped installing click), but a cap at the # next minor would turn away each new typer until we release, so it is 1.0. dependencies = [ - # 3.15 leaks 100 to 160 bytes of memory on every dataset read, write and - # `.shape`, so streaming used memory in proportion to its blocks. + # 3.15 leaks 32 to 160 bytes of memory on every dataset read, write and + # `.shape`, so streaming used memory in proportion to its blocks. 3.16 + # brings HDF5 2.0, whose 8 MiB chunk cache `open_store` sets back to 1 MiB. "h5py>=3.16.0,<4", # Zarr v2 compressors (storage/__init__.py). "numcodecs>=0.14", diff --git a/src/adata/commands/recast.py b/src/adata/commands/recast.py index 1cbc5d1..2914669 100644 --- a/src/adata/commands/recast.py +++ b/src/adata/commands/recast.py @@ -145,6 +145,11 @@ def recast_store( src, dst, plans, chunk=chunk, chunk_rows=chunk_rows, in_memory=in_memory, console=console, + # Dense to CSC learns its nnz only once it has read the + # matrix, so it checks again then, for its scratch. + room=lambda nbytes: check_free_space( + dst_path, nbytes, check=space_check + ), ) ensure_anndata_skeleton(dst) @@ -183,24 +188,25 @@ def _extra(plan: Plan) -> float: def _scratch(plan: Plan, in_memory: bool) -> float: - """Bytes of scratch a streaming transpose writes beside the output. + """Bytes of scratch a conversion writes beside the output. - Every entry passes through a bucket as an int64 major, an int64 minor - and its value; a CSC matrix being densified is also sorted into a CSR - first. A dense matrix sparsified to CSC goes through the same, but its - nnz is not known until it is read, so it is not counted. + A streaming transpose passes every entry through a bucket as an int64 + major, an int64 minor and its value (`--in-memory` does not). A CSC + matrix being densified is also sorted into a CSR first, either way. A + dense matrix sparsified to CSC goes through the same, but its nnz is not + known until it is read, so `sparsify` checks for that itself. """ src = plan.source - if in_memory or not src.sparse: + if not src.sparse: return 0.0 transposed = plan.layout in spec.SPARSE_TYPES and plan.layout != src.kind densified_csc = plan.layout == "dense" and src.kind == spec.CSC_MATRIX - if not (transposed or densified_csc): - return 0.0 - entry = 16 + src.dtype.itemsize + needed = 0.0 + if (transposed or densified_csc) and not in_memory: + needed += src.nnz * (16 + src.dtype.itemsize) if densified_csc: - entry += 8 + src.dtype.itemsize - return float(src.nnz * entry) + needed += src.nnz * (8 + src.dtype.itemsize) + 8 * (src.shape[0] + 1) + return needed def _on_disk(obj: Any) -> int: diff --git a/src/adata/core/recast.py b/src/adata/core/recast.py index c874225..c34c19a 100644 --- a/src/adata/core/recast.py +++ b/src/adata/core/recast.py @@ -24,7 +24,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any, Dict, Iterator, Optional, Tuple +from typing import Any, Callable, Dict, Iterator, Optional, Tuple import numpy as np from rich.console import Console @@ -206,6 +206,8 @@ def check_index_dtype( index_dtype: np.dtype, pointer_dtype: np.dtype, shape: Tuple[int, int], + *, + minor: Optional[int] = None, ) -> None: """Refuse index dtypes that cannot address this matrix. @@ -217,12 +219,12 @@ def check_index_dtype( A dense source has no nnz to read: pass `group=None` and only the coordinates are checked here, and `sparsify` checks the offsets as it - counts them. + counts them. `minor`, the extent of the output's minor axis, is what the + coordinates reach (columns for CSR); without it, the larger dimension. """ nnz = None if group is None else int(group["indices"].shape[0]) - check_index_range( - "indices", index_dtype, max(int(shape[0]), int(shape[1])), shape, nnz - ) + reach = int(minor) if minor is not None else max(int(shape[0]), int(shape[1])) + check_index_range("indices", index_dtype, reach, shape, nnz) if nnz is not None: check_index_range("indptr", pointer_dtype, nnz, shape, nnz) @@ -359,6 +361,7 @@ def _write_sparse_arrays( index_dtype: np.dtype, pointer_dtype: np.dtype, source: Any, + codecs_from: Any = None, ) -> None: for name, values, dtype, template in ( ("data", data, data_dtype, source["data"]), @@ -366,7 +369,7 @@ def _write_sparse_arrays( ("indptr", indptr, pointer_dtype, source["indptr"]), ): cast = values.astype(dtype, copy=False) - dataset = _sparse_dataset(group, name, dtype, cast.size, template) + dataset = _output_array(group, name, dtype, cast.size, template, codecs_from) if cast.size: dataset[:] = cast @@ -418,11 +421,13 @@ def transpose_sparse_in_memory( data_dtype: np.dtype, index_dtype: np.dtype, pointer_dtype: np.dtype = np.dtype("int64"), + codecs_from: Any = None, ) -> None: """Swap CSR<->CSC by loading the matrix and sorting it once. Faster than streaming whenever the matrix fits, and the whole matrix is - what it needs -- so it is opt-in, never the default. + what it needs -- so it is opt-in, never the default. `codecs_from` as for + `transpose_sparse_streaming`. """ target = ( spec.CSC_MATRIX if src.kind == spec.CSR_MATRIX else spec.CSR_MATRIX @@ -452,6 +457,7 @@ def transpose_sparse_in_memory( index_dtype=index_dtype, pointer_dtype=pointer_dtype, source=src.obj, + codecs_from=codecs_from, ) @@ -472,6 +478,7 @@ def transpose_sparse_streaming( console: Optional[Console] = None, compression: Optional[str] = None, keep_layout: bool = True, + codecs_from: Any = None, ) -> None: """Swap CSR<->CSC without loading the matrix. @@ -489,7 +496,9 @@ def transpose_sparse_streaming( Costs about two extra passes over nnz, which is the price of not holding the matrix. `--in-memory` is there for when you would rather pay in RAM. `keep_layout=False` writes plain arrays rather than the source's chunks - and codecs, for a transpose into scratch. + and codecs, for a transpose into scratch. `codecs_from`, a dense matrix, + gives the output that matrix's codecs and element-sized chunks instead, + for a transpose out of a scratch built from it. """ target = ( spec.CSC_MATRIX if src.kind == spec.CSR_MATRIX else spec.CSR_MATRIX @@ -518,11 +527,12 @@ def transpose_sparse_streaming( # written by slice rather than grown block by block. group = _new_sparse_group(src, dst_parent, name, target, src.shape) templates = (source_data, source_indices, src.obj["indptr"]) if keep_layout else (None,) * 3 - out_data = _sparse_dataset( - group, "data", data_dtype, nnz, templates[0], compression=compression + out_data = _output_array( + group, "data", data_dtype, nnz, templates[0], codecs_from, compression=compression ) - out_indices = _sparse_dataset( - group, "indices", index_dtype, nnz, templates[1], compression=compression + out_indices = _output_array( + group, "indices", index_dtype, nnz, templates[1], codecs_from, + compression=compression, ) written = 0 @@ -554,8 +564,8 @@ def transpose_sparse_streaming( written += count cast_indptr = out_indptr.astype(pointer_dtype, copy=False) - dataset = _sparse_dataset( - group, "indptr", pointer_dtype, cast_indptr.size, templates[2], + dataset = _output_array( + group, "indptr", pointer_dtype, cast_indptr.size, templates[2], codecs_from, compression=compression, ) dataset[:] = cast_indptr @@ -716,6 +726,7 @@ def sparsify( in_memory: bool = False, console: Optional[Console] = None, compression: Optional[str] = None, + room: Optional[Callable[[int], None]] = None, ) -> Tuple[int, float]: """Write a dense array out as CSR or CSC, `chunk_rows` rows at a time. @@ -726,28 +737,37 @@ def sparsify( CSC goes through a CSR in scratch and the streaming transpose. It read `chunk_rows` columns of every row at a time, n_obs x --chunk values (F030), and each of those reads decoded every chunk of a row-chunked - source, once per block of columns. + source, once per block of columns. `room(nbytes)` is called before the + transpose with what it will write, nnz being known only then, so that a + free-space check can refuse first. """ n_rows, n_cols = src.shape if enc == spec.CSC_MATRIX: with scratch_group(dst_parent, f"__{name}_sparsify_scratch__") as scratch: + # Plain arrays: scratch has no chunk cache, so a compressed one + # decoded and re-encoded a chunk on every append, 30 to 100 + # times slower. The output takes the source's codecs instead. nnz = _sparsify_rows( src, scratch, "csr", data_dtype=data_dtype, index_dtype=np.dtype("int64"), pointer_dtype=np.dtype("int64"), - chunk_rows=chunk_rows, compression=compression, + chunk_rows=chunk_rows, compression=None, template=None, + offsets_within=pointer_dtype, ) - # The scratch's offsets are int64; the output's must hold nnz too. - check_index_range("indptr", pointer_dtype, nnz, src.shape, nnz) + if room is not None: + # nnz is known only now: the buckets and the output next. + room(nnz * (16 + 2 * np.dtype(data_dtype).itemsize + + np.dtype(index_dtype).itemsize)) _transpose( describe(scratch["csr"]), dst_parent, name, in_memory=in_memory, chunk=chunk, data_dtype=data_dtype, index_dtype=index_dtype, pointer_dtype=pointer_dtype, compression=compression, + codecs_from=src.obj, ) else: nnz = _sparsify_rows( src, dst_parent, name, data_dtype=data_dtype, index_dtype=index_dtype, pointer_dtype=pointer_dtype, chunk_rows=chunk_rows, - compression=compression, + compression=compression, template=src.obj, ) density = nnz / max(1, n_rows * n_cols) @@ -769,8 +789,16 @@ def _sparsify_rows( pointer_dtype: np.dtype, chunk_rows: int, compression: Optional[str], + template: Any, + offsets_within: Optional[np.dtype] = None, ) -> int: - """Write a dense matrix out as CSR, `chunk_rows` rows at a time; its nnz.""" + """Write a dense matrix out as CSR, `chunk_rows` rows at a time; its nnz. + + `template` gives the arrays its codecs (None for plain ones). nnz is + refused as soon as it passes `offsets_within` (by default + `pointer_dtype`), so a CSR written to scratch on its way to narrower + offsets is not finished first. + """ from adata.core.subset import _append n_rows, n_cols = src.shape @@ -781,15 +809,16 @@ def _sparsify_rows( set_shape_attr(group, src.shape) out_data = _growable_from( - group, "data", data_dtype, src.obj, compression=compression + group, "data", data_dtype, template, compression=compression ) out_indices = _growable_from( - group, "indices", index_dtype, src.obj, compression=compression + group, "indices", index_dtype, template, compression=compression ) # Per-row counts in one int64 array: a Python list of int objects took # 36 to 40 bytes a row against 8. counts = np.zeros(n_rows, dtype=np.int64) - pointer_limit = int(np.iinfo(pointer_dtype).max) + offsets_within = np.dtype(offsets_within or pointer_dtype) + pointer_limit = int(np.iinfo(offsets_within).max) nnz = 0 for lo in range(0, n_rows, chunk_rows): @@ -803,7 +832,7 @@ def _sparsify_rows( nnz += int(rows.size) if nnz > pointer_limit: # Offsets cast with a plain astype wrapped negative (F013). - check_index_range("indptr", pointer_dtype, nnz, src.shape, None) + check_index_range("indptr", offsets_within, nnz, src.shape, None) counts[lo:hi] = np.bincount(rows, minlength=hi - lo) _append(out_indices, cols.astype(index_dtype, copy=False)) _append(out_data, block[rows, cols]) @@ -834,11 +863,48 @@ def _growable_from( `_growable_like` takes a template's first chunk dimension as the element chunk, which for a dense matrix is a row height: a source chunked 16 rows tall gave `data` 16-element chunks, 281,188 of them for 4.5M - nonzeros and a file three times the size (F078). Its fill value and - scale-offset filter belong to the template's dtype, not this one's. + nonzeros and a file three times the size (F078). With no template, a + plain array. + """ + kw = _codecs(group, dtype, template, compression) + if is_zarr_group(group): + return group.create_array(name, shape=(0,), dtype=dtype, chunks=(ELEMENT_CHUNK,), **kw) + return group.create_dataset( + name, shape=(0,), maxshape=(None,), dtype=dtype, chunks=(ELEMENT_CHUNK,), **kw + ) + + +def _output_array( + group: Any, + name: str, + dtype: np.dtype, + n: int, + template: Any, + codecs_from: Any, + *, + compression: Optional[str] = None, +) -> Any: + """A 1-D array of `n` elements: laid out like `template`, or, given + `codecs_from` (a dense matrix), with its codecs and element chunks.""" + if codecs_from is None: + return _sparse_dataset(group, name, dtype, n, template, compression=compression) + backend = "zarr" if is_zarr_group(group) else "hdf5" + kw = _codecs(group, dtype, codecs_from, compression) + kw = fit_chunks({**kw, "chunks": (ELEMENT_CHUNK,)}, (n,), backend=backend) + return create_dataset(group, name, shape=(n,), dtype=dtype, **kw) + + +def _codecs(group: Any, dtype: np.dtype, template: Any, compression: Optional[str]) -> dict: + """Creation kwargs for a 1-D array: `template`'s codecs, not its layout. + + Its chunks and shards are a grid of the template's own shape, and its + fill value and scale-offset filter belong to the template's dtype. """ backend = "zarr" if is_zarr_group(group) else "hdf5" - kw = dataset_create_kwargs(template, target_backend=backend, dst_parent=group) + kw = ( + dataset_create_kwargs(template, target_backend=backend, dst_parent=group) + if template is not None else {} + ) for key in ("chunks", "shards", "fillvalue", "scaleoffset"): kw.pop(key, None) if compression is not None: @@ -846,11 +912,7 @@ def _growable_from( kw, compression, shape=(0,), dtype=dtype, backend=backend, zarr_format=zarr_format_of(group), ) - if backend == "zarr": - return group.create_array(name, shape=(0,), dtype=dtype, chunks=(ELEMENT_CHUNK,), **kw) - return group.create_dataset( - name, shape=(0,), maxshape=(None,), dtype=dtype, chunks=(ELEMENT_CHUNK,), **kw - ) + return kw def cast_dense( @@ -943,6 +1005,8 @@ def plan_conversion( else: idx_dtype = ptr_dtype = np.dtype("int64") + # The refusals that need no pass over the values come first, so that a + # range or growth refusal does not wait for the whole matrix to be read. if target_layout in spec.SPARSE_TYPES and data_dtype == np.dtype("float16"): # scipy has no float16 sparse matrix, so anndata, scanpy and this # tool's own readers fail on one. It was written, with a message @@ -953,17 +1017,6 @@ def plan_conversion( "--layout dense." ) - report: Optional[CastReport] = None - if dtype is not None and data_dtype != src.dtype: - values = src.obj["data"] if src.sparse else src.obj - report = check_cast(values, data_dtype, chunk=chunk, rows=chunk_rows) - message = report.describe(src.dtype, data_dtype) - if not report.lossless and not force: - raise ValueError(f"{name}: {message}. Pass --force to convert anyway.") - if console is not None: - colour = "dim" if report.lossless else "yellow" - console.print(f"[{colour}]{name}: {message}[/]") - # Always, not only when asked: an inferred dtype can be too narrow too, # and a silently overflowed offset is indistinguishable from corruption. # Nor does --force waive it: it accepts lossy values and a large store, @@ -971,7 +1024,8 @@ def plan_conversion( # are checked here and its offsets as `sparsify` counts them. if target_layout in spec.SPARSE_TYPES: check_index_dtype( - src.obj if src.sparse else None, idx_dtype, ptr_dtype, src.shape + src.obj if src.sparse else None, idx_dtype, ptr_dtype, src.shape, + minor=src.shape[1] if target_layout == spec.CSR_MATRIX else src.shape[0], ) if target_layout == "dense" and src.sparse: @@ -980,6 +1034,17 @@ def plan_conversion( _stored_bytes(src.obj), projected, force=force, what=f"{name} as dense" ) + report: Optional[CastReport] = None + if dtype is not None and data_dtype != src.dtype: + values = src.obj["data"] if src.sparse else src.obj + report = check_cast(values, data_dtype, chunk=chunk, rows=chunk_rows) + message = report.describe(src.dtype, data_dtype) + if not report.lossless and not force: + raise ValueError(f"{name}: {message}. Pass --force to convert anyway.") + if console is not None: + colour = "dim" if report.lossless else "yellow" + console.print(f"[{colour}]{name}: {message}[/]") + return Plan(src, target_layout, data_dtype, idx_dtype, ptr_dtype, report) @@ -997,6 +1062,7 @@ def convert_matrix( in_memory: bool = False, force: bool = False, console: Optional[Console] = None, + room: Optional[Callable[[int], None]] = None, ) -> None: """Write `obj` into `dst_parent` under `name`, converted. @@ -1039,6 +1105,7 @@ def convert_matrix( chunk=chunk, in_memory=in_memory, console=console, + room=room, ) return diff --git a/src/adata/storage/__init__.py b/src/adata/storage/__init__.py index 06bb980..f6dd919 100644 --- a/src/adata/storage/__init__.py +++ b/src/adata/storage/__init__.py @@ -118,6 +118,13 @@ def detect_backend(path: Path) -> str: return "hdf5" +#: Bytes of chunk cache HDF5 keeps for each dataset of a store. HDF5 2.0 +#: (h5py 3.16) raised its default from 1 MiB to 8 MiB, and a writer holds one +#: per dataset it fills: `split` holds many outputs open at once, so memory +#: followed the number of outputs, eight times faster. 1 MiB, as before. +HDF5_CHUNK_CACHE = 1 << 20 + + def open_store( path: Path, mode: str, @@ -162,7 +169,7 @@ def open_store( zarr_format=zarr_format_of(root), consolidate=writable, ) - root = h5py.File(path, mode) + root = h5py.File(path, mode, rdcc_nbytes=HDF5_CHUNK_CACHE) if _is_writable_mode(mode): ensure_anndata_root_attrs(root) elif require_anndata: diff --git a/tests/perf_counters.py b/tests/perf_counters.py index 1145f26..4428e1c 100644 --- a/tests/perf_counters.py +++ b/tests/perf_counters.py @@ -34,8 +34,7 @@ Decoded and encoded bytes are a model worked out from the selection, the chunk grid and the codecs, not a measurement of the library, which is what -keeps them deterministic. It follows HDF5's per-dataset chunk cache, and -counts an HDF5 chunk with no filter as read or written in place; see +keeps them deterministic. It follows HDF5's per-dataset chunk cache; see `_ChunkCache`. Memory has two instruments. `count_allocations` (tracemalloc) is exact and @@ -275,14 +274,16 @@ def _axis_indices(key: Any, extent: int, size: int) -> np.ndarray: class _ChunkCache: """Which chunks a backend still holds decoded, dataset by dataset. - HDF5 keeps a chunk cache per dataset (1 MiB by default, none at all for a - chunk bigger than that), and serves a read or a write from it without - touching the file: appending a block at a time to a growing dataset - encodes each chunk once, when it leaves the cache, not once a block. - Modelled as least recently used, at the dataset's own cache size. A - chunk with no filter has nothing to decode, and HDF5 reads or writes - just the part asked for. Zarr caches nothing, so every chunk a selection - touches is decoded and every shard it writes is rewritten. + HDF5 keeps a chunk cache per dataset (1 MiB in the stores `open_store` + opens, 8 MiB by default since HDF5 2.0, none for a chunk bigger than the + cache) and serves a read or a write from it without touching the file: + appending a block at a time to a growing dataset encodes each chunk + once, when it leaves the cache, not once a block. Modelled as least + recently used, at the dataset's own cache size. A chunk is read into the + cache whole, filtered or not; only an unfiltered chunk too big for the + cache is read and written in place, just the part asked for. Zarr caches + nothing, so every chunk a selection touches is decoded and every shard + it writes is rewritten. """ def __init__(self) -> None: @@ -290,7 +291,7 @@ def __init__(self) -> None: @staticmethod def _unfiltered_hdf5(array: Any) -> bool: - """An HDF5 chunk with no filter is read and written in place, in part.""" + """Whether an HDF5 dataset's chunks have no filter.""" plist = getattr(getattr(array, "id", None), "get_create_plist", None) if plist is None: return False # Zarr @@ -317,13 +318,13 @@ def cost(self, array: Any, key: Any, returned: int, *, write: bool) -> int: shape = tuple(int(n) for n in array.shape) except Exception: # pragma: no cover - an object without a layout return returned - if not grid or not shape or self._unfiltered_hdf5(array): - # Contiguous, a scalar, or chunks with no codec: what was asked - # for is what moved. - return returned + if not grid or not shape: + return returned # contiguous, or a scalar: what was asked is what moved unit = _unit_bytes(array, grid) - touched = chunks_touched(key, shape, grid) capacity = self._capacity(array, unit) + if capacity == 0 and self._unfiltered_hdf5(array): + return returned # bypasses the cache, in place + touched = chunks_touched(key, shape, grid) if capacity == 0 or touched > _MODELLED_CHUNKS: return touched * unit @@ -348,7 +349,7 @@ def cost(self, array: Any, key: Any, returned: int, *, write: bool) -> int: def count_io() -> Iterator[IOCounts]: """Count backend reads and writes for the duration of the block. - Patches these seams, verified against h5py 3.15.1 and zarr 3.1.5: + Patches these seams, verified against h5py 3.16.0 and zarr 3.1.5: ====================================== ================================= ``h5py.Dataset.__getitem__`` HDF5 reads diff --git a/tests/test_performance.py b/tests/test_performance.py index 12709bd..5f83c4d 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -212,14 +212,16 @@ def test_the_byte_counters_see_whole_chunks_decoded_and_rewritten(tmp_path): with h5py.File(tmp_path / "c.h5", "r") as handle, count_io() as io: handle["gz"][5:6] # one row: the whole compressed chunk is decoded handle["gz"][6:7] # the same chunk again, from HDF5's chunk cache - handle["raw"][5:6] # no filter: HDF5 reads just the row - assert (io.read_bytes, io.largest_read) == (3 * row, row) - assert io.decoded_bytes == chunk + row + handle["raw"][5:6] # no filter, but the chunk fits: read into the cache whole + handle["raw"][6:7] + assert (io.read_bytes, io.largest_read) == (4 * row, row) + assert io.decoded_bytes == 2 * chunk with h5py.File(tmp_path / "c.h5", "r", rdcc_nbytes=0) as handle, count_io() as io: handle["gz"][5:6] handle["gz"][6:7] - assert io.decoded_bytes == 2 * chunk, "with no cache, each read decodes the chunk" + handle["raw"][5:6] # no cache and no filter: just the row, in place + assert io.decoded_bytes == 2 * chunk + row, "with no cache, each read decodes the chunk" array = zarr.create_array( str(tmp_path / "s.zarr"), shape=(64, 64), chunks=(8, 8), shards=(32, 64), @@ -236,7 +238,7 @@ def test_the_byte_counters_see_whole_chunks_decoded_and_rewritten(tmp_path): def test_h5py_calls_keep_no_memory_once_they_return(tmp_path): - """h5py 3.15 kept 100 to 160 bytes of every dataset read, write and + """h5py 3.15 kept 32 to 160 bytes of every dataset read, write and `.shape` for good: 48 MB of resident memory after 300,000 writes. Every streaming loop makes calls in proportion to its blocks, so memory grew with the input, and in the transpose, whose bucket writes grow with the @@ -275,6 +277,23 @@ def work() -> None: ) +@pytest.mark.parametrize("mode", ["w", "r"]) +def test_stores_give_each_hdf5_dataset_a_one_mebibyte_chunk_cache(tmp_path, mode): + """HDF5 2.0, which h5py 3.16 brings, raised the chunk cache from 1 MiB to + 8 MiB a dataset. A writer fills one for each dataset it grows, and + `split` holds many outputs open, so memory followed the outputs eight + times faster: eight groups of a 200,000 x 2,000 CSR took 530 MB of + resident memory against 429 MB on h5py 3.15. Stores keep 1 MiB.""" + from adata.storage import HDF5_CHUNK_CACHE, open_store + + path = tmp_path / "cache.h5ad" + with open_store(path, "w") as store: + store.root.create_dataset("x", shape=(10,), dtype="f4", chunks=(5,)) + with open_store(path, mode if mode == "r" else "a") as store: + cache = store.root["x"].id.get_access_plist().get_chunk_cache()[1] + assert cache == HDF5_CHUNK_CACHE == 1 << 20 + + @pytest.mark.slow def test_peak_rss_sees_a_command_that_holds_the_matrix(tmp_path): """F146: memory was measured only with tracemalloc, which cannot see what @@ -2072,14 +2091,14 @@ def record(*args, _name=name, _original=original, **kwargs): # rows adds a few hundredths of it; one holding the matrix adds all of it. #: Rows in a dense block, and entries in a sparse read, while the grid runs. -GRID_ROWS, GRID_ENTRIES = 32, 16_384 +GRID_ROWS, GRID_ENTRIES = 32, 32_768 #: Sizes along the grown axis, and the other axis's fixed size. Small enough -#: that a transpose needs no more than a few dozen buckets (see -#: `test_a_transpose_buckets_hold_about_a_chunk_on_a_skewed_matrix`). +#: that a transpose needs no more than a few dozen buckets, each of which +#: costs some memory of its own; at this chunk, a few hundredths of the input. GRID_SIZES = {"n_obs": (1024, 8192), "n_var": (128, 1024)} GRID_FIXED = {"n_obs": 1024, "n_var": 4096} -GRID_DENSITY = 0.05 +GRID_DENSITY = 0.1 def _grid_shape(axis: str, n: int) -> tuple: @@ -2116,17 +2135,11 @@ def _in_memory_bytes(matrix) -> int: #: Known gaps, each owned by another batch, and whether the fix must flip it. _F032 = "F032: sparse subset blocks are 4,096 majors, whatever --chunk (B6b, #97)" -_CONCAT = ( - "concat keeps every obs name by design, and its row pointers in Python " - "lists until B3b (#116, F132); together a third of this fixture's input" -) GRID_GAPS = { ("subset", "csr", "n_obs"): (_F032, True), ("subset", "csr", "n_var"): (_F032, True), ("subset", "csc", "n_obs"): (_F032, True), ("subset", "csc", "n_var"): (_F032, True), - ("concat", "csr", "n_obs"): (_CONCAT, False), - ("concat", "csc", "n_obs"): (_CONCAT, False), } @@ -2163,9 +2176,14 @@ def measure(n: int) -> int: layout=layout, matrix=matrix(n)) if command == "concat": shutil.copy(source, source.with_name(f"b{source.name}")) - with count_allocations() as peak: - run(source, tmp_path / f"g{n}" / "out.h5ad") - return peak[0] + peaks = [] + # Twice, keeping the smaller, so a one-time growth of the process's + # own structures cannot land in one measurement and decide it. + for attempt in range(2): + with count_allocations() as peak: + run(source, tmp_path / f"g{n}" / f"out{attempt}.h5ad") + peaks.append(peak[0]) + return min(peaks) copies = 2 if command == "concat" else 1 assert_memory_is_a_fraction_of_the_input( @@ -2223,7 +2241,9 @@ def measure(n: int) -> int: with count_io() as io: _recast(source, out, dtype="float64", chunk_rows=256) assert io.read_bytes >= n * 16 * 4, "the matrix was read" - return _largest_matrix_read(io) + largest = _largest_matrix_read(io) + assert largest > 0, "no read of X was seen: the names no longer match" + return largest assert_independent_of( measure, what=f"the largest read of recast --dtype ({fmt})", axis="n_obs", @@ -2254,7 +2274,7 @@ def measure(n: int) -> int: peaks.append(peak[0]) # No read larger than a chunk of entries or a block of rows. bound = GRID_ENTRIES * 8 + rows * 64 * 8 - assert io.read_bytes > 0 and _largest_matrix_read(io) <= bound, io + assert 0 < _largest_matrix_read(io) <= bound, io return min(peaks) # Per row added, so the fixed cost of the store cancels. A block of @@ -2292,13 +2312,41 @@ def test_a_transpose_reads_about_a_chunk_on_a_skewed_matrix(tmp_path, fmt, sourc _recast(store, tmp_path / f"o{n}{store.suffix}", layout=target, chunk=chunk) assert io.read_bytes >= n * 8 * 4, "the dense genes were read" largest = _largest_matrix_read(io, elements=True) - assert largest <= bound, ( + assert 0 < largest <= bound, ( f"a skewed {source} transpose of {n:,} cells ({fmt}) read " f"{largest:,} entries at once, against {bound:,} for a chunk of " f"{chunk:,}: a run of genes in every cell is being read whole" ) +def test_a_compressed_dense_matrix_to_csc_does_not_rewrite_its_scratch(tmp_path): + """F030's route through a CSR in scratch gave that scratch the source's + gzip, and scratch has no chunk cache, so every block's append decoded + and re-encoded a whole 65,536-element chunk: 2,471 MB encoded for 14.5 MB + written, 30 to 100 times slower at --chunk 16. The scratch is plain now, + and the output still takes the source's codecs.""" + import h5py + + from tests.perf_fixtures import counts_matrix + + dense = counts_matrix(20_000, 64, "dense", density=0.1) + source = tmp_path / "gz.h5ad" + ad.AnnData(X=dense).write_h5ad(source, compression="gzip") + out = tmp_path / "csc.h5ad" + + with count_io() as io: + _recast(source, out, layout="csc", chunk_rows=16, chunk=GRID_ENTRIES) + + assert io.written_bytes > 0 + assert io.encoded_bytes <= 4 * io.written_bytes, ( + f"{io.encoded_bytes / 1e6:.0f} MB of chunks encoded for " + f"{io.written_bytes / 1e6:.1f} MB written: something is rewritten per block" + ) + with h5py.File(out, "r") as handle: + for name in ("data", "indices"): + assert handle[f"X/{name}"].compression == "gzip", name + + @pytest.mark.slow def test_a_dense_cast_adds_a_block_of_resident_memory_not_the_matrix(tmp_path): """F029 as the operating system sees it, through `peak_rss`: at 64 MB of diff --git a/tests/test_recast_streaming.py b/tests/test_recast_streaming.py index 4b4173a..72d7a4d 100644 --- a/tests/test_recast_streaming.py +++ b/tests/test_recast_streaming.py @@ -273,17 +273,45 @@ def test_a_dense_source_whose_nonzeros_overflow_the_offsets_is_refused( dtype's range the offsets wrapped and the command still succeeded. nnz is only known by reading, so this one is refused while writing; staging then removes the half-written output.""" + from tests.perf_counters import count_io + source_path = _put(new_store, np.ones((300, 200), dtype="float32")) out = _output(source_path) - with pytest.raises(ValueError, match="cannot hold this matrix's offsets"): + with count_io() as io, pytest.raises(ValueError, match="cannot hold this matrix's offsets"): recast_store(source_path, ["X"], out, QUIET, layout=layout, indices_dtype="int16", chunk_rows=16) assert not out.exists() + # As soon as the count passes the range, CSC's scratch included: not + # after writing all of it. + assert io.read_bytes < 300 * 200 * 4, f"read {io.read_bytes:,} bytes first" leftovers = [p.name for p in source_path.parent.iterdir() if "adata-" in p.name] assert leftovers == [], leftovers +@pytest.mark.parametrize("layout,refused", [("csr", False), ("csc", True)]) +def test_the_coordinates_checked_are_those_of_the_output_minor_axis( + new_store, int16_indices, layout, refused +): + """With F013's check always running, it must check the right thing: CSR + indices are column coordinates and CSC ones row coordinates. Checking + the larger dimension refused int16 indices for a CSR matrix with more + rows than int16 holds but two columns.""" + rows = np.arange(0, 40_000, 800) + matrix = sparse.csr_matrix( + (np.ones(rows.size), (rows, np.zeros(rows.size, dtype=int))), shape=(40_000, 2) + ) + source_path = _put(new_store, matrix) + + if refused: + with pytest.raises(ValueError, match="cannot hold this matrix's coordinates"): + _recast(source_path, layout=layout, indices_dtype="int16") + assert not _output(source_path).exists() + return + out = _recast(source_path, layout=layout, indices_dtype="int16") + _same(_read(out).X, matrix) + + def test_offsets_that_fit_are_written_in_the_narrow_dtype(new_store, int16_indices): """The other side of F013's check: a dense source whose nonzeros fit the narrow dtype is converted, with both arrays in it.""" @@ -461,38 +489,63 @@ def test_a_sparse_matrix_from_a_dense_one_has_element_sized_chunks(new_store, la @pytest.mark.parametrize( - "layout,source,scratch", - [("csc", "csr", True), ("dense", "csc", True), ("dense", "csr", False), ("csr", "csr", False)], + "layout,source,streamed,in_memory", + [ + # bucket entries: an int64 major, an int64 minor and a float64 value + ("csc", "csr", 24, 0), + # and the CSR the CSC goes through first, with its indptr + ("dense", "csc", 24 + 16, 16), + ("dense", "csr", 0, 0), + ("csr", "csr", 0, 0), + ], ) -def test_the_free_space_estimate_counts_the_transpose_scratch( - tmp_path, monkeypatch, layout, source, scratch +def test_the_free_space_estimate_counts_the_scratch( + tmp_path, layout, source, streamed, in_memory ): """A streaming transpose writes every entry to scratch beside the output, - as an int64 major, an int64 minor and its value, and a CSC matrix being - densified goes through a CSR first. The estimate counted only the output, - so a disk with room for that ran out part way through the scratch.""" + and a CSC matrix being densified goes through a CSR first, in memory or + not. The estimate counted only the output, so a disk with room for that + could fill part way through the scratch.""" import adata.commands.recast as command + from adata.storage import open_store matrix = _matrix(n_obs=50, n_var=40) - value = _as(matrix, source) path = tmp_path / "in.h5ad" - ad.AnnData(X=value).write_h5ad(path) + ad.AnnData(X=_as(matrix, source)).write_h5ad(path) + + with open_store(path, "r") as store: + plan = core.plan_conversion(store.root["X"], "X", layout=layout, force=True) + rows = 8 * (matrix.shape[0] + 1) if layout == "dense" and source == "csc" else 0 + for flag, per_entry in ((False, streamed), (True, in_memory)): + expected = matrix.nnz * per_entry + (rows if per_entry else 0) + assert command._scratch(plan, flag) == expected, (flag, per_entry) + + +def test_dense_to_csc_checks_free_space_again_once_it_knows_its_nonzeros( + tmp_path, monkeypatch +): + """Dense to CSC goes through a CSR in scratch and the transpose's buckets, + sized by nnz, which is known only once the matrix has been read. The + up-front estimate cannot count them, so `sparsify` asks again before + transposing, and a disk too small for them is refused there.""" + import adata.commands.recast as command + + dense = _matrix(n_obs=60, n_var=30).toarray() + path = tmp_path / "in.h5ad" + ad.AnnData(X=dense).write_h5ad(path) asked = [] - monkeypatch.setattr( - command, "check_free_space", - lambda output, needed, check=True: asked.append(needed()), - ) - def estimate(in_memory: bool) -> float: - asked.clear() - out = tmp_path / f"o{len(asked)}{in_memory}{layout}.h5ad" - recast_store(path, ["X"], out, QUIET, layout=layout, force=True, - dtype="float32", in_memory=in_memory) - return asked[0] + def check(output, needed, check=True): + asked.append(needed() if callable(needed) else needed) + if len(asked) == 2: + raise OSError("Not enough free space") - extra = estimate(False) - estimate(True) - if not scratch: - assert extra == 0 - return - per_entry = 16 + 8 + (8 + 8 if layout == "dense" else 0) # float64 source - assert extra == pytest.approx(matrix.nnz * per_entry), extra + monkeypatch.setattr(command, "check_free_space", check) + out = tmp_path / "out.h5ad" + with pytest.raises(OSError, match="Not enough free space"): + recast_store(path, ["X"], out, QUIET, layout="csc", chunk_rows=7) + + nnz = int(np.count_nonzero(dense)) + assert asked[1] == nnz * (16 + 8 + 8 + 8), asked + assert not out.exists() + assert [p.name for p in tmp_path.iterdir() if "adata-" in p.name] == [] From 7288efa65bbb4e4b0c518e7e7ba7f518282a27a3 Mon Sep 17 00:00:00 2001 From: Aljes Date: Mon, 2 Nov 2026 11:12:08 +0000 Subject: [PATCH 4/5] RSS on Linux from VmHWM; cut the new tests' CI time ru_maxrss on Linux starts at the parent's peak (the kernel folds the replaced memory map into it at exec), so on CI every command added nothing: peak_rss reads VmHWM there. The new tests cost about 350 s of the unit job. Grid stores are built once and shared, each size measured once against a share of a half (bounded paths reach 0.21), with --chunk 256 rows along n_obs and 32 along n_var. The per-finding guards run on h5ad and sharded Zarr v3, at smaller sizes, and the recast correctness tests use real chunks where the blocks do not matter. Co-Authored-By: Claude Opus 5.5 --- docs/reference/testing.md | 17 +++-- tests/perf_counters.py | 23 +++++-- tests/test_performance.py | 115 +++++++++++++++++++-------------- tests/test_recast_streaming.py | 6 +- 4 files changed, 100 insertions(+), 61 deletions(-) diff --git a/docs/reference/testing.md b/docs/reference/testing.md index d58033c..79f8ed8 100644 --- a/docs/reference/testing.md +++ b/docs/reference/testing.md @@ -167,8 +167,9 @@ model against numbers worked out by hand. buffers but not what HDF5, blosc or zarr's codecs allocate in C. `peak_rss(argv)` runs `adata ` in a fresh interpreter and returns how far the command raised the operating system's high-water mark (`ru_maxrss`, read -once the CLI is imported and again when the command ends). A subprocess, -because the high-water mark only rises. It has a few megabytes of noise, so +once the CLI is imported and again when the command ends; on Linux `VmHWM`, +since `ru_maxrss` there starts at the parent's peak). A subprocess, because +the high-water mark only rises. It has a few megabytes of noise, so guard it with generous ratios against inputs of tens of megabytes, as `test_a_dense_cast_adds_a_block_of_resident_memory_not_the_matrix` does; `test_peak_rss_sees_a_command_that_holds_the_matrix` checks that it sees @@ -220,7 +221,7 @@ accept a 64x increase in a command that is supposed to read nothing. |---|---|---| | `assert_grows_linearly` | no worse than linear in one axis | most commands | | `assert_grows_slower_than_input` | grows at least N times slower than the input | streaming at a fixed chunk | -| `assert_memory_is_a_fraction_of_the_input` | adds at most a share (a quarter) of what the input adds, in bytes | memory, along either axis | +| `assert_memory_is_a_fraction_of_the_input` | adds at most a share (half) of what the input adds, in bytes | memory, along either axis | | `assert_independent_of` | does not grow at all | inspection, grouping work per row, the largest read | **Inspection is free, and that is an exact number.** `view` and `ls` reach only @@ -264,12 +265,16 @@ A growth ratio cannot be the claim along n_var: widen a matrix and a block of `--chunk` rows widens with it, as fast as the input, so a ratio cannot tell it from the whole matrix. Peak memory added per byte of input added can: a block of 32 rows of a matrix thousands of rows tall adds a few hundredths, and a path -holding the matrix adds a whole byte and more for each copy. +holding the matrix adds a whole byte and more for each copy. The limit is half; +bounded paths came to at most 0.21. Growing n_obs, `--chunk` is 256 rows, so a +block of `--chunk` columns over every row (CSC densify before F030) is plain. +Each fixture is built once and shared by the commands. The known gap is marked `xfail` with its owner: sparse `subset` (F032, B6b), strict, so the fix has to flip it. -The recast guards beside it pin each block in bytes, on every format: the cast +The recast guards beside it pin each block in bytes, on h5ad and sharded Zarr +v3 (Zarr v2 is covered by the correctness tests): the cast check reads `--chunk` rows of a dense matrix, CSC densify and sparsify read blocks set by the chunk, and a skewed transpose never reads more than about two buckets of entries at once. @@ -341,7 +346,7 @@ should scale with, scale only that, and hold everything else fixed. | `convert` CellBender (plain, `--min-cell-probability`, `--raw` full and filtered) | n_barcodes | linear; **flat** write calls | | `recast --dtype`, `--layout csc`, `--layout dense` | nnz, n_obs | linear; the streaming transpose's **peak memory** is bounded | | `recast` (every dtype and layout change), `subset`, `concat` | n_obs, n_var; CSR, CSC, dense | **peak memory** a small share of the input | -| `recast --dtype` (dense), CSC `--layout dense`, dense `--layout csc`, a skewed transpose | n_obs; h5ad, Zarr v2, sharded Zarr v3 | the **largest read** set by the chunk | +| `recast --dtype` (dense), CSC `--layout dense`, dense `--layout csc`, a skewed transpose | n_obs; h5ad, sharded Zarr v3 | the **largest read** set by the chunk | The `slow` marker is on the guards that build a 65,536-row store or run `peak_rss`; they still gate merges, and `-m "not slow"` skips them locally. `perf` is on the diff --git a/tests/perf_counters.py b/tests/perf_counters.py index 4428e1c..c0c811b 100644 --- a/tests/perf_counters.py +++ b/tests/perf_counters.py @@ -537,10 +537,23 @@ def count_allocations() -> Iterator[List[int]]: #: Run by `peak_rss` in a fresh interpreter. The high-water mark is read once #: the CLI and everything it imports is loaded, and again after the command, #: so the difference is what the command itself added. +#: +#: On Linux that is `VmHWM`, the peak of the process's own memory map, not +#: `ru_maxrss`: at exec the kernel folds the memory map being replaced -- the +#: parent's, after fork or vfork -- into `ru_maxrss`, so a child of a large +#: pytest process starts at the parent's peak and its command added nothing. +#: macOS keeps no such history, and has no /proc. _RSS_RUNNER = r""" import json, resource, sys def high_water(): + try: + with open("/proc/self/status") as status: + for line in status: + if line.startswith("VmHWM:"): + return int(line.split()[1]) * 1024 # kB + except OSError: + pass value = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss return value if sys.platform == "darwin" else value * 1024 # Linux: KiB @@ -581,7 +594,8 @@ def peak_rss(argv: List[str], *, cwd: Any = None) -> PeakRSS: blosc and zarr's codecs decode into, and anything else allocated in C, none of which tracemalloc can see. A subprocess because the high-water mark only rises: in this process it would hold whatever the test suite - reached before. Not exact -- the allocator keeps a few megabytes it has + reached before (and on Linux, see `_RSS_RUNNER`, `ru_maxrss` would carry + it into the child too). Not exact -- the allocator keeps a few megabytes it has freed -- so guard it with ratios that leave room, against inputs large enough (tens of megabytes) that holding one would be unmistakable. """ @@ -885,9 +899,10 @@ def assert_grows_slower_than_input( #: The most memory a streaming path may add per byte of input added. A path #: holding a block of `--chunk` rows of a matrix many blocks tall adds a few #: hundredths; one holding the matrix adds a whole byte, and more for every -#: copy it makes. A quarter leaves room for the per-row arrays a path keeps -#: by design (an indptr, a row's count) on a narrow fixture. -INPUT_SHARE = 0.25 +#: copy it makes. Half leaves room on both sides: what a bounded path keeps +#: by design (an indptr, a row's count, concat's obs names, the buckets of a +#: transpose) came to at most 0.21 on the test fixtures. +INPUT_SHARE = 0.5 def assert_memory_is_a_fraction_of_the_input( diff --git a/tests/test_performance.py b/tests/test_performance.py index 5f83c4d..429dbfe 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -249,7 +249,7 @@ def test_h5py_calls_keep_no_memory_once_they_return(tmp_path): import h5py - calls = 20_000 + calls = 5_000 with h5py.File(tmp_path / "calls.h5", "w") as handle: dataset = handle.create_dataset("x", shape=(100,), dtype="int64") values = np.arange(3) @@ -271,7 +271,7 @@ def work() -> None: finally: tracemalloc.stop() - assert kept <= calls, ( + assert kept <= 8 * calls, ( f"h5py {h5py.__version__} kept {kept / calls:.0f} bytes a call " f"({kept / 1e6:.1f} MB over {calls:,} calls of a read, a write and .shape)" ) @@ -2090,8 +2090,13 @@ def record(*args, _name=name, _original=original, **kwargs): # memory added per byte of input added. A path holding a block of --chunk # rows adds a few hundredths of it; one holding the matrix adds all of it. -#: Rows in a dense block, and entries in a sparse read, while the grid runs. -GRID_ROWS, GRID_ENTRIES = 32, 32_768 +#: Rows in a dense block (--chunk) along each axis, and entries in a sparse +#: read. Growing n_obs, a block of rows is the same at both sizes and cancels, +#: while a block of --chunk columns over every row grows with it: 256 makes +#: that plain. Growing n_var, a row block widens with the input, so it is +#: kept to 32 rows of thousands. +GRID_ROWS = {"n_obs": 256, "n_var": 32} +GRID_ENTRIES = 32_768 #: Sizes along the grown axis, and the other axis's fixed size. Small enough #: that a transpose needs no more than a few dozen buckets, each of which @@ -2111,25 +2116,52 @@ def _in_memory_bytes(matrix) -> int: return int(matrix.nbytes) -#: What each command does to a store, at the grid's chunk sizes. +#: Grid stores, built once and shared by every command: (path, input bytes). +_GRID_STORES: dict = {} + + +@pytest.fixture(scope="module") +def grid_dir(tmp_path_factory): + yield tmp_path_factory.mktemp("grid") + _GRID_STORES.clear() + + +def _grid_store(directory: Path, layout: str, axis: str, n: int) -> tuple: + """The grid's store of `layout` at size `n` along `axis`, and a copy of it + for concat, built once: no command writes to its source.""" + import shutil + + from tests.perf_fixtures import counts_matrix, matrix_store + + key = (layout, axis, n) + if key not in _GRID_STORES: + matrix = counts_matrix(*_grid_shape(axis, n), layout, density=GRID_DENSITY) + path = matrix_store(directory / f"{layout}-{axis}-{n}", n_obs=matrix.shape[0], + n_var=matrix.shape[1], layout=layout, matrix=matrix) + shutil.copy(path, path.with_name(f"b{path.name}")) + _GRID_STORES[key] = (path, _in_memory_bytes(matrix)) + return _GRID_STORES[key] + + +#: What each command does to a store, given the grid's rows per block. GRID_COMMANDS = { - "recast --dtype": lambda s, o: _recast( - s, o, dtype="float64", chunk_rows=GRID_ROWS, chunk=GRID_ENTRIES + "recast --dtype": lambda s, o, rows: _recast( + s, o, dtype="float64", chunk_rows=rows, chunk=GRID_ENTRIES ), - "recast --layout csr": lambda s, o: _recast( - s, o, layout="csr", force=True, chunk_rows=GRID_ROWS, chunk=GRID_ENTRIES + "recast --layout csr": lambda s, o, rows: _recast( + s, o, layout="csr", force=True, chunk_rows=rows, chunk=GRID_ENTRIES ), - "recast --layout csc": lambda s, o: _recast( - s, o, layout="csc", force=True, chunk_rows=GRID_ROWS, chunk=GRID_ENTRIES + "recast --layout csc": lambda s, o, rows: _recast( + s, o, layout="csc", force=True, chunk_rows=rows, chunk=GRID_ENTRIES ), - "recast --layout dense": lambda s, o: _recast( - s, o, layout="dense", force=True, chunk_rows=GRID_ROWS, chunk=GRID_ENTRIES + "recast --layout dense": lambda s, o, rows: _recast( + s, o, layout="dense", force=True, chunk_rows=rows, chunk=GRID_ENTRIES ), - "subset": lambda s, o: subset_h5ad( - s, o, None, None, console=QUIET, obs_query="col >= 0", chunk_rows=GRID_ROWS + "subset": lambda s, o, rows: subset_h5ad( + s, o, None, None, console=QUIET, obs_query="col >= 0", chunk_rows=rows ), - "concat": lambda s, o: concat_on_disk( - [s, s.with_name(f"b{s.name}")], o, QUIET, chunk_rows=GRID_ROWS, index_unique="-" + "concat": lambda s, o, rows: concat_on_disk( + [s, s.with_name(f"b{s.name}")], o, QUIET, chunk_rows=rows, index_unique="-" ), } @@ -2156,39 +2188,20 @@ def _grid_cases(): @pytest.mark.parametrize("command,layout,axis", list(_grid_cases())) -def test_peak_memory_is_a_small_share_of_the_input(tmp_path, command, layout, axis): +def test_peak_memory_is_a_small_share_of_the_input(grid_dir, tmp_path, command, layout, axis): """F091: every layout, both axes, against what the input itself costs.""" - import shutil - - from tests.perf_fixtures import counts_matrix, matrix_store - run = GRID_COMMANDS[command] - matrices = {} - - def matrix(n: int): - if n not in matrices: - matrices[n] = counts_matrix(*_grid_shape(axis, n), layout, density=GRID_DENSITY) - return matrices[n] + copies = 2 if command == "concat" else 1 def measure(n: int) -> int: - n_obs, n_var = _grid_shape(axis, n) - source = matrix_store(tmp_path / f"g{n}", n_obs=n_obs, n_var=n_var, - layout=layout, matrix=matrix(n)) - if command == "concat": - shutil.copy(source, source.with_name(f"b{source.name}")) - peaks = [] - # Twice, keeping the smaller, so a one-time growth of the process's - # own structures cannot land in one measurement and decide it. - for attempt in range(2): - with count_allocations() as peak: - run(source, tmp_path / f"g{n}" / f"out{attempt}.h5ad") - peaks.append(peak[0]) - return min(peaks) + source, _ = _grid_store(grid_dir, layout, axis, n) + with count_allocations() as peak: + run(source, tmp_path / f"out{n}.h5ad", GRID_ROWS[axis]) + return peak[0] - copies = 2 if command == "concat" else 1 assert_memory_is_a_fraction_of_the_input( measure, - lambda n: copies * _in_memory_bytes(matrix(n)), + lambda n: copies * _grid_store(grid_dir, layout, axis, n)[1], what=f"{command} of a {layout} X", axis=axis, sizes=GRID_SIZES[axis], @@ -2203,7 +2216,9 @@ def measure(n: int) -> int: # set by --chunk, not by the input. Every format, since anndata 0.13's # sharded Zarr v3 lays a matrix out unlike anything else. -STORE_FORMATS = ("h5ad", "zarr2", "zarr3") +#: h5ad and anndata's sharded Zarr v3, whose layout is the one unlike the +#: others; Zarr v2 is covered by the correctness tests, every path of them. +STORE_FORMATS = ("h5ad", "zarr3") def _largest_matrix_read(io, *, elements: bool = False) -> int: @@ -2247,7 +2262,7 @@ def measure(n: int) -> int: assert_independent_of( measure, what=f"the largest read of recast --dtype ({fmt})", axis="n_obs", - sizes=(1024, 16384), + sizes=(1024, 8192), ) @@ -2261,6 +2276,7 @@ def test_csc_densify_and_sparsify_read_blocks_set_by_the_chunk(tmp_path, fmt, la from tests.perf_fixtures import matrix_store rows = 256 # --chunk + entries = 4096 # below nnz at either size, so the blocks are full at both def measure(n: int) -> int: store = matrix_store(tmp_path, n_obs=n, n_var=64, layout=source, fmt=fmt) @@ -2270,17 +2286,17 @@ def measure(n: int) -> int: for attempt in range(2): with count_allocations() as peak, count_io() as io: _recast(store, tmp_path / f"o{n}-{attempt}{store.suffix}", layout=layout, - force=True, chunk_rows=rows, chunk=GRID_ENTRIES) + force=True, chunk_rows=rows, chunk=entries) peaks.append(peak[0]) # No read larger than a chunk of entries or a block of rows. - bound = GRID_ENTRIES * 8 + rows * 64 * 8 + bound = entries * 8 + rows * 64 * 8 assert 0 < _largest_matrix_read(io) <= bound, io return min(peaks) # Per row added, so the fixed cost of the store cancels. A block of # --chunk columns over every row holds `rows` float32 values a row; what # is left is the per-row index arrays (an indptr, a count), a few int64. - small, large = 1024, 8192 + small, large = 1024, 4096 per_row = (measure(large) - measure(small)) / (large - small) assert per_row <= rows * 4 / 4, ( f"recast {source} -> {layout} ({fmt}) held {per_row:.0f} bytes more per " @@ -2304,7 +2320,8 @@ def test_a_transpose_reads_about_a_chunk_on_a_skewed_matrix(tmp_path, fmt, sourc # majors, read whole, up to one bucket plus one major that fits in one. bound = 2 * max(chunk, MIN_BUCKET_ENTRIES) - for n in (1024, 16384): + # Above two buckets' worth, so a gene read whole is past the bound. + for n in (12_288,): store = matrix_store(tmp_path, n_obs=n, n_var=64, layout=source, fmt=fmt, matrix=skewed_matrix(n, 64, layout=source), tag="skew") target = "csr" if source == "csc" else "csc" diff --git a/tests/test_recast_streaming.py b/tests/test_recast_streaming.py index 72d7a4d..a82cffd 100644 --- a/tests/test_recast_streaming.py +++ b/tests/test_recast_streaming.py @@ -390,7 +390,9 @@ def test_anndata_reads_back_every_dtype_and_layout_recast_writes( if dtype == "float16" and layout != "dense": pytest.skip("refused, see test_float16_is_refused_for_a_sparse_output") - out = _recast(source_path, dtype=dtype, layout=layout, force=True) + # Blocks do not matter here; the dtype and layout do. + out = _recast(source_path, dtype=dtype, layout=layout, force=True, + chunk=4096, chunk_rows=64) if layout == "dense": expected = matrix.toarray().astype(dtype) @@ -470,7 +472,7 @@ def test_a_sparse_matrix_from_a_dense_one_has_element_sized_chunks(new_store, la ew.write_dataframe_header(root, "var", [f"g{i}" for i in range(50)], []) ew.write_dense(root, "X", dense, chunks=(4, 50)) - out = _recast(path, layout=layout, chunk_rows=64) + out = _recast(path, layout=layout, chunk_rows=64, chunk=4096) from adata.storage import open_store From 1bec201db3f6db8ae857d8bb475e624750b81c1e Mon Sep 17 00:00:00 2001 From: Aljes Date: Mon, 2 Nov 2026 14:47:58 +0000 Subject: [PATCH 5/5] Run two Zarr v3 guard cases on h5ad only, for CI time The CSC per-row guard's Zarr v3 cases took 38 s of CI and measure the algorithm, not the store; the skewed CSR case reads only the transpose's own scratch. The sharded reads stay covered by the dense-cast guard and the skewed CSC case. Co-Authored-By: Claude Opus 5.5 --- docs/reference/testing.md | 3 ++- tests/test_performance.py | 11 ++++++++--- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/docs/reference/testing.md b/docs/reference/testing.md index a901708..812cb43 100644 --- a/docs/reference/testing.md +++ b/docs/reference/testing.md @@ -366,7 +366,8 @@ should scale with, scale only that, and hold everything else fixed. | `convert` CellBender (plain, `--min-cell-probability`, `--raw` full and filtered) | n_barcodes | linear; **flat** write calls | | `recast --dtype`, `--layout csc`, `--layout dense` | nnz, n_obs | linear; the streaming transpose's **peak memory** is bounded | | `recast` (every dtype and layout change), `subset`, `concat` | n_obs, n_var; CSR, CSC, dense | **peak memory** a small share of the input | -| `recast --dtype` (dense), CSC `--layout dense`, dense `--layout csc`, a skewed transpose | n_obs; h5ad, sharded Zarr v3 | the **largest read** set by the chunk | +| `recast --dtype` (dense), a skewed transpose | n_obs; h5ad, sharded Zarr v3 | the **largest read** set by the chunk | +| CSC `--layout dense`, dense `--layout csc` | n_obs; h5ad | the **largest read** set by the chunk, and memory per row | The `slow` marker is on the guards that build a 65,536-row store or run `peak_rss`; they still gate merges, and `-m "not slow"` skips them locally. `perf` is on the diff --git a/tests/test_performance.py b/tests/test_performance.py index f43c932..4efa844 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -2520,7 +2520,9 @@ def measure(n: int) -> int: ) -@pytest.mark.parametrize("fmt", STORE_FORMATS) +# h5ad only: what a row costs is the algorithm's, and Zarr v3's reads are +# the dense-cast and skewed guards' to check. On Zarr this took 38 s of CI. +@pytest.mark.parametrize("fmt", ["h5ad"]) @pytest.mark.parametrize("layout,source", [("dense", "csc"), ("csc", "dense")]) def test_csc_densify_and_sparsify_read_blocks_set_by_the_chunk(tmp_path, fmt, layout, source): """F030: densifying a CSC matrix allocated n_rows x --chunk per block, @@ -2559,8 +2561,11 @@ def measure(n: int) -> int: ) -@pytest.mark.parametrize("fmt", STORE_FORMATS) -@pytest.mark.parametrize("source", ["csc", "csr"]) +# CSC's pass 2 reads the source, so it runs on sharded Zarr v3 too; CSR's +# pass 3 reads the transpose's own scratch, the same on either. +@pytest.mark.parametrize( + "fmt,source", [("h5ad", "csc"), ("zarr3", "csc"), ("h5ad", "csr")] +) def test_a_transpose_reads_about_a_chunk_on_a_skewed_matrix(tmp_path, fmt, source): """F077: pass 2 of the streaming transpose stepped by the average nnz per major. Eight adjacent genes in every cell made one step many times the