From 5a35496469b0212da62d81d7f17e13a34a73dc1f Mon Sep 17 00:00:00 2001 From: Aljes Date: Thu, 24 Sep 2026 15:33:36 +0100 Subject: [PATCH 1/3] Keep sparse index dtypes and compression, and split in one pass Every sparse output of subset, split and concat widened `indices` and `indptr` to int64 and dropped the source's compression, because the growable datasets they wrote into forwarded no storage settings. Outputs came out about twice the size anndata writes: 309.7 MB against 146.6 MB for subset in the 0.6.0 benchmark. They now keep the source's dtypes and layout, through the helpers convert already used, moved into subset.py. concat widens to int64 only when the combined matrix needs it. split ran one subset per group. With interleaved groups each subset read the span of its own rows, which was nearly all of X, so X was read k times. The matrix writers now take a list of outputs and share each source block among them; split_store makes one split_h5ad call, batched by the file-descriptor limit. Rows are gathered in one vectorised step rather than a Python loop. Dataframe columns were read through an h5py fancy index, which cost 0.38 s per column per group on 50,000 cells. They are now read as contiguous blocks and selected in memory. ci tier, locally: split 56.7 s -> 1.6 s (anndata 2.3 s), outputs within 0.3% of anndata's size for split, subset and concat. Co-Authored-By: Claude Opus 5.5 (1M context) --- CHANGELOG.md | 31 ++ src/adata/commands/split.py | 33 +- src/adata/core/concat.py | 53 ++- src/adata/core/convert.py | 36 +- src/adata/core/subset.py | 880 ++++++++++++++++++++++++----------- tests/test_performance.py | 42 +- tests/test_subset_fan_out.py | 352 ++++++++++++++ 7 files changed, 1086 insertions(+), 341 deletions(-) create mode 100644 tests/test_subset_fan_out.py diff --git a/CHANGELOG.md b/CHANGELOG.md index d043cfc..b7f1a6e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,37 @@ Notable changes to `adata-cli`. Versions are `MAJOR.MINOR.PATCH`; tags carry no `v` prefix. +## Unreleased + +### Fixed + +- **`subset`, `split` and `concat` wrote sparse matrices about twice as large + as they needed to be.** Every sparse output widened `indices` and `indptr` + to int64, whatever the source used, and was written without the source's + compression, so with an lzf input the `data` and `indices` came out + uncompressed. Outputs now keep the source's index dtypes and storage + settings. `concat` widens to int64 only when the combined matrix could not + be addressed otherwise. `convert` already did this; the code next to it + did not. + +- **`split` read the whole matrix once per group.** It ran a separate subset + for each group. When groups were interleaved through the file, as samples + usually are, each subset read nearly all of X, so a split into k groups + read and decompressed X k times: 56.7 s against 7.8 s for a hand-written + anndata loop in the 0.6.0 benchmark. `split` now reads each matrix once + and shares every block out among all the outputs. It holds up to a quarter + of the file-descriptor limit open at once, and makes one pass per batch + beyond that. `subset` shares the same writer, which now gathers rows in + one vectorised step instead of a Python loop over rows, and skips blocks + that hold none of the selected rows. A new guard in + `tests/test_performance.py` checks that split's reads of X do not grow + with the number of groups. + + Dataframe columns are no longer read through an h5py fancy index, which + cost 0.38 s per column per group on 50,000 cells. They are now read as + contiguous blocks and selected in memory. With both changes, the ci-tier + benchmark on a laptop splits in 1.6 s against anndata's 2.3 s. + ## 0.6.0 Adds `adata convert`, and fixes four quadratic paths that made `concat` and diff --git a/src/adata/commands/split.py b/src/adata/commands/split.py index 95c2409..cd1ec58 100644 --- a/src/adata/commands/split.py +++ b/src/adata/commands/split.py @@ -3,7 +3,7 @@ Requested in issue #2, modelled on cellgeni/scraft's `split_h5ad` but streaming: the source is never loaded into memory, only the column being split on is read, and each output is produced by the same subset machinery used by -`adata subset`. +`adata subset`, fanned out so every matrix is read once for all groups. """ from __future__ import annotations @@ -17,7 +17,7 @@ from rich.console import Console from adata.core.select import group_indices -from adata.core.subset import subset_h5ad +from adata.core.subset import split_h5ad from adata.storage import detect_backend, open_store @@ -112,19 +112,22 @@ def split_store( output_dir.mkdir(parents=True, exist_ok=True) - for label, out_path, _ in planned: - indices = np.sort(groups[label]) - subset_h5ad( - file=file, - output=out_path, - obs_file=None, - var_file=None, - chunk_rows=chunk_rows, - console=console, - obs_indices=indices if axis == "obs" else None, - var_indices=indices if axis == "var" else None, - zarr_format=zarr_format, - ) + # One call for every group, so each matrix is read once rather than once + # per group. + split_h5ad( + file, + [ + ( + out_path, + np.sort(groups[label]) if axis == "obs" else None, + np.sort(groups[label]) if axis == "var" else None, + ) + for label, out_path, _ in planned + ], + chunk_rows=chunk_rows, + console=console, + zarr_format=zarr_format, + ) if manifest: _write_manifest(file, output_dir, column, planned, console) diff --git a/src/adata/core/concat.py b/src/adata/core/concat.py index 39aa7ea..c9c67ea 100644 --- a/src/adata/core/concat.py +++ b/src/adata/core/concat.py @@ -389,6 +389,37 @@ def _matrix_kind(obj: Any) -> str: return "other" +def _concat_index_dtypes( + sources: List[Any], n_minor: int +) -> Tuple[np.dtype, np.dtype]: + """The `indices` and `indptr` dtypes for a concatenation of `sources`. + + The widest the inputs used, so int32 inputs give an int32 output rather + than the int64 every output used to get. That is widened to int64 only + when the result could not be addressed otherwise: `indices` must reach + the output's minor dimension, and `indptr` its total nonzero count, + which is known before writing as the sum of the inputs'. + """ + limit = np.iinfo(np.int32).max + indices = np.result_type(*[s["indices"].dtype for s in sources]) + indptr = np.result_type(*[s["indptr"].dtype for s in sources]) + nnz = sum(int(s["data"].shape[0]) for s in sources) + if n_minor > limit: + indices = np.result_type(indices, np.int64) + if nnz > limit: + indptr = np.result_type(indptr, np.int64) + return np.dtype(indices), np.dtype(indptr) + + +def _write_indptr( + group: Any, indptr: List[int], dtype: np.dtype, template: Any +) -> None: + from adata.core.subset import _sparse_dataset + + values = np.asarray(indptr, dtype=np.int64).astype(dtype, copy=False) + _sparse_dataset(group, "indptr", dtype, values.size, template)[:] = values + + def _concat_sparse( dst_parent: Any, name: str, @@ -403,7 +434,7 @@ def _concat_sparse( Absent columns need no fill: a sparse matrix's zeros are implicit, so dropped entries simply do not appear in the output. """ - from adata.core.subset import _append, _growable + from adata.core.subset import _append, _growable_like group = dst_parent.create_group(name) spec.set_encoding(group, spec.CSR_MATRIX) @@ -412,8 +443,11 @@ def _concat_sparse( set_shape_attr(group, (n_rows, n_cols)) dtype = np.result_type(*[s["data"].dtype for s in sources]) - out_data = _growable(group, "data", dtype) - out_indices = _growable(group, "indices", np.int64) + index_dtype, pointer_dtype = _concat_index_dtypes(sources, n_cols) + out_data = _growable_like(group, "data", dtype, sources[0]["data"]) + out_indices = _growable_like( + group, "indices", index_dtype, sources[0]["indices"] + ) indptr = [0] nnz = 0 @@ -449,7 +483,7 @@ def _concat_sparse( _append(out_indices, np.concatenate(kept_idx)) _append(out_data, np.concatenate(kept_data).astype(dtype)) - create_dataset(group, "indptr", data=np.asarray(indptr, dtype=np.int64)) + _write_indptr(group, indptr, pointer_dtype, sources[0]["indptr"]) def _concat_dense( @@ -509,7 +543,7 @@ def _concat_csc( input's row offset added. Because every input's column is already sorted and the offsets increase, the result is sorted without a re-sort. """ - from adata.core.subset import _append, _growable + from adata.core.subset import _append, _growable_like from adata.elements.write import set_shape_attr group = dst_parent.create_group(name) @@ -517,8 +551,11 @@ def _concat_csc( set_shape_attr(group, (n_rows, n_cols)) dtype = np.result_type(*[s["data"].dtype for s in sources]) - out_data = _growable(group, "data", dtype) - out_indices = _growable(group, "indices", np.int64) + index_dtype, pointer_dtype = _concat_index_dtypes(sources, n_rows) + out_data = _growable_like(group, "data", dtype, sources[0]["data"]) + out_indices = _growable_like( + group, "indices", index_dtype, sources[0]["indices"] + ) indptr = [0] nnz = 0 @@ -547,7 +584,7 @@ def _concat_csc( nnz += int(sum(len(r) for r in rows)) indptr.append(nnz) - create_dataset(group, "indptr", data=np.asarray(indptr, dtype=np.int64)) + _write_indptr(group, indptr, pointer_dtype, sources[0]["indptr"]) def check_matrix_encodings(roots: List[Any], console: Console) -> None: diff --git a/src/adata/core/convert.py b/src/adata/core/convert.py index fb3cfd3..4aadaa2 100644 --- a/src/adata/core/convert.py +++ b/src/adata/core/convert.py @@ -31,6 +31,7 @@ from adata.elements import spec from adata.elements.write import set_shape_attr +from adata.core.subset import _growable_like, _sparse_dataset from adata.storage import ( copy_attrs, create_dataset, @@ -298,41 +299,6 @@ def _new_sparse_group( return group -def _sparse_dataset( - group: Any, name: str, dtype: np.dtype, n: int, template: Any -) -> Any: - """A 1-D dataset of `n` elements laid out like `template`. - - Forwarding compression and chunking matters more here than anywhere - else: the point of a dtype change is usually to make the file smaller, - and creating the output with a fixed 65,536-element chunk made a - 2,400-nonzero matrix allocate 786 KB of mostly empty chunk -- eight - times the source, from a conversion asked for to halve it. - """ - from adata.core.subset import _clamp_chunks - - backend = "zarr" if is_zarr_group(group) else "hdf5" - kw = dataset_create_kwargs(template, target_backend=backend, dst_parent=group) - kw = _clamp_chunks(kw, max(1, n)) - if "chunks" not in kw and n: - kw["chunks"] = (min(n, 1 << 16),) - return create_dataset(group, name, shape=(n,), dtype=dtype, **kw) - - -def _growable_like(group: Any, name: str, dtype: np.dtype, template: Any) -> Any: - """Like `_sparse_dataset`, but extensible for a size not yet known.""" - backend = "zarr" if is_zarr_group(group) else "hdf5" - kw = dataset_create_kwargs(template, target_backend=backend, dst_parent=group) - kw.pop("shards", None) - chunks = kw.pop("chunks", None) - step = int(chunks[0]) if chunks else 1 << 16 - 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,), **kw - ) - - def _write_sparse_arrays( group: Any, data: np.ndarray, diff --git a/src/adata/core/subset.py b/src/adata/core/subset.py index 09e7767..002bb48 100644 --- a/src/adata/core/subset.py +++ b/src/adata/core/subset.py @@ -2,6 +2,8 @@ from __future__ import annotations +from contextlib import ExitStack +from dataclasses import dataclass from pathlib import Path import shutil from typing import Optional, Set, Tuple, List, Dict, Any @@ -112,15 +114,37 @@ def indices_from_name_set( return np.asarray(found_indices, dtype=np.int64), remaining +#: Rows `_take_rows` reads from HDF5 in one contiguous slice. +TAKE_BLOCK_ROWS = 1 << 20 + + def _take_rows(obj: Any, indices: Optional[np.ndarray]) -> Any: - """Read the selected rows of a dataset, handling both backends' indexing.""" + """Read the selected rows of a dataset, handling both backends' indexing. + + On HDF5 the rows are read as contiguous blocks and picked out in memory, + not by passing `indices` to h5py. h5py turns a fancy index into one + hyperslab per row, which made reading 6,250 rows of a 50,000-row column + take 0.38 s against 5 ms for the whole column. `split` does that for every + column of every group, so it came to most of split's run time. + """ if indices is None: return obj[...] if is_zarr_array(obj): if obj.ndim == 1: return obj.oindex[indices] return obj.oindex[(indices,) + (slice(None),) * (obj.ndim - 1)] - return obj[indices, ...] + _check_sorted(indices, "row") + if indices.size == 0: + return obj[0:0, ...] + + parts = [] + first, last = int(indices[0]), int(indices[-1]) + 1 + for start in range(first, last, TAKE_BLOCK_ROWS): + end = min(start + TAKE_BLOCK_ROWS, last) + lo, hi = np.searchsorted(indices, [start, end], side="left") + if hi > lo: + parts.append(np.asarray(obj[start:end, ...])[indices[lo:hi] - start]) + return parts[0] if len(parts) == 1 else np.concatenate(parts) def _copy_rows( @@ -223,6 +247,60 @@ def subset_axis_group( copy_tree(obj, dst, key) +def _minor_remap(keep: Optional[np.ndarray], size: int) -> Optional[np.ndarray]: + """Build a lookup from old minor index to new, with -1 for dropped entries. + + A dense lookup table costs one int64 per column of the source, which is + negligible beside the matrix itself and turns the remap into a single + vectorised gather rather than a per-entry dict lookup. + """ + if keep is None: + return None + remap = np.full(size, -1, dtype=np.int64) + remap[keep] = np.arange(len(keep), dtype=np.int64) + return remap + + +#: One output of a matrix write: the parent to create it in, and the obs and +#: var indices it keeps (None for all). Several of these share one read of +#: the source, which is what lets `split` make a single pass. +MatrixTarget = Tuple[Any, Optional[np.ndarray], Optional[np.ndarray]] + + +def _check_sorted(indices: Optional[np.ndarray], what: str) -> None: + """Refuse a selection the streaming writers cannot honour. + + The writers walk the source once, in order, so the rows they emit come + out in source order. A selection that is not strictly increasing would + be silently reordered or deduplicated, so it is refused instead. + h5py's fancy indexing has the same requirement, so every selection the + commands produce already meets it. + """ + if indices is not None and indices.size > 1 and np.any(np.diff(indices) <= 0): + 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 subset_dense_matrix( src: Any, dst_parent: Any, @@ -232,56 +310,146 @@ def subset_dense_matrix( *, chunk_rows: int = 1024, ) -> None: - if src.ndim != 2: - copy_tree(src, dst_parent, name) - return - - n_obs, n_var = src.shape - out_obs = len(obs_idx) if obs_idx is not None else n_obs - out_var = len(var_idx) if var_idx is not None else n_var + fan_out_dense_matrix(src, name, [(dst_parent, obs_idx, var_idx)], chunk_rows=chunk_rows) - target_backend = _target_backend(dst_parent) - kw = dataset_create_kwargs( - src, target_backend=target_backend, dst_parent=dst_parent - ) - kw = _clamp_chunks(kw, out_obs, out_var) - dst = create_dataset( - dst_parent, - name, - shape=(out_obs, out_var), - dtype=src.dtype, - **kw, - ) - copy_attrs(src.attrs, dst.attrs, target_backend=_target_backend(dst_parent)) +def fan_out_dense_matrix( + src: Any, + name: str, + targets: List[MatrixTarget], + *, + chunk_rows: int = 1024, +) -> None: + """Write one row/column subset of a dense matrix per target, in one pass. - for out_start in range(0, out_obs, chunk_rows): - out_end = min(out_start + chunk_rows, out_obs) + 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. + """ + if src.ndim != 2: + for dst_parent, _, _ in targets: + copy_tree(src, dst_parent, name) + return - if obs_idx is None: - block = src[out_start:out_end, :] + n_obs, n_var = src.shape + outputs = [] + for dst_parent, obs_idx, var_idx in targets: + _check_sorted(obs_idx, "obs") + out_obs = len(obs_idx) if obs_idx is not None else n_obs + out_var = len(var_idx) if var_idx is not None else n_var + + target_backend = _target_backend(dst_parent) + kw = dataset_create_kwargs( + src, target_backend=target_backend, dst_parent=dst_parent + ) + kw = _clamp_chunks(kw, out_obs, out_var) + dst = create_dataset( + dst_parent, name, shape=(out_obs, out_var), dtype=src.dtype, **kw + ) + copy_attrs(src.attrs, dst.attrs, target_backend=target_backend) + outputs.append([dst, _Cursor(obs_idx), var_idx, 0]) + + 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] + 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) + + +class _SparseOutput: + """One target's share of a streamed sparse subset.""" + + def __init__( + self, + src: Any, + dst_parent: Any, + name: str, + enc: str, + out_shape: Tuple[int, int], + major_idx: Optional[np.ndarray], + remap: Optional[np.ndarray], + ) -> None: + self.group = dst_parent.create_group(name) + copy_attrs(src.attrs, self.group.attrs, target_backend=_target_backend(dst_parent)) + spec.set_encoding(self.group, enc) + set_shape_attr(self.group, out_shape) + + # The source's own dtypes, not int64: a subset has no more nonzeros + # than its source and no coordinate beyond the source's dimensions, + # 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"]) + self.indices = _growable_like( + self.group, "indices", src["indices"].dtype, src["indices"] + ) + self.cursor = _Cursor(major_idx) + self.remap = remap + self.counts: List[np.ndarray] = [] + + def take( + self, + rows: np.ndarray, + indptr: np.ndarray, + lo: 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] + total = int(lengths.sum()) + + if total and self.cursor.selection is None: + # Every row of the block: its entries are one contiguous run. + begin = int(starts[0]) + positions = slice(begin, begin + total) else: - rows = obs_idx[out_start:out_end] - block = src[rows, :] - - if var_idx is not None: - block = block[:, var_idx] - - dst[out_start:out_end, :] = block - - -def _minor_remap(keep: Optional[np.ndarray], size: int) -> Optional[np.ndarray]: - """Build a lookup from old minor index to new, with -1 for dropped entries. + # Each row's run, laid end to end: the offset of an entry within + # its row plus the start of that row in the block. + ends = np.cumsum(lengths) + positions = np.repeat(starts - (ends - lengths), lengths) + np.arange( + total, dtype=np.int64 + ) - A dense lookup table costs one int32 per column of the source, which is - negligible beside the matrix itself and turns the remap into a single - vectorised gather rather than a per-entry dict lookup. - """ - if keep is None: - return None - remap = np.full(size, -1, dtype=np.int64) - remap[keep] = np.arange(len(keep), dtype=np.int64) - return remap + minor = block_indices[positions] + values = block_data[positions] + counts = lengths + if self.remap is not None: + mapped = self.remap[minor] + keep = mapped >= 0 + owner = np.repeat(np.arange(len(rows), dtype=np.int64), lengths) + counts = np.bincount(owner[keep], minlength=len(rows)) + minor, values = mapped[keep], values[keep] + + self.counts.append(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))) + dtype = self.indptr_template.dtype + dataset = _sparse_dataset( + self.group, "indptr", dtype, indptr.size, self.indptr_template + ) + dataset[:] = indptr.astype(dtype, copy=False) def subset_sparse_matrix_group( @@ -293,12 +461,30 @@ def subset_sparse_matrix_group( *, chunk_major: int = 4096, ) -> 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 major axis at a time.""" + fan_out_sparse_matrix( + src, name, [(dst_parent, obs_idx, var_idx)], chunk_major=chunk_major + ) + + +def fan_out_sparse_matrix( + src: Any, + name: str, + targets: List[MatrixTarget], + *, + chunk_major: int = 4096, +) -> 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. - Only the slice of `data`/`indices` spanned by the current block is read, so - peak memory is set by `chunk_major` rather than by the matrix. The output - datasets are grown as each block is appended, since the final nnz is not - known until the pass completes. + 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. """ enc = _decode_attr(src.attrs.get("encoding-type", b"")) if enc not in spec.SPARSE_TYPES: @@ -311,67 +497,58 @@ def subset_sparse_matrix_group( data_ds, indices_ds = src["data"], src["indices"] indptr = np.asarray(src["indptr"][...], dtype=np.int64) + csr = enc == spec.CSR_MATRIX + n_major, n_minor = (n_rows, n_cols) if csr else (n_cols, n_rows) + + outputs: List[_SparseOutput] = [] + for dst_parent, obs_idx, var_idx in targets: + major_idx, minor_keep = (obs_idx, var_idx) if csr else (var_idx, obs_idx) + _check_sorted(major_idx, "obs" if csr else "var") + out_shape = ( + len(obs_idx) if obs_idx is not None else n_rows, + len(var_idx) if var_idx is not None else n_cols, + ) + outputs.append( + _SparseOutput( + src, + dst_parent, + name, + enc, + out_shape, + major_idx, + _minor_remap(minor_keep, n_minor), + ) + ) - if enc == spec.CSR_MATRIX: - major_idx, minor_keep, n_minor = obs_idx, var_idx, n_cols - out_rows = len(obs_idx) if obs_idx is not None else n_rows - out_cols = len(var_idx) if var_idx is not None else n_cols - else: - major_idx, minor_keep, n_minor = var_idx, obs_idx, n_rows - out_rows = len(obs_idx) if obs_idx is not None else n_rows - out_cols = len(var_idx) if var_idx is not None else n_cols - - n_major = n_rows if enc == spec.CSR_MATRIX else n_cols - majors = major_idx if major_idx is not None else np.arange(n_major, dtype=np.int64) - remap = _minor_remap(minor_keep, n_minor) - - group = dst_parent.create_group(name) - copy_attrs(src.attrs, group.attrs, target_backend=_target_backend(dst_parent)) - spec.set_encoding(group, enc) - set_shape_attr(group, (out_rows, out_cols)) - - out_data = _growable(group, "data", data_ds.dtype) - out_indices = _growable(group, "indices", np.int64) - out_indptr = [0] - nnz = 0 - - for block_start in range(0, len(majors), chunk_major): - block = majors[block_start : block_start + chunk_major] - # One contiguous read covers the whole block's entries. - lo, hi = int(indptr[block].min()), int(indptr[block + 1].max()) + 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]) if hi > lo: - block_indices = np.asarray(indices_ds[lo:hi], dtype=np.int64) + block_indices = np.asarray(indices_ds[lo:hi]) block_data = np.asarray(data_ds[lo:hi]) else: - block_indices = np.empty(0, dtype=np.int64) + block_indices = np.empty(0, dtype=indices_ds.dtype) block_data = np.empty(0, dtype=data_ds.dtype) - kept_indices: List[np.ndarray] = [] - kept_data: List[np.ndarray] = [] - for m in block: - sl = slice(int(indptr[m]) - lo, int(indptr[m + 1]) - lo) - minor = block_indices[sl] - values = block_data[sl] - if remap is not None: - mapped = remap[minor] - keep = mapped >= 0 - minor, values = mapped[keep], values[keep] - kept_indices.append(minor) - kept_data.append(values) - nnz += len(minor) - out_indptr.append(nnz) - - if kept_indices: - _append(out_indices, np.concatenate(kept_indices)) - _append(out_data, np.concatenate(kept_data)) - - create_dataset( - group, "indptr", data=np.asarray(out_indptr, dtype=np.int64) - ) + for out, rows in zip(outputs, wanted): + if rows.size: + out.take(rows, indptr, lo, block_indices, block_data) + + for out in outputs: + out.finish() def _growable(group: Any, name: str, dtype: Any) -> Any: - """Create an empty 1-D dataset that can be extended as blocks arrive.""" + """Create an empty 1-D scratch dataset that can be extended as blocks arrive. + + For intermediate buffers only: it forwards no storage settings. Output + arrays go through `_growable_like`, which lays them out like the source. + """ if is_zarr_group(group): return group.create_array(name, shape=(0,), dtype=dtype, chunks=(65536,)) return group.create_dataset( @@ -379,6 +556,44 @@ def _growable(group: Any, name: str, dtype: Any) -> Any: ) +def _sparse_dataset( + group: Any, name: str, dtype: np.dtype, n: int, template: Any +) -> Any: + """A 1-D dataset of `n` elements laid out like `template`. + + Forwarding compression and chunking matters more here than anywhere + else: the point of a dtype change is usually to make the file smaller, + and creating the output with a fixed 65,536-element chunk made a + 2,400-nonzero matrix allocate 786 KB of mostly empty chunk -- eight + times the source, from a conversion asked for to halve it. + """ + backend = _target_backend(group) + kw = dataset_create_kwargs(template, target_backend=backend, dst_parent=group) + kw = _clamp_chunks(kw, max(1, n)) + if "chunks" not in kw and n: + kw["chunks"] = (min(n, 1 << 16),) + return create_dataset(group, name, shape=(n,), dtype=dtype, **kw) + + +def _growable_like(group: Any, name: str, dtype: np.dtype, template: Any) -> Any: + """Like `_sparse_dataset`, but extensible for a size not yet known. + + 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. + """ + backend = _target_backend(group) + kw = dataset_create_kwargs(template, target_backend=backend, dst_parent=group) + kw.pop("shards", None) + chunks = kw.pop("chunks", None) + step = int(chunks[0]) if chunks else 1 << 16 + 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,), **kw + ) + + def _append(ds: Any, values: np.ndarray) -> None: """Append a block to a growable 1-D dataset.""" if values.size == 0: @@ -397,21 +612,37 @@ def subset_matrix_entry( *, chunk_rows: int, entry_label: str, +) -> None: + fan_out_matrix_entry( + obj, + name, + [(dst_parent, obs_idx, var_idx)], + chunk_rows=chunk_rows, + entry_label=entry_label, + ) + + +def fan_out_matrix_entry( + obj: Any, + name: str, + targets: List[MatrixTarget], + *, + chunk_rows: int, + entry_label: str, ) -> None: if is_dataset(obj): - subset_dense_matrix( - obj, dst_parent, name, obs_idx, var_idx, chunk_rows=chunk_rows - ) + fan_out_dense_matrix(obj, name, targets, chunk_rows=chunk_rows) return if is_group(obj): enc = _decode_attr(obj.attrs.get("encoding-type", b"")) if enc in spec.SPARSE_TYPES: - subset_sparse_matrix_group(obj, dst_parent, name, obs_idx, var_idx) + fan_out_sparse_matrix(obj, name, targets) return if enc == spec.DATAFRAME: # obsm/varm may hold a dataframe; it is row-aligned like obs/var. - subset_axis_group(obj, dst_parent.create_group(name), obs_idx) + for dst_parent, row_idx, _ in targets: + subset_axis_group(obj, dst_parent.create_group(name), row_idx) return raise ValueError(f"Unsupported {entry_label} encoding type: {enc}") @@ -423,6 +654,20 @@ def subset_matrix_entry( ) +@dataclass +class Target: + """One output store of a subset, with the rows it keeps on each axis. + + `var_keep` is the names behind `var_idx`, for matching `raw/var`, which + has its own var axis. + """ + + root: Any + obs_idx: Optional[np.ndarray] + var_idx: Optional[np.ndarray] + var_keep: Optional[Set[str]] = None + + def subset_raw_group( src_raw: Any, dst: Any, @@ -432,59 +677,215 @@ def subset_raw_group( chunk_rows: int, console: Console, ) -> None: - """Subset a `raw/` group, which carries its own var axis. + _fan_out_raw( + src_raw, + [Target(dst, obs_idx, None, var_keep)], + chunk_rows=chunk_rows, + console=console, + ) + + +def _fan_out_raw( + src_raw: Any, + targets: List[Target], + *, + chunk_rows: int, + console: Console, +) -> None: + """Subset a `raw/` group into every target, reading its matrices once. `raw` typically holds more genes than the main object, so its var names are matched independently rather than reusing the outer var indices -- using those would select the wrong columns entirely. """ - raw_dst = dst.create_group("raw") - copy_attrs(src_raw.attrs, raw_dst.attrs, target_backend=_target_backend(dst)) - spec.set_encoding(raw_dst, spec.RAW) - - raw_var_idx: Optional[np.ndarray] = None - if var_keep is not None and "var" in src_raw: - raw_var_names, _ = resolve_index(src_raw["var"], "var") - raw_var_idx, missing = indices_from_name_set(raw_var_names, var_keep) - console.print( - f"[green]Selected {len(raw_var_idx)} raw/var " - f"(of {element_len(raw_var_names)})[/]" + raw_dsts: List[Any] = [] + raw_var_idxs: List[Optional[np.ndarray]] = [] + for target in targets: + raw_dst = target.root.create_group("raw") + copy_attrs( + src_raw.attrs, raw_dst.attrs, target_backend=_target_backend(target.root) ) - if missing: + spec.set_encoding(raw_dst, spec.RAW) + raw_dsts.append(raw_dst) + + raw_var_idx: Optional[np.ndarray] = None + if target.var_keep is not None and "var" in src_raw: + raw_var_names, _ = resolve_index(src_raw["var"], "var") + raw_var_idx, missing = indices_from_name_set( + raw_var_names, target.var_keep + ) console.print( - f"[yellow]Warning: {len(missing)} var names not found in raw/var[/]" + f"[green]Selected {len(raw_var_idx)} raw/var " + f"(of {element_len(raw_var_names)})[/]" ) + if missing: + console.print( + f"[yellow]Warning: {len(missing)} var names not found in raw/var[/]" + ) + raw_var_idxs.append(raw_var_idx) if "var" in src_raw: - subset_axis_group(src_raw["var"], raw_dst.create_group("var"), raw_var_idx) + for raw_dst, raw_var_idx in zip(raw_dsts, raw_var_idxs): + subset_axis_group(src_raw["var"], raw_dst.create_group("var"), raw_var_idx) if "X" in src_raw: - subset_matrix_entry( + fan_out_matrix_entry( src_raw["X"], - raw_dst, "X", - obs_idx, - raw_var_idx, + [ + (raw_dst, target.obs_idx, raw_var_idx) + for raw_dst, target, raw_var_idx in zip(raw_dsts, targets, raw_var_idxs) + ], chunk_rows=chunk_rows, entry_label="raw/X", ) if "varm" in src_raw: - varm_dst = _ensure_group(raw_dst, "varm") + varm_dsts = [_ensure_group(raw_dst, "varm") for raw_dst in raw_dsts] for key in src_raw["varm"].keys(): - subset_matrix_entry( + fan_out_matrix_entry( src_raw["varm"][key], - varm_dst, key, - raw_var_idx, - None, + [ + (varm_dst, raw_var_idx, None) + for varm_dst, raw_var_idx in zip(varm_dsts, raw_var_idxs) + ], chunk_rows=chunk_rows, entry_label=f"raw/varm:{key}", ) for key in src_raw.keys(): if key not in ("X", "var", "varm"): - copy_tree(src_raw[key], raw_dst, key) + for raw_dst in raw_dsts: + copy_tree(src_raw[key], raw_dst, key) + + +#: Mapping groups whose entries are matrices, and the axes each entry is +#: aligned to, as the (rows, columns) a target keeps. +_MATRIX_MAPPINGS = { + "layers": ("obs", "var"), + "obsm": ("obs", None), + "varm": ("var", None), + "obsp": ("obs", "obs"), + "varp": ("var", "var"), +} + +_TASK_LABELS = {"layers": "layer"} + + +def _axis_idx(target: Target, axis: Optional[str]) -> Optional[np.ndarray]: + if axis is None: + return None + return target.obs_idx if axis == "obs" else target.var_idx + + +def _write_targets( + src: Any, + targets: List[Target], + *, + chunk_rows: int, + console: Console, +) -> None: + """Write every element of `src`, narrowed, into each target's store. + + Matrices are read once and shared out among the targets; the rest -- the + dataframes, `uns`, anything unrecognised -- is small or row-proportional, + and is written target by target with the single-output helpers. + """ + tasks: List[str] = [] + if "obs" in src: + tasks.append("obs") + if "var" in src: + tasks.append("var") + if "X" in src: + tasks.append("X") + for mapping in _MATRIX_MAPPINGS: + if mapping in src: + label = _TASK_LABELS.get(mapping, mapping) + tasks.extend(f"{label}:{k}" for k in src[mapping].keys()) + if "uns" in src: + tasks.append("uns") + # anndata writes a placeholder `raw` even when there is none, and + # on Zarr that placeholder is an array rather than a group. + if "raw" in src and is_group(src["raw"]): + tasks.append("raw") + elif "raw" in src: + tasks.append("copy:raw") + + passthrough = [k for k in src.keys() if k not in HANDLED_KEYS] + if passthrough: + console.print( + "[yellow]Copying unrecognised top-level " + f"{'keys' if len(passthrough) > 1 else 'key'} verbatim: " + f"{', '.join(sorted(passthrough))}[/]" + ) + tasks.extend(f"copy:{k}" for k in passthrough) + + labels = {v: k for k, v in _TASK_LABELS.items()} + + with Progress( + SpinnerColumn(finished_text="[green]✓[/]"), + TextColumn("[progress.description]{task.description}"), + console=console, + transient=False, + ) as progress: + for task in tasks: + task_id = progress.add_task(f"[cyan]Subsetting {task}...[/]", total=None) + if task in ("obs", "var"): + for target in targets: + subset_axis_group( + src[task], + target.root.create_group(task), + _axis_idx(target, task), + ) + elif task == "X": + 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) + elif is_group(X): + fan_out_sparse_matrix(X, "X", pairs) + else: + for target in targets: + copy_tree(X, target.root, "X") + elif task == "uns": + for target in targets: + copy_tree(src["uns"], target.root, "uns") + elif task == "raw": + _fan_out_raw( + src["raw"], targets, chunk_rows=chunk_rows, console=console + ) + elif task.startswith("copy:"): + key = task.split(":", 1)[1] + for target in targets: + copy_tree(src[key], target.root, key) + else: + label, key = task.split(":", 1) + mapping = labels.get(label, label) + rows, cols = _MATRIX_MAPPINGS[mapping] + fan_out_matrix_entry( + src[mapping][key], + key, + [ + ( + _ensure_group(t.root, mapping), + _axis_idx(t, rows), + _axis_idx(t, cols), + ) + for t in targets + ], + chunk_rows=chunk_rows, + entry_label=task, + ) + progress.update( + task_id, + description=f"[green]Subsetting {task}[/]", + completed=1, + total=1, + ) + + for target in targets: + _ensure_optional_anndata_groups(target.root) def _select_axis( @@ -602,156 +1003,12 @@ def subset_h5ad( status.stop() - tasks: List[str] = [] - if "obs" in src: - tasks.append("obs") - if "var" in src: - tasks.append("var") - if "X" in src: - tasks.append("X") - if "layers" in src: - tasks.extend([f"layer:{k}" for k in src["layers"].keys()]) - if "obsm" in src: - tasks.extend([f"obsm:{k}" for k in src["obsm"].keys()]) - if "varm" in src: - tasks.extend([f"varm:{k}" for k in src["varm"].keys()]) - if "obsp" in src: - tasks.extend([f"obsp:{k}" for k in src["obsp"].keys()]) - if "varp" in src: - tasks.extend([f"varp:{k}" for k in src["varp"].keys()]) - if "uns" in src: - tasks.append("uns") - # anndata writes a placeholder `raw` even when there is none, and - # on Zarr that placeholder is an array rather than a group. - if "raw" in src and is_group(src["raw"]): - tasks.append("raw") - elif "raw" in src: - tasks.append("copy:raw") - - passthrough = [ - k for k in src.keys() if k not in HANDLED_KEYS - ] - if passthrough: - console.print( - "[yellow]Copying unrecognised top-level " - f"{'keys' if len(passthrough) > 1 else 'key'} verbatim: " - f"{', '.join(sorted(passthrough))}[/]" - ) - tasks.extend(f"copy:{k}" for k in passthrough) - - with Progress( - SpinnerColumn(finished_text="[green]✓[/]"), - TextColumn("[progress.description]{task.description}"), + _write_targets( + src, + [Target(dst, obs_idx, var_idx, var_keep)], + chunk_rows=chunk_rows, console=console, - transient=False, - ) as progress: - for task in tasks: - task_id = progress.add_task( - f"[cyan]Subsetting {task}...[/]", total=None - ) - if task == "obs": - obs_dst = dst.create_group("obs") - subset_axis_group(src["obs"], obs_dst, obs_idx) - elif task == "var": - var_dst = dst.create_group("var") - subset_axis_group(src["var"], var_dst, var_idx) - elif task == "X": - X = src["X"] - if is_dataset(X): - subset_dense_matrix( - X, dst, "X", obs_idx, var_idx, chunk_rows=chunk_rows - ) - elif is_group(X): - subset_sparse_matrix_group(X, dst, "X", obs_idx, var_idx) - else: - copy_tree(X, dst, "X") - elif task.startswith("layer:"): - key = task.split(":", 1)[1] - layer_src = src["layers"][key] - layers_dst = _ensure_group(dst, "layers") - subset_matrix_entry( - layer_src, - layers_dst, - key, - obs_idx, - var_idx, - chunk_rows=chunk_rows, - entry_label=f"layer:{key}", - ) - elif task.startswith("obsm:"): - key = task.split(":", 1)[1] - obsm_dst = _ensure_group(dst, "obsm") - obsm_obj = src["obsm"][key] - subset_matrix_entry( - obsm_obj, - obsm_dst, - key, - obs_idx, - None, - chunk_rows=chunk_rows, - entry_label=f"obsm:{key}", - ) - elif task.startswith("varm:"): - key = task.split(":", 1)[1] - varm_dst = _ensure_group(dst, "varm") - varm_obj = src["varm"][key] - subset_matrix_entry( - varm_obj, - varm_dst, - key, - var_idx, - None, - chunk_rows=chunk_rows, - entry_label=f"varm:{key}", - ) - elif task.startswith("obsp:"): - key = task.split(":", 1)[1] - obsp_dst = _ensure_group(dst, "obsp") - obsp_obj = src["obsp"][key] - subset_matrix_entry( - obsp_obj, - obsp_dst, - key, - obs_idx, - obs_idx, - chunk_rows=chunk_rows, - entry_label=f"obsp:{key}", - ) - elif task.startswith("varp:"): - key = task.split(":", 1)[1] - varp_dst = _ensure_group(dst, "varp") - varp_obj = src["varp"][key] - subset_matrix_entry( - varp_obj, - varp_dst, - key, - var_idx, - var_idx, - chunk_rows=chunk_rows, - entry_label=f"varp:{key}", - ) - elif task == "uns": - copy_tree(src["uns"], dst, "uns") - elif task == "raw": - subset_raw_group( - src["raw"], - dst, - obs_idx, - var_keep, - chunk_rows=chunk_rows, - console=console, - ) - elif task.startswith("copy:"): - key = task.split(":", 1)[1] - copy_tree(src[key], dst, key) - progress.update( - task_id, - description=f"[green]Subsetting {task}[/]", - completed=1, - total=1, - ) - - _ensure_optional_anndata_groups(dst) + ) finally: status.stop() @@ -765,3 +1022,62 @@ def subset_h5ad( shutil.move(str(dst_path), str(file)) else: dst_path.replace(file) + + +def _open_output_limit() -> int: + """How many output stores `split_h5ad` may hold open at once. + + Each open HDF5 file costs a descriptor, and the soft limit is 256 on + macOS and 1024 on most Linux systems by default. A quarter of it leaves + room for everything else the process has open. + """ + try: + import resource + + soft, _ = resource.getrlimit(resource.RLIMIT_NOFILE) + except (ImportError, OSError, ValueError): # pragma: no cover - Windows + return 64 + if soft == resource.RLIM_INFINITY: + return 1024 + return max(8, min(1024, soft // 4)) + + +#: A split into more groups than this makes one pass per batch of this many, +#: rather than running out of file descriptors. +MAX_OPEN_OUTPUTS = _open_output_limit() + + +def split_h5ad( + file: Path, + outputs: List[Tuple[Path, Optional[np.ndarray], Optional[np.ndarray]]], + *, + chunk_rows: int = 1024, + console: Console, + zarr_format: Optional[int] = None, +) -> None: + """Write one subset of `file` per ``(path, obs_idx, var_idx)`` in `outputs`. + + Equivalent to calling :func:`subset_h5ad` once per output, but every + matrix is read once for all of them (per batch of `MAX_OPEN_OUTPUTS`) + rather than once per output. + """ + if zarr_format is None and detect_backend(file) == "zarr": + 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] + with ExitStack() as stack: + src = stack.enter_context(open_store(file, "r")).root + targets = [] + for path, obs_idx, var_idx in batch: + dst = stack.enter_context( + open_store(path, "w", zarr_format=zarr_format) + ).root + var_keep = ( + set(read_names(src, "var", var_idx)) + if var_idx is not None and "var" in src + else None + ) + targets.append(Target(dst, obs_idx, var_idx, var_keep)) + _write_targets(src, targets, chunk_rows=chunk_rows, console=console) diff --git a/tests/test_performance.py b/tests/test_performance.py index f272055..c2fabb3 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -464,14 +464,21 @@ def measure(n: int) -> int: assert_grows_linearly(measure, what="h5ad to zarr", axis="n_obs") -def test_split_reads_grow_linearly_in_group_count(tmp_path): +def test_split_reads_grow_linearly_in_group_count(tmp_path, monkeypatch): """Splitting into k groups must stay linear in k at fixed store size. n_obs is held constant deliberately. Scaling rows alongside groups would make the total legitimately quadratic -- k passes over 4k rows -- and the guard would be measuring the fixture, not the implementation. + + The output batch is raised above the largest k: batching adds one pass + per batch, a step the linearity check would misread as curvature. + `test_split_in_batches_matches_split_in_one` covers the batches. """ from adata.commands.split import split_store + from adata.core import subset + + monkeypatch.setattr(subset, "MAX_OPEN_OUTPUTS", 1024) n_obs = 512 @@ -493,6 +500,39 @@ def measure(n: int) -> int: ) +def test_split_reads_the_matrix_once_whatever_the_group_count(tmp_path, monkeypatch): + """X is read once for all groups, not once per group. + + Linear growth in k, which the guard above accepts, is exactly what split + used to do: it subset once per group, and with interleaved groups each + subset read the span of its own rows -- nearly all of X. On the release + benchmark that was 56.7 s against 7.8 s for a naive anndata loop. The + other elements are still written per group, so only X is counted here. + """ + from adata.commands.split import split_store + from adata.core import subset + + monkeypatch.setattr(subset, "MAX_OPEN_OUTPUTS", 1024) + n_obs = 256 + + def measure(n: int) -> int: + source = _store( + tmp_path / f"x{n}.h5ad", name="x", n_obs=n_obs, n_var=64, n_categories=n + ) + with count_io() as io: + split_store(source, "ct", tmp_path / f"ox{n}", QUIET, manifest=False) + read = sum( + v for k, v in io.by_name.items() if k.endswith(("X/data", "X/indices")) + ) + # Every nonzero at least once: the counters can see these reads. + assert read >= 2 * n_obs * 64, io.by_name + return read + + assert_independent_of( + measure, what="split X reads", axis="n_groups", sizes=(2, 128) + ) + + # --------------------------------------------------------------------------- # the constants themselves diff --git a/tests/test_subset_fan_out.py b/tests/test_subset_fan_out.py new file mode 100644 index 0000000..0668ca8 --- /dev/null +++ b/tests/test_subset_fan_out.py @@ -0,0 +1,352 @@ +"""The streamed sparse writers: dtypes, storage settings, and one-pass split. + +Two defects shared this code. Every sparse output of subset, split and +concat widened `indices`/`indptr` to int64 and dropped the source's +compression, so outputs came out about twice the size anndata writes; and +split subset once per group, reading all of X once per group whenever the +groups were interleaved. These tests pin the fixes against anndata and +scipy rather than against the code itself. +""" + +from __future__ import annotations + +from pathlib import Path + +import h5py +import numpy as np +import pytest +from rich.console import Console + +from adata.commands.split import split_store +from adata.core import subset as subset_mod +from adata.core.concat import _concat_index_dtypes, concat_on_disk +from adata.core.subset import ( + fan_out_sparse_matrix, + split_h5ad, + subset_h5ad, + subset_sparse_matrix_group, +) +from adata.storage import open_store + +ad = pytest.importorskip("anndata") +pd = pytest.importorskip("pandas") +sparse = pytest.importorskip("scipy.sparse") + +QUIET = Console(quiet=True) + + +def _anndata(n_obs=40, n_var=12, *, layout="csr", seed=0, raw=True): + rng = np.random.default_rng(seed) + dense = rng.poisson(0.6, (n_obs, n_var)).astype("float32") + X = sparse.csr_matrix(dense) if layout == "csr" else ( + sparse.csc_matrix(dense) if layout == "csc" else dense + ) + obs = pd.DataFrame( + { + # Interleaved, the case that made split read X once per group. + "group": pd.Categorical([f"g{i % 3}" for i in range(n_obs)]), + "score": np.arange(n_obs, dtype="float32"), + }, + index=[f"c{i}" for i in range(n_obs)], + ) + var = pd.DataFrame( + {"kind": pd.Categorical([f"k{i % 2}" for i in range(n_var)])}, + index=[f"v{i}" for i in range(n_var)], + ) + obj = ad.AnnData(X=X, obs=obs, var=var) + obj.layers["counts"] = sparse.csc_matrix(dense * 2) + obj.obsm["X_pca"] = rng.normal(size=(n_obs, 3)).astype("float32") + obj.varm["loadings"] = rng.normal(size=(n_var, 2)).astype("float32") + obj.obsp["conn"] = sparse.csr_matrix( + rng.random((n_obs, n_obs)).astype("float32") * (rng.random((n_obs, n_obs)) < 0.2) + ) + if raw: + obj.raw = obj + return obj + + +def _write(obj, path: Path) -> Path: + if path.suffix == ".zarr": + obj.write_zarr(path) + else: + obj.write_h5ad(path, compression="lzf") + return path + + +def _dense(m): + return m.toarray() if sparse.issparse(m) else np.asarray(m) + + +def _assert_same(a, b): + assert list(a.obs_names) == list(b.obs_names) + assert list(a.var_names) == list(b.var_names) + np.testing.assert_array_equal(_dense(a.X), _dense(b.X)) + np.testing.assert_array_equal(_dense(a.layers["counts"]), _dense(b.layers["counts"])) + np.testing.assert_array_equal(a.obsm["X_pca"], b.obsm["X_pca"]) + np.testing.assert_array_equal(a.varm["loadings"], b.varm["loadings"]) + np.testing.assert_array_equal(_dense(a.obsp["conn"]), _dense(b.obsp["conn"])) + pd.testing.assert_frame_equal(a.obs, b.obs) + if a.raw is not None or b.raw is not None: + np.testing.assert_array_equal(_dense(a.raw.X), _dense(b.raw.X)) + assert list(a.raw.var_names) == list(b.raw.var_names) + + +# --------------------------------------------------------------------------- +# dtypes and storage settings + + +@pytest.mark.parametrize("suffix", [".h5ad", ".zarr"]) +@pytest.mark.parametrize("layout", ["csr", "csc"]) +def test_subset_keeps_the_source_index_dtypes(tmp_path, suffix, layout): + src = _write(_anndata(layout=layout), tmp_path / f"src{suffix}") + out = tmp_path / f"out{suffix}" + subset_h5ad(src, out, None, None, console=QUIET, obs_indices=np.arange(0, 40, 3)) + + with open_store(src, "r") as s, open_store(out, "r") as o: + for key in ("indices", "indptr", "data"): + assert o.root["X"][key].dtype == s.root["X"][key].dtype, key + # scipy writes int32 for a matrix this small; the defect was int64. + assert o.root["X"]["indices"].dtype == np.int32 + + +def test_subset_keeps_int64_indices_int64(tmp_path): + obj = _anndata(raw=False) + obj.X = sparse.csr_matrix(obj.X) + obj.X.indices = obj.X.indices.astype(np.int64) + obj.X.indptr = obj.X.indptr.astype(np.int64) + src = _write(obj, tmp_path / "wide.h5ad") + with h5py.File(src, "r") as f: + assert f["X/indices"].dtype == np.int64, "fixture must really be int64" + + out = tmp_path / "out.h5ad" + subset_h5ad(src, out, None, None, console=QUIET, obs_indices=np.arange(5)) + with h5py.File(out, "r") as f: + assert f["X/indices"].dtype == np.int64 + assert f["X/indptr"].dtype == np.int64 + + +def test_subset_keeps_the_source_compression(tmp_path): + src = _write(_anndata(), tmp_path / "src.h5ad") + out = tmp_path / "out.h5ad" + subset_h5ad(src, out, None, None, console=QUIET, obs_indices=np.arange(0, 40, 2)) + + with h5py.File(src, "r") as s, h5py.File(out, "r") as o: + for key in ("data", "indices", "indptr"): + assert s[f"X/{key}"].compression == "lzf" + assert o[f"X/{key}"].compression == "lzf", key + + +# --------------------------------------------------------------------------- +# the vectorised gather, against scipy + + +@pytest.mark.parametrize("layout", ["csr", "csc"]) +@pytest.mark.parametrize("seed", range(6)) +def test_sparse_subset_matches_scipy(tmp_path, layout, seed): + rng = np.random.default_rng(seed) + n_obs, n_var = 57, 23 + dense = rng.poisson(0.5, (n_obs, n_var)).astype("float32") + dense[rng.random(n_obs) < 0.3] = 0 # empty rows + matrix = sparse.csr_matrix(dense) if layout == "csr" else sparse.csc_matrix(dense) + src = _write(ad.AnnData(X=matrix), tmp_path / "m.h5ad") + + def pick(n): + choice = rng.integers(0, 4) + if choice == 0: + return None + if choice == 1: + return np.array([], dtype=np.int64) + return np.sort(rng.choice(n, size=rng.integers(1, n), replace=False)) + + targets = [(pick(n_obs), pick(n_var)) for _ in range(4)] + out = tmp_path / "out.h5ad" + with h5py.File(src, "r") as s, h5py.File(out, "w") as o: + fan_out_sparse_matrix( + s["X"], + "X", + [(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, + ) + for i, (ob, va) in enumerate(targets): + expected = dense + if ob is not None: + expected = expected[ob] + if va is not None: + expected = expected[:, va] + with h5py.File(out, "r") as o: + g = o[f"t{i}/X"] + cls = sparse.csr_matrix if layout == "csr" else sparse.csc_matrix + got = cls( + (g["data"][...], g["indices"][...], g["indptr"][...]), + shape=tuple(g.attrs["shape"]), + ) + np.testing.assert_array_equal(got.toarray(), expected) + assert got.has_sorted_indices + + +def test_unsorted_selection_is_refused(tmp_path): + src = _write(_anndata(raw=False), tmp_path / "s.h5ad") + with h5py.File(src, "r") as s, h5py.File(tmp_path / "o.h5ad", "w") as o: + with pytest.raises(ValueError, match="sorted and unique"): + subset_sparse_matrix_group(s["X"], o, "X", np.array([3, 1]), None) + + +# --------------------------------------------------------------------------- +# split in one pass + + +@pytest.mark.parametrize("suffix", [".h5ad", ".zarr"]) +@pytest.mark.parametrize("layout", ["csr", "csc", "dense"]) +@pytest.mark.parametrize("axis", ["obs", "var"]) +def test_split_matches_subset_per_group(tmp_path, suffix, layout, axis): + src = _write(_anndata(layout=layout), tmp_path / f"src{suffix}") + column = "group" if axis == "obs" else "kind" + planned = split_store( + src, column, tmp_path / "parts", QUIET, axis=axis, manifest=False + ) + assert len(planned) == (3 if axis == "obs" else 2) + + source = ad.read_h5ad(src) if suffix == ".h5ad" else ad.read_zarr(src) + frame = source.obs if axis == "obs" else source.var + for label, path, count in planned: + indices = np.flatnonzero(frame[column].astype(str).to_numpy() == label) + expected_path = tmp_path / f"expected-{label}{suffix}" + subset_h5ad( + src, + expected_path, + None, + None, + console=QUIET, + obs_indices=indices if axis == "obs" else None, + var_indices=indices if axis == "var" else None, + ) + read = ad.read_h5ad if suffix == ".h5ad" else ad.read_zarr + got, expected = read(path), read(expected_path) + assert (got.n_obs if axis == "obs" else got.n_vars) == count + _assert_same(got, expected) + # And against anndata's own slicing, not just our other code path. + want = source[indices] if axis == "obs" else source[:, indices] + np.testing.assert_array_equal(_dense(got.X), _dense(want.X)) + + +def test_split_in_batches_matches_split_in_one(tmp_path, monkeypatch): + src = _write(_anndata(n_obs=50), tmp_path / "src.h5ad") + source = ad.read_h5ad(src) + groups = [np.arange(i, 50, 5) for i in range(5)] + + whole = [(tmp_path / f"w{i}.h5ad", g, None) for i, g in enumerate(groups)] + split_h5ad(src, whole, console=QUIET) + + monkeypatch.setattr(subset_mod, "MAX_OPEN_OUTPUTS", 2) + batched = [(tmp_path / f"b{i}.h5ad", g, None) for i, g in enumerate(groups)] + split_h5ad(src, batched, console=QUIET) + + for (w, g, _), (b, _, _) in zip(whole, batched): + _assert_same(ad.read_h5ad(b), ad.read_h5ad(w)) + np.testing.assert_array_equal(_dense(ad.read_h5ad(b).X), _dense(source[g].X)) + + +def test_split_outputs_keep_dtypes_and_compression(tmp_path): + src = _write(_anndata(), tmp_path / "src.h5ad") + planned = split_store(src, "group", tmp_path / "parts", QUIET, manifest=False) + for _, path, _ in planned: + with h5py.File(path, "r") as f: + assert f["X/indices"].dtype == np.int32 + assert f["X/indptr"].dtype == np.int32 + assert f["X/data"].compression == "lzf" + + +# --------------------------------------------------------------------------- +# concat + + +def test_concat_keeps_int32_indices_and_compression(tmp_path): + parts = [ + _write(_anndata(n_obs=10, seed=s, raw=False), tmp_path / f"p{s}.h5ad") + for s in range(3) + ] + for s, p in enumerate(parts): + obj = ad.read_h5ad(p) + obj.obs_names = [f"p{s}-{n}" for n in obj.obs_names] + del obj.obsp["conn"] + obj.write_h5ad(p, compression="lzf") + out = tmp_path / "all.h5ad" + concat_on_disk(parts, out, QUIET) + + with h5py.File(out, "r") as f: + for element in ("X", "layers/counts"): + assert f[f"{element}/indices"].dtype == np.int32, element + assert f[f"{element}/indptr"].dtype == np.int32, element + assert f[f"{element}/data"].compression == "lzf", element + + expected = ad.concat([ad.read_h5ad(p) for p in parts]) + got = ad.read_h5ad(out) + np.testing.assert_array_equal(_dense(got.X), _dense(expected.X)) + np.testing.assert_array_equal( + _dense(got.layers["counts"]), _dense(expected.layers["counts"]) + ) + + +class _Fake: + """Just enough of a sparse group for `_concat_index_dtypes`.""" + + def __init__(self, nnz, indices=np.int32, indptr=np.int32): + self._d = { + "data": np.empty(0, dtype=np.float32), + "indices": np.empty(0, dtype=indices), + "indptr": np.empty(0, dtype=indptr), + } + self._nnz = nnz + + def __getitem__(self, key): + if key == "data": + return type("D", (), {"shape": (self._nnz,)})() + return self._d[key] + + +def test_concat_index_dtypes_widen_only_when_needed(): + limit = int(np.iinfo(np.int32).max) + small = [_Fake(10), _Fake(20)] + assert _concat_index_dtypes(small, 1000) == (np.int32, np.int32) + + # Each input fits int32 on its own; together their nonzeros do not. + big = [_Fake(limit // 2 + 1), _Fake(limit // 2 + 1)] + assert _concat_index_dtypes(big, 1000) == (np.int32, np.int64) + + assert _concat_index_dtypes(small, limit + 1) == (np.int64, np.int32) + + # A wider input is never narrowed. + mixed = [_Fake(10), _Fake(10, indices=np.int64, indptr=np.int64)] + assert _concat_index_dtypes(mixed, 1000) == (np.int64, np.int64) + + +# --------------------------------------------------------------------------- +# row reads for dataframes and obsm + + +@pytest.mark.parametrize("block", [1, 3, 1 << 20]) +def test_take_rows_reads_blocks_not_a_fancy_index(tmp_path, monkeypatch, block): + monkeypatch.setattr(subset_mod, "TAKE_BLOCK_ROWS", block) + rng = np.random.default_rng(0) + numbers = rng.normal(size=(30, 2)) + strings = np.array([f"s{i}" for i in range(30)], dtype=object) + with h5py.File(tmp_path / "t.h5", "w") as f: + f["n"] = numbers + f.create_dataset("s", data=strings, dtype=h5py.string_dtype()) + for indices in ( + np.array([], dtype=np.int64), + np.array([7]), + np.array([0, 1, 2, 29]), + np.sort(rng.choice(30, 11, replace=False)), + ): + np.testing.assert_array_equal( + subset_mod._take_rows(f["n"], indices), numbers[indices] + ) + got = subset_mod._take_rows(f["s"], indices) + assert [x.decode() if isinstance(x, bytes) else x for x in got] == list( + strings[indices] + ) + with pytest.raises(ValueError, match="sorted and unique"): + subset_mod._take_rows(f["n"], np.array([2, 1])) From d0b59e740cf0e436716ebd9171675cae98536636 Mon Sep 17 00:00:00 2001 From: Aljes Date: Thu, 24 Sep 2026 15:33:53 +0100 Subject: [PATCH 2/3] Release 0.6.1 A patch: two fixes to existing commands and no new surface. The one behavioural change is that a row selection passed to the subset writers must be sorted and unique; every command already produced one, and h5py already required it. Co-Authored-By: Claude Opus 5.5 (1M context) --- CHANGELOG.md | 7 ++++++- pyproject.toml | 2 +- uv.lock | 2 +- 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b7f1a6e..6a9dc72 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,7 +3,12 @@ Notable changes to `adata-cli`. Versions are `MAJOR.MINOR.PATCH`; tags carry no `v` prefix. -## Unreleased +## 0.6.1 + +Makes the outputs of `subset`, `split` and `concat` the size anndata writes +(they were about twice that), and makes `split` read each matrix once rather +than once per group. In 0.6.0 a hand-written anndata loop split the ci-tier +benchmark 7x faster than `split` did; `split` is now the faster of the two. ### Fixed diff --git a/pyproject.toml b/pyproject.toml index c4de0fa..02f04fb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ # so the distribution is published as `pyadata-cli`. The import package # and the command are both still `adata`. name = "pyadata-cli" -version = "0.6.0" +version = "0.6.1" description = "Streaming CLI for exploring and editing large AnnData .h5ad and .zarr stores" readme = "README.md" requires-python = ">=3.12" diff --git a/uv.lock b/uv.lock index dd5144a..3e9201d 100644 --- a/uv.lock +++ b/uv.lock @@ -486,7 +486,7 @@ wheels = [ [[package]] name = "pyadata-cli" -version = "0.6.0" +version = "0.6.1" source = { editable = "." } dependencies = [ { name = "h5py" }, From 91243671a8e846c454030c3236251b790fdc4d94 Mon Sep 17 00:00:00 2001 From: Aljes Date: Thu, 24 Sep 2026 15:48:12 +0100 Subject: [PATCH 3/3] Widen concat's index dtypes against their own limits, not int32's `_concat_index_dtypes` only compared the combined nnz and the minor dimension with int32's maximum, so inputs with a narrower index dtype kept it past its range: two int16 `indptr`s of 20,000 nonzeros each gave an int16 output whose 40,000 offset wrapped silently. Each is now checked against the chosen dtype's own limit and widened to the narrowest of int32/int64 that holds it. Found by Codex review on #19. Co-Authored-By: Claude Opus 5.5 (1M context) --- CHANGELOG.md | 2 +- src/adata/core/concat.py | 28 ++++++++++++++++++---------- tests/test_subset_fan_out.py | 35 +++++++++++++++++++++++++++++++++++ 3 files changed, 54 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6a9dc72..fbb5f87 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,7 +17,7 @@ benchmark 7x faster than `split` did; `split` is now the faster of the two. to int64, whatever the source used, and was written without the source's compression, so with an lzf input the `data` and `indices` came out uncompressed. Outputs now keep the source's index dtypes and storage - settings. `concat` widens to int64 only when the combined matrix could not + settings. `concat` widens only when the combined matrix could not be addressed otherwise. `convert` already did this; the code next to it did not. diff --git a/src/adata/core/concat.py b/src/adata/core/concat.py index c9c67ea..d2bebfe 100644 --- a/src/adata/core/concat.py +++ b/src/adata/core/concat.py @@ -389,26 +389,34 @@ def _matrix_kind(obj: Any) -> str: return "other" +def _widen_to_hold(dtype: np.dtype, largest: int) -> np.dtype: + """`dtype`, or the narrowest of int32/int64 wider than it that holds `largest`.""" + dtype = np.dtype(dtype) + if largest <= np.iinfo(dtype).max: + return dtype + for candidate in (np.int32, np.int64): + if largest <= np.iinfo(candidate).max: + return np.dtype(np.result_type(dtype, candidate)) + return np.dtype(np.int64) + + def _concat_index_dtypes( sources: List[Any], n_minor: int ) -> Tuple[np.dtype, np.dtype]: """The `indices` and `indptr` dtypes for a concatenation of `sources`. The widest the inputs used, so int32 inputs give an int32 output rather - than the int64 every output used to get. That is widened to int64 only - when the result could not be addressed otherwise: `indices` must reach - the output's minor dimension, and `indptr` its total nonzero count, - which is known before writing as the sum of the inputs'. + than the int64 every output used to get. That is widened only when the + result could not be addressed otherwise: `indices` must reach the + output's minor dimension, and `indptr` its total nonzero count, which is + known before writing as the sum of the inputs'. Both are checked against + the chosen dtype's own limit, not int32's, since the spec allows + narrower index arrays and an int16 `indptr` would otherwise wrap. """ - limit = np.iinfo(np.int32).max indices = np.result_type(*[s["indices"].dtype for s in sources]) indptr = np.result_type(*[s["indptr"].dtype for s in sources]) nnz = sum(int(s["data"].shape[0]) for s in sources) - if n_minor > limit: - indices = np.result_type(indices, np.int64) - if nnz > limit: - indptr = np.result_type(indptr, np.int64) - return np.dtype(indices), np.dtype(indptr) + return _widen_to_hold(indices, n_minor), _widen_to_hold(indptr, nnz) def _write_indptr( diff --git a/tests/test_subset_fan_out.py b/tests/test_subset_fan_out.py index 0668ca8..0528cce 100644 --- a/tests/test_subset_fan_out.py +++ b/tests/test_subset_fan_out.py @@ -317,6 +317,13 @@ def test_concat_index_dtypes_widen_only_when_needed(): assert _concat_index_dtypes(small, limit + 1) == (np.int64, np.int32) + # Narrower than int32: checked against its own limit, not int32's. Two + # int16 indptrs of 20,000 nonzeros each need offsets up to 40,000. + narrow = [_Fake(20_000, np.int16, np.int16), _Fake(20_000, np.int16, np.int16)] + assert _concat_index_dtypes(narrow, 100) == (np.int16, np.int32) + assert _concat_index_dtypes(narrow[:1], 40_000) == (np.int32, np.int16) + assert _concat_index_dtypes(narrow[:1], 100) == (np.int16, np.int16) + # A wider input is never narrowed. mixed = [_Fake(10), _Fake(10, indices=np.int64, indptr=np.int64)] assert _concat_index_dtypes(mixed, 1000) == (np.int64, np.int64) @@ -350,3 +357,31 @@ def test_take_rows_reads_blocks_not_a_fancy_index(tmp_path, monkeypatch, block): ) with pytest.raises(ValueError, match="sorted and unique"): subset_mod._take_rows(f["n"], np.array([2, 1])) + + +def test_concat_widens_a_narrow_indptr_rather_than_wrapping_it(tmp_path): + """End to end: int16 inputs whose combined nnz overflows int16.""" + parts = [] + for s in range(2): + rng = np.random.default_rng(s) + dense = (rng.random((200, 200)) < 0.6).astype("float32") # ~24,000 nnz + path = tmp_path / f"n{s}.h5ad" + ad.AnnData( + X=sparse.csr_matrix(dense), + obs=pd.DataFrame(index=[f"n{s}-{i}" for i in range(200)]), + var=pd.DataFrame(index=[f"g{i}" for i in range(200)]), + ).write_h5ad(path) + with h5py.File(path, "r+") as f: + for key in ("indices", "indptr"): + values = f[f"X/{key}"][...] + del f[f"X/{key}"] + f[f"X/{key}"] = values.astype(np.int16) + parts.append(path) + + out = tmp_path / "n.h5ad" + concat_on_disk(parts, out, QUIET) + with h5py.File(out, "r") as f: + assert f["X/indices"].dtype == np.int16 + assert f["X/indptr"].dtype == np.int32 + expected = sparse.vstack([ad.read_h5ad(p).X for p in parts]).toarray() + np.testing.assert_array_equal(ad.read_h5ad(out).X.toarray(), expected)