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 dfa9b56..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`. @@ -191,7 +193,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..f7193fe 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. @@ -164,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/guides/pipelines.md b/docs/guides/pipelines.md index 91afae0..d651ab5 100644 --- a/docs/guides/pipelines.md +++ b/docs/guides/pipelines.md @@ -123,8 +123,21 @@ 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 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 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 770aee7..b482125 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 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` Concatenate along the obs axis, streaming each input in turn. @@ -621,6 +626,39 @@ 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 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), 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. +- 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..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 @@ -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,10 +372,19 @@ 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`, `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 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; @@ -547,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/docs/troubleshooting.md b/docs/troubleshooting.md index 52d429b..474756f 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 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 + +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 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` A loom has no sparse matrices, so a sparse store's zeros are all written. 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/src/adata/cli.py b/src/adata/cli.py index b97532c..a62c1c6 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..e1f567e 100644 --- a/src/adata/core/concat.py +++ b/src/adata/core/concat.py @@ -53,6 +53,8 @@ ) from adata.elements.strings import is_string_dtype from adata.storage import ( + RowWriter, + row_bands, copy_tree, create_dataset, dataset_create_kwargs, @@ -877,6 +879,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 +906,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 +947,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 @@ -949,14 +957,12 @@ 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] - 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,11 @@ 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] + for start, end in row_bands(source, chunk_rows): + writer.write(source[start:end, ...]) + writer.flush() return True @@ -1416,8 +1423,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 +1512,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..fb4308a 100644 --- a/src/adata/core/recast.py +++ b/src/adata/core/recast.py @@ -40,6 +40,9 @@ ) from adata.core.subset import _sparse_dataset from adata.storage import ( + RowWriter, + backend_chunks, + row_bands, copy_attrs, create_dataset, dataset_create_kwargs, @@ -162,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 @@ -394,17 +398,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 +536,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 +560,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( @@ -681,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) @@ -694,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]) @@ -709,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( @@ -808,12 +821,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) @@ -821,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. @@ -837,6 +851,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 +950,17 @@ 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) - 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 - ) + # 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, 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 e604c94..2d3d427 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,22 @@ resolve_index, ) from adata.storage import ( + TARGET_READ_BYTES, + WRITE_UNIT_BYTES, + RowWriter, create_dataset, copy_attrs, + backend_chunks, 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 +399,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 +420,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 +443,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 +464,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 +490,69 @@ 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 + 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( 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 +567,9 @@ def __init__( out_shape: Tuple[int, int], major_idx: Optional[np.ndarray], 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)) @@ -513,29 +581,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), size=size, + ) 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), size=size, ) - 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 +638,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 +660,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: @@ -621,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. @@ -641,17 +794,18 @@ def fan_out_sparse_matrix( out_shape, major_idx, _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()), ) ) - 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]) + 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] + ): if hi > lo: block_indices = np.asarray(indices_ds[lo:hi]) block_data = np.asarray(data_ds[lo:hi]) @@ -659,9 +813,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 +854,10 @@ 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, + unit_bytes=WRITE_UNIT_BYTES // 2) if compression is not None: kw = with_compression( kw, compression, shape=(n,), dtype=dtype, backend=backend, @@ -717,12 +875,22 @@ 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, + 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 + 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 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) @@ -731,20 +899,34 @@ 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]),) + 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): - 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 +958,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 +1021,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 +1069,7 @@ def _fan_out_raw( ], chunk_rows=chunk_rows, entry_label="raw/X", + share=share, ) if "varm" in src_raw: @@ -899,6 +1084,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 +1166,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 +1193,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 +1206,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 +1244,7 @@ def _write_targets( ], chunk_rows=chunk_rows, entry_label=task, + share=share, ) progress.update( task_id, @@ -1301,6 +1492,16 @@ def _open_output_limit() -> int: #: rather than running out of file descriptors. MAX_OPEN_OUTPUTS = _open_output_limit() +#: 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( file: Path, @@ -1320,14 +1521,19 @@ 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 = 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: 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/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 21e94b3..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,12 +125,116 @@ 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 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). + + 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` 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. + """ + import weakref + + 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._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) + + def _lookup(self, key: str, byte_range: Any) -> tuple: + if not isinstance(byte_range, SuffixByteRequest): + return None, 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 or slot in _SHARD_INDEXES: + return + _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) + 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 +246,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 +265,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 +285,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 +608,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 +670,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 +856,94 @@ 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 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 = 128 * 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 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], + *, + 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 +952,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 +979,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 +1204,224 @@ 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 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. + + 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 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 + 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.buffer: Optional[np.ndarray] = None + self.filled = 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.filled + + def write(self, values: Any) -> None: + values = np.asarray(values) + if not values.shape or values.shape[0] == 0: + return + 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.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 + elif self.cols is not None: + self.ds[self.written:end, self.cols, ...] = values + else: + self.ds[self.written:end, ...] = values + self.written = end + + def copy_dataset( src: Any, dst_group: Any, name: str, *, compression: Optional[str] = None, skip_attrs: Sequence[str] = (), @@ -1008,18 +1460,41 @@ 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, ) + 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( + 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 +1505,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 +1633,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..b94cff8 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) @@ -558,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 f00f75e..a142f2e 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,73 @@ 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", 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. `dense_chunks` gives a dense matrix other inner chunks, in the + same shards.""" + 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 = dense_chunks + 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 +197,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..e46efb2 --- /dev/null +++ b/tests/test_block_layout.py @@ -0,0 +1,444 @@ +"""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) +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] + 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) + + +# 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", ["h5ad", "zarr31"]) +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) + + +#: 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, + 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, directory, "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(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(stores, 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 = _sharded_counts(stores, layout, stem="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) + # 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" + 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(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(stores, "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(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(stores, "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(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 + from zarr.storage import LocalStore + + import adata.storage as storage + from adata.storage import open_store + + _, src = _sharded_counts(stores, "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" + + +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.""" + import gc + + import adata.storage as storage + from adata.storage import open_store + + _, src = _sharded_counts(stores, "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 4efa844..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") @@ -2420,13 +2424,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(): @@ -2499,25 +2500,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 @@ -2641,3 +2646,345 @@ 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 = 4096, 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 + + +#: 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 + + 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: + 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( + 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 + (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(sharded_dir, 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("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. 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: + 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"{command} 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(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(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: + 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.parametrize("layout,n_obs,n_var,density", [ + ("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 +): + """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 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 << 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, + 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 (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 + 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(320) + assert many <= few + 32e6, ( + 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" + ) 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"): 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