From 3db4a92fc7c694b3180f17e35ea29ce064732492 Mon Sep 17 00:00:00 2001 From: Aljes Date: Mon, 2 Nov 2026 16:06:49 +0000 Subject: [PATCH 1/5] Align reads and writes with chunks and Zarr v3 shards (B6b) Outputs keep the source's chunks and shards, clamped to the output, and are written a whole shard or chunk at a time through storage.RowWriter (F031, F033, F036, F060). Reads take whole source chunks, in column tiles where a chunk height across the matrix is over the read budget (F061, F081). Sparse subset and split read about --chunk rows' entries a block, a major larger than that in parts (F032). Shard indexes are kept for the life of a store opened for reading (F062). split holds at most 128 HDF5 outputs open a pass, with no chunk cache (F079). Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 9 +- docs/developers.md | 8 +- docs/guides/pipelines.md | 14 +- docs/reference/commands.md | 37 ++- docs/reference/testing.md | 21 +- docs/troubleshooting.md | 16 +- src/adata/cli.py | 4 +- src/adata/core/concat.py | 32 ++- src/adata/core/csr_build.py | 2 + src/adata/core/recast.py | 67 ++--- src/adata/core/subset.py | 358 ++++++++++++++++++++------- src/adata/storage/__init__.py | 444 +++++++++++++++++++++++++++++++--- tests/perf_counters.py | 10 +- tests/perf_fixtures.py | 76 +++++- tests/test_block_layout.py | 390 +++++++++++++++++++++++++++++ tests/test_performance.py | 279 ++++++++++++++++++++- tests/test_storage.py | 35 ++- tests/test_subset_fan_out.py | 7 +- 18 files changed, 1611 insertions(+), 198 deletions(-) create mode 100644 tests/test_block_layout.py diff --git a/AGENTS.md b/AGENTS.md index dfa9b56..1ca6313 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -191,7 +191,14 @@ These are easy to break without any test noticing straight away: - Commands read the source's version from `Store.zarr_format`. - `concat` uses the version its Zarr inputs share. - **Outputs keep the source's index dtypes, chunking and compression.** - Changing them is `recast`'s job, or that of an explicit flag. + Changing them is `recast`'s job, or that of an explicit flag. Zarr v3 + shards are kept too: `fit_chunks` clamps them with the chunk, and given + `itemsize` cuts a unit too large to hold along its rows. +- **Write a streamed output through `storage.RowWriter`**, and `flush()` it + after the last block. A write to part of a Zarr shard rewrites the whole + shard, and to part of a chunk Zarr keeps no cache, so writing a `--chunk` + block at a time straight into the array costs a unit per block. + `_growable_like` returns one for a sparse matrix's `data` and `indices`. ## Tests diff --git a/docs/developers.md b/docs/developers.md index e739246..66253ab 100644 --- a/docs/developers.md +++ b/docs/developers.md @@ -136,7 +136,13 @@ before changing anything that writes: `open_store`; without it zarr-python writes v3, upgrading a v2 input. `concat` uses the version its Zarr inputs share. - **Outputs keep the source's index dtypes, chunking and compression.** - Changing them is `recast`'s job, or an explicit flag's. + Changing them is `recast`'s job, or an explicit flag's. Zarr v3 shards are + kept too: `fit_chunks` clamps them with the chunk, and given `itemsize` + cuts a unit too large to hold along its rows. +- **A streamed output is written through `storage.RowWriter`**, flushed after + the last block. A write to part of a Zarr shard rewrites the whole shard, + and Zarr keeps no chunk cache, so a `--chunk` block written straight into + the array costs a whole unit a block. - **Memory is set by `--chunk`.** Loading a whole matrix or column is a regression unless it is behind an explicit `--in-memory`. `--chunk` is defined once in `cli.py` (`_chunk_option`, with `min=1`); reuse it. diff --git a/docs/guides/pipelines.md b/docs/guides/pipelines.md index 91afae0..7216b80 100644 --- a/docs/guides/pipelines.md +++ b/docs/guides/pipelines.md @@ -123,8 +123,20 @@ most data: | `import array` | 1,048,576 | elements | A dense block holds `--chunk` × n_var values, so lower it for very wide dense -matrices. It must be at least 1; zero or a negative value is a usage error. Sparse matrices stream in blocks too. obs and var columns are read +matrices. It must be at least 1; zero or a negative value is a usage error. Sparse matrices stream in blocks too: in `subset` and `split`, about as +many entries a block as `--chunk` rows hold. obs and var columns are read whole, which is what grows with the number of cells. + +Three things are set by the store's layout rather than by `--chunk`, each +bounded: + +- A read takes whole chunks of the source, up to 32 MiB. +- An output is written a whole shard or chunk at a time, so its rows are + held until one is full: at most 64 MiB, and 256 MiB across the outputs + `split` writes at once. +- `split` holds up to 128 HDF5 outputs open, about 0.6 MB each. + +See [Writing outputs](../reference/commands.md#writing-outputs). [Benchmarks](../reference/benchmarks.md) gives peak memory per command on a 50,000 × 20,000 store, measured on every release: between 70 and 330 MB in 0.16.1. A request of 1 to 2 GB leaves room for larger stores. diff --git a/docs/reference/commands.md b/docs/reference/commands.md index f8af9a9..58d1747 100644 --- a/docs/reference/commands.md +++ b/docs/reference/commands.md @@ -87,7 +87,7 @@ adata subset data.h5ad --inplace --obs barcodes.txt | `--inplace` | Replace the source (written to a temporary path first) | | `--obs` / `--var` | File of names to keep, one per line | | `--obs-query`, `-q` / `--var-query` | Keep rows matching an [expression](query.md) | -| `--chunk`, `-C` | Row chunk size for dense matrices (at least 1) | +| `--chunk`, `-C` | Rows per block, of a dense or a sparse matrix (at least 1) | | `--zarr-format` | Zarr version to write; defaults to the source's | | `--no-space-check` | Write even if the disk looks too small; see [Writing outputs](#writing-outputs) | @@ -285,7 +285,7 @@ adata split data.h5ad --by cluster -o out/ --min-size 50 | `--min-size` | Skip groups smaller than this | | `--dry-run` | Print the plan without writing | | `--manifest / --no-manifest` | Write a CSV listing the outputs (default on) | -| `--chunk`, `-C` | Row chunk size for dense matrices (at least 1) | +| `--chunk`, `-C` | Rows per block, of a dense or a sparse matrix (at least 1) | | `--zarr-format` | Zarr version to write; defaults to the source's | | `--no-space-check` | Write even if the disk looks too small; see [Writing outputs](#writing-outputs) | @@ -304,6 +304,11 @@ all of them are complete, and the manifest is written last, so a split that fails leaves no partial stores. An output that would land on the input is refused, on a case-insensitive filesystem too (a label `A` beside `a.h5ad`). +Each matrix is read once for up to 128 HDF5 outputs (or as many Zarr +outputs as a quarter of the open-file limit allows): an open HDF5 file costs +about 0.6 MB however little it holds, so with more groups than that, the +matrices are read once per 128. + ## `concat` Concatenate along the obs axis, streaming each input in turn. @@ -607,6 +612,34 @@ skips it. **`--chunk` is at least 1** wherever it appears. A zero or negative value is a usage error (exit 2), where it used to write an all-zero dense matrix. +**Outputs keep the source's chunks and shards**, and are read and written a +whole one at a time: + +- A Zarr v3 output keeps the source's shards and the inner chunks inside + them, each clamped to the output. Until 0.17.6, `subset`, `split` and + `concat` dropped the shard of a sparse matrix and kept its inner chunk: + under anndata 0.13 that is 245 elements, a file each. +- Each shard, or each chunk of a Zarr or chunked HDF5 output, is written + once, whole. A shard is rewritten in full for a write of any part of it, + so writing one a `--chunk` block at a time took hundreds of times longer. + To write a shard whole, the command holds its rows: a shard, or chunk, + whose rows across the matrix come to more than 64 MiB is cut along its + rows to fit (anndata aims at 1 GB shards). `split` shares 256 MiB among + the outputs it writes at once, so with many groups their shards and + chunks are cut further. +- Zarr v2 has no shards: written from a sharded store, a shard becomes the + chunk. An HDF5 output of a sharded store lets HDF5 choose its chunks, + rather than taking the shard's inner chunk. +- A dense matrix is read in whole chunks of the source: `--chunk` rows + rounded up to whole chunk heights, up to 32 MiB a read (more if + `--chunk` asks for more). Where one chunk height across every column is + over that, as for chunks of (n_obs, 8), a read takes a band of columns, + whole chunks at a time. +- A sparse matrix in `subset` and `split` is read about as many entries a + block as `--chunk` rows hold on average, in either layout; a column of a + CSC matrix holding more, a gene expressed in every cell, is read in parts. +- A Zarr v3 shard's index is read once per command, not once a read. + ## Zarr format selection New Zarr stores follow the source store's version, so a v2 input is not diff --git a/docs/reference/testing.md b/docs/reference/testing.md index 812cb43..afbf7d6 100644 --- a/docs/reference/testing.md +++ b/docs/reference/testing.md @@ -169,6 +169,7 @@ shard rewritten on every block, so `IOCounts` also counts bytes: | `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 | +| `written_by_name`, `encoded_by_name` | the two above, per array | a writer's own output apart from its scratch | | `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 @@ -276,10 +277,13 @@ 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 +default sharded Zarr v3 -- 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. +layout, growing n_obs and then n_var. zarr 3.4 gives anndata's shards inner +chunks of about a megabyte; `zarr31_layout` rewrites a store's X as zarr 3.1 +laid it out, inner chunks of a kilobyte in shards of the whole array, which +is what users' stores hold and where a read or write that splits a unit +costs most. 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 @@ -290,8 +294,8 @@ 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. +Sparse `subset` was the known gap, read 4,096 majors a block whatever +`--chunk` said (F032); 0.17.6 closed it, and `GRID_GAPS` is empty. 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 @@ -368,6 +372,13 @@ should scale with, scale only that, and hold everything else fixed. | `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), 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 | +| `subset`, `split`, `concat`, `recast`, `convert` into anndata 0.13's shards | Zarr v3 as zarr 3.1 wrote it | each shard **encoded once** | +| dense `subset` | tall gzip chunks, zarr's own chunks | each chunk **decoded once** | +| `split` into chunked HDF5 | eight groups | each output chunk **encoded once** | +| `subset`, `concat` of a sharded store | Zarr v3 as zarr 3.1 wrote it | each shard index **fetched once** | +| sparse `subset` | n_obs; CSR, CSC, a gene in every cell | the **largest read** about `--chunk` rows of entries | +| `convert` of a matrix chunked (n_obs, 8) | n_obs | the **largest read** a tile within the budget | +| `split` | 128 against 512 groups | **resident memory** flat past 128 HDF5 outputs | 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/docs/troubleshooting.md b/docs/troubleshooting.md index 52d429b..c76df14 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -338,9 +338,23 @@ store it wrote; for a store another tool changed, run Lower `--chunk`: a dense block holds `--chunk` × n_var values, 1,024 rows by default. Two things are held whole by design: obs and var columns, and a -pycisTopic pickle, which cannot be read in part. See +pycisTopic pickle, which cannot be read in part. Writing an output a whole +shard or chunk at a time holds at most 64 MiB more (256 MiB across +`split`'s outputs). See [Pipelines and containers](guides/pipelines.md#memory-and---chunk). +### It is slow on a Zarr store anndata 0.13 wrote + +anndata 0.13 writes Zarr v3 sharded: a whole array in one shard file, +inner chunks of a kilobyte (with zarr 3.1) inside. Until 0.17.6, `subset`, +`split`, `concat`, `recast` and `convert` rewrote a whole shard for every +block they wrote and read its index again for every block they read, +hundreds of times slower than the same data as `.h5ad`. They write each +shard once now, keep each index once read, and keep the source's shards on +output. To cap what writing a shard whole holds, a shard whose rows come to +more than 64 MiB is cut along its rows. See +[Writing outputs](reference/commands.md#writing-outputs). + ### `A loom holds its matrices dense: 54.0 GB to compress and write. It is written, but may take a while` A loom has no sparse matrices, so a sparse store's zeros are all written. diff --git a/src/adata/cli.py b/src/adata/cli.py index 93281ef..0f1f580 100644 --- a/src/adata/cli.py +++ b/src/adata/cli.py @@ -344,7 +344,7 @@ def subset( help="Keep var matching an expression, e.g. \"highly_variable == True\"", ), chunk_rows: int = _chunk_option( - 1024, "Row chunk size for dense matrices", "--chunk-rows", "-r" + 1024, "Rows per block, of a dense or a sparse matrix", "--chunk-rows", "-r" ), zarr_format: Optional[int] = typer.Option( None, @@ -1163,7 +1163,7 @@ def split( dry_run: bool = typer.Option( False, "--dry-run", help="Show what would be written without writing it" ), - chunk_rows: int = _chunk_option(1024, "Row chunk size for dense matrices"), + chunk_rows: int = _chunk_option(1024, "Rows per block, of a dense or a sparse matrix"), zarr_format: Optional[int] = typer.Option( None, "--zarr-format", diff --git a/src/adata/core/concat.py b/src/adata/core/concat.py index 2cb539a..86ea384 100644 --- a/src/adata/core/concat.py +++ b/src/adata/core/concat.py @@ -53,6 +53,7 @@ ) from adata.elements.strings import is_string_dtype from adata.storage import ( + RowWriter, copy_tree, create_dataset, dataset_create_kwargs, @@ -877,6 +878,8 @@ def _concat_sparse( # rather than hold both: 8 bytes a row more at the peak. del src_indptr + out_indices.flush() + out_data.flush() _write_indptr(group, indptr, pointer_dtype, inputs.get(0, path)["indptr"]) @@ -902,8 +905,10 @@ def _layout_like(src: Any, dst_parent: Any, shape: Tuple[int, ...], dtype: Any) compression, the chunks fitted to its own shape, as every copied dataset's are kept.""" backend = "zarr" if is_zarr_group(dst_parent) else "hdf5" + # A unit cut to what `RowWriter` holds, for the writers that fill it a + # block of rows at a time. kw = fit_chunks(dataset_create_kwargs(src, target_backend=backend, dst_parent=dst_parent), - shape, backend=backend) + shape, backend=backend, itemsize=np.dtype(dtype).itemsize or None) if np.dtype(dtype) != np.dtype(src.dtype): # The source's fill value is of its own dtype, not the output's. kw.pop("fillvalue", None) @@ -941,7 +946,9 @@ def _concat_dense( **_layout_like(inputs.get(0, path), dst_parent, shape, dtype)) spec.set_encoding(out, spec.ARRAY) - row_offset = 0 + # Whole shards (or chunks) at a time: a block of --chunk rows into the + # first input's shard rewrote it once a block (F036). + writer = RowWriter(out) for i, col_map in enumerate(col_maps): source = inputs.get(i, path) keep = col_map >= 0 @@ -954,9 +961,8 @@ def _concat_dense( block = (np.full((end - start, n_cols), fill_value, dtype=dtype) if gaps else np.empty((end - start, n_cols), dtype=dtype)) block[:, targets] = np.asarray(source[start:end, :])[:, keep] - out[row_offset + start : row_offset + end, :] = block - - row_offset += n_src_rows + writer.write(block) + writer.flush() @contextmanager @@ -1141,6 +1147,8 @@ def _concat_csc( _append(out_indices, rows.astype(index_dtype, copy=False)) _append(out_data, values) + out_indices.flush() + out_data.flush() _write_indptr(group, indptr, pointer_dtype, inputs.get(0, path)["indptr"]) @@ -1312,12 +1320,12 @@ def _concat_obsm( **_layout_like(inputs.get(0, path), dst_parent, shape, dtype)) spec.set_encoding(out, spec.ARRAY) - offset = 0 + writer = RowWriter(out) for source in inputs.each(path): for start in range(0, source.shape[0], chunk_rows): end = min(start + chunk_rows, source.shape[0]) - out[offset + start : offset + end, ...] = source[start:end, ...] - offset += source.shape[0] + writer.write(source[start:end, ...]) + writer.flush() return True @@ -1416,8 +1424,10 @@ def _write_varm_array(parent: Any, name: str, obj: Any, take: np.ndarray, block_ out = create_dataset(parent, name, shape=shape, dtype=obj.dtype, **_layout_like(obj, parent, shape, obj.dtype)) spec.set_encoding(out, spec.STRING_ARRAY if text else spec.ARRAY) + writer = RowWriter(out) for start in range(0, len(take), block_rows): - out[start:start + block_rows, ...] = aligned_rows(obj, take[start:start + block_rows]) + writer.write(aligned_rows(obj, take[start:start + block_rows])) + writer.flush() def _entries_on(obj: Any, col_map: np.ndarray, entries: int) -> Iterator[Tuple[np.ndarray, np.ndarray, np.ndarray]]: @@ -1503,8 +1513,10 @@ def _write_pairwise( out = create_dataset(parent, name, shape=(n, n), dtype=obj.dtype, **_layout_like(obj, parent, (n, n), obj.dtype)) spec.set_encoding(out, spec.STRING_ARRAY if is_string_dtype(obj.dtype) else spec.ARRAY) + writer = RowWriter(out) for start in range(0, n, block_rows): - out[start:start + block_rows, :] = aligned_rows(obj, take[start:start + block_rows])[:, take] + writer.write(aligned_rows(obj, take[start:start + block_rows])[:, take]) + writer.flush() return enc = spec.encoding_type(obj) if enc not in spec.SPARSE_TYPES: diff --git a/src/adata/core/csr_build.py b/src/adata/core/csr_build.py index b14e320..8ab92dd 100644 --- a/src/adata/core/csr_build.py +++ b/src/adata/core/csr_build.py @@ -333,6 +333,8 @@ def append(self, rows: np.ndarray, cols: np.ndarray, values: np.ndarray) -> None def finish(self) -> int: from adata.core.subset import _sparse_dataset + self.data.flush() + self.indices.flush() dtype = np.int32 if self.nnz <= INT32_MAX else np.int64 template = self.templates[2] if template is not None and np.dtype(template.dtype).itemsize > np.dtype(dtype).itemsize: diff --git a/src/adata/core/recast.py b/src/adata/core/recast.py index c34c19a..92f3ccc 100644 --- a/src/adata/core/recast.py +++ b/src/adata/core/recast.py @@ -40,6 +40,7 @@ ) from adata.core.subset import _sparse_dataset from adata.storage import ( + RowWriter, copy_attrs, create_dataset, dataset_create_kwargs, @@ -394,17 +395,22 @@ def cast_sparse( source_data, source_indices = src.obj["data"], src.obj["indices"] nnz = int(source_data.shape[0]) - out_data = _sparse_dataset(group, "data", data_dtype, nnz, source_data) - out_indices = _sparse_dataset(group, "indices", index_dtype, nnz, source_indices) + # Written whole shards at a time: a block of `chunk` entries into the + # source's shard, a whole array under anndata 0.13, rewrote it once a + # block (F036). + out_data = RowWriter(_sparse_dataset(group, "data", data_dtype, nnz, source_data)) + out_indices = RowWriter( + _sparse_dataset(group, "indices", index_dtype, nnz, source_indices) + ) for start in range(0, nnz, chunk): end = min(start + chunk, nnz) - out_data[start:end] = np.asarray(source_data[start:end]).astype( - data_dtype, copy=False - ) - out_indices[start:end] = np.asarray(source_indices[start:end]).astype( - index_dtype, copy=False + out_data.write(np.asarray(source_data[start:end]).astype(data_dtype, copy=False)) + out_indices.write( + np.asarray(source_indices[start:end]).astype(index_dtype, copy=False) ) + out_data.flush() + out_indices.flush() indptr = np.asarray(src.obj["indptr"][...]).astype(pointer_dtype, copy=False) out_indptr = _sparse_dataset( @@ -527,14 +533,13 @@ 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 = _output_array( + out_data = RowWriter(_output_array( group, "data", data_dtype, nnz, templates[0], codecs_from, compression=compression - ) - out_indices = _output_array( + )) + out_indices = RowWriter(_output_array( group, "indices", index_dtype, nnz, templates[1], codecs_from, compression=compression, - ) - written = 0 + )) with scratch_group(dst_parent, f"__{name}_transpose_scratch__") as scratch: sorter = BucketSorter(scratch, bounds, src.dtype) @@ -552,16 +557,13 @@ def transpose_sparse_streaming( # The input's minor axis is the output's major one. sorter.add(minor, _majors_of(indptr, start, end), values) - # Pass 3: each bucket in turn, in output order. + # Pass 3: each bucket in turn, in output order, written a whole + # shard at a time: a bucket at a time rewrote the shard (F036). 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 - ) - out_data[written : written + count] = values.astype( - data_dtype, copy=False - ) - written += count + out_indices.write(new_minor.astype(index_dtype, copy=False)) + out_data.write(values.astype(data_dtype, copy=False)) + out_indices.flush() + out_data.flush() cast_indptr = out_indptr.astype(pointer_dtype, copy=False) dataset = _output_array( @@ -808,12 +810,14 @@ def _sparsify_rows( spec.set_encoding(group, spec.CSR_MATRIX) set_shape_attr(group, src.shape) - out_data = _growable_from( + # Appended to whole chunks at a time: Zarr keeps no cache, so a block + # appended to part of a chunk decoded and encoded it again. + out_data = RowWriter(_growable_from( group, "data", data_dtype, template, compression=compression - ) - out_indices = _growable_from( + ), grow=True) + out_indices = RowWriter(_growable_from( group, "indices", index_dtype, template, compression=compression - ) + ), grow=True) # 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) @@ -837,6 +841,8 @@ def _sparsify_rows( _append(out_indices, cols.astype(index_dtype, copy=False)) _append(out_data, block[rows, cols]) + out_indices.flush() + out_data.flush() indptr = np.concatenate(([0], np.cumsum(counts))) kw = {} if compression is not None: @@ -934,15 +940,18 @@ def cast_dense( name, shape=(n_rows, n_cols), dtype=data_dtype, - **fit_chunks(kw, (n_rows, n_cols), backend=backend), + **fit_chunks(kw, (n_rows, n_cols), backend=backend, + itemsize=np.dtype(data_dtype).itemsize), ) copy_attrs(src.obj.attrs, dst.attrs, target_backend=backend) + # A whole shard at a time: a block of --chunk rows rewrote the source's + # shard once a block, 537 MB encoded for 4.4 MB written (F036). + writer = RowWriter(dst) for lo in range(0, n_rows, chunk_rows): hi = min(lo + chunk_rows, n_rows) - dst[lo:hi, :] = np.asarray(src.obj[lo:hi, :]).astype( - data_dtype, copy=False - ) + writer.write(np.asarray(src.obj[lo:hi, :]).astype(data_dtype, copy=False)) + writer.flush() # --------------------------------------------------------------------------- diff --git a/src/adata/core/subset.py b/src/adata/core/subset.py index e604c94..20d1f5b 100644 --- a/src/adata/core/subset.py +++ b/src/adata/core/subset.py @@ -6,7 +6,7 @@ from contextlib import ExitStack from dataclasses import dataclass from pathlib import Path -from typing import Optional, Set, Tuple, List, Dict, Any +from typing import Optional, Set, Tuple, List, Dict, Any, Iterator import numpy as np from rich.console import Console @@ -30,13 +30,21 @@ resolve_index, ) from adata.storage import ( + TARGET_READ_BYTES, + WRITE_UNIT_BYTES, + RowWriter, create_dataset, copy_attrs, copy_tree, + cut_unit, dataset_create_kwargs, + decode_chunks, detect_backend, fit_chunks, fit_fill, + read_blocks, + tile_shards, + write_share, is_dataset, is_group, is_legacy_categorical, @@ -390,25 +398,13 @@ def _check_sorted(indices: Optional[np.ndarray], what: str) -> None: raise ValueError(f"{what} indices must be sorted and unique.") -class _Cursor: - """Hand out, block by block, the rows of one target's major selection. - - The blocks arrive in increasing order and the selection is sorted, so - each block's rows are the next slice of it: finding them is one - `searchsorted`, not a scan. - """ - - def __init__(self, selection: Optional[np.ndarray]) -> None: - self.selection = selection - self.pos = 0 - - def rows(self, lo: int, hi: int) -> np.ndarray: - if self.selection is None: - return np.arange(lo, hi, dtype=np.int64) - end = int(np.searchsorted(self.selection, hi, side="left")) - rows = self.selection[self.pos : end] - self.pos = end - return rows +def _selected(selection: Optional[np.ndarray], lo: int, hi: int) -> Tuple[int, np.ndarray]: + """The rows of a sorted selection in `[lo, hi)`, and where the first of + them lands in the output: two `searchsorted`s, not a scan.""" + if selection is None: + return lo, np.arange(lo, hi, dtype=np.int64) + a, b = np.searchsorted(selection, [lo, hi], side="left") + return int(a), selection[int(a) : int(b)] def subset_dense_matrix( @@ -423,12 +419,21 @@ def subset_dense_matrix( fan_out_dense_matrix(src, name, [(dst_parent, obs_idx, var_idx)], chunk_rows=chunk_rows) +def _itemsize(dtype: Any) -> int: + """Bytes an element holds in memory, a guess for text.""" + from adata.storage import VLEN_ELEMENT_BYTES + + size = int(getattr(dtype, "itemsize", 0) or 0) + return size if size and getattr(dtype, "kind", "") not in ("O", "T") else VLEN_ELEMENT_BYTES + + def fan_out_dense_matrix( src: Any, name: str, targets: List[MatrixTarget], *, chunk_rows: int = 1024, + share: Optional[int] = None, ) -> None: """Write one row/column subset of a dense array per target, in one pass. @@ -437,9 +442,20 @@ def fan_out_dense_matrix( axes whole, as anndata's subsetting keeps them. Only a 0-d value has no rows, and is copied as it is. - The source is read a block of `chunk_rows` rows at a time, each block - once whatever the number of targets, and only the span between the - first and last row any target keeps in it. + The source is read in bands of `chunk_rows` rows rounded to whole + chunks of it (`read_blocks`), each once whatever the number of targets, + and only the span between the first and last row any target keeps in + it. A band of --chunk rows split a chunk 50,000 rows tall across fifty + reads, each decoding all of it (F061). Where one chunk height across + the array is over the read budget -- a chunk of (n_obs, 8) -- a band is + read a tile of whole chunks at a time. + + Each output keeps the source's chunks and shards, and is written a + whole unit at a time (`RowWriter`): a block of rows wrote part of a + shard, which zarr rewrites whole, once a block (F031), and `split`'s + outputs, given a few rows each a block, decoded and encoded each chunk + tens of times (F060). A unit's rows are cut to `share` bytes (by + default `WRITE_UNIT_BYTES`), the most each output may hold. """ if src.ndim == 0: for dst_parent, _, _ in targets: @@ -447,6 +463,17 @@ def fan_out_dense_matrix( return n_obs = src.shape[0] + itemsize = _itemsize(src.dtype) + row_bytes = itemsize * int(np.prod(src.shape[1:], dtype=np.int64)) + band, width = read_blocks( + src.shape, decode_chunks(src), itemsize, rows=chunk_rows, + budget=max(chunk_rows * row_bytes, TARGET_READ_BYTES), + ) + tiles = ( + [(0, src.shape[1] if src.ndim > 1 else 0)] if width is None + else [(left, min(left + width, src.shape[1])) for left in range(0, src.shape[1], width)] + ) + outputs = [] for dst_parent, obs_idx, var_idx in targets: _check_sorted(obs_idx, "obs") @@ -462,32 +489,59 @@ def fan_out_dense_matrix( kw = dataset_create_kwargs( src, target_backend=target_backend, dst_parent=dst_parent ) - kw = fit_chunks(kw, shape, backend=target_backend) + budget = share if share is not None else WRITE_UNIT_BYTES + kw = fit_chunks(kw, shape, backend=target_backend, itemsize=itemsize, unit_bytes=budget) + if width is not None and kw.get("shards"): + kw["shards"] = tile_shards(kw["chunks"], kw["shards"], width, rows=False) dst = create_dataset( dst_parent, name, shape=tuple(shape), dtype=src.dtype, **kw ) copy_attrs(src.attrs, dst.attrs, target_backend=target_backend) - outputs.append([dst, _Cursor(obs_idx), var_idx, 0]) + # A writer per tile of columns, each buffering its columns' rows to + # whole units: together they hold one unit of rows across the output. + writers = [] + for left, right in tiles: + if width is None: + writers.append(((None, var_idx), RowWriter(dst, budget=budget))) + continue + p0, p1 = ( + (left, right) if var_idx is None + else np.searchsorted(var_idx, [left, right], side="left") + ) + if p1 > p0: + picked = None if var_idx is None else var_idx[int(p0):int(p1)] - left + writers.append(((left, picked), RowWriter(dst, cols=(p0, p1), budget=budget))) + else: + writers.append(None) + outputs.append((obs_idx, writers)) - for start in range(0, n_obs, chunk_rows): - end = min(start + chunk_rows, n_obs) - wanted = [out[1].rows(start, end) for out in outputs] + for start in range(0, n_obs, band): + end = min(start + band, n_obs) + wanted = [_selected(obs_idx, start, end)[1] for obs_idx, _ in outputs] nonempty = [rows for rows in wanted if rows.size] if not nonempty: continue first = min(int(rows[0]) for rows in nonempty) last = max(int(rows[-1]) for rows in nonempty) - block = np.asarray(src[first : last + 1, ...]) - for out, rows in zip(outputs, wanted): - if not rows.size: - continue - dst, _, var_idx, written = out - part = block[rows - first] - if var_idx is not None: - part = part[:, var_idx] - dst[written : written + len(rows), ...] = part - out[3] = written + len(rows) + for t, (left, right) in enumerate(tiles): + block = np.asarray( + src[first : last + 1, ...] if width is None + else src[first : last + 1, left:right, ...] + ) + for (_, writers), rows in zip(outputs, wanted): + if not rows.size or writers[t] is None: + continue + (_, picked), writer = writers[t] + part = block[rows - first] + if picked is not None: + part = part[:, picked] + writer.write(part) + + for _, writers in outputs: + for entry in writers: + if entry is not None: + entry[1].flush() class _SparseOutput: @@ -502,6 +556,8 @@ def __init__( out_shape: Tuple[int, int], major_idx: Optional[np.ndarray], remap: Optional[np.ndarray], + n_major: int, + share: Optional[int] = None, ) -> None: self.group = dst_parent.create_group(name) copy_attrs(src.attrs, self.group.attrs, target_backend=_target_backend(dst_parent)) @@ -513,29 +569,43 @@ def __init__( # so whatever held the source holds the subset. Widening them made # every output about half as large again as it needed to be. self.indptr_template = src["indptr"] - self.data = _growable_like(self.group, "data", src["data"].dtype, src["data"]) + data_dtype, index_dtype = np.dtype(src["data"].dtype), np.dtype(src["indices"].dtype) + share = WRITE_UNIT_BYTES if share is None else share + # The two arrays share the output's budget in proportion to their + # widths: whole units of both, at once, are what it holds. + width = data_dtype.itemsize + index_dtype.itemsize + self.data = _growable_like( + self.group, "data", data_dtype, src["data"], + share=max(1, share * data_dtype.itemsize // width), + ) self.indices = _growable_like( - self.group, "indices", src["indices"].dtype, src["indices"] + self.group, "indices", index_dtype, src["indices"], + share=max(1, share * index_dtype.itemsize // width), ) - self.cursor = _Cursor(major_idx) + self.selection = major_idx self.remap = remap - self.counts: List[np.ndarray] = [] + # Entries per output major, added to as blocks arrive: a major split + # across blocks adds to its count twice. + self.counts = np.zeros(len(major_idx) if major_idx is not None else n_major, np.int64) def take( self, + at: int, rows: np.ndarray, indptr: np.ndarray, lo: int, + hi: int, block_indices: np.ndarray, block_data: np.ndarray, ) -> None: - """Append the entries of `rows`, whose data starts at offset `lo`.""" - starts = indptr[rows] - lo - lengths = indptr[rows + 1] - indptr[rows] + """Append the entries of `rows` in `[lo, hi)`, the span the block + holds from `lo`; `at` is the first row's position in the output.""" + starts = np.maximum(indptr[rows], lo) - lo + lengths = np.minimum(indptr[rows + 1], hi) - lo - starts total = int(lengths.sum()) - if total and self.cursor.selection is None: - # Every row of the block: its entries are one contiguous run. + if total and (self.selection is None or len(rows) == 1): + # Every row of the block, or one: its entries are one run. begin = int(starts[0]) positions = slice(begin, begin + total) else: @@ -556,15 +626,14 @@ def take( counts = np.bincount(owner[keep], minlength=len(rows)) minor, values = mapped[keep], values[keep] - self.counts.append(counts) + self.counts[at : at + len(rows)] += counts _append(self.indices, minor.astype(self.indices.dtype, copy=False)) _append(self.data, values) def finish(self) -> None: - counts = ( - np.concatenate(self.counts) if self.counts else np.empty(0, dtype=np.int64) - ) - indptr = np.concatenate(([0], np.cumsum(counts, dtype=np.int64))) + self.data.flush() + self.indices.flush() + indptr = np.concatenate(([0], np.cumsum(self.counts, dtype=np.int64))) dtype = self.indptr_template.dtype dataset = _sparse_dataset( self.group, "indptr", dtype, indptr.size, self.indptr_template @@ -579,32 +648,103 @@ def subset_sparse_matrix_group( obs_idx: Optional[np.ndarray], var_idx: Optional[np.ndarray], *, - chunk_major: int = 4096, + chunk_rows: int = 1024, ) -> None: - """Subset a CSR/CSC matrix, streaming it a block of major axis at a time.""" + """Subset a CSR/CSC matrix, streaming it a block of entries at a time.""" fan_out_sparse_matrix( - src, name, [(dst_parent, obs_idx, var_idx)], chunk_major=chunk_major + src, name, [(dst_parent, obs_idx, var_idx)], chunk_rows=chunk_rows ) +def sparse_block_entries(n_rows: int, nnz: int, chunk_rows: int) -> int: + """Entries a block of a sparse matrix holds: as many as `chunk_rows` + rows hold on average, in either layout. + + Blocks were 4,096 majors, whatever --chunk said. Of a CSC matrix that is + 4,096 genes across every cell, a block growing with n_obs: 270M entries, + several GB, at a million cells (F032). + """ + per_row = -(-int(nnz) // max(1, int(n_rows))) if nnz else 1 + return max(1, int(chunk_rows) * max(1, per_row)) + + +def _sparse_blocks( + indptr: np.ndarray, + budget: int, + selections: List[Optional[np.ndarray]], +) -> Iterator[Tuple[int, int, int, int]]: + """`(first, stop, lo, hi)` for each block: majors `[first, stop)`, of + which entries `[lo, hi)` are read. + + A block takes whole majors while they fit `budget` entries; a major + holding more is read in parts of `budget`, a block each (then `stop` + is `first + 1`). Majors no selection keeps are skipped, and a block is + read only from the first to the last major some selection keeps in it. + """ + n_major = indptr.size - 1 + whole = any(sel is None for sel in selections) + # Every major some selection keeps, once: one search a block, not one + # for each of `split`'s outputs. + kept = None if whole else np.unique(np.concatenate( + [np.asarray(sel, dtype=np.int64) for sel in selections] or [np.empty(0, np.int64)] + )) + + def next_kept(m: int) -> Optional[int]: + if whole: + return m if m < n_major else None + i = int(np.searchsorted(kept, m, side="left")) + return int(kept[i]) if i < kept.size else None + + def last_kept(stop: int) -> int: + if whole: + return stop - 1 + return int(kept[int(np.searchsorted(kept, stop, side="left")) - 1]) + + m = next_kept(0) + at = None # within major m, when it is being read in parts + while m is not None and m < n_major: + start = int(indptr[m]) if at is None else at + if at is None: + stop = int(np.searchsorted(indptr, start + budget, side="right")) - 1 + stop = min(stop, n_major) + if stop > m: + hi = int(indptr[last_kept(stop) + 1]) + yield m, stop, start, hi + m = next_kept(stop) + continue + # Major m holds more than a block: a part of it. + end = min(start + budget, int(indptr[m + 1])) + yield m, m + 1, start, end + if end >= int(indptr[m + 1]): + m, at = next_kept(m + 1), None + else: + at = end + + def fan_out_sparse_matrix( src: Any, name: str, targets: List[MatrixTarget], *, - chunk_major: int = 4096, + chunk_rows: int = 1024, + share: Optional[int] = None, ) -> None: """Write one subset of a CSR/CSC matrix per target, reading it once. - The source's major axis is walked in contiguous blocks of `chunk_major`, - and each block's `data`/`indices` are read once -- only the span between - the first and last row any target keeps -- then shared out among the - targets. Peak memory is set by the block, not the matrix. + The source is read a block of entries at a time -- about as many as + `chunk_rows` rows hold (`sparse_block_entries`), whole majors where + they fit and a major too large for one block in parts -- and each block + is read once, only from the first to the last major any target keeps, + then shared out among the targets. Peak memory is set by the block, + not the matrix, in either layout and along either axis. This is what makes `split` one pass. An earlier version subset once per target, reading the span of each target's own selected rows; when the groups were interleaved that span was nearly the whole matrix, so a split into k groups read and decompressed X k times. + + Each output's `data` and `indices` keep the source's chunks and shards, + and are written a whole unit at a time (`RowWriter`), cut to `share`. """ enc = _decode_attr(src.attrs.get("encoding-type", b"")) if enc not in spec.SPARSE_TYPES: @@ -641,17 +781,15 @@ def fan_out_sparse_matrix( out_shape, major_idx, _minor_remap(minor_keep, n_minor), + n_major, + share, ) ) - for start in range(0, n_major, chunk_major): - end = min(start + chunk_major, n_major) - wanted = [out.cursor.rows(start, end) for out in outputs] - nonempty = [rows for rows in wanted if rows.size] - if not nonempty: - continue - lo = int(indptr[min(int(rows[0]) for rows in nonempty)]) - hi = int(indptr[max(int(rows[-1]) for rows in nonempty) + 1]) + budget = sparse_block_entries(n_rows, int(indptr[-1]) if indptr.size else 0, chunk_rows) + for first, stop, lo, hi in _sparse_blocks( + indptr, budget, [out.selection for out in outputs] + ): if hi > lo: block_indices = np.asarray(indices_ds[lo:hi]) block_data = np.asarray(data_ds[lo:hi]) @@ -659,9 +797,10 @@ def fan_out_sparse_matrix( block_indices = np.empty(0, dtype=indices_ds.dtype) block_data = np.empty(0, dtype=data_ds.dtype) - for out, rows in zip(outputs, wanted): + for out in outputs: + at, rows = _selected(out.selection, first, stop) if rows.size: - out.take(rows, indptr, lo, block_indices, block_data) + out.take(at, rows, indptr, lo, hi, block_indices, block_data) for out in outputs: out.finish() @@ -699,7 +838,9 @@ def _sparse_dataset( """ backend = _target_backend(group) kw = dataset_create_kwargs(template, target_backend=backend, dst_parent=group) - kw = fit_chunks(kw, (n,), backend=backend) + # A unit cut to what `RowWriter` holds: writers fill these a block at + # a time, and anndata's shard of a whole array may be a gigabyte. + kw = fit_chunks(kw, (n,), backend=backend, itemsize=np.dtype(dtype).itemsize) if compression is not None: kw = with_compression( kw, compression, shape=(n,), dtype=dtype, backend=backend, @@ -717,12 +858,18 @@ def _growable_like( template: Any, *, compression: Optional[str] = None, -) -> Any: - """Like `_sparse_dataset`, but extensible for a size not yet known. + share: Optional[int] = None, +) -> RowWriter: + """Like `_sparse_dataset`, but extensible for a size not yet known: a + `RowWriter` that grows it, appended to with `_append`. Flush it once the + last entries are in. The sparse writers of subset and concat once used a bare growable dataset here, which dropped the source's compression: their `data` and - `indices` were written uncompressed whatever the input. + `indices` were written uncompressed whatever the input. They dropped its + shard too, and kept the inner chunk: 245 elements, a file each, under + anndata 0.13 (F033). The shard is kept now, cut to `share` bytes (by + default `WRITE_UNIT_BYTES`), and written whole. """ backend = _target_backend(group) kw = dataset_create_kwargs(template, target_backend=backend, dst_parent=group) @@ -731,20 +878,28 @@ def _growable_like( kw, compression, shape=(0,), dtype=dtype, backend=backend, zarr_format=zarr_format_of(group), ) - kw.pop("shards", None) chunks = kw.pop("chunks", None) - step = int(chunks[0]) if chunks else 1 << 16 + kw["chunks"] = (int(chunks[0]),) if chunks else (1 << 16,) + if kw.get("shards"): + kw["shards"] = (int(kw["shards"][0]),) + budget = WRITE_UNIT_BYTES if share is None else share + kw = cut_unit(kw, (0,), np.dtype(dtype).itemsize, budget) if is_zarr_group(group): - return group.create_array(name, shape=(0,), dtype=dtype, chunks=(step,), **kw) - return group.create_dataset( - name, shape=(0,), maxshape=(None,), dtype=dtype, chunks=(step,), **fit_fill(kw, dtype) - ) + ds = group.create_array(name, shape=(0,), dtype=dtype, **kw) + else: + ds = group.create_dataset( + name, shape=(0,), maxshape=(None,), dtype=dtype, **fit_fill(kw, dtype) + ) + return RowWriter(ds, grow=True, budget=budget) def _append(ds: Any, values: np.ndarray) -> None: - """Append a block to a growable 1-D dataset.""" + """Append a block to a growable 1-D dataset, or to its `RowWriter`.""" if values.size == 0: return + if isinstance(ds, RowWriter): + ds.write(values) + return start = ds.shape[0] ds.resize((start + values.size,)) ds[start:] = values @@ -776,15 +931,16 @@ def fan_out_matrix_entry( *, chunk_rows: int, entry_label: str, + share: Optional[int] = None, ) -> None: if is_dataset(obj): - fan_out_dense_matrix(obj, name, targets, chunk_rows=chunk_rows) + fan_out_dense_matrix(obj, name, targets, chunk_rows=chunk_rows, share=share) return if is_group(obj): enc = spec.encoding_type(obj) if enc in spec.SPARSE_TYPES: - fan_out_sparse_matrix(obj, name, targets) + fan_out_sparse_matrix(obj, name, targets, chunk_rows=chunk_rows, share=share) return if enc == spec.DATAFRAME: # obsm/varm may hold a dataframe; it is row-aligned like obs/var. @@ -838,6 +994,7 @@ def _fan_out_raw( *, chunk_rows: int, console: Console, + share: Optional[int] = None, ) -> None: """Subset a `raw/` group into every target, reading its matrices once. @@ -885,6 +1042,7 @@ def _fan_out_raw( ], chunk_rows=chunk_rows, entry_label="raw/X", + share=share, ) if "varm" in src_raw: @@ -899,6 +1057,7 @@ def _fan_out_raw( ], chunk_rows=chunk_rows, entry_label=f"raw/varm:{key}", + share=share, ) for key in src_raw.keys(): @@ -980,6 +1139,9 @@ def _write_targets( labels = {v: k for k, v in _TASK_LABELS.items()} pruned: List[Dict[str, Pruned]] = [{} for _ in targets] + # What each output may hold to write whole units, its part of what + # one pass holds across them all (F079). + share = write_share(len(targets)) tasks = [t for t in tasks if labels.get(t.split(":", 1)[0], t.split(":", 1)[0]) not in skip and not (t.startswith("copy:") and t.split(":", 1)[1] in skip)] @@ -1004,9 +1166,9 @@ def _write_targets( X = src["X"] pairs = [(t.root, t.obs_idx, t.var_idx) for t in targets] if is_dataset(X): - fan_out_dense_matrix(X, "X", pairs, chunk_rows=chunk_rows) + fan_out_dense_matrix(X, "X", pairs, chunk_rows=chunk_rows, share=share) elif is_group(X): - fan_out_sparse_matrix(X, "X", pairs) + fan_out_sparse_matrix(X, "X", pairs, chunk_rows=chunk_rows, share=share) else: for target in targets: copy_tree(X, target.root, "X") @@ -1017,7 +1179,8 @@ def _write_targets( prune_colors(target.root["uns"], narrowed.get(axis, {})) elif task == "raw": _fan_out_raw( - src["raw"], targets, chunk_rows=chunk_rows, console=console + src["raw"], targets, chunk_rows=chunk_rows, console=console, + share=share, ) elif task.startswith("copy:"): key = task.split(":", 1)[1] @@ -1054,6 +1217,7 @@ def _write_targets( ], chunk_rows=chunk_rows, entry_label=task, + share=share, ) progress.update( task_id, @@ -1301,6 +1465,15 @@ def _open_output_limit() -> int: #: rather than running out of file descriptors. MAX_OPEN_OUTPUTS = _open_output_limit() +#: How many HDF5 outputs a pass holds open, whatever the descriptor limit +#: allows. An open HDF5 file costs about 0.6 MB however it is opened -- the +#: library's own per-file structures; a chunk cache or a smaller metadata +#: cache changes nothing -- so with a container's high limit, 1,024 outputs +#: a pass took RSS to about 1 GB, following the number of groups rather than +#: --chunk (F079). 128 holds that near 80 MB, for a pass more per 128 +#: groups. Zarr outputs hold no descriptor and no such cost. +MAX_OPEN_HDF5_OUTPUTS = 128 + def split_h5ad( file: Path, @@ -1320,14 +1493,21 @@ def split_h5ad( with open_store(file, "r") as probe: zarr_format = probe.zarr_format - for start in range(0, len(outputs), MAX_OPEN_OUTPUTS): - batch = outputs[start : start + MAX_OPEN_OUTPUTS] + per_pass = MAX_OPEN_OUTPUTS + if any(detect_backend(Path(path)) == "hdf5" for path, _, _ in outputs): + per_pass = min(per_pass, MAX_OPEN_HDF5_OUTPUTS) + for start in range(0, len(outputs), per_pass): + batch = outputs[start : start + per_pass] with ExitStack() as stack: src = stack.enter_context(open_store(file, "r")).root targets = [] for path, obs_idx, var_idx in batch: + # No chunk cache: every chunked array is written whole units + # at a time (`RowWriter`), and an HDF5 cache for each dataset + # of each of up to 1,024 open outputs, holding the chunks it + # was given until closed, took RSS to about 1 GB (F079). dst = stack.enter_context( - open_store(path, "w", zarr_format=zarr_format) + open_store(path, "w", zarr_format=zarr_format, chunk_cache=0) ).root var_keep = ( set(read_names(src, "var", var_idx)) diff --git a/src/adata/storage/__init__.py b/src/adata/storage/__init__.py index 21e94b3..31df0cb 100644 --- a/src/adata/storage/__init__.py +++ b/src/adata/storage/__init__.py @@ -125,12 +125,91 @@ def detect_backend(path: Path) -> str: #: followed the number of outputs, eight times faster. 1 MiB, as before. HDF5_CHUNK_CACHE = 1 << 20 +#: Bytes of Zarr shard indexes a store opened for reading keeps (see +#: `_shard_index_cache`): 16 bytes an inner chunk, so the index of a 1 GB +#: shard of 1 KB chunks, as anndata 0.13 wrote with zarr 3.1, is 16 MB. +SHARD_INDEX_CACHE_BYTES = 64 * 1024 * 1024 + + +def _shard_index_cache(store: Any) -> Any: + """`store`, keeping the shard indexes read through it (F062). + + zarr-python reads a shard's index again for every read of part of the + shard: 16 bytes an inner chunk, about 1 MB for a 64 MB shard of anndata + 0.13's 1 KB chunks. A command streaming such an array a block at a + time read the whole index once a block, so the bytes read grew with + the square of the array: subset of an anndata-default store read 4.2 + times the store, concat 27.7 times. A store opened for reading cannot + change underneath, so each index is read once and kept, least recently + used first out, up to `SHARD_INDEX_CACHE_BYTES`. + + The index is the one read zarr makes by suffix (its default + `index_location`, and anndata's). Inner chunks are read by range, and + passed through. + """ + from collections import OrderedDict + + try: + from zarr.abc.store import SuffixByteRequest + from zarr.storage import WrapperStore + except ImportError: # pragma: no cover - a zarr without either + return store + + class ShardIndexCache(WrapperStore): + def __init__(self, inner: Any) -> None: + super().__init__(inner) + self._held: "OrderedDict[tuple, Any]" = OrderedDict() + self._held_bytes = 0 + + def _with_store(self, inner: Any) -> Any: + return type(self)(inner) + + def _lookup(self, key: str, byte_range: Any) -> tuple: + if not isinstance(byte_range, SuffixByteRequest): + return None, None + slot = (key, int(byte_range.suffix)) + if slot in self._held: + self._held.move_to_end(slot) + return slot, self._held[slot] + return slot, None + + def _keep(self, slot: tuple, value: Any) -> None: + size = len(value) + if size > SHARD_INDEX_CACHE_BYTES: + return + self._held[slot] = value + self._held_bytes += size + while self._held_bytes > SHARD_INDEX_CACHE_BYTES: + _, dropped = self._held.popitem(last=False) + self._held_bytes -= len(dropped) + + async def get(self, key: str, prototype: Any, byte_range: Any = None) -> Any: + slot, value = self._lookup(key, byte_range) + if value is not None: + return value + value = await self._store.get(key, prototype, byte_range) + if slot is not None and value is not None: + self._keep(slot, value) + return value + + def get_sync(self, key: str, *, prototype: Any = None, byte_range: Any = None) -> Any: + slot, value = self._lookup(key, byte_range) + if value is not None: + return value + value = super().get_sync(key, prototype=prototype, byte_range=byte_range) + if slot is not None and value is not None: + self._keep(slot, value) + return value + + return ShardIndexCache(store) + def open_store( path: Path, mode: str, zarr_format: Optional[int] = None, require_anndata: bool = True, + chunk_cache: int = HDF5_CHUNK_CACHE, ) -> Store: """Open a store, auto-detecting the backend from `path`. @@ -142,6 +221,10 @@ def open_store( Set `require_anndata=False` for format-agnostic commands such as `ls`, which are expected to open plain HDF5 and Zarr stores and should not warn about a missing AnnData root. + + `chunk_cache` is the bytes of chunk cache an HDF5 store keeps for each + dataset; a writer holding many outputs at once, which writes whole + chunks through `RowWriter`, needs none. """ path = Path(path) backend = detect_backend(path) @@ -157,8 +240,15 @@ def open_store( # snapshot, so members added through it are invisible even to the # handle that created them. kwargs["use_consolidated"] = False + target: Any = str(path) + elif path.is_dir(): + from zarr.storage import LocalStore + + target = _shard_index_cache(LocalStore(str(path), read_only=True)) + else: + target = str(path) - root = zarr.open_group(str(path), mode=mode, **kwargs) + root = zarr.open_group(target, mode=mode, **kwargs) if writable: ensure_anndata_root_attrs(root) elif require_anndata: @@ -170,7 +260,7 @@ def open_store( zarr_format=zarr_format_of(root), consolidate=writable, ) - root = h5py.File(path, mode, rdcc_nbytes=HDF5_CHUNK_CACHE) + root = h5py.File(path, mode, rdcc_nbytes=chunk_cache) if mode not in ("w", "w-", "x"): # Anything opened for its contents, read or appended to (--inplace). try: @@ -493,7 +583,12 @@ def dataset_create_kwargs( kw_target = zarr_format_of(dst_parent) kw: dict = {} chunks = getattr(src, "chunks", None) - if chunks is not None: + shards = getattr(src, "shards", None) if is_zarr_array(src) else None + if chunks is not None and not (shards and target_backend == "hdf5"): + # A shard's inner chunk is a unit of compression inside a file, a + # kilobyte under anndata 0.13 on zarr 3.1: as the chunk of an HDF5 + # dataset it made 16,327 chunks where anndata's own h5ad has 512, + # and a file twice the size (F033). HDF5 picks its own instead. kw["chunks"] = chunks if target_backend == "hdf5" and is_hdf5_dataset(src): if src.compression is not None: @@ -550,9 +645,16 @@ def dataset_create_kwargs( if filters and not _is_string_src(src): kw["filters"] = filters - shards = src.shards if shards is not None and target_format == 3: kw["shards"] = shards + elif shards is not None and target_format == 2: + # v2 has no shards: the shard itself is the chunk, as F033 asks, + # rather than thousands of inner-chunk files. Cut along rows to + # a write unit, as a shard would be (`fit_chunks`). + kw["chunks"] = cut_unit( + {"chunks": tuple(shards)}, tuple(shards), + int(src.dtype.itemsize) or VLEN_ELEMENT_BYTES, WRITE_UNIT_BYTES, + )["chunks"] fill_value = src.fill_value # A record's fill is a value of its fields, which change type when @@ -729,7 +831,71 @@ def _zarr_is_compressed(src: Any) -> bool: return bool(src.compressors) -def fit_chunks(kw: dict, shape: Sequence[int], *, backend: str) -> dict: +#: The most one output holds to write a whole unit -- a shard, or a chunk +#: that a write of part of would decode and encode again -- at once +#: (`RowWriter`). A unit whose rows across the array's width come to more is +#: cut along its rows when the output is created (`fit_chunks`): anndata +#: 0.13 aims at 1 GB shards, which a writer filling them from blocks of +#: --chunk rows would otherwise hold, or rewrite once a block. +WRITE_UNIT_BYTES = 64 * 1024 * 1024 + +#: The most one pass writing several outputs at once (`split`) holds to +#: write whole units, across all of them: each output's units are cut to a +#: share of it (`write_share`). +WRITE_BUFFER_BYTES = 256 * 1024 * 1024 + + +def write_share(outputs: int) -> int: + """Bytes each of `outputs` written in one pass may hold for a unit.""" + return max(1, min(WRITE_UNIT_BYTES, WRITE_BUFFER_BYTES // max(1, outputs))) + + +def _prod(values: Sequence[int]) -> int: + total = 1 + for v in values: + total *= max(1, int(v)) + return total + + +def cut_unit(kw: dict, shape: Sequence[int], itemsize: int, budget: int) -> dict: + """`kw` with its write unit -- the shard, or else the chunk -- cut along + its rows so that the unit's rows across `shape`'s width fit `budget`. + + A shard keeps a whole number of its inner chunks, and its inner chunk is + cut too only when one inner chunk's rows are already over the budget. + """ + chunks = kw.get("chunks") + if not isinstance(chunks, (tuple, list)) or not chunks: + return kw + shards = kw.get("shards") + shards = list(shards) if isinstance(shards, (tuple, list)) and shards else None + chunks = list(chunks) + row = max(1, int(itemsize)) * _prod(shape[1:]) + unit = shards[0] if shards else chunks[0] + if unit * row <= budget: + return kw + limit = max(1, budget // row) + if shards and chunks[0] <= limit: + shards[0] = (limit // chunks[0]) * chunks[0] + elif shards: + chunks[0] = shards[0] = limit + else: + chunks[0] = limit + kw = dict(kw) + kw["chunks"] = tuple(chunks) + if shards: + kw["shards"] = tuple(shards) + return kw + + +def fit_chunks( + kw: dict, + shape: Sequence[int], + *, + backend: str, + itemsize: Optional[int] = None, + unit_bytes: Optional[int] = None, +) -> dict: """Make a forwarded chunk shape legal for a dataset of `shape`. h5py rejects a chunk larger than the dataset, so a chunked source copied @@ -738,10 +904,16 @@ def fit_chunks(kw: dict, shape: Sequence[int], *, backend: str) -> dict: its own chunking, which it does whenever compression asks for one. Zarr accepts a chunk larger than its array, so there it only has to shrink. - Zarr additionally requires a shard to be a whole number of chunks, so a - clamped chunk invalidates the source's shard geometry. Sharding is a - storage-layout choice rather than data, so it is dropped and left to the - backend rather than recomputed into something arbitrary. + A Zarr shard is kept, clamped to the array as its chunk is, in whole + chunks. Dropping it once a chunk was clamped left anndata 0.13's 1 KB + inner chunks as the output's chunks, a file each: 4,083 files for one + array of 4M nonzeros, where the source has 3 (F033). A shard larger + than the whole output is no use either: a writer has to hold all of it. + + Given the dtype's `itemsize`, a write unit (the shard, or the chunk) + whose rows across the array come to more than `unit_bytes` -- by default + `WRITE_UNIT_BYTES` -- is cut along its rows (`_cut_rows`), for an output + written a block at a time through `RowWriter`. """ chunks = kw.get("chunks") if not isinstance(chunks, (tuple, list)) or not chunks: @@ -759,12 +931,26 @@ def fit_chunks(kw: dict, shape: Sequence[int], *, backend: str) -> dict: min(c, dims[i]) if i < len(dims) and dims[i] > 0 else c for i, c in enumerate(original) ) - if clamped == original: - return kw - - kw = dict(kw) - kw["chunks"] = clamped - kw.pop("shards", None) + shards = kw.get("shards") + fitted_shards = None + if isinstance(shards, (tuple, list)) and shards: + fitted = [] + for i, (s, c) in enumerate(zip(shards, clamped)): + s = int(s) + if i < len(dims) and dims[i] > 0: + s = min(s, -(-dims[i] // c) * c) + fitted.append(max(c, (s // c) * c)) + fitted_shards = tuple(fitted) + + if clamped != original or (fitted_shards is not None and fitted_shards != tuple(shards)): + kw = dict(kw) + kw["chunks"] = clamped + if fitted_shards is not None: + kw["shards"] = fitted_shards + if itemsize is not None: + kw = cut_unit( + kw, shape, itemsize, WRITE_UNIT_BYTES if unit_bytes is None else unit_bytes + ) return kw @@ -970,6 +1156,172 @@ def _chunk_step( return min(step, n_rows) +def decode_chunks(array: Any) -> Optional[tuple]: + """The unit a read of `array` decodes whole: its chunk, the inner chunk + of a sharded Zarr array, or None for a contiguous one.""" + chunks = getattr(array, "chunks", None) + return tuple(int(c) for c in chunks) if chunks else None + + +def read_blocks( + shape: Sequence[int], + chunks: Optional[Sequence[int]], + itemsize: int, + *, + rows: int, + budget: int, +) -> tuple: + """How to read an array in whole chunks: `(band, width)`. + + `band` is the rows per read, `rows` rounded to whole chunk heights, and + as many of them as `budget` bytes allow when `rows` is more. A read of + part of a chunk decodes all of it, so blocks of --chunk rows decoded a + chunk 50,000 rows tall fifty times at --chunk 1024 (F061). `width` is + None to read whole rows; when one chunk height across every column is + over `budget` -- chunks of (n_obs, 8), or zarr's own guess for a wide + matrix -- it is the columns per read, a whole number of chunks, and a + band is one chunk height (F081). A single chunk over `budget` is read + whole: reading it in parts would decode it once a part. + """ + if not shape: + return 1, None + n_rows = max(1, int(shape[0])) + height = int(chunks[0]) if chunks else 1 + row = max(1, int(itemsize)) * _prod(shape[1:]) + band = -(-max(1, int(rows)) // height) * height + if band * row > budget: + band = max(height, (budget // row) // height * height) + band = min(band, -(-n_rows // height) * height) + if height * row <= budget or len(shape) < 2 or not chunks or len(chunks) < 2: + return band, None + across = int(chunks[1]) + column = height * max(1, int(itemsize)) * _prod(shape[2:]) + width = max(across, (budget // column) // across * across) + return height, (None if width >= int(shape[1]) else width) + + +def tile_shards( + chunks: Sequence[int], shards: Sequence[int], width: int, *, rows: bool = True +) -> tuple: + """A shard written whole by tiles of one chunk height and `width` + columns (`read_blocks`): one chunk tall (unless `rows` is False, for a + writer that buffers its rows), and a whole number of chunks that divides + the tile's width. A taller or wider shard was rewritten once for each + tile that wrote part of it.""" + shards = [int(s) for s in shards] + if rows: + shards[0] = int(chunks[0]) + across = min(shards[1], width) + if width % across: + across = int(chunks[1]) + shards[1] = across + return tuple(shards) + + +def write_unit_rows(ds: Any) -> int: + """Rows a write to `ds` should cover, in multiples of, for each unit it + stores -- a Zarr shard, or a chunk -- to be encoded once. + + A sharded Zarr array rewrites a whole shard for a write of any part of + it, and Zarr keeps no cache, so a write of part of a chunk decodes and + encodes the chunk again. HDF5's chunk cache can hold a part-written chunk + for the next write, but only while every chunk across the array fits + in it, and not at all in a store opened without one. A contiguous + dataset, or a scalar, has no unit. + """ + shape = getattr(ds, "shape", None) + if not shape: + return 1 + if is_zarr_array(ds): + shards = getattr(ds, "shards", None) + if shards: + return int(shards[0]) + chunks = getattr(ds, "chunks", None) + return int(chunks[0]) if chunks else 1 + + +class RowWriter: + """Write the rows of `ds` in order, each unit of it once. + + Rows are buffered until they reach the end of a unit (`write_unit_rows`) + and written in whole units, so a writer filling an output a block of + --chunk rows at a time no longer rewrites a shard -- 1 GB, by anndata + 0.13's default -- once a block (F031, F036), nor decodes and encodes a + chunk once a block of `split`'s many outputs (F060). It holds less than + one unit of rows, plus the block it was given: `fit_chunks` cuts a unit + to fit its budget when the output is created. + + `grow` extends a 1-D array of unknown length (a sparse matrix's `data` + or `indices`) as it goes. `cols`, a `(start, stop)`, writes only those + columns, for a writer given a band of columns at a time. A unit whose + rows (in those columns) come to more than `budget` -- one the backend + chose, where `fit_chunks` cut none -- is not held: rows are written as + they come, as they were before. Call `flush` after the last rows: what + is still buffered is written then, and nothing before. + """ + + def __init__( + self, ds: Any, *, grow: bool = False, start: int = 0, + cols: Optional[tuple] = None, budget: Optional[int] = None, + ) -> None: + self.ds = ds + self.grow = grow + self.cols = None if cols is None else slice(int(cols[0]), int(cols[1])) + self.unit = max(1, write_unit_rows(ds)) + shape = tuple(ds.shape) + width = (int(cols[1]) - int(cols[0])) if cols is not None else ( + shape[1] if len(shape) > 1 else 1 + ) + row = (int(getattr(ds.dtype, "itemsize", 0) or VLEN_ELEMENT_BYTES) + * max(1, width) * _prod(shape[2:])) + if self.unit * row > (WRITE_UNIT_BYTES if budget is None else budget): + self.unit = 1 + self.written = int(start) + self.pending: list = [] + self.pending_rows = 0 + + @property + def dtype(self) -> Any: + return self.ds.dtype + + @property + def rows(self) -> int: + """Rows given so far, written or not.""" + return self.written + self.pending_rows + + def write(self, values: Any) -> None: + values = np.asarray(values) + if not values.shape or values.shape[0] == 0: + return + end = ((self.rows + values.shape[0]) // self.unit) * self.unit + # Kept past this call, rows are copied: the caller may reuse its + # buffer for the next block. + self.pending.append(values if end > self.written else values.copy()) + self.pending_rows += values.shape[0] + if end > self.written: + self._commit(end) + + def flush(self) -> None: + if self.pending_rows: + self._commit(self.rows) + + def _commit(self, end: int) -> None: + values = self.pending[0] if len(self.pending) == 1 else np.concatenate(self.pending) + count = end - self.written + if self.grow and self.ds.shape[0] < end: + self.ds.resize((end,) + tuple(self.ds.shape[1:])) + if values.ndim == 1: + self.ds[self.written:end] = values[:count] + elif self.cols is not None: + self.ds[self.written:end, self.cols, ...] = values[:count] + else: + self.ds[self.written:end, ...] = values[:count] + rest = values[count:] + self.pending = [rest.copy()] if rest.shape[0] else [] + self.pending_rows = int(rest.shape[0]) + self.written = end + + def copy_dataset( src: Any, dst_group: Any, name: str, *, compression: Optional[str] = None, skip_attrs: Sequence[str] = (), @@ -1008,18 +1360,39 @@ def copy_dataset( shape = tuple(src.shape) zformat = zarr_format_of(dst_group) + element_bytes = ( + _sample_element_bytes(src, shape[0]) if shape and _is_vlen(src.dtype) + else VLEN_ELEMENT_BYTES + ) + itemsize = element_bytes if _is_vlen(src.dtype) else int(src.dtype.itemsize) + chunks = decode_chunks(src) + band, width = ( + read_blocks( + shape, chunks, itemsize, + rows=_chunk_step(shape, chunks, src.dtype, element_bytes), + budget=TARGET_READ_BYTES, + ) + if shape else (1, None) + ) kw = dataset_create_kwargs(src, target_backend=target_backend, zarr_format=zformat) if compression is not None: kw = with_compression( kw, compression, shape=shape, dtype=src.dtype, backend=target_backend, zarr_format=zformat, ) + # Tiles are written where they belong, whole chunks at a time, so only + # whole rows need a unit cut to what `RowWriter` may hold. + kw = fit_chunks( + kw, shape, backend=target_backend, itemsize=None if width is not None else itemsize + ) + if width is not None and kw.get("shards"): + kw["shards"] = tile_shards(kw["chunks"], kw["shards"], width) ds = create_dataset( dst_group, name, shape=shape, dtype=target_dtype(src.dtype, target_backend, zformat), - **fit_chunks(kw, shape, backend=target_backend), + **kw, ) attrs = src.attrs if not skip_attrs else { k: v for k, v in src.attrs.items() if k not in skip_attrs @@ -1030,17 +1403,23 @@ def copy_dataset( ds[()] = src[()] return ds - element_bytes = ( - _sample_element_bytes(src, shape[0]) if _is_vlen(src.dtype) - else VLEN_ELEMENT_BYTES - ) - step = _chunk_step(shape, getattr(src, "chunks", None), src.dtype, element_bytes) - for start in range(0, shape[0], step): - end = min(start + step, shape[0]) - if len(shape) == 1: - ds[start:end] = src[start:end] - else: - ds[start:end, ...] = src[start:end, ...] + if width is not None: + # One chunk height across every column is over the budget: whole + # chunks a band of columns at a time, written where they belong. + # Read whole rows, a column-chunked (n_obs, 8) array was read in + # one piece, the whole matrix (F081). + for start in range(0, shape[0], band): + end = min(start + band, shape[0]) + for left in range(0, shape[1], width): + right = min(left + width, shape[1]) + ds[start:end, left:right, ...] = src[start:end, left:right, ...] + return ds + + writer = RowWriter(ds) + for start in range(0, shape[0], band): + end = min(start + band, shape[0]) + writer.write(src[start:end] if len(shape) == 1 else src[start:end, ...]) + writer.flush() return ds @@ -1152,12 +1531,13 @@ def read(block: tuple) -> np.ndarray: k: v for k, v in src.attrs.items() if k not in skip_attrs } copy_attrs(attrs, ds.attrs, target_backend=target_backend) + if shape == (): + ds[()] = as_records(read(blocks[0]), dtype) + return ds + writer = RowWriter(ds) for block in blocks: - values = as_records(read(block), dtype) - if shape == (): - ds[()] = values - else: - ds[block[0]:block[1]] = values + writer.write(as_records(read(block), dtype)) + writer.flush() return ds diff --git a/tests/perf_counters.py b/tests/perf_counters.py index c294775..f9d136a 100644 --- a/tests/perf_counters.py +++ b/tests/perf_counters.py @@ -112,6 +112,10 @@ class IOCounts: #: 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) + #: Bytes written, and whole units encoded, per array: a writer's own + #: output apart from its scratch. + written_by_name: Dict[str, int] = field(default_factory=dict) + encoded_by_name: Dict[str, int] = field(default_factory=dict) @property def elements(self) -> int: @@ -412,9 +416,13 @@ def _read(array: Any, key: Any, result: Any, name: str) -> int: def _write(array: Any, key: Any, value: Any) -> None: given = _nbytes(value) + encoded = cache.cost(array, key, given, write=True) counts.written_bytes += given counts.largest_write = max(counts.largest_write, given) - counts.encoded_bytes += cache.cost(array, key, given, write=True) + counts.encoded_bytes += encoded + name = str(getattr(array, "name", None) or getattr(array, "path", None) or "/") + counts.written_by_name[name] = counts.written_by_name.get(name, 0) + given + counts.encoded_by_name[name] = counts.encoded_by_name.get(name, 0) + encoded def h5_wrapper(self: Any, key: Any) -> Any: result = h5_get(self, key) diff --git a/tests/perf_fixtures.py b/tests/perf_fixtures.py index f00f75e..433df17 100644 --- a/tests/perf_fixtures.py +++ b/tests/perf_fixtures.py @@ -77,12 +77,12 @@ def skewed_matrix(n_obs: int, n_var: int = 64, *, dense_genes: Tuple[int, int] = return matrix.tocsc() if layout == "csc" else matrix -def write_anndata(obj, path: Path, fmt: str) -> Path: +def write_anndata(obj, path: Path, fmt: str, compression: Optional[str] = None) -> Path: """Write `obj` as anndata would in `fmt`.""" import anndata as ad if fmt == "h5ad": - obj.write_h5ad(path) + obj.write_h5ad(path, compression=compression) elif fmt == "zarr2": with ad.settings.override(zarr_write_format=2): obj.write_zarr(path) @@ -106,11 +106,15 @@ def matrix_store( matrix=None, dense_chunks: Optional[Tuple[int, int]] = None, tag: str = "m", + groups: int = 0, + compression: Optional[str] = None, ) -> 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). + tool might be (anndata writes it contiguous), with `compression` (an + h5ad's). `groups` adds a categorical obs column `grp` of that many + labels, interleaved, for `split`. """ import anndata as ad import pandas as pd @@ -119,22 +123,70 @@ def matrix_store( if matrix is None: matrix = counts_matrix(n_obs, n_var, layout, density=density, dtype=dtype) n_obs, n_var = matrix.shape + columns = {"col": np.arange(n_obs, dtype="int32")} + if groups: + columns["grp"] = pd.Categorical([f"s{i % groups}" for i in range(n_obs)]) 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)], - ), + obs=pd.DataFrame(columns, 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) + write_anndata(obj, path, fmt, compression) if dense_chunks is not None: - rechunk_dense(path, dense_chunks) + rechunk_dense(path, dense_chunks, compression=compression) return path -def rechunk_dense(path: Path, chunks: Tuple[int, int], key: str = "X") -> None: +#: anndata 0.13's layout as zarr 3.1 chose it (`auto_shard_zarr_v3`): inner +#: chunks of about a kilobyte, in shards of the whole array. zarr 3.4 +#: chooses inner chunks of about a megabyte, so a store written today does +#: not show what the users' stores do (F031, F033, F062). +ZARR31_DENSE_CHUNK = (16, 8) +ZARR31_ELEMENT_CHUNK = 245 + + +def zarr31_layout(path: Path, key: str = "X") -> Path: + """Rewrite `key` of a Zarr v3 store in anndata 0.13's layout under zarr + 3.1: a dense matrix in inner chunks of (16, 8) and shards of 4,096 + rows across every column, a sparse one's `data` and `indices` in chunks + of 245 elements and a shard of the whole array. Values and attributes + are kept.""" + import zarr + + def whole(n: int, c: int) -> int: + return max(c, -(-n // c) * c) + + root = zarr.open_group(str(path), mode="r+", use_consolidated=False) + node = root[key] + names = [None] if isinstance(node, zarr.Array) else ["data", "indices"] + for name in names: + parent = root if name is None else node + member = key if name is None else name + array = parent[member] + values, attrs = array[...], dict(array.attrs) + if array.ndim == 2: + chunks = ZARR31_DENSE_CHUNK + shards = (whole(min(array.shape[0], 4096), chunks[0]), + whole(array.shape[1], chunks[1])) + else: + chunks = (ZARR31_ELEMENT_CHUNK,) + shards = (whole(array.shape[0], chunks[0]),) + compressors = array.compressors + del parent[member] + rewritten = parent.create_array( + member, shape=values.shape, dtype=values.dtype, chunks=chunks, + shards=shards, compressors=compressors, + ) + rewritten[...] = values + rewritten.attrs.update(attrs) + zarr.consolidate_metadata(str(path)) + return path + + +def rechunk_dense( + path: Path, chunks: Tuple[int, int], key: str = "X", compression: Optional[str] = None +) -> None: """Rewrite a dense h5ad matrix with `chunks`, keeping its attributes.""" import h5py @@ -142,6 +194,8 @@ def rechunk_dense(path: Path, chunks: Tuple[int, int], key: str = "X") -> None: values = handle[key][...] attrs = dict(handle[key].attrs) del handle[key] - dataset = handle.create_dataset(key, data=values, chunks=chunks) + dataset = handle.create_dataset( + key, data=values, chunks=chunks, compression=compression + ) for name, value in attrs.items(): dataset.attrs[name] = value diff --git a/tests/test_block_layout.py b/tests/test_block_layout.py new file mode 100644 index 0000000..393ec07 --- /dev/null +++ b/tests/test_block_layout.py @@ -0,0 +1,390 @@ +"""Blocks and shards: outputs keep their source's chunks and shards, are +written a whole unit at a time, and still equal anndata's own subsetting +(B6b, #97). + +The cost of each of these is guarded in `test_performance.py`; here the +outputs are checked against anndata, on h5ad, Zarr v2, Zarr v3 as anndata +writes it today, and Zarr v3 in anndata 0.13's layout under zarr 3.1 +(`zarr31`: inner chunks of a kilobyte inside shards of a whole array), at +blocks small enough that every boundary is crossed: a major read in parts, +a unit cut to a sliver of a split's write budget, a matrix read in tiles. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import numpy as np +import pytest + +from tests.conftest import backend_path +from tests.perf_counters import count_io +from tests.perf_fixtures import counts_matrix, zarr31_layout +from tests.test_subset_oracle import ( + assert_matches_anndata, + ok, + read_anndata, + rich_anndata, + write_anndata, +) + +ad = pytest.importorskip("anndata") +pd = pytest.importorskip("pandas") +sparse = pytest.importorskip("scipy.sparse") +zarr = pytest.importorskip("zarr") + +FORMATS = ("h5ad", "zarr2", "zarr3", "zarr31") + + +def write_store(obj: Any, tmp_path: Path, fmt: str, stem: str = "src") -> Path: + """`obj` as anndata writes it in `fmt`; `zarr31` then rewrites X.""" + path = write_anndata(obj, tmp_path, "zarr3" if fmt == "zarr31" else fmt, stem) + return zarr31_layout(path) if fmt == "zarr31" else path + + +def _names(tmp_path: Path, stem: str, prefix: str, picked) -> Path: + path = tmp_path / f"{stem}.txt" + path.write_text("".join(f"{prefix}{i}\n" for i in picked)) + return path + + +# --------------------------------------------------------------------------- +# anndata as the oracle, at blocks smaller than every unit + + +@pytest.mark.parametrize("layout", ["csr", "csc", "dense"]) +@pytest.mark.parametrize("fmt", FORMATS) +@pytest.mark.parametrize("chunk", [1, 3]) +def test_a_subset_in_tiny_blocks_matches_anndata(tmp_path, fmt, layout, chunk): + """At --chunk 1 a sparse block holds a row's worth of entries, so + majors of the CSC layers are read in parts (F032), and every output + unit is filled over many blocks before it is written (F031).""" + source = rich_anndata(layout, seed=3) + src = write_store(source, tmp_path, fmt) + obs = [0, 1, 2, 5, 9, 10, 11, 17, 23, 28, 29] + var = [0, 3, 4, 8] + out = backend_path(tmp_path, "h5ad" if fmt == "h5ad" else "zarr3", "out") + + ok("subset", src, "-o", out, "--obs", _names(tmp_path, "c", "c", obs), + "--var", _names(tmp_path, "g", "v", var), "-C", chunk) + assert_matches_anndata(read_anndata(out), source, obs, var) + + +@pytest.mark.parametrize("layout", ["csr", "csc", "dense"]) +@pytest.mark.parametrize("fmt", FORMATS) +def test_split_in_tiny_blocks_and_a_tiny_write_budget_matches_anndata( + tmp_path, monkeypatch, fmt, layout +): + """Each output's units are cut to its share of the pass's write budget + (F079); here the budget is a few hundred bytes, so units are a row or + two, and every output still equals anndata's slice.""" + import adata.storage as storage + + monkeypatch.setattr(storage, "WRITE_BUFFER_BYTES", 512) + source = rich_anndata(layout, seed=4) + src = write_store(source, tmp_path, fmt) + parts = tmp_path / "parts" + + ok("split", src, "--by", "group", "-o", parts, "--no-manifest", "-C", 2) + labels = source.obs["group"].astype(str).to_numpy() + for label in set(labels): + got = read_anndata(parts / f"{label}{src.suffix}") + assert_matches_anndata(got, source, obs=np.flatnonzero(labels == label)) + + +def test_a_matrix_read_in_tiles_matches_anndata(tmp_path, monkeypatch): + """F061, F081: where one chunk height across the matrix is over the read + budget -- chunks of (n_obs, 2) -- a band is read in tiles of whole + chunks, each output column band written by its own writer. Made to + happen at this size by a read budget of a few bytes.""" + import h5py + + import adata.storage as storage + from adata.core import subset + + monkeypatch.setattr(storage, "TARGET_READ_BYTES", 64) + monkeypatch.setattr(subset, "TARGET_READ_BYTES", 64) + source = rich_anndata("dense", seed=5) + src = write_store(source, tmp_path, "h5ad") + with h5py.File(src, "a") as handle: + for key in ("X", "layers/dense"): + values, attrs = handle[key][...], dict(handle[key].attrs) + del handle[key] + ds = handle.create_dataset(key, data=values, chunks=(values.shape[0], 2), + compression="gzip") + ds.attrs.update(attrs) + obs, var = [1, 2, 3, 8, 13, 21, 22, 27], [0, 2, 3, 6, 7] + out = tmp_path / "out.h5ad" + ok("subset", src, "-o", out, "--obs", _names(tmp_path, "c", "c", obs), + "--var", _names(tmp_path, "g", "v", var), "-C", 1) + assert_matches_anndata(read_anndata(out), source, obs, var) + + parts = tmp_path / "parts" + ok("split", src, "--by", "group", "-o", parts, "--no-manifest", "-C", 1) + labels = source.obs["group"].astype(str).to_numpy() + for label in set(labels): + assert_matches_anndata(read_anndata(parts / f"{label}.h5ad"), source, + obs=np.flatnonzero(labels == label)) + + zout = tmp_path / "copy.zarr" + ok("convert", src, "-o", zout) + got = read_anndata(zout) + np.testing.assert_array_equal(got.X, source.X) + np.testing.assert_array_equal(got.layers["dense"], source.layers["dense"]) + + +# --------------------------------------------------------------------------- +# outputs keep their source's layout (F033) + + +def _sharded_counts(tmp_path: Path, layout: str, n_obs: int = 3000, n_var: int = 40, + stem: str = "src") -> tuple: + matrix = counts_matrix(n_obs, n_var, layout, density=0.2) + obj = ad.AnnData( + X=matrix, + obs=pd.DataFrame({"n": np.arange(n_obs, dtype=np.int64), + "group": pd.Categorical([f"g{i % 3}" for i in range(n_obs)])}, + index=[f"c{i}" for i in range(n_obs)]), + var=pd.DataFrame(index=[f"v{j}" for j in range(n_var)]), + ) + return obj, write_store(obj, tmp_path, "zarr31", stem) + + +def _arrays(group: Any) -> dict: + x = group["X"] + if isinstance(x, zarr.Array): + return {"X": x} + return {f"X/{k}": x[k] for k in ("data", "indices")} + + +@pytest.mark.parametrize("layout", ["csr", "csc", "dense"]) +@pytest.mark.parametrize("command", ["subset", "split", "concat", "recast"]) +def test_outputs_keep_the_sources_inner_chunks_and_shards(tmp_path, command, layout): + """F033: outputs of anndata 0.13's layout dropped the shard and kept the + 245-element inner chunk as their own, a file each: 4,083 files for one + array the source holds in 3. Each output array has the source's inner + chunk and shard now, the shard clamped to the output, in whole chunks.""" + obj, src = _sharded_counts(tmp_path, layout) + source = _arrays(zarr.open_group(str(src), mode="r")) + if command == "subset": + outputs = [tmp_path / "out.zarr"] + ok("subset", src, "-o", outputs[0], "-q", "n >= 100", "-C", 64) + elif command == "split": + ok("split", src, "--by", "group", "-o", tmp_path / "parts", "--no-manifest", "-C", 64) + outputs = sorted((tmp_path / "parts").iterdir()) + elif command == "concat": + other = write_store(obj, tmp_path, "zarr31", "other") + outputs = [tmp_path / "out.zarr"] + ok("concat", src, other, "-o", outputs[0], "--index-unique", "-", "-C", 64) + else: + outputs = [tmp_path / "out.zarr"] + ok("recast", src, "X", "-o", outputs[0], "--dtype", "float64", "-C", 64) + + for path in outputs: + group = zarr.open_group(str(path), mode="r") + for name, array in _arrays(group).items(): + want = source[name] + assert array.chunks == want.chunks, (path.name, name) + shards = tuple( + min(s, max(c, -(-n // c) * c)) if array.ndim > 1 or name == "X" else s + for s, c, n in zip(want.shards, want.chunks, array.shape) + ) + assert array.shards == shards, (path.name, name, array.shards, shards) + files = [f for f in (path / name).rglob("*") + if f.is_file() and f.name not in ("zarr.json", ".zarray", ".zattrs")] + assert len(files) <= 4, f"{name} of {path.name} is {len(files)} files" + got = read_anndata(path) + rows = got.obs["n"].to_numpy() % len(obj) + want_x = obj.X[rows] + got_x = got.X + if sparse.issparse(want_x): + want_x, got_x = want_x.toarray(), got_x.toarray() + np.testing.assert_array_equal(got_x, want_x) + + +def test_converting_a_sharded_store_to_h5ad_lets_hdf5_choose_its_chunks(tmp_path): + """F033: a shard's 245-element inner chunk became the HDF5 chunk, 16,327 + chunks where anndata's own h5ad has 512, and a file twice the size. HDF5 + picks the chunks now, and X takes what anndata's own takes.""" + import h5py + + obj, src = _sharded_counts(tmp_path, "csr", n_obs=20_000, n_var=100) + out = tmp_path / "out.h5ad" + ok("convert", src, "-o", out) + theirs = tmp_path / "theirs.h5ad" + ad.read_zarr(src).write_h5ad(theirs, compression="gzip") + + def stored(path: Path) -> int: + with h5py.File(path, "r") as handle: + for name in ("data", "indices"): + assert handle[f"X/{name}"].chunks[0] > 4096, handle[f"X/{name}"].chunks + return sum(handle[f"X/{n}"].id.get_storage_size() for n in ("data", "indices")) + + assert stored(out) <= 1.1 * stored(theirs), (stored(out), stored(theirs)) + np.testing.assert_array_equal(read_anndata(out).X.toarray(), obj.X.toarray()) + + +def test_a_sharded_store_written_as_zarr_v2_takes_its_shard_as_the_chunk(tmp_path): + """F033: Zarr v2 has no shards, so `--zarr-format 2` kept only the inner + chunk, a 245-element file each. The shard is the chunk now.""" + obj, src = _sharded_counts(tmp_path, "csr") + out = tmp_path / "v2.zarr" + ok("convert", src, "-o", out, "--zarr-format", "2") + source = zarr.open_group(str(src), mode="r") + group = zarr.open_group(str(out), mode="r") + for name in ("data", "indices"): + array, shards = group[f"X/{name}"], source[f"X/{name}"].shards + # The shard, clamped to the array as any chunk is. + assert array.chunks == tuple(min(s, n) for s, n in zip(shards, array.shape)) + np.testing.assert_array_equal(read_anndata(out).X.toarray(), obj.X.toarray()) + + +# --------------------------------------------------------------------------- +# the pieces + + +@pytest.mark.parametrize("sharded", [False, True]) +def test_a_row_writer_writes_each_unit_once_whatever_the_blocks(new_store, backend, sharded): + """Blocks of every size, none of them a unit, land where they belong, and + each unit (chunk, or shard) is written once.""" + from adata.storage import RowWriter, create_dataset + + if sharded and backend != "zarr3": + pytest.skip("only Zarr v3 shards") + path, opener = new_store() + values = np.arange(103 * 6, dtype="float32").reshape(103, 6) + kw = {"chunks": (7, 3)} + if sharded: + kw["shards"] = (14, 6) + with opener("a") as root: + ds = create_dataset(root["uns"], "m", shape=values.shape, dtype=values.dtype, **kw) + writer = RowWriter(ds) + with count_io() as io: + at = 0 + for size in (1, 5, 2, 13, 30, 4, 48): + writer.write(values[at:at + size]) + at += size + writer.flush() + np.testing.assert_array_equal(ds[...], values) + unit = 14 if sharded else 7 + # The last unit is part-filled, and counted whole. + assert io.encoded_bytes <= io.written_bytes + unit * 6 * 4, io + + +def test_a_row_writer_grows_an_array_and_writes_columns(new_store): + from adata.storage import RowWriter, create_dataset + + path, opener = new_store() + with opener("a") as root: + grown = create_dataset(root["uns"], "g", shape=(0,), dtype="int64", chunks=(10,), + **({} if "zarr" in str(path) else {"maxshape": (None,)})) + writer = RowWriter(grown, grow=True) + for start in range(0, 57, 4): + writer.write(np.arange(start, min(start + 4, 57))) + assert grown.shape[0] % 10 == 0, "grown a whole unit at a time" + writer.flush() + np.testing.assert_array_equal(grown[...], np.arange(57)) + + wide = create_dataset(root["uns"], "w", shape=(20, 9), dtype="int32", chunks=(4, 3)) + left, right = RowWriter(wide, cols=(0, 4)), RowWriter(wide, cols=(4, 9)) + values = np.arange(180, dtype="int32").reshape(20, 9) + for start in range(0, 20, 3): + left.write(values[start:start + 3, :4]) + right.write(values[start:start + 3, 4:]) + left.flush(), right.flush() + np.testing.assert_array_equal(wide[...], values) + + +def test_a_row_writer_does_not_hold_a_unit_over_its_budget(new_store): + """A unit the backend chose, too large to hold, is written as rows come.""" + from adata.storage import RowWriter, create_dataset + + path, opener = new_store() + with opener("a") as root: + ds = create_dataset(root["uns"], "m", shape=(100, 4), dtype="float64", chunks=(50, 4)) + assert RowWriter(ds).unit == 50 + assert RowWriter(ds, budget=50 * 4 * 8 - 1).unit == 1 + + +def test_read_blocks_round_to_whole_chunks_within_the_budget(): + from adata.storage import read_blocks + + # Contiguous: --chunk rows, as before. + assert read_blocks((1000, 10), None, 4, rows=64, budget=10**6) == (64, None) + # Rounded up to whole chunk heights. + assert read_blocks((1000, 10), (50, 10), 4, rows=64, budget=10**6) == (100, None) + # A chunk taller than --chunk is read whole, within the budget. + assert read_blocks((100_000, 10), (5000, 10), 4, rows=64, budget=10**6) == (5000, None) + # More than the budget allows: fewer whole chunk heights. + assert read_blocks((100_000, 10), (50, 10), 4, rows=10_000, budget=50 * 40 * 3) == (150, None) + # One chunk height across the matrix is over the budget: tiles of whole chunks. + assert read_blocks((100_000, 64), (100_000, 8), 4, rows=64, budget=100_000 * 4 * 16) == ( + 100_000, 16) + # Never more rows than the array, in whole chunks. + assert read_blocks((10, 4), (4, 4), 4, rows=1000, budget=10**6) == (12, None) + + +@pytest.mark.parametrize("seed", range(20)) +def test_sparse_blocks_cover_every_kept_major_once_within_the_budget(seed): + """F032's blocks: whole majors while they fit, a major over the budget in + parts, majors no target keeps skipped.""" + from adata.core.subset import _selected, _sparse_blocks + + rng = np.random.default_rng(seed) + n_major = int(rng.integers(1, 40)) + sizes = rng.integers(0, 12, n_major) * (rng.random(n_major) < 0.8) + indptr = np.concatenate(([0], np.cumsum(sizes))) + budget = int(rng.integers(1, 15)) + selections = [None if rng.random() < 0.3 else + np.sort(rng.choice(n_major, int(rng.integers(0, n_major + 1)), replace=False)) + for _ in range(int(rng.integers(1, 4)))] + if all(s is not None and s.size == 0 for s in selections): + selections[0] = None + + got = {i: np.zeros(n_major, dtype=np.int64) for i in range(len(selections))} + last = -1 + for first, stop, lo, hi in _sparse_blocks(indptr, budget, selections): + assert hi - lo <= budget + assert lo >= last, "blocks arrive in order" + last = hi + for i, sel in enumerate(selections): + _, rows = _selected(sel, first, stop) + for r in rows: + got[i][r] += max(0, min(indptr[r + 1], hi) - max(indptr[r], lo)) + for i, sel in enumerate(selections): + kept = np.arange(n_major) if sel is None else sel + np.testing.assert_array_equal(got[i][kept], sizes[kept]) + + +def test_a_shard_index_is_read_once_and_the_cache_is_bounded(tmp_path, monkeypatch): + """F062: each read of part of a shard fetched its index; a store opened + for reading keeps the indexes now, up to `SHARD_INDEX_CACHE_BYTES`.""" + from zarr.abc.store import SuffixByteRequest + from zarr.storage import LocalStore + + import adata.storage as storage + from adata.storage import open_store + + _, src = _sharded_counts(tmp_path, "csr") + fetched = [] + original = LocalStore.get + + async def get(self, key, prototype, byte_range=None): + if isinstance(byte_range, SuffixByteRequest): + fetched.append(key) + return await original(self, key, prototype, byte_range) + + monkeypatch.setattr(LocalStore, "get", get) + + def read_in_blocks() -> int: + fetched.clear() + with open_store(src, "r") as store: + data = store.root["X/data"] + for start in range(0, data.shape[0], 1000): + data[start:start + 1000] + return len(fetched) + + assert read_in_blocks() == 1 + monkeypatch.setattr(storage, "SHARD_INDEX_CACHE_BYTES", 0) + assert read_in_blocks() > 10, "a cache too small for an index keeps none" diff --git a/tests/test_performance.py b/tests/test_performance.py index 4efa844..d79e888 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -2420,13 +2420,10 @@ def _grid_store(directory: Path, layout: str, axis: str, n: int) -> tuple: } #: 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)" -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), -} +#: Sparse subset was one (F032): its blocks were 4,096 majors, whatever +#: --chunk said, 3 to 5 times the input added per byte; a block holds about +#: --chunk rows' entries now, in either layout (B6b, #97). +GRID_GAPS: dict = {} def _grid_cases(): @@ -2641,3 +2638,271 @@ def measure(n: int) -> int: measure, lambda n: n * 256 * 4, what="recast --dtype of a dense X (RSS)", axis="n_obs", sizes=(8192, 65536), ) + + +# --------------------------------------------------------------------------- +# blocks and shards: reads and writes line up with the store's own units +# +# anndata 0.13 writes Zarr v3 sharded: under zarr 3.1, inner chunks of about +# a kilobyte inside shards of a whole array (`zarr31_layout`). A read of +# part of a shard fetched the shard's index again; a write of part of one +# rewrote all of it; a read of part of a chunk decoded all of it. Each guard +# names what it measures in whole units, so a block that splits a unit shows +# as bytes decoded or encoded many times over, without a clock. + +#: Rows of the stores below, and --chunk: blocks far smaller than a unit. +UNIT_ROWS, UNIT_CHUNK = 8192, 256 + + +def _x_amplification(io, *, encoded: bool = True) -> float: + """Bytes of X's own arrays encoded per byte written to them: 1.0 when + every shard or chunk is written once. Not obs, var or a scratch.""" + def of_x(name: str) -> bool: + parts = name.strip("/").split("/") + return "X" in parts and not any(p.startswith("__") for p in parts) + + names = [n for n in io.written_by_name if of_x(n)] + written = sum(io.written_by_name[n] for n in names) + assert written > 0, f"no write to X was seen: {sorted(io.written_by_name)}" + return sum(io.encoded_by_name[n] for n in names) / written + + +def _sharded(tmp_path: Path, layout: str, n_obs: int = UNIT_ROWS, n_var: int = 64, + groups: int = 0, tag: str = "s") -> Path: + from tests.perf_fixtures import matrix_store, zarr31_layout + + return zarr31_layout(matrix_store(tmp_path, n_obs=n_obs, n_var=n_var, layout=layout, + fmt="zarr3", groups=groups, tag=tag)) + + +def _split(source: Path, out: Path, chunk_rows: int = UNIT_CHUNK) -> None: + from adata.commands.split import split_store + + split_store(source, "grp", out, QUIET, manifest=False, chunk_rows=chunk_rows) + + +def _twice(source: Path) -> List[Path]: + """`source` and a copy of it, for concat, whose inputs need names of + their own.""" + import shutil + + copy = source.with_name(f"b{source.name}") + if not copy.exists(): + (shutil.copytree if source.is_dir() else shutil.copy)(source, copy) + return [source, copy] + + +def _subset(source: Path, out: Path, chunk_rows: int = UNIT_CHUNK, **kw) -> None: + subset_h5ad(source, out, None, None, console=QUIET, + obs_query=kw.pop("obs_query", "col >= 0"), chunk_rows=chunk_rows, **kw) + + +@pytest.mark.parametrize("command,layout", [ + ("subset", "dense"), + ("split", "dense"), + ("concat", "dense"), + ("recast --dtype", "dense"), + ("recast --indices-dtype", "csr"), + ("recast --layout csc", "csr"), + ("convert", "dense"), +]) +def test_a_sharded_output_is_written_a_whole_shard_at_a_time(tmp_path, monkeypatch, command, layout): + """F031, F036: an output keeping anndata 0.13's shard was written a + --chunk block at a time, and zarr rewrites a whole shard for a write of + part of it: 15 times the bytes for subset and concat, 121 for split + (eight groups, each given a few rows a block), 15 for recast --dtype, + 138 for a transpose. Each shard is written once now (`RowWriter`). + + A transpose's own scratch is another matter (a follow-up); only X's + arrays are counted. `convert` copies in reads of `TARGET_READ_BYTES`, + made small here so that a read is smaller than a shard.""" + import adata.storage as storage + + source = _sharded(tmp_path, layout, groups=8 if command == "split" else 0) + out = tmp_path / "out.zarr" + with count_io() as io: + if command == "subset": + _subset(source, out) + elif command == "split": + _split(source, tmp_path / "parts") + elif command == "concat": + concat_on_disk(_twice(source), out, QUIET, chunk_rows=UNIT_CHUNK, + index_unique="-") + elif command == "recast --dtype": + _recast(source, out, dtype="float64", chunk_rows=UNIT_CHUNK) + elif command == "recast --indices-dtype": + _recast(source, out, indices_dtype="int64", chunk=4096) + elif command == "recast --layout csc": + _recast(source, out, layout="csc", chunk=4096) + else: + from adata.commands.convert import convert_store + + monkeypatch.setattr(storage, "TARGET_READ_BYTES", 64 * 1024) + convert_store(source, out, QUIET) + amplification = _x_amplification(io) + assert amplification <= 1.25, ( + f"{command} of a sharded {layout} X encoded {amplification:.1f} times the " + "bytes it wrote to X: shards are being rewritten once a block" + ) + + +@pytest.mark.parametrize("fmt,chunks", [("h5ad", (UNIT_ROWS, 8)), ("zarr2", None)], + ids=["tall-chunked-h5ad", "zarr2-default"]) +def test_a_dense_subset_decodes_each_source_chunk_once(tmp_path, fmt, chunks): + """F061: blocks of --chunk rows were not aligned to the source's chunk + height, so every chunk was decoded once a block: 15 times the matrix + for gzip chunks of (n_obs, 8) at 16,384 rows (60 at 65,536), 7 for + zarr's own chunks (8192, 16), 29 at 262,144 rows. Blocks are whole + chunk heights now, within the read budget.""" + from tests.perf_fixtures import matrix_store + + source = matrix_store(tmp_path, n_obs=4 * UNIT_ROWS if fmt == "zarr2" else UNIT_ROWS, + n_var=64, layout="dense", fmt=fmt, density=0.5, + dense_chunks=chunks, compression="gzip" if chunks else None) + with count_io() as io: + _subset(source, tmp_path / f"o{source.suffix}", chunk_rows=1024) + x = [n for n in io.by_name if n.strip("/") == "X"] + assert io.read_bytes >= 4 * UNIT_ROWS * 64, "the matrix was read" + assert io.read_amplification <= 1.25, ( + f"subset decoded {io.read_amplification:.1f} times the bytes it read " + f"({fmt}, X read as {x}): a block splits the source's chunks" + ) + + +def test_split_writes_each_output_chunk_once(tmp_path): + """F060: split's outputs keep the source's chunk height, and each got a + share of every --chunk block: 32 rows a block into chunks 512 rows tall, + across 128 chunks of columns, more than HDF5's chunk cache holds, so + each chunk was decoded and encoded 16 times. Each output is written + whole chunk heights at a time now.""" + from tests.perf_fixtures import matrix_store + + source = matrix_store(tmp_path, n_obs=4096, n_var=1024, layout="dense", density=0.5, + groups=8, dense_chunks=(512, 8), compression="gzip") + with count_io() as io: + _split(source, tmp_path / "parts") + amplification = _x_amplification(io) + assert amplification <= 1.25, ( + f"split encoded {amplification:.1f} times the bytes it wrote to X" + ) + + +@pytest.mark.parametrize("command,layout", [("subset", "dense"), ("concat", "csr")]) +def test_a_sharded_source_is_read_with_its_shard_index_once(tmp_path, command, layout): + """F062: zarr reads a shard's index again for every read of part of the + shard, 16 bytes an inner chunk, so a store read a block at a time read + the whole index once a block: subset fetched 2.3 times the dense X's + bytes from the store, concat several times. Indexes are kept for the + life of the open store now.""" + source = _sharded(tmp_path, layout, n_obs=4 * UNIT_ROWS if layout == "csr" else UNIT_ROWS) + stored = sum(f.stat().st_size for f in source.rglob("*") if f.is_file()) + inputs = _twice(source) if command == "concat" else [source] + with count_io() as io: + if command == "subset": + _subset(source, tmp_path / "o.h5ad") + else: + concat_on_disk(inputs, tmp_path / "o.h5ad", QUIET, + chunk_rows=UNIT_CHUNK, index_unique="-") + fetched = io.store_get_bytes / (len(inputs) * stored) + assert fetched <= 1.3, ( + f"{command} fetched {fetched:.2f} times the store's bytes on disk: shard " + "indexes are being read again for every block" + ) + + +@pytest.mark.parametrize("layout", ["csr", "csc"]) +def test_a_sparse_subset_reads_a_block_set_by_the_chunk(tmp_path, layout): + """F032: sparse subset read 4,096 majors a block, whatever --chunk said: + of a CSC matrix, 4,096 genes across every cell, a block growing with + n_obs (26,214 entries at 4,096 cells, 209,715 at 32,768, at --chunk + 64). A block holds about --chunk rows' entries now.""" + from tests.perf_fixtures import matrix_store + + rows = 64 + + def measure(n: int) -> int: + source = matrix_store(tmp_path, n_obs=n, n_var=64, layout=layout) + with count_io() as io: + _subset(source, tmp_path / f"o{n}.h5ad", chunk_rows=rows) + largest = _largest_matrix_read(io, elements=True) + # --chunk rows of entries, at the matrix's mean per row (6.4). + bound = rows * 7 + assert 0 < largest <= bound, ( + f"a {layout} subset of {n:,} cells read {largest:,} entries at once " + f"against {bound:,} for --chunk {rows}" + ) + return largest + + assert_independent_of(measure, what=f"the largest read of a {layout} subset", + axis="n_obs", sizes=(4096, 32768)) + + +def test_a_gene_in_every_cell_is_read_in_parts(tmp_path): + """F032, its other half: a major holding more entries than a block -- + a gene expressed in every cell, read whole -- grew with n_obs however + small the block. It is read in parts of a block now.""" + from tests.perf_fixtures import matrix_store, skewed_matrix + + n = 16384 + matrix = skewed_matrix(n, 64, layout="csc") + source = matrix_store(tmp_path, n_obs=n, n_var=64, layout="csc", matrix=matrix, tag="skew") + with count_io() as io: + _subset(source, tmp_path / "o.h5ad", chunk_rows=16) + # --chunk rows of entries, at the matrix's mean per row. + bound = 16 * -(-matrix.nnz // n) + assert bound < n, "a dense gene holds more than a block" + largest = _largest_matrix_read(io, elements=True) + assert 0 < largest <= bound, f"read {largest:,} entries at once against {bound:,}" + + +def test_copying_a_column_chunked_matrix_reads_tiles_not_the_matrix(tmp_path, monkeypatch): + """F081: `copy_dataset` read whole rows, at least one chunk height, so a + matrix chunked (n_obs, 8) was read in one piece, the whole matrix, + whatever the read budget. Past the budget it reads tiles of whole + chunks now: within the budget, or one chunk where a chunk alone is + over it. `TARGET_READ_BYTES` is made small rather than the matrix + large.""" + import h5py + import zarr + + import adata.storage as storage + from adata.commands.convert import convert_store + from tests.perf_fixtures import matrix_store + + budget = 256 * 1024 + monkeypatch.setattr(storage, "TARGET_READ_BYTES", budget) + for n in (4096, 16384): + source = matrix_store(tmp_path, n_obs=n, n_var=64, layout="dense", density=0.5, + dense_chunks=(n, 8)) + out = tmp_path / f"o{n}.zarr" + with count_io() as io: + convert_store(source, out, QUIET) + largest, chunk = _largest_matrix_read(io), n * 8 * 4 + assert 0 < largest <= max(budget, chunk), ( + f"convert of a (n_obs, 8)-chunked X of {n:,} rows read {largest:,} bytes " + f"at once against a budget of {budget:,} and chunks of {chunk:,}" + ) + with h5py.File(source, "r") as handle: + assert np.array_equal(zarr.open_group(str(out))["X"][...], handle["X"][...]) + + +@pytest.mark.slow +def test_split_memory_stops_growing_with_the_group_count(tmp_path): + """F079: split held up to 1,024 HDF5 outputs open in a pass, about 0.6 MB + each however they are opened (HDF5's own structures for a file), so + RSS followed the group count, not --chunk: 136 MB at 128 groups, 423 MB + at 512. A pass holds at most 128 HDF5 outputs open now.""" + from tests.perf_counters import peak_rss + from tests.perf_fixtures import matrix_store + + def measure(groups: int) -> int: + source = matrix_store(tmp_path / str(groups), n_obs=16384, n_var=64, layout="csr", + groups=groups) + return peak_rss(["split", source, "--by", "grp", "-o", tmp_path / f"p{groups}", + "--no-manifest"]).added + + few, many = measure(128), measure(512) + assert many <= few + 32e6, ( + f"split into 512 groups added {many / 1e6:.0f} MB of RSS against " + f"{few / 1e6:.0f} MB into 128: memory follows the number of groups" + ) diff --git a/tests/test_storage.py b/tests/test_storage.py index 81a11eb..b52d8d8 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -463,13 +463,40 @@ def test_subsetting_a_sharded_v3_store(temp_dir, target): assert np.array_equal(store.root["X"][...], expected) -def test_fitting_a_chunk_drops_an_incompatible_shard(): +def test_fitting_a_chunk_keeps_its_shard_in_whole_chunks(): + """F033: a clamped chunk dropped the shard, so an output of anndata + 0.13's sharded layout kept only its 1 KB inner chunks, a file each. + The shard is clamped with the chunk now, to a whole number of them.""" kept = fit_chunks({"chunks": (4, 2), "shards": (8, 4)}, (100, 100), backend="zarr") assert kept["shards"] == (8, 4), "an unclamped chunk keeps its shard" - dropped = fit_chunks({"chunks": (4, 2), "shards": (8, 4)}, (3, 4), backend="zarr") - assert dropped["chunks"] == (3, 2) - assert "shards" not in dropped + clamped = fit_chunks({"chunks": (4, 2), "shards": (8, 4)}, (3, 4), backend="zarr") + assert clamped["chunks"] == (3, 2) + assert clamped["shards"] == (3, 4) + + # A shard larger than the array is clamped to it, in whole chunks. + fitted = fit_chunks({"chunks": (4, 2), "shards": (64, 4)}, (10, 4), backend="zarr") + assert fitted["shards"] == (12, 4) + + +def test_fitting_cuts_a_unit_too_large_to_hold_along_its_rows(): + """A shard whose rows across the array come to more than the write + budget is cut to fewer whole inner chunks; one inner chunk too many + for it is cut too, and so is an unsharded chunk.""" + kw = {"chunks": (16, 8), "shards": (4096, 64)} + cut = fit_chunks(kw, (8192, 64), backend="zarr", itemsize=4, unit_bytes=64 * 4 * 160) + assert cut["chunks"] == (16, 8) + assert cut["shards"] == (160, 64) + + tight = fit_chunks(kw, (8192, 64), backend="zarr", itemsize=4, unit_bytes=64 * 4 * 10) + assert tight["chunks"] == (10, 8) and tight["shards"] == (10, 64) + + plain = fit_chunks({"chunks": (1000, 8)}, (8192, 64), backend="hdf5", itemsize=4, + unit_bytes=64 * 4 * 100) + assert plain["chunks"] == (100, 8) + + within = fit_chunks(kw, (8192, 64), backend="zarr", itemsize=4) + assert within == kw, "a unit within the budget is kept as it is" def test_fitting_handles_one_dimensional_chunks(): diff --git a/tests/test_subset_fan_out.py b/tests/test_subset_fan_out.py index 0528cce..e82a4d5 100644 --- a/tests/test_subset_fan_out.py +++ b/tests/test_subset_fan_out.py @@ -140,9 +140,12 @@ def test_subset_keeps_the_source_compression(tmp_path): # the vectorised gather, against scipy +@pytest.mark.parametrize("chunk_rows", [1, 5]) @pytest.mark.parametrize("layout", ["csr", "csc"]) @pytest.mark.parametrize("seed", range(6)) -def test_sparse_subset_matches_scipy(tmp_path, layout, seed): +def test_sparse_subset_matches_scipy(tmp_path, layout, seed, chunk_rows): + """At --chunk 1 a block holds fewer entries than many majors do, so + majors are read in parts across blocks (F032's split majors).""" rng = np.random.default_rng(seed) n_obs, n_var = 57, 23 dense = rng.poisson(0.5, (n_obs, n_var)).astype("float32") @@ -167,7 +170,7 @@ def pick(n): [(o.create_group(f"t{i}"), ob, va) for i, (ob, va) in enumerate(targets)], # Small blocks, so selections straddle block boundaries and some # blocks hold nothing a target wants. - chunk_major=5, + chunk_rows=chunk_rows, ) for i, (ob, va) in enumerate(targets): expected = dense From 0755190c7a55a988e1d3043450ce1bcb233277c7 Mon Sep 17 00:00:00 2001 From: Aljes Date: Mon, 2 Nov 2026 16:09:54 +0000 Subject: [PATCH 2/5] Trim the new block and shard tests for CI time Co-Authored-By: Claude Opus 5.5 --- tests/test_block_layout.py | 14 ++++++++------ tests/test_performance.py | 10 +++++----- 2 files changed, 13 insertions(+), 11 deletions(-) diff --git a/tests/test_block_layout.py b/tests/test_block_layout.py index 393ec07..dfd1c7d 100644 --- a/tests/test_block_layout.py +++ b/tests/test_block_layout.py @@ -55,11 +55,11 @@ def _names(tmp_path: Path, stem: str, prefix: str, picked) -> Path: @pytest.mark.parametrize("layout", ["csr", "csc", "dense"]) @pytest.mark.parametrize("fmt", FORMATS) -@pytest.mark.parametrize("chunk", [1, 3]) -def test_a_subset_in_tiny_blocks_matches_anndata(tmp_path, fmt, layout, chunk): - """At --chunk 1 a sparse block holds a row's worth of entries, so - majors of the CSC layers are read in parts (F032), and every output - unit is filled over many blocks before it is written (F031).""" +def test_a_subset_in_tiny_blocks_matches_anndata(tmp_path, fmt, layout): + """At --chunk 2 a sparse block holds two rows' worth of entries, fewer + than a column of a CSC layer, which is read in parts (F032); every + output unit is filled over many blocks before it is written (F031).""" + chunk = 2 source = rich_anndata(layout, seed=3) src = write_store(source, tmp_path, fmt) obs = [0, 1, 2, 5, 9, 10, 11, 17, 23, 28, 29] @@ -71,8 +71,10 @@ def test_a_subset_in_tiny_blocks_matches_anndata(tmp_path, fmt, layout, chunk): assert_matches_anndata(read_anndata(out), source, obs, var) +# h5ad's chunks and anndata 0.13's shards: the other formats' units are +# cut by the same code, and `test_subset_oracle` splits them. @pytest.mark.parametrize("layout", ["csr", "csc", "dense"]) -@pytest.mark.parametrize("fmt", FORMATS) +@pytest.mark.parametrize("fmt", ["h5ad", "zarr31"]) def test_split_in_tiny_blocks_and_a_tiny_write_budget_matches_anndata( tmp_path, monkeypatch, fmt, layout ): diff --git a/tests/test_performance.py b/tests/test_performance.py index d79e888..050b158 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -2651,7 +2651,7 @@ def measure(n: int) -> int: # as bytes decoded or encoded many times over, without a clock. #: Rows of the stores below, and --chunk: blocks far smaller than a unit. -UNIT_ROWS, UNIT_CHUNK = 8192, 256 +UNIT_ROWS, UNIT_CHUNK = 4096, 256 def _x_amplification(io, *, encoded: bool = True) -> float: @@ -2746,9 +2746,9 @@ def test_a_sharded_output_is_written_a_whole_shard_at_a_time(tmp_path, monkeypat ) -@pytest.mark.parametrize("fmt,chunks", [("h5ad", (UNIT_ROWS, 8)), ("zarr2", None)], +@pytest.mark.parametrize("fmt,n_obs,chunks", [("h5ad", 8192, (8192, 8)), ("zarr2", 32768, None)], ids=["tall-chunked-h5ad", "zarr2-default"]) -def test_a_dense_subset_decodes_each_source_chunk_once(tmp_path, fmt, chunks): +def test_a_dense_subset_decodes_each_source_chunk_once(tmp_path, fmt, n_obs, chunks): """F061: blocks of --chunk rows were not aligned to the source's chunk height, so every chunk was decoded once a block: 15 times the matrix for gzip chunks of (n_obs, 8) at 16,384 rows (60 at 65,536), 7 for @@ -2756,13 +2756,13 @@ def test_a_dense_subset_decodes_each_source_chunk_once(tmp_path, fmt, chunks): chunk heights now, within the read budget.""" from tests.perf_fixtures import matrix_store - source = matrix_store(tmp_path, n_obs=4 * UNIT_ROWS if fmt == "zarr2" else UNIT_ROWS, + source = matrix_store(tmp_path, n_obs=n_obs, n_var=64, layout="dense", fmt=fmt, density=0.5, dense_chunks=chunks, compression="gzip" if chunks else None) with count_io() as io: _subset(source, tmp_path / f"o{source.suffix}", chunk_rows=1024) x = [n for n in io.by_name if n.strip("/") == "X"] - assert io.read_bytes >= 4 * UNIT_ROWS * 64, "the matrix was read" + assert io.read_bytes >= n_obs * 64 * 4, "the matrix was read" assert io.read_amplification <= 1.25, ( f"subset decoded {io.read_amplification:.1f} times the bytes it read " f"({fmt}, X read as {x}): a block splits the source's chunks" From e48bd1cb64abad1b60f3a465425ad7bd1848aa2a Mon Sep 17 00:00:00 2001 From: Aljes Date: Mon, 2 Nov 2026 18:55:55 +0000 Subject: [PATCH 3/5] Address review: unit memory, a shared index cache, aligned reads everywhere RowWriter fills one preallocated unit in place and writes whole units from the block itself; the unit budget is 8 MiB per output, data and indices sharing it, and 128 MiB across split's outputs. A guard checks peak allocation against the unit. The shard-index cache is one 64 MiB budget for the process, its entries dropped with their store. concat, recast (and its cast check) and export read whole source chunks too (row_bands). recast --layout dense writes through RowWriter, and a Zarr output with no source chunks takes zarr's own, cut to the budget. Tile mode no longer cuts chunks when one band covers every row. split writes 128 outputs a pass in any format. Growable outputs clamp their shard to the entries they can hold. The RSS runner raises the open-file limit. Co-Authored-By: Claude Opus 5.5 --- docs/guides/pipelines.md | 7 +- docs/reference/commands.md | 27 +++-- docs/reference/testing.md | 5 +- docs/troubleshooting.md | 6 +- src/adata/core/concat.py | 7 +- src/adata/core/recast.py | 25 +++-- src/adata/core/subset.py | 60 +++++++--- src/adata/formats/array.py | 14 ++- src/adata/formats/mtx.py | 14 ++- src/adata/storage/__init__.py | 202 +++++++++++++++++++++++++--------- tests/perf_counters.py | 9 ++ tests/perf_fixtures.py | 9 +- tests/test_block_layout.py | 43 +++++++- tests/test_performance.py | 72 ++++++++++-- 14 files changed, 376 insertions(+), 124 deletions(-) diff --git a/docs/guides/pipelines.md b/docs/guides/pipelines.md index 7216b80..d651ab5 100644 --- a/docs/guides/pipelines.md +++ b/docs/guides/pipelines.md @@ -132,9 +132,10 @@ bounded: - A read takes whole chunks of the source, up to 32 MiB. - An output is written a whole shard or chunk at a time, so its rows are - held until one is full: at most 64 MiB, and 256 MiB across the outputs - `split` writes at once. -- `split` holds up to 128 HDF5 outputs open, about 0.6 MB each. + held until one is full: at most 8 MiB, and 128 MiB across the outputs + `split` writes at once. Encoding a shard holds about four times its + size again, one shard at a time. +- `split` writes up to 128 outputs a pass, an HDF5 one about 0.6 MB open. See [Writing outputs](../reference/commands.md#writing-outputs). [Benchmarks](../reference/benchmarks.md) gives peak memory per command on a diff --git a/docs/reference/commands.md b/docs/reference/commands.md index 58d1747..cc6142a 100644 --- a/docs/reference/commands.md +++ b/docs/reference/commands.md @@ -304,10 +304,10 @@ all of them are complete, and the manifest is written last, so a split that fails leaves no partial stores. An output that would land on the input is refused, on a case-insensitive filesystem too (a label `A` beside `a.h5ad`). -Each matrix is read once for up to 128 HDF5 outputs (or as many Zarr -outputs as a quarter of the open-file limit allows): an open HDF5 file costs -about 0.6 MB however little it holds, so with more groups than that, the -matrices are read once per 128. +Each matrix is read once for up to 128 outputs: an open HDF5 file costs +about 0.6 MB however little it holds, and each output holds rows to write +its shards and chunks whole. With more groups than that, the matrices are +read once per 128. ## `concat` @@ -623,18 +623,23 @@ whole one at a time: once, whole. A shard is rewritten in full for a write of any part of it, so writing one a `--chunk` block at a time took hundreds of times longer. To write a shard whole, the command holds its rows: a shard, or chunk, - whose rows across the matrix come to more than 64 MiB is cut along its - rows to fit (anndata aims at 1 GB shards). `split` shares 256 MiB among - the outputs it writes at once, so with many groups their shards and - chunks are cut further. + whose rows across the matrix come to more than 8 MiB is cut along its + rows to fit (anndata aims at 1 GB shards). While zarr encodes a unit it + holds about four times that again, so this costs up to about 40 MB. + `split` shares 128 MiB among the outputs it writes at once, so with many + groups their shards and chunks are cut further. A Zarr output of a store + with no chunks of its own (a contiguous h5ad X) takes zarr's own chunks, + cut the same way. - Zarr v2 has no shards: written from a sharded store, a shard becomes the chunk. An HDF5 output of a sharded store lets HDF5 choose its chunks, rather than taking the shard's inner chunk. - A dense matrix is read in whole chunks of the source: `--chunk` rows rounded up to whole chunk heights, up to 32 MiB a read (more if - `--chunk` asks for more). Where one chunk height across every column is - over that, as for chunks of (n_obs, 8), a read takes a band of columns, - whole chunks at a time. + `--chunk` asks for more), in `subset`, `split`, `concat`, `recast`, + `convert` and `export array` alike. Where one chunk height across every + column is over that, as for chunks of (n_obs, 8), `subset`, `split` and + `convert` read a band of columns, whole chunks at a time; `concat`, + `recast` and `export` read `--chunk` rows, as before. - A sparse matrix in `subset` and `split` is read about as many entries a block as `--chunk` rows hold on average, in either layout; a column of a CSC matrix holding more, a gene expressed in every cell, is read in parts. diff --git a/docs/reference/testing.md b/docs/reference/testing.md index afbf7d6..f6b616e 100644 --- a/docs/reference/testing.md +++ b/docs/reference/testing.md @@ -373,12 +373,13 @@ should scale with, scale only that, and hold everything else fixed. | `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 | | `subset`, `split`, `concat`, `recast`, `convert` into anndata 0.13's shards | Zarr v3 as zarr 3.1 wrote it | each shard **encoded once** | -| dense `subset` | tall gzip chunks, zarr's own chunks | each chunk **decoded once** | +| dense `subset`, `concat`, `recast`, `export array` | tall gzip chunks, zarr's own chunks | each chunk **decoded once** | +| `subset` into a sharded output | 1 MiB units; dense, CSR | **peak allocation** a few units, not the shard | | `split` into chunked HDF5 | eight groups | each output chunk **encoded once** | | `subset`, `concat` of a sharded store | Zarr v3 as zarr 3.1 wrote it | each shard index **fetched once** | | sparse `subset` | n_obs; CSR, CSC, a gene in every cell | the **largest read** about `--chunk` rows of entries | | `convert` of a matrix chunked (n_obs, 8) | n_obs | the **largest read** a tile within the budget | -| `split` | 128 against 512 groups | **resident memory** flat past 128 HDF5 outputs | +| `split` | 128 against 512 groups | **resident memory** flat past 128 outputs a pass | 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/docs/troubleshooting.md b/docs/troubleshooting.md index c76df14..474756f 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -339,8 +339,8 @@ store it wrote; for a store another tool changed, run Lower `--chunk`: a dense block holds `--chunk` × n_var values, 1,024 rows by default. Two things are held whole by design: obs and var columns, and a pycisTopic pickle, which cannot be read in part. Writing an output a whole -shard or chunk at a time holds at most 64 MiB more (256 MiB across -`split`'s outputs). See +shard or chunk at a time holds at most 8 MiB more (128 MiB across +`split`'s outputs), and about four times a shard while it is encoded. See [Pipelines and containers](guides/pipelines.md#memory-and---chunk). ### It is slow on a Zarr store anndata 0.13 wrote @@ -352,7 +352,7 @@ block they wrote and read its index again for every block they read, hundreds of times slower than the same data as `.h5ad`. They write each shard once now, keep each index once read, and keep the source's shards on output. To cap what writing a shard whole holds, a shard whose rows come to -more than 64 MiB is cut along its rows. See +more than 8 MiB is cut along its rows. See [Writing outputs](reference/commands.md#writing-outputs). ### `A loom holds its matrices dense: 54.0 GB to compress and write. It is written, but may take a while` diff --git a/src/adata/core/concat.py b/src/adata/core/concat.py index 86ea384..e1f567e 100644 --- a/src/adata/core/concat.py +++ b/src/adata/core/concat.py @@ -54,6 +54,7 @@ from adata.elements.strings import is_string_dtype from adata.storage import ( RowWriter, + row_bands, copy_tree, create_dataset, dataset_create_kwargs, @@ -956,8 +957,7 @@ def _concat_dense( gaps = targets.size < n_cols n_src_rows = source.shape[0] - for start in range(0, n_src_rows, chunk_rows): - end = min(start + chunk_rows, n_src_rows) + for start, end in row_bands(source, chunk_rows): block = (np.full((end - start, n_cols), fill_value, dtype=dtype) if gaps else np.empty((end - start, n_cols), dtype=dtype)) block[:, targets] = np.asarray(source[start:end, :])[:, keep] @@ -1322,8 +1322,7 @@ def _concat_obsm( writer = RowWriter(out) for source in inputs.each(path): - for start in range(0, source.shape[0], chunk_rows): - end = min(start + chunk_rows, source.shape[0]) + for start, end in row_bands(source, chunk_rows): writer.write(source[start:end, ...]) writer.flush() return True diff --git a/src/adata/core/recast.py b/src/adata/core/recast.py index 92f3ccc..fb4308a 100644 --- a/src/adata/core/recast.py +++ b/src/adata/core/recast.py @@ -41,6 +41,8 @@ from adata.core.subset import _sparse_dataset from adata.storage import ( RowWriter, + backend_chunks, + row_bands, copy_attrs, create_dataset, dataset_create_kwargs, @@ -163,8 +165,9 @@ def check_cast( np.iinfo(target) if target.kind in "iu" else None ) - for start in range(0, n, step): - block = np.asarray(dataset[start : min(start + step, n)]) + # In whole chunks of it: the check reads every value, as the cast does. + for start, end in row_bands(dataset, step) if shape else (): + block = np.asarray(dataset[start:end]) report.total += block.size if block.size == 0: continue @@ -683,8 +686,14 @@ def _densify_rows( """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" + # zarr's own chunks, cut to what `RowWriter` holds, so each is written + # once: --chunk rows into chunks 12,500 tall encoded 13 times the matrix. + kw = fit_chunks( + backend_chunks(dst_parent, {}, (n_rows, n_cols), data_dtype), (n_rows, n_cols), + backend=backend, itemsize=np.dtype(data_dtype).itemsize, + ) dst = create_dataset( - dst_parent, name, shape=(n_rows, n_cols), dtype=data_dtype + dst_parent, name, shape=(n_rows, n_cols), dtype=data_dtype, **kw ) copy_attrs(attrs, dst.attrs, target_backend=backend) spec.set_encoding(dst, spec.ARRAY) @@ -696,6 +705,7 @@ def _densify_rows( indptr = np.asarray(src.obj["indptr"][...], dtype=np.int64) indices, values = src.obj["indices"], src.obj["data"] + writer = RowWriter(dst) for lo in range(0, n_rows, chunk_rows): hi = min(lo + chunk_rows, n_rows) start, end = int(indptr[lo]), int(indptr[hi]) @@ -711,7 +721,8 @@ def _densify_rows( # Plain assignment keeps whichever came last, so a # non-canonical input silently changed value on densifying. np.add.at(block, (major, minor), data) - dst[lo:hi, :] = block + writer.write(block) + writer.flush() def sparsify( @@ -825,8 +836,7 @@ def _sparsify_rows( pointer_limit = int(np.iinfo(offsets_within).max) nnz = 0 - for lo in range(0, n_rows, chunk_rows): - hi = min(lo + chunk_rows, n_rows) + for lo, hi in row_bands(src.obj, chunk_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. @@ -948,8 +958,7 @@ def cast_dense( # A whole shard at a time: a block of --chunk rows rewrote the source's # shard once a block, 537 MB encoded for 4.4 MB written (F036). writer = RowWriter(dst) - for lo in range(0, n_rows, chunk_rows): - hi = min(lo + chunk_rows, n_rows) + for lo, hi in row_bands(src.obj, chunk_rows): writer.write(np.asarray(src.obj[lo:hi, :]).astype(data_dtype, copy=False)) writer.flush() diff --git a/src/adata/core/subset.py b/src/adata/core/subset.py index 20d1f5b..2d3d427 100644 --- a/src/adata/core/subset.py +++ b/src/adata/core/subset.py @@ -35,6 +35,7 @@ RowWriter, create_dataset, copy_attrs, + backend_chunks, copy_tree, cut_unit, dataset_create_kwargs, @@ -490,7 +491,17 @@ def fan_out_dense_matrix( src, target_backend=target_backend, dst_parent=dst_parent ) budget = share if share is not None else WRITE_UNIT_BYTES - kw = fit_chunks(kw, shape, backend=target_backend, itemsize=itemsize, unit_bytes=budget) + if src.dtype.kind not in ("O", "T"): + kw = backend_chunks(dst_parent, kw, shape, src.dtype) + # One band of tiles covering every row (chunks of (n_obs, 8)) gives + # each tile writer whole columns of chunks at once: nothing is held, + # and the chunks need no cut. Otherwise a unit's rows across the + # output are what its writers hold together. + kw = fit_chunks( + kw, shape, backend=target_backend, + itemsize=None if width is not None and band >= n_obs else itemsize, + unit_bytes=budget, + ) if width is not None and kw.get("shards"): kw["shards"] = tile_shards(kw["chunks"], kw["shards"], width, rows=False) dst = create_dataset( @@ -558,6 +569,7 @@ def __init__( remap: Optional[np.ndarray], n_major: int, share: Optional[int] = None, + size: Optional[int] = None, ) -> None: self.group = dst_parent.create_group(name) copy_attrs(src.attrs, self.group.attrs, target_backend=_target_backend(dst_parent)) @@ -576,11 +588,11 @@ def __init__( width = data_dtype.itemsize + index_dtype.itemsize self.data = _growable_like( self.group, "data", data_dtype, src["data"], - share=max(1, share * data_dtype.itemsize // width), + share=max(1, share * data_dtype.itemsize // width), size=size, ) self.indices = _growable_like( self.group, "indices", index_dtype, src["indices"], - share=max(1, share * index_dtype.itemsize // width), + share=max(1, share * index_dtype.itemsize // width), size=size, ) self.selection = major_idx self.remap = remap @@ -761,6 +773,7 @@ def fan_out_sparse_matrix( n_major, n_minor = (n_rows, n_cols) if csr else (n_cols, n_rows) outputs: List[_SparseOutput] = [] + lengths = np.diff(indptr) for target in targets: # A fourth element names this target's matrix, where it differs # from `name`: one read can then fill X and a layer alike. @@ -783,9 +796,12 @@ def fan_out_sparse_matrix( _minor_remap(minor_keep, n_minor), n_major, share, + # The most entries it can hold: its majors' entries. + int(lengths.sum()) if major_idx is None else int(lengths[major_idx].sum()), ) ) + del lengths budget = sparse_block_entries(n_rows, int(indptr[-1]) if indptr.size else 0, chunk_rows) for first, stop, lo, hi in _sparse_blocks( indptr, budget, [out.selection for out in outputs] @@ -840,7 +856,8 @@ def _sparse_dataset( kw = dataset_create_kwargs(template, target_backend=backend, dst_parent=group) # A unit cut to what `RowWriter` holds: writers fill these a block at # a time, and anndata's shard of a whole array may be a gigabyte. - kw = fit_chunks(kw, (n,), backend=backend, itemsize=np.dtype(dtype).itemsize) + kw = fit_chunks(kw, (n,), backend=backend, itemsize=np.dtype(dtype).itemsize, + unit_bytes=WRITE_UNIT_BYTES // 2) if compression is not None: kw = with_compression( kw, compression, shape=(n,), dtype=dtype, backend=backend, @@ -859,6 +876,7 @@ def _growable_like( *, compression: Optional[str] = None, share: Optional[int] = None, + size: Optional[int] = None, ) -> RowWriter: """Like `_sparse_dataset`, but extensible for a size not yet known: a `RowWriter` that grows it, appended to with `_append`. Flush it once the @@ -869,7 +887,10 @@ def _growable_like( `indices` were written uncompressed whatever the input. They dropped its shard too, and kept the inner chunk: 245 elements, a file each, under anndata 0.13 (F033). The shard is kept now, cut to `share` bytes (by - default `WRITE_UNIT_BYTES`), and written whole. + default half of `WRITE_UNIT_BYTES`: `data` and `indices` come as a + pair), and written whole. `size`, an upper bound on the entries it will + hold, clamps the shard (and the chunk) to it: an output of a few + thousand entries carried the source's shard of millions, and its index. """ backend = _target_backend(group) kw = dataset_create_kwargs(template, target_backend=backend, dst_parent=group) @@ -882,7 +903,13 @@ def _growable_like( kw["chunks"] = (int(chunks[0]),) if chunks else (1 << 16,) if kw.get("shards"): kw["shards"] = (int(kw["shards"][0]),) - budget = WRITE_UNIT_BYTES if share is None else share + if size is not None and size > 0: + step = min(kw["chunks"][0], int(size)) + kw["chunks"] = (step,) + if kw.get("shards"): + kw["shards"] = (max(step, min(kw["shards"][0], -(-int(size) // step) * step)),) + kw["shards"] = ((kw["shards"][0] // step) * step,) + budget = WRITE_UNIT_BYTES // 2 if share is None else share kw = cut_unit(kw, (0,), np.dtype(dtype).itemsize, budget) if is_zarr_group(group): ds = group.create_array(name, shape=(0,), dtype=dtype, **kw) @@ -1465,14 +1492,15 @@ def _open_output_limit() -> int: #: rather than running out of file descriptors. MAX_OPEN_OUTPUTS = _open_output_limit() -#: How many HDF5 outputs a pass holds open, whatever the descriptor limit -#: allows. An open HDF5 file costs about 0.6 MB however it is opened -- the -#: library's own per-file structures; a chunk cache or a smaller metadata -#: cache changes nothing -- so with a container's high limit, 1,024 outputs -#: a pass took RSS to about 1 GB, following the number of groups rather than -#: --chunk (F079). 128 holds that near 80 MB, for a pass more per 128 -#: groups. Zarr outputs hold no descriptor and no such cost. -MAX_OPEN_HDF5_OUTPUTS = 128 +#: How many outputs a pass writes, whatever the descriptor limit allows. An +#: open HDF5 file costs about 0.6 MB however it is opened -- the library's +#: own per-file structures; a chunk cache or a smaller metadata cache +#: changes nothing -- so with a container's high limit, 1,024 outputs a pass +#: took RSS to about 1 GB, following the number of groups rather than +#: --chunk (F079). And every output's units are cut to its share of +#: `WRITE_BUFFER_BYTES`, a row or two each at 1,024 of them. 128 holds the +#: files near 80 MB and a share at 1 MiB, for a pass more per 128 groups. +MAX_OUTPUTS_PER_PASS = 128 def split_h5ad( @@ -1493,9 +1521,7 @@ def split_h5ad( with open_store(file, "r") as probe: zarr_format = probe.zarr_format - per_pass = MAX_OPEN_OUTPUTS - if any(detect_backend(Path(path)) == "hdf5" for path, _, _ in outputs): - per_pass = min(per_pass, MAX_OPEN_HDF5_OUTPUTS) + per_pass = min(MAX_OPEN_OUTPUTS, MAX_OUTPUTS_PER_PASS) for start in range(0, len(outputs), per_pass): batch = outputs[start : start + per_pass] with ExitStack() as stack: diff --git a/src/adata/formats/array.py b/src/adata/formats/array.py index fa5bdb2..7804657 100644 --- a/src/adata/formats/array.py +++ b/src/adata/formats/array.py @@ -24,18 +24,24 @@ from adata.elements.write import write_dense, write_mapping, write_scalar from adata.formats.common import _get_encoding_type, _resolve, stdout_reader_may_stop from adata.formats.validate import check_target, validate_dimensions -from adata.storage import create_dataset, is_dataset, is_group, is_zarr_group +from adata.storage import create_dataset, is_dataset, is_group, is_zarr_group, row_bands from adata.elements.strings import is_string_dtype #: Elements of a `.npy` copied into the store at once by `import array`. IMPORT_CHUNK_ELEMENTS = 1 << 20 -def _row_blocks(shape: Tuple[int, ...], chunk_elements: int) -> Iterator[Tuple[int, int]]: - """[start, end) row ranges of about `chunk_elements` elements each.""" +def _row_blocks( + shape: Tuple[int, ...], chunk_elements: int, array: Any = None +) -> Iterator[Tuple[int, int]]: + """[start, end) row ranges of about `chunk_elements` elements each; of + `array`, rounded to whole chunks of it (`storage.row_bands`).""" n = int(shape[0]) row = int(np.prod(shape[1:])) if len(shape) > 1 else 1 step = max(1, int(chunk_elements) // max(1, row)) + if array is not None: + yield from row_bands(array, step) + return for start in range(0, n, step): yield start, min(start + step, n) @@ -184,7 +190,7 @@ def convert(start: Optional[int], end: Optional[int]) -> np.ndarray: fh.write(convert(None, None).tobytes()) else: n = shape[0] - for start, end in _row_blocks(shape, chunk_elements): + for start, end in _row_blocks(shape, chunk_elements, values): if progress: progress.update(f"[magenta]Exporting {obj}: {start}-{end} of {n}...[/]") fh.write(convert(start, end).tobytes()) diff --git a/src/adata/formats/mtx.py b/src/adata/formats/mtx.py index fd7bfa4..5d84fc1 100644 --- a/src/adata/formats/mtx.py +++ b/src/adata/formats/mtx.py @@ -471,9 +471,10 @@ def iter_matrix_entries( yield (cols, rows, values) if transpose else (rows, cols, values) return - n_rows = matrix.shape[0] - for lo in range(0, n_rows, chunk_rows): - block = np.asarray(matrix.obj[lo : min(lo + chunk_rows, n_rows)]) + from adata.storage import row_bands + + for lo, hi in row_bands(matrix.obj, chunk_rows): + block = np.asarray(matrix.obj[lo:hi]) r, c = np.nonzero(block) rows = r.astype(np.int64) + lo cols = c.astype(np.int64) @@ -499,10 +500,11 @@ def count_entries(matrix: Matrix, *, chunk_rows: int = 1024) -> int: """The entries a matrix writes: its stored nonzeros, or a dense one's nonzeros.""" if matrix.sparse: return matrix.nnz + from adata.storage import row_bands + total = 0 - n_rows = matrix.shape[0] - for lo in range(0, n_rows, chunk_rows): - total += int(np.count_nonzero(np.asarray(matrix.obj[lo : min(lo + chunk_rows, n_rows)]))) + for lo, hi in row_bands(matrix.obj, chunk_rows): + total += int(np.count_nonzero(np.asarray(matrix.obj[lo:hi]))) return total diff --git a/src/adata/storage/__init__.py b/src/adata/storage/__init__.py index 31df0cb..e5d78c4 100644 --- a/src/adata/storage/__init__.py +++ b/src/adata/storage/__init__.py @@ -2,7 +2,7 @@ from dataclasses import dataclass from pathlib import Path -from typing import Any, Iterable, Optional, Sequence +from typing import Any, Iterable, Iterator, Optional, Sequence import shutil import warnings @@ -125,11 +125,33 @@ def detect_backend(path: Path) -> str: #: followed the number of outputs, eight times faster. 1 MiB, as before. HDF5_CHUNK_CACHE = 1 << 20 -#: Bytes of Zarr shard indexes a store opened for reading keeps (see -#: `_shard_index_cache`): 16 bytes an inner chunk, so the index of a 1 GB -#: shard of 1 KB chunks, as anndata 0.13 wrote with zarr 3.1, is 16 MB. +#: Bytes of Zarr shard indexes kept, across every store opened for reading +#: (see `_shard_index_cache`): 16 bytes an inner chunk, so the index of a +#: 1 GB shard of 1 KB chunks, as anndata 0.13 wrote with zarr 3.1, is 16 MB. +#: One budget for the process: concat holds up to 64 inputs open. SHARD_INDEX_CACHE_BYTES = 64 * 1024 * 1024 +#: The shard indexes kept, least recently used first: `(store, key, +#: length)` to bytes, `store` a number each wrapped store takes once. +_SHARD_INDEXES: "dict" = {} +_SHARD_INDEX_STATE = {"bytes": 0, "next": 0} + + +def _drop_shard_index(slot: tuple) -> None: + value = _SHARD_INDEXES.pop(slot, None) + if value is not None: + _SHARD_INDEX_STATE["bytes"] -= len(value) + + +def _forget_shard_indexes(store_id: int) -> None: + """Drop what a store kept, once nothing can read through it again. + + A finalizer, so it can run in the middle of anything, zarr's own thread + included: it walks a copy of the keys, and pops what is still there. + """ + for slot in [slot for slot in list(_SHARD_INDEXES) if slot[0] == store_id]: + _drop_shard_index(slot) + def _shard_index_cache(store: Any) -> Any: """`store`, keeping the shard indexes read through it (F062). @@ -141,13 +163,16 @@ def _shard_index_cache(store: Any) -> Any: the square of the array: subset of an anndata-default store read 4.2 times the store, concat 27.7 times. A store opened for reading cannot change underneath, so each index is read once and kept, least recently - used first out, up to `SHARD_INDEX_CACHE_BYTES`. + used first out, up to `SHARD_INDEX_CACHE_BYTES` across every store + (not each: concat holds up to 64 open). Each wrapped store keys its own + entries, which go when it does, so a store written again at the same + path is never read through an older one's indexes. The index is the one read zarr makes by suffix (its default `index_location`, and anndata's). Inner chunks are read by range, and passed through. """ - from collections import OrderedDict + import weakref try: from zarr.abc.store import SuffixByteRequest @@ -158,8 +183,9 @@ def _shard_index_cache(store: Any) -> Any: class ShardIndexCache(WrapperStore): def __init__(self, inner: Any) -> None: super().__init__(inner) - self._held: "OrderedDict[tuple, Any]" = OrderedDict() - self._held_bytes = 0 + self._id = _SHARD_INDEX_STATE["next"] + _SHARD_INDEX_STATE["next"] += 1 + weakref.finalize(self, _forget_shard_indexes, self._id) def _with_store(self, inner: Any) -> Any: return type(self)(inner) @@ -167,21 +193,20 @@ def _with_store(self, inner: Any) -> Any: def _lookup(self, key: str, byte_range: Any) -> tuple: if not isinstance(byte_range, SuffixByteRequest): return None, None - slot = (key, int(byte_range.suffix)) - if slot in self._held: - self._held.move_to_end(slot) - return slot, self._held[slot] - return slot, None + slot = (self._id, key, int(byte_range.suffix)) + value = _SHARD_INDEXES.pop(slot, None) + if value is not None: + _SHARD_INDEXES[slot] = value # now the most recently used + return slot, value def _keep(self, slot: tuple, value: Any) -> None: size = len(value) - if size > SHARD_INDEX_CACHE_BYTES: + if size > SHARD_INDEX_CACHE_BYTES or slot in _SHARD_INDEXES: return - self._held[slot] = value - self._held_bytes += size - while self._held_bytes > SHARD_INDEX_CACHE_BYTES: - _, dropped = self._held.popitem(last=False) - self._held_bytes -= len(dropped) + _SHARD_INDEXES[slot] = value + _SHARD_INDEX_STATE["bytes"] += size + while _SHARD_INDEX_STATE["bytes"] > SHARD_INDEX_CACHE_BYTES and _SHARD_INDEXES: + _drop_shard_index(next(iter(list(_SHARD_INDEXES)))) async def get(self, key: str, prototype: Any, byte_range: Any = None) -> Any: slot, value = self._lookup(key, byte_range) @@ -833,16 +858,19 @@ def _zarr_is_compressed(src: Any) -> bool: #: The most one output holds to write a whole unit -- a shard, or a chunk #: that a write of part of would decode and encode again -- at once -#: (`RowWriter`). A unit whose rows across the array's width come to more is -#: cut along its rows when the output is created (`fit_chunks`): anndata -#: 0.13 aims at 1 GB shards, which a writer filling them from blocks of -#: --chunk rows would otherwise hold, or rewrite once a block. -WRITE_UNIT_BYTES = 64 * 1024 * 1024 +#: (`RowWriter`), a sparse matrix's `data` and `indices` together. A unit +#: whose rows across the array's width come to more is cut along its rows +#: when the output is created (`fit_chunks`): anndata 0.13 aims at 1 GB +#: shards, which a writer filling them from blocks of --chunk rows would +#: otherwise hold, or rewrite once a block. While zarr encodes a unit it +#: holds about four times as much again, so 8 MiB costs about 40 MB at most +#: (64 MiB cost about 300 MB). +WRITE_UNIT_BYTES = 8 * 1024 * 1024 #: The most one pass writing several outputs at once (`split`) holds to #: write whole units, across all of them: each output's units are cut to a #: share of it (`write_share`). -WRITE_BUFFER_BYTES = 256 * 1024 * 1024 +WRITE_BUFFER_BYTES = 128 * 1024 * 1024 def write_share(outputs: int) -> int: @@ -888,6 +916,26 @@ def cut_unit(kw: dict, shape: Sequence[int], itemsize: int, budget: int) -> dict return kw +def backend_chunks(parent: Any, kw: dict, shape: Sequence[int], dtype: Any) -> dict: + """`kw` with the chunks a Zarr array would take by default made explicit, + where nothing gave it any (a contiguous h5ad source), so `fit_chunks` + can cut them. zarr's own choice for a wide matrix can span thousands of + rows across every column, more than a writer may hold to write it whole. + zarr is asked on an array of the same shape in memory.""" + if not is_zarr_group(parent) or kw.get("chunks") or not shape or 0 in tuple(shape): + return kw + try: + from zarr.storage import MemoryStore + + probe = zarr.create_array( + store=MemoryStore(), shape=tuple(int(n) for n in shape), dtype=np.dtype(dtype), + zarr_format=zarr_format_of(parent) or 3, + ) + return {**kw, "chunks": tuple(int(c) for c in probe.chunks)} + except Exception: # pragma: no cover - a dtype zarr cannot probe + return kw + + def fit_chunks( kw: dict, shape: Sequence[int], @@ -1218,6 +1266,31 @@ def tile_shards( return tuple(shards) +def row_bands(array: Any, rows: int) -> Iterator[tuple]: + """`[start, end)` row ranges for reading `array` whole rows at a time: + `rows` rounded to whole chunk heights of it, within the larger of `rows` + rows and `TARGET_READ_BYTES` (`read_blocks`). Where a chunk height + across the array is over that, `rows` as they are, as before: a reader + of whole rows cannot take tiles. + + Blocks of --chunk rows split a chunk taller than --chunk across reads, + each decoding all of it: concat and recast decoded a gzip matrix in + chunks of (n_obs, 8) 15 times over, `export array` 11 (F061).""" + shape = tuple(array.shape) + if not shape: + return + rows = max(1, int(rows)) + dtype = getattr(array, "dtype", None) + itemsize = VLEN_ELEMENT_BYTES if _is_vlen(dtype) else int(dtype.itemsize) + band, width = read_blocks( + shape, decode_chunks(array), itemsize, rows=rows, + budget=max(rows * itemsize * _prod(shape[1:]), TARGET_READ_BYTES), + ) + step = rows if width is not None else band + for start in range(0, int(shape[0]), step): + yield start, min(start + step, int(shape[0])) + + def write_unit_rows(ds: Any) -> int: """Rows a write to `ds` should cover, in multiples of, for each unit it stores -- a Zarr shard, or a chunk -- to be encoded once. @@ -1247,9 +1320,14 @@ class RowWriter: and written in whole units, so a writer filling an output a block of --chunk rows at a time no longer rewrites a shard -- 1 GB, by anndata 0.13's default -- once a block (F031, F036), nor decodes and encodes a - chunk once a block of `split`'s many outputs (F060). It holds less than - one unit of rows, plus the block it was given: `fit_chunks` cuts a unit - to fit its budget when the output is created. + chunk once a block of `split`'s many outputs (F060). + + It holds one unit's rows, in a buffer filled in place, plus the block it + is given; whole units within a block are written from the block itself. + `fit_chunks` cuts a unit to fit its budget when the output is created. + Gathering a list of blocks and joining them, as it first did, held the + unit twice over at each write, and zarr holds the shard it encodes on + top of that: a 64 MiB unit cost about 300 MB. `grow` extends a 1-D array of unknown length (a sparse matrix's `data` or `indices`) as it goes. `cols`, a `(start, stop)`, writes only those @@ -1277,8 +1355,8 @@ def __init__( if self.unit * row > (WRITE_UNIT_BYTES if budget is None else budget): self.unit = 1 self.written = int(start) - self.pending: list = [] - self.pending_rows = 0 + self.buffer: Optional[np.ndarray] = None + self.filled = 0 @property def dtype(self) -> Any: @@ -1287,38 +1365,60 @@ def dtype(self) -> Any: @property def rows(self) -> int: """Rows given so far, written or not.""" - return self.written + self.pending_rows + return self.written + self.filled def write(self, values: Any) -> None: values = np.asarray(values) if not values.shape or values.shape[0] == 0: return - end = ((self.rows + values.shape[0]) // self.unit) * self.unit - # Kept past this call, rows are copied: the caller may reuse its - # buffer for the next block. - self.pending.append(values if end > self.written else values.copy()) - self.pending_rows += values.shape[0] - if end > self.written: - self._commit(end) + if self.unit == 1: + self._put(values) + return + at = 0 + if self.filled: + # Complete the unit already begun. + at = self._buffer(values[: self.unit - self.filled]) + if self.filled == self.unit: + self._put(self.buffer) + self.filled = 0 + whole = (values.shape[0] - at) // self.unit * self.unit + if whole: + # Whole units in the block go from the block itself. + self._put(values[at : at + whole]) + at += whole + if at < values.shape[0]: + self._buffer(values[at:]) def flush(self) -> None: - if self.pending_rows: - self._commit(self.rows) - - def _commit(self, end: int) -> None: - values = self.pending[0] if len(self.pending) == 1 else np.concatenate(self.pending) - count = end - self.written + if self.filled: + self._put(self.buffer[: self.filled]) + self.filled = 0 + self.buffer = None + + def _buffer(self, values: np.ndarray) -> int: + """Copy `values` into the unit's buffer; how many rows that took.""" + buffer = self.buffer + if buffer is None or buffer.shape[1:] != values.shape[1:]: + buffer = np.empty((self.unit,) + values.shape[1:], dtype=values.dtype) + elif buffer.dtype != values.dtype: + # Text read a block at a time can come wider in a later block. + buffer = buffer.astype(np.result_type(buffer.dtype, values.dtype)) + self.buffer = buffer + n = values.shape[0] + buffer[self.filled : self.filled + n] = values + self.filled += n + return n + + def _put(self, values: np.ndarray) -> None: + end = self.written + values.shape[0] if self.grow and self.ds.shape[0] < end: self.ds.resize((end,) + tuple(self.ds.shape[1:])) if values.ndim == 1: - self.ds[self.written:end] = values[:count] + self.ds[self.written:end] = values elif self.cols is not None: - self.ds[self.written:end, self.cols, ...] = values[:count] + self.ds[self.written:end, self.cols, ...] = values else: - self.ds[self.written:end, ...] = values[:count] - rest = values[count:] - self.pending = [rest.copy()] if rest.shape[0] else [] - self.pending_rows = int(rest.shape[0]) + self.ds[self.written:end, ...] = values self.written = end @@ -1380,6 +1480,8 @@ def copy_dataset( kw, compression, shape=shape, dtype=src.dtype, backend=target_backend, zarr_format=zformat, ) + if not _is_vlen(src.dtype): + kw = backend_chunks(dst_group, kw, shape, src.dtype) # Tiles are written where they belong, whole chunks at a time, so only # whole rows need a unit cut to what `RowWriter` may hold. kw = fit_chunks( diff --git a/tests/perf_counters.py b/tests/perf_counters.py index f9d136a..b94cff8 100644 --- a/tests/perf_counters.py +++ b/tests/perf_counters.py @@ -566,6 +566,15 @@ def high_water(): return value if sys.platform == "darwin" else value * 1024 # Linux: KiB argv, report = json.loads(sys.argv[1]), sys.argv[2] +# The open-file limit a cluster or a container gives, not macOS's 256: the +# limits it sets (outputs open at once) are what a guard measures. +soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) +want = 65536 if hard == resource.RLIM_INFINITY else min(hard, 65536) +if soft != resource.RLIM_INFINITY and soft < want: + try: + resource.setrlimit(resource.RLIMIT_NOFILE, (want, hard)) + except (ValueError, OSError): + pass from adata.cli import app before = high_water() diff --git a/tests/perf_fixtures.py b/tests/perf_fixtures.py index 433df17..a142f2e 100644 --- a/tests/perf_fixtures.py +++ b/tests/perf_fixtures.py @@ -146,12 +146,15 @@ def matrix_store( ZARR31_ELEMENT_CHUNK = 245 -def zarr31_layout(path: Path, key: str = "X") -> Path: +def zarr31_layout( + path: Path, key: str = "X", dense_chunks: Tuple[int, int] = ZARR31_DENSE_CHUNK +) -> Path: """Rewrite `key` of a Zarr v3 store in anndata 0.13's layout under zarr 3.1: a dense matrix in inner chunks of (16, 8) and shards of 4,096 rows across every column, a sparse one's `data` and `indices` in chunks of 245 elements and a shard of the whole array. Values and attributes - are kept.""" + are kept. `dense_chunks` gives a dense matrix other inner chunks, in the + same shards.""" import zarr def whole(n: int, c: int) -> int: @@ -166,7 +169,7 @@ def whole(n: int, c: int) -> int: array = parent[member] values, attrs = array[...], dict(array.attrs) if array.ndim == 2: - chunks = ZARR31_DENSE_CHUNK + chunks = dense_chunks shards = (whole(min(array.shape[0], 4096), chunks[0]), whole(array.shape[1], chunks[1])) else: diff --git a/tests/test_block_layout.py b/tests/test_block_layout.py index dfd1c7d..d667e25 100644 --- a/tests/test_block_layout.py +++ b/tests/test_block_layout.py @@ -188,11 +188,16 @@ def test_outputs_keep_the_sources_inner_chunks_and_shards(tmp_path, command, lay for name, array in _arrays(group).items(): want = source[name] assert array.chunks == want.chunks, (path.name, name) - shards = tuple( - min(s, max(c, -(-n // c) * c)) if array.ndim > 1 or name == "X" else s - for s, c, n in zip(want.shards, want.chunks, array.shape) - ) - assert array.shards == shards, (path.name, name, array.shards, shards) + # The source's, clamped to the output in whole chunks: a dense + # matrix to its shape, a sparse one's arrays to the most entries + # its majors hold, which a selection of minors can leave above + # the entries it keeps (CSC, selecting cells). + for s, c, n, got in zip(want.shards, want.chunks, array.shape, array.shards): + fitted = min(s, max(c, -(-n // c) * c)) + if array.ndim == 1 and layout == "csc" and command in ("subset", "split"): + assert fitted <= got <= s and got % c == 0, (path.name, name, got) + else: + assert got == fitted, (path.name, name, got, fitted) files = [f for f in (path / name).rglob("*") if f.is_file() and f.name not in ("zarr.json", ".zarray", ".zattrs")] assert len(files) <= 4, f"{name} of {path.name} is {len(files)} files" @@ -390,3 +395,31 @@ def read_in_blocks() -> int: assert read_in_blocks() == 1 monkeypatch.setattr(storage, "SHARD_INDEX_CACHE_BYTES", 0) assert read_in_blocks() > 10, "a cache too small for an index keeps none" + + +def test_the_shard_index_cache_is_one_budget_across_stores(tmp_path, monkeypatch): + """From the review: each open store kept up to 64 MiB of indexes, and + concat holds up to 64 inputs open, so the cache grew with the number of + inputs. One budget is shared now, and a store's indexes go with it.""" + import gc + + import adata.storage as storage + from adata.storage import open_store + + _, src = _sharded_counts(tmp_path, "csr") + with open_store(src, "r") as store: + store.root["X/data"][0:10] + one = storage._SHARD_INDEX_STATE["bytes"] + assert one > 0 + del store + gc.collect() + assert storage._SHARD_INDEX_STATE["bytes"] == 0, "a closed store's indexes stay" + + monkeypatch.setattr(storage, "SHARD_INDEX_CACHE_BYTES", 2 * one) + stores = [open_store(src, "r") for _ in range(5)] + for store in stores: + store.root["X/data"][0:10] + store.root["X/indices"][0:10] + assert 0 < storage._SHARD_INDEX_STATE["bytes"] <= 2 * one + for store in stores: + store.close() diff --git a/tests/test_performance.py b/tests/test_performance.py index 050b158..b3cad0b 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -2746,25 +2746,44 @@ def test_a_sharded_output_is_written_a_whole_shard_at_a_time(tmp_path, monkeypat ) -@pytest.mark.parametrize("fmt,n_obs,chunks", [("h5ad", 8192, (8192, 8)), ("zarr2", 32768, None)], - ids=["tall-chunked-h5ad", "zarr2-default"]) -def test_a_dense_subset_decodes_each_source_chunk_once(tmp_path, fmt, n_obs, chunks): +@pytest.mark.parametrize("command,fmt,n_obs,chunks", [ + ("subset", "h5ad", 8192, (8192, 8)), + ("subset", "zarr2", 32768, None), + ("concat", "h5ad", 8192, (8192, 8)), + ("recast --dtype", "h5ad", 8192, (8192, 8)), + ("recast --layout csr", "h5ad", 8192, (8192, 8)), + ("export array", "h5ad", 8192, (8192, 8)), +], ids=["subset-tall-h5ad", "subset-zarr2-default", "concat-tall-h5ad", + "recast-dtype-tall-h5ad", "recast-csr-tall-h5ad", "export-array-tall-h5ad"]) +def test_a_dense_read_decodes_each_source_chunk_once(tmp_path, command, fmt, n_obs, chunks): """F061: blocks of --chunk rows were not aligned to the source's chunk height, so every chunk was decoded once a block: 15 times the matrix for gzip chunks of (n_obs, 8) at 16,384 rows (60 at 65,536), 7 for - zarr's own chunks (8192, 16), 29 at 262,144 rows. Blocks are whole - chunk heights now, within the read budget.""" + zarr's own chunks (8192, 16), 29 at 262,144 rows. subset, concat, + recast and `export array` alike. Blocks are whole chunk heights now, + within the read budget.""" + from adata.commands.export import export_npy from tests.perf_fixtures import matrix_store source = matrix_store(tmp_path, n_obs=n_obs, n_var=64, layout="dense", fmt=fmt, density=0.5, dense_chunks=chunks, compression="gzip" if chunks else None) + out = tmp_path / f"o{source.suffix}" with count_io() as io: - _subset(source, tmp_path / f"o{source.suffix}", chunk_rows=1024) + if command == "subset": + _subset(source, out, chunk_rows=1024) + elif command == "concat": + concat_on_disk(_twice(source), out, QUIET, chunk_rows=1024, index_unique="-") + elif command == "recast --dtype": + _recast(source, out, dtype="float64", chunk_rows=1024) + elif command == "recast --layout csr": + _recast(source, out, layout="csr", chunk_rows=1024) + else: + export_npy(source, "X", tmp_path / "x.npy", 64 * 1024, QUIET) x = [n for n in io.by_name if n.strip("/") == "X"] assert io.read_bytes >= n_obs * 64 * 4, "the matrix was read" assert io.read_amplification <= 1.25, ( - f"subset decoded {io.read_amplification:.1f} times the bytes it read " + f"{command} decoded {io.read_amplification:.1f} times the bytes it read " f"({fmt}, X read as {x}): a block splits the source's chunks" ) @@ -2886,12 +2905,49 @@ def test_copying_a_column_chunked_matrix_reads_tiles_not_the_matrix(tmp_path, mo assert np.array_equal(zarr.open_group(str(out))["X"][...], handle["X"][...]) +@pytest.mark.parametrize("layout,n_obs,n_var,density", [ + ("dense", 4096, 2048, 0.5), ("csr", 2048, 2048, 0.25), +]) +def test_writing_whole_units_holds_a_few_units_not_the_shard( + tmp_path, monkeypatch, layout, n_obs, n_var, density +): + """F031, F036, and the review of their fix: writing a shard whole holds + its rows, so memory is set by the unit (`WRITE_UNIT_BYTES`), not by the + source's shard of the whole array. Gathering blocks in a list and + joining them held each unit twice over, and a sparse output gave its + `data` and `indices` a budget each: 335 MB to subset a 16M-entry CSR at + 64 MiB units, against 50 MB before. Here the unit is 1 MiB, and the + source a shard of the whole array, 32 MB dense (inner chunks 16 rows + tall, so a read is a few rows) or 8 MB of entries in anndata's own + layout: main held the shard or, for CSR, the matrix (F032).""" + import adata.core.subset as subset + import adata.storage as storage + from tests.perf_fixtures import matrix_store, zarr31_layout + + unit = 1 << 20 + monkeypatch.setattr(storage, "WRITE_UNIT_BYTES", unit, raising=False) + monkeypatch.setattr(subset, "WRITE_UNIT_BYTES", unit, raising=False) + source = matrix_store(tmp_path, n_obs=n_obs, n_var=n_var, layout=layout, + fmt="zarr3", density=density) + if layout == "dense": + zarr31_layout(source, dense_chunks=(16, 256)) + with count_allocations() as peak: + _subset(source, tmp_path / "o.zarr", chunk_rows=64) + # The unit, zarr's own copies while it encodes one, and a block. + assert peak[0] <= 10 * unit, ( + f"subset of a sharded {layout} X peaked at {peak[0] / 1e6:.1f} MB with " + f"{unit / 1e6:.1f} MB units: more than a few units are held" + ) + + @pytest.mark.slow def test_split_memory_stops_growing_with_the_group_count(tmp_path): """F079: split held up to 1,024 HDF5 outputs open in a pass, about 0.6 MB each however they are opened (HDF5's own structures for a file), so RSS followed the group count, not --chunk: 136 MB at 128 groups, 423 MB - at 512. A pass holds at most 128 HDF5 outputs open now.""" + at 512. A pass writes at most 128 outputs now. `peak_rss` raises the + open-file limit, as a cluster sets it, so macOS's 256 does not cap the + outputs first.""" from tests.perf_counters import peak_rss from tests.perf_fixtures import matrix_store From 8ddfb29d6321c5f6b4f5b9c29c13e8bf8c3f69a3 Mon Sep 17 00:00:00 2001 From: Aljes Date: Mon, 2 Nov 2026 19:05:31 +0000 Subject: [PATCH 4/5] Restate two guards for reads in whole source chunks The F029 cast-check guard bounds the largest read by --chunk rows or one chunk height of the source, whichever is taller; the F013 overflow test gives its source chunks 16 rows tall, so it still shows the refusal after the first block. Co-Authored-By: Claude Opus 5.5 --- tests/test_performance.py | 28 ++++++++++++++++------------ tests/test_recast_streaming.py | 11 +++++++++++ 2 files changed, 27 insertions(+), 12 deletions(-) diff --git a/tests/test_performance.py b/tests/test_performance.py index b3cad0b..077434c 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -2496,25 +2496,29 @@ 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.""" + matrix at --chunk 100. The largest read is --chunk rows now, or one + chunk height of the source where that is taller, since a read of part + of a chunk decodes all of it (F061): anndata's Zarr v3 chunks grow with + the matrix, to about a megabyte, so at this size they are half of it.""" + import zarr + 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. + rows = 256 + for n in (1024, 8192): source = matrix_store(tmp_path, n_obs=n, n_var=16, layout="dense", fmt=fmt) + height = (zarr.open_group(str(source), mode="r")["X"].chunks[0] + if fmt != "h5ad" else 1) out = tmp_path / f"o{n}{source.suffix}" with count_io() as io: - _recast(source, out, dtype="float64", chunk_rows=256) + _recast(source, out, dtype="float64", chunk_rows=rows) assert io.read_bytes >= n * 16 * 4, "the matrix was read" 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", - sizes=(1024, 8192), - ) + bound = -(-rows // height) * height * 16 * 4 + assert 0 < largest <= bound, ( + f"recast --dtype ({fmt}) of {n:,} rows read {largest:,} bytes of X at " + f"once against {bound:,} for --chunk {rows} in chunks {height} tall" + ) # h5ad only: what a row costs is the algorithm's, and Zarr v3's reads are diff --git a/tests/test_recast_streaming.py b/tests/test_recast_streaming.py index 6c3ab51..3758ed6 100644 --- a/tests/test_recast_streaming.py +++ b/tests/test_recast_streaming.py @@ -275,7 +275,18 @@ def test_a_dense_source_whose_nonzeros_overflow_the_offsets_is_refused( removes the half-written output.""" from tests.perf_counters import count_io + from adata.storage import copy_attrs, create_dataset, open_store + source_path = _put(new_store, np.ones((300, 200), dtype="float32")) + # Chunks 16 rows tall, as --chunk: reads take whole chunks of the source + # (F061), so a source held in one chunk is read in one piece. + with open_store(source_path, "a") as store: + root = store.root + values, attrs = root["X"][...], dict(root["X"].attrs) + del root["X"] + rechunked = create_dataset(root, "X", data=values, chunks=(16, 200)) + copy_attrs(attrs, rechunked.attrs, + target_backend="zarr" if source_path.suffix == ".zarr" else "hdf5") out = _output(source_path) with count_io() as io, pytest.raises(ValueError, match="cannot hold this matrix's offsets"): From 72c66c585a524f411c27f74bea547063dcec6d31 Mon Sep 17 00:00:00 2001 From: Aljes Date: Tue, 3 Nov 2026 17:50:07 +0000 Subject: [PATCH 5/5] Run the cost guards as their own CI job; cheaper B6b guards The unit job ran past its 20-minute timeout: coverage's trace function roughly doubled the suite on CI, and the cost guards are a fifth of it. The coverage job now measures through sys.monitoring (COVERAGE_CORE= sysmon: the same line coverage, 95.6%) and deselects `perf`; a new parallel job runs the guards, marked `perf` as a module, on both Pythons. publish.yml calls tests.yml whole, so a release needs both. B6b's guards share their sharded stores (built once per module), and the unit-memory and split RSS guards use half the data; each still fails on main. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/tests.yml | 69 ++++++++++++++++++++++++++++++++++++- AGENTS.md | 8 +++-- docs/developers.md | 4 +-- docs/reference/testing.md | 13 +++---- pytest.ini | 2 +- tests/test_block_layout.py | 45 +++++++++++++++++------- tests/test_performance.py | 50 +++++++++++++++++++-------- 7 files changed, 151 insertions(+), 40 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 9b685f3..1d9b1be 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -40,6 +40,11 @@ jobs: # before the one setup-python puts on PATH, so up to 0.17.4 the py3.13 # leg ran the suite on 3.12. UV_PYTHON outranks the file. UV_PYTHON: ${{ matrix.python-version }} + # Line coverage through sys.monitoring (3.12+) rather than a trace + # function: the same numbers, without the trace function's cost on + # every line, which had doubled this job's time and taken it past + # its timeout. + COVERAGE_CORE: sysmon steps: - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0 @@ -72,9 +77,10 @@ jobs: exit 1 fi + # The cost guards (`perf`) run in the job below, at the same time. - name: Run tests with coverage run: | - uv run pytest -v -W default tests/ -m "not integration" \ + uv run pytest -v -W default tests/ -m "not integration and not perf" \ --cov=adata \ --cov-report=term-missing \ --cov-report=xml \ @@ -107,6 +113,67 @@ jobs: files: coverage.xml fail_ci_if_error: false + # The cost guards: what an operation reads, writes and holds, counted + # rather than timed (tests/test_performance.py, and the few `perf` tests + # beside the code they guard). A job of their own, beside the unit suite, + # because together they ran past the unit job's timeout; publish.yml runs + # this workflow whole, so a release still needs both. + perf: + runs-on: ubuntu-latest + timeout-minutes: 20 + + permissions: + contents: read + + strategy: + fail-fast: false + matrix: + python-version: ["3.12", "3.13"] + + name: cost guards (py${{ matrix.python-version }}) + + env: + UV_PYTHON: ${{ matrix.python-version }} + + steps: + - uses: actions/checkout@11d5960a326750d5838078e36cf38b85af677262 # v4.4.0 + with: + persist-credentials: false + + - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 + with: + python-version: ${{ matrix.python-version }} + + - name: Set up uv + uses: astral-sh/setup-uv@caf0cab7a618c569241d31dcd442f54681755d39 # v3.2.4 + with: + version: "0.13.1" + enable-cache: false + + - name: Install dependencies (frozen) + run: uv sync --extra dev --frozen + + - name: The guards run on the Python this leg names + run: | + actual=$(uv run --frozen python -c 'import sys; print("%d.%d" % sys.version_info[:2])') + if [ "$actual" != "$UV_PYTHON" ]; then + echo "::error::the py$UV_PYTHON leg would run the guards on Python $actual" + exit 1 + fi + + - name: Run the cost guards + run: | + uv run pytest -v -W default tests/ -m "perf and not integration" \ + --junitxml="perf-results-py$UV_PYTHON.xml" + + - name: Upload the guards' results + uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 + if: always() + with: + name: perf-py${{ matrix.python-version }} + path: perf-results-py${{ matrix.python-version }}.xml + retention-days: 30 + # What `pip install pyadata-cli` gets today, not what uv.lock pins. 0.7.0 # to 0.17.2 imported click, which typer had stopped installing in 0.26, so # every `adata convert` on a fresh install crashed while the locked suite diff --git a/AGENTS.md b/AGENTS.md index 1ca6313..66e14c2 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -35,8 +35,8 @@ uv sync --extra dev # set up uv run adata --help # run the CLI from source uv run pytest -m "not integration" # unit suite, ~2.5 min; what CI gates on uv run pytest tests/test_recast.py::test_inplace_replaces_the_source # one test -uv run pytest -m "not integration" --cov=adata --cov-fail-under=90 # CI's coverage floor -uv run pytest -m perf # tracing complexity guards only +uv run pytest -m "not integration and not perf" --cov=adata --cov-fail-under=90 # CI's coverage job +uv run pytest -m perf # the cost guards, CI's own job ADATA_REQUIRE_VERSION_FIXTURES=1 uv run pytest tests/test_anndata_versions.py -m integration uv run python -m benchmarks.run --tier smoke --case concat-inner --out results.json uv run python -m benchmarks.report results.json @@ -46,7 +46,9 @@ Test markers: - `integration` tests build uv environments with real anndata releases (0.7.8 to 0.13) and the oracles (scanpy, PyTables and others), all pinned exactly in `tests/reference_stores.py`, so they need the network on the first run. -- `perf` tests are slow, but CI runs them in the main job. +- `perf` marks the cost guards (all of `tests/test_performance.py`, and a + few beside the code they guard). CI runs them as a job of their own, + beside the unit suite's coverage job, and a release needs both. CI also installs the package from a fresh resolution, without the lock, imports every module on the base dependencies alone, and runs `convert`. diff --git a/docs/developers.md b/docs/developers.md index 66253ab..f7193fe 100644 --- a/docs/developers.md +++ b/docs/developers.md @@ -170,8 +170,8 @@ before changing anything that writes: page has the front matter the sidebar needs. ```bash -uv run pytest -m "not integration" --cov=adata --cov-fail-under=90 # CI's coverage floor -uv run pytest -m perf # the complexity guards +uv run pytest -m "not integration and not perf" --cov=adata --cov-fail-under=90 # CI's coverage job +uv run pytest -m perf # the cost guards, CI's own job ADATA_REQUIRE_VERSION_FIXTURES=1 uv run pytest tests/test_anndata_versions.py -m integration ``` diff --git a/docs/reference/testing.md b/docs/reference/testing.md index f6b616e..6a426cf 100644 --- a/docs/reference/testing.md +++ b/docs/reference/testing.md @@ -12,13 +12,13 @@ redirect_from: uv sync --extra dev uv run pytest # everything uv run pytest -m "not integration" # no environment building, ~2.5 min -uv run pytest -m perf # just the tracing complexity guards +uv run pytest -m perf # just the cost guards uv run pytest -m integration # anndata releases, scanpy and PyTables ``` CI runs the unit suite on Python 3.12 and 3.13 with a **90% coverage floor**, -and the anndata compatibility suite and the 10x external-reader suite as -separate jobs. A fresh-install job installs the package alone from a fresh +and the cost guards (`-m perf`), the anndata compatibility suite and the 10x +external-reader suite as separate jobs. A fresh-install job installs the package alone from a fresh resolution, without `uv.lock`, as `pip install pyadata-cli` would. It imports every module on the base dependencies alone, runs `adata convert` with and without click beside typer, and runs the docs-accuracy suite with no @@ -382,8 +382,9 @@ should scale with, scale only that, and hold everything else fixed. | `split` | 128 against 512 groups | **resident memory** flat past 128 outputs a pass | 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. +they still gate merges, and `-m "not slow"` skips them locally. `perf` is on +every guard in the file, and on the few beside the code they guard; CI runs +them as a job of their own, beside the unit suite's coverage job. 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; @@ -559,7 +560,7 @@ ADATA_SKIP_VERSION_FIXTURES=1 uv run pytest ## Coverage ```bash -uv run pytest -m "not integration" --cov=adata --cov-report=term-missing +uv run pytest -m "not integration and not perf" --cov=adata --cov-report=term-missing ``` The floor is 90%. What remains uncovered is mostly defensive `except` branches diff --git a/pytest.ini b/pytest.ini index a116cbf..8b381bd 100644 --- a/pytest.ini +++ b/pytest.ini @@ -6,5 +6,5 @@ python_functions = test_* addopts = -v --strict-markers --tb=short --timeout=300 --timeout-method=thread markers = slow: marks tests as slow (deselect with '-m "not slow"') - perf: complexity guards that trace Python execution (10-50x slower) + perf: cost guards (test_performance.py, and a few beside their code); CI runs them as a job of their own integration: builds environments with uv and needs the network on first run diff --git a/tests/test_block_layout.py b/tests/test_block_layout.py index d667e25..e46efb2 100644 --- a/tests/test_block_layout.py +++ b/tests/test_block_layout.py @@ -140,8 +140,27 @@ def test_a_matrix_read_in_tiles_matches_anndata(tmp_path, monkeypatch): # outputs keep their source's layout (F033) -def _sharded_counts(tmp_path: Path, layout: str, n_obs: int = 3000, n_var: int = 40, +#: Sharded stores, each built once per module: no test writes to its source. +_STORES: dict = {} + + +@pytest.fixture(scope="module") +def stores(tmp_path_factory): + yield tmp_path_factory.mktemp("stores") + _STORES.clear() + + +def _sharded_counts(directory: Path, layout: str, n_obs: int = 3000, n_var: int = 40, stem: str = "src") -> tuple: + key = (layout, n_obs, n_var, stem) + if key not in _STORES: + _STORES[key] = _build_counts(directory / "-".join(map(str, key)), layout, n_obs, + n_var, stem) + return _STORES[key] + + +def _build_counts(directory: Path, layout: str, n_obs: int, n_var: int, stem: str) -> tuple: + directory.mkdir(parents=True, exist_ok=True) matrix = counts_matrix(n_obs, n_var, layout, density=0.2) obj = ad.AnnData( X=matrix, @@ -150,7 +169,7 @@ def _sharded_counts(tmp_path: Path, layout: str, n_obs: int = 3000, n_var: int = index=[f"c{i}" for i in range(n_obs)]), var=pd.DataFrame(index=[f"v{j}" for j in range(n_var)]), ) - return obj, write_store(obj, tmp_path, "zarr31", stem) + return obj, write_store(obj, directory, "zarr31", stem) def _arrays(group: Any) -> dict: @@ -162,12 +181,12 @@ def _arrays(group: Any) -> dict: @pytest.mark.parametrize("layout", ["csr", "csc", "dense"]) @pytest.mark.parametrize("command", ["subset", "split", "concat", "recast"]) -def test_outputs_keep_the_sources_inner_chunks_and_shards(tmp_path, command, layout): +def test_outputs_keep_the_sources_inner_chunks_and_shards(stores, tmp_path, command, layout): """F033: outputs of anndata 0.13's layout dropped the shard and kept the 245-element inner chunk as their own, a file each: 4,083 files for one array the source holds in 3. Each output array has the source's inner chunk and shard now, the shard clamped to the output, in whole chunks.""" - obj, src = _sharded_counts(tmp_path, layout) + obj, src = _sharded_counts(stores, layout) source = _arrays(zarr.open_group(str(src), mode="r")) if command == "subset": outputs = [tmp_path / "out.zarr"] @@ -176,7 +195,7 @@ def test_outputs_keep_the_sources_inner_chunks_and_shards(tmp_path, command, lay ok("split", src, "--by", "group", "-o", tmp_path / "parts", "--no-manifest", "-C", 64) outputs = sorted((tmp_path / "parts").iterdir()) elif command == "concat": - other = write_store(obj, tmp_path, "zarr31", "other") + _, other = _sharded_counts(stores, layout, stem="other") outputs = [tmp_path / "out.zarr"] ok("concat", src, other, "-o", outputs[0], "--index-unique", "-", "-C", 64) else: @@ -210,13 +229,13 @@ def test_outputs_keep_the_sources_inner_chunks_and_shards(tmp_path, command, lay np.testing.assert_array_equal(got_x, want_x) -def test_converting_a_sharded_store_to_h5ad_lets_hdf5_choose_its_chunks(tmp_path): +def test_converting_a_sharded_store_to_h5ad_lets_hdf5_choose_its_chunks(stores, tmp_path): """F033: a shard's 245-element inner chunk became the HDF5 chunk, 16,327 chunks where anndata's own h5ad has 512, and a file twice the size. HDF5 picks the chunks now, and X takes what anndata's own takes.""" import h5py - obj, src = _sharded_counts(tmp_path, "csr", n_obs=20_000, n_var=100) + obj, src = _sharded_counts(stores, "csr", n_obs=20_000, n_var=100) out = tmp_path / "out.h5ad" ok("convert", src, "-o", out) theirs = tmp_path / "theirs.h5ad" @@ -232,10 +251,10 @@ def stored(path: Path) -> int: np.testing.assert_array_equal(read_anndata(out).X.toarray(), obj.X.toarray()) -def test_a_sharded_store_written_as_zarr_v2_takes_its_shard_as_the_chunk(tmp_path): +def test_a_sharded_store_written_as_zarr_v2_takes_its_shard_as_the_chunk(stores, tmp_path): """F033: Zarr v2 has no shards, so `--zarr-format 2` kept only the inner chunk, a 245-element file each. The shard is the chunk now.""" - obj, src = _sharded_counts(tmp_path, "csr") + obj, src = _sharded_counts(stores, "csr") out = tmp_path / "v2.zarr" ok("convert", src, "-o", out, "--zarr-format", "2") source = zarr.open_group(str(src), mode="r") @@ -364,7 +383,7 @@ def test_sparse_blocks_cover_every_kept_major_once_within_the_budget(seed): np.testing.assert_array_equal(got[i][kept], sizes[kept]) -def test_a_shard_index_is_read_once_and_the_cache_is_bounded(tmp_path, monkeypatch): +def test_a_shard_index_is_read_once_and_the_cache_is_bounded(stores, monkeypatch): """F062: each read of part of a shard fetched its index; a store opened for reading keeps the indexes now, up to `SHARD_INDEX_CACHE_BYTES`.""" from zarr.abc.store import SuffixByteRequest @@ -373,7 +392,7 @@ def test_a_shard_index_is_read_once_and_the_cache_is_bounded(tmp_path, monkeypat import adata.storage as storage from adata.storage import open_store - _, src = _sharded_counts(tmp_path, "csr") + _, src = _sharded_counts(stores, "csr") fetched = [] original = LocalStore.get @@ -397,7 +416,7 @@ def read_in_blocks() -> int: assert read_in_blocks() > 10, "a cache too small for an index keeps none" -def test_the_shard_index_cache_is_one_budget_across_stores(tmp_path, monkeypatch): +def test_the_shard_index_cache_is_one_budget_across_stores(stores, monkeypatch): """From the review: each open store kept up to 64 MiB of indexes, and concat holds up to 64 inputs open, so the cache grew with the number of inputs. One budget is shared now, and a store's indexes go with it.""" @@ -406,7 +425,7 @@ def test_the_shard_index_cache_is_one_budget_across_stores(tmp_path, monkeypatch import adata.storage as storage from adata.storage import open_store - _, src = _sharded_counts(tmp_path, "csr") + _, src = _sharded_counts(stores, "csr") with open_store(src, "r") as store: store.root["X/data"][0:10] one = storage._SHARD_INDEX_STATE["bytes"] diff --git a/tests/test_performance.py b/tests/test_performance.py index 077434c..018cc89 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -67,6 +67,10 @@ count_lines, ) +#: Every guard here is a `perf` test: CI runs them as a job of their own, +#: beside the unit suite and its coverage, so neither waits on the other. +pytestmark = pytest.mark.perf + ad = pytest.importorskip("anndata", reason="anndata builds the fixtures") pd = pytest.importorskip("pandas") sparse = pytest.importorskip("scipy.sparse") @@ -2671,12 +2675,27 @@ def of_x(name: str) -> bool: return sum(io.encoded_by_name[n] for n in names) / written -def _sharded(tmp_path: Path, layout: str, n_obs: int = UNIT_ROWS, n_var: int = 64, +#: Sharded stores, built once per module and shared: no guard writes to its +#: source, and building one is most of a guard's time. +_SHARDED: dict = {} + + +@pytest.fixture(scope="module") +def sharded_dir(tmp_path_factory): + yield tmp_path_factory.mktemp("sharded") + _SHARDED.clear() + + +def _sharded(directory: Path, layout: str, n_obs: int = UNIT_ROWS, n_var: int = 64, groups: int = 0, tag: str = "s") -> Path: from tests.perf_fixtures import matrix_store, zarr31_layout - return zarr31_layout(matrix_store(tmp_path, n_obs=n_obs, n_var=n_var, layout=layout, - fmt="zarr3", groups=groups, tag=tag)) + key = (layout, n_obs, n_var, groups, tag) + if key not in _SHARDED: + _SHARDED[key] = zarr31_layout(matrix_store( + directory / "-".join(map(str, key)), n_obs=n_obs, n_var=n_var, layout=layout, + fmt="zarr3", groups=groups, tag=tag)) + return _SHARDED[key] def _split(source: Path, out: Path, chunk_rows: int = UNIT_CHUNK) -> None: @@ -2710,7 +2729,9 @@ def _subset(source: Path, out: Path, chunk_rows: int = UNIT_CHUNK, **kw) -> None ("recast --layout csc", "csr"), ("convert", "dense"), ]) -def test_a_sharded_output_is_written_a_whole_shard_at_a_time(tmp_path, monkeypatch, command, layout): +def test_a_sharded_output_is_written_a_whole_shard_at_a_time( + sharded_dir, tmp_path, monkeypatch, command, layout +): """F031, F036: an output keeping anndata 0.13's shard was written a --chunk block at a time, and zarr rewrites a whole shard for a write of part of it: 15 times the bytes for subset and concat, 121 for split @@ -2722,7 +2743,7 @@ def test_a_sharded_output_is_written_a_whole_shard_at_a_time(tmp_path, monkeypat made small here so that a read is smaller than a shard.""" import adata.storage as storage - source = _sharded(tmp_path, layout, groups=8 if command == "split" else 0) + source = _sharded(sharded_dir, layout, groups=8 if command == "split" else 0) out = tmp_path / "out.zarr" with count_io() as io: if command == "subset": @@ -2811,13 +2832,13 @@ def test_split_writes_each_output_chunk_once(tmp_path): @pytest.mark.parametrize("command,layout", [("subset", "dense"), ("concat", "csr")]) -def test_a_sharded_source_is_read_with_its_shard_index_once(tmp_path, command, layout): +def test_a_sharded_source_is_read_with_its_shard_index_once(sharded_dir, tmp_path, command, layout): """F062: zarr reads a shard's index again for every read of part of the shard, 16 bytes an inner chunk, so a store read a block at a time read the whole index once a block: subset fetched 2.3 times the dense X's bytes from the store, concat several times. Indexes are kept for the life of the open store now.""" - source = _sharded(tmp_path, layout, n_obs=4 * UNIT_ROWS if layout == "csr" else UNIT_ROWS) + source = _sharded(sharded_dir, layout, n_obs=4 * UNIT_ROWS if layout == "csr" else UNIT_ROWS) stored = sum(f.stat().st_size for f in source.rglob("*") if f.is_file()) inputs = _twice(source) if command == "concat" else [source] with count_io() as io: @@ -2910,7 +2931,7 @@ def test_copying_a_column_chunked_matrix_reads_tiles_not_the_matrix(tmp_path, mo @pytest.mark.parametrize("layout,n_obs,n_var,density", [ - ("dense", 4096, 2048, 0.5), ("csr", 2048, 2048, 0.25), + ("dense", 4096, 1024, 0.5), ("csr", 2048, 2048, 0.25), ]) def test_writing_whole_units_holds_a_few_units_not_the_shard( tmp_path, monkeypatch, layout, n_obs, n_var, density @@ -2920,15 +2941,15 @@ def test_writing_whole_units_holds_a_few_units_not_the_shard( source's shard of the whole array. Gathering blocks in a list and joining them held each unit twice over, and a sparse output gave its `data` and `indices` a budget each: 335 MB to subset a 16M-entry CSR at - 64 MiB units, against 50 MB before. Here the unit is 1 MiB, and the - source a shard of the whole array, 32 MB dense (inner chunks 16 rows + 64 MiB units, against 50 MB before. Here the unit is 512 KiB, and the + source a shard of the whole array, 16 MB dense (inner chunks 16 rows tall, so a read is a few rows) or 8 MB of entries in anndata's own layout: main held the shard or, for CSR, the matrix (F032).""" import adata.core.subset as subset import adata.storage as storage from tests.perf_fixtures import matrix_store, zarr31_layout - unit = 1 << 20 + unit = 1 << 19 monkeypatch.setattr(storage, "WRITE_UNIT_BYTES", unit, raising=False) monkeypatch.setattr(subset, "WRITE_UNIT_BYTES", unit, raising=False) source = matrix_store(tmp_path, n_obs=n_obs, n_var=n_var, layout=layout, @@ -2949,7 +2970,8 @@ def test_split_memory_stops_growing_with_the_group_count(tmp_path): """F079: split held up to 1,024 HDF5 outputs open in a pass, about 0.6 MB each however they are opened (HDF5's own structures for a file), so RSS followed the group count, not --chunk: 136 MB at 128 groups, 423 MB - at 512. A pass writes at most 128 outputs now. `peak_rss` raises the + at 512 (here, 138 against 320 MB at 320). A pass writes at most 128 + outputs now. `peak_rss` raises the open-file limit, as a cluster sets it, so macOS's 256 does not cap the outputs first.""" from tests.perf_counters import peak_rss @@ -2961,8 +2983,8 @@ def measure(groups: int) -> int: return peak_rss(["split", source, "--by", "grp", "-o", tmp_path / f"p{groups}", "--no-manifest"]).added - few, many = measure(128), measure(512) + few, many = measure(128), measure(320) assert many <= few + 32e6, ( - f"split into 512 groups added {many / 1e6:.0f} MB of RSS against " + f"split into 320 groups added {many / 1e6:.0f} MB of RSS against " f"{few / 1e6:.0f} MB into 128: memory follows the number of groups" )