diff --git a/CHANGELOG.md b/CHANGELOG.md index d043cfc..fbb5f87 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,42 @@ Notable changes to `adata-cli`. Versions are `MAJOR.MINOR.PATCH`; tags carry no `v` prefix. +## 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 + +- **`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 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/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/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..d2bebfe 100644 --- a/src/adata/core/concat.py +++ b/src/adata/core/concat.py @@ -389,6 +389,45 @@ 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 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. + """ + 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) + return _widen_to_hold(indices, n_minor), _widen_to_hold(indptr, nnz) + + +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 +442,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 +451,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 +491,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 +551,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 +559,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 +592,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..0528cce --- /dev/null +++ b/tests/test_subset_fan_out.py @@ -0,0 +1,387 @@ +"""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) + + # 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) + + +# --------------------------------------------------------------------------- +# 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])) + + +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) 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" },