From c95fd31ae08bf8e26fc26392192d35332f1e50d4 Mon Sep 17 00:00:00 2001 From: Aljes Date: Fri, 30 Oct 2026 14:48:38 +0000 Subject: [PATCH 1/2] B3b: concat fidelity and scale concat now opens its inputs as it needs them, at most a quarter of the file descriptor limit at once, closing the most recently used to make room, so any number of inputs works and memory no longer grows with them (F076). Every element is looked up again on each pass rather than held. CSC inputs are read once each, in blocks of entries sized by --chunk, and their columns put together through buckets in an HDF5 scratch file beside the output, instead of 2 x n_var x n_inputs reads (F011). A column bigger than a block is streamed. The free-space check counts the scratch, for a Zarr output too, and moves into concat_on_disk to do so. Row pointers are one int64 array, not a Python list, in concat and in recast's sparsify (F132). Dense X, layers and obsm arrays keep the first input's chunking and compression (F075). --fill-value widens an integer matrix only when an outer join fills cells it cannot hold (F116). varp follows --merge as varm does, sparse entries compared by value (F118). An outer join's var and merged categories are ordered as anndata.concat orders them (D2; F113, F114), and ordered categoricals whose orders conflict are written unordered with a warning. A matrix whose rows are not its input's obs is refused before writing. Closes #93. Part of #87. Co-Authored-By: Claude Opus 5.5 --- docs/reference/commands.md | 47 +- src/adata/cli.py | 7 +- src/adata/commands/concat.py | 10 +- src/adata/core/concat.py | 1371 +++++++++++++++++++--------- src/adata/core/recast.py | 10 +- tests/test_concat_axis_elements.py | 147 +++ tests/test_concat_correctness.py | 256 +++++- tests/test_concat_scale.py | 361 ++++++++ tests/test_concat_versions.py | 30 + tests/test_performance.py | 216 ++++- 10 files changed, 1995 insertions(+), 460 deletions(-) create mode 100644 tests/test_concat_scale.py diff --git a/docs/reference/commands.md b/docs/reference/commands.md index 6ed3a31..1704828 100644 --- a/docs/reference/commands.md +++ b/docs/reference/commands.md @@ -287,7 +287,7 @@ adata concat a.h5ad b.h5ad -o m.h5ad --keys a,b --index-unique - --uns-merge sam | `--keys` | Names for the inputs; defaults to their filenames | | `--index-unique` | Delimiter for suffixing obs names with their key | | `--merge` / `--uns-merge` | `drop` (default), `same`, `unique`, `first`, `only` | -| `--fill-value` | Value for dense cells introduced by an outer join | +| `--fill-value` | Value for dense cells introduced by an outer join; an integer matrix that cannot hold it becomes float64 | | `--chunk`, `-C` | Row chunk size while streaming (at least 1) | | `--zarr-format` | Zarr version to write; defaults to the inputs' when they agree, otherwise v3 | | `--no-space-check` | Write even if the disk looks too small; see [Writing outputs](#writing-outputs) | @@ -299,11 +299,21 @@ them; make them unique first (anndata's `var_names_make_unique()`). are refused under `--join inner`; inputs with no variables at all concatenate to n_obs × 0. +var is in the order `anndata.concat` gives it. Under `--join inner` that is +the first input's order. Under `--join outer` it is the inputs' own order +when every input has the same variables in the same order, and the union +sorted otherwise; an input with no variables has no say. + obs columns keep their dtypes: categoricals union their category sets (numeric categories stay numeric, and anndata 0.7's categoricals are read by label), and nullable columns keep their masks. Under `--join outer` a column -missing from one input is padded rather than dropped. Missing values stay -missing: +missing from one input is padded rather than dropped. Merged categories are +ordered as anndata orders them: the first input's order, if its categories +already include every input's; otherwise the union in natural order (`t2` +before `t10`, numbers by value). An ordered categorical stays ordered only +if every input has the same categories in the same order; otherwise it is +written unordered, with a warning, as anndata keeps no order either. +Missing values stay missing: - Text with missing values (from a categorical, a nullable column, or an input without the column) is written as a categorical whose missing rows @@ -324,15 +334,28 @@ obsm entries present in every input are stacked: arrays as they are, sparse matrices as X is, and dataframes (a SCope import's `RegulonsAUC`) as obs is, their columns following `--join`. Like X, an entry must be stored the same way in every input (CSR, CSC or dense); concat refuses a mix before writing, -and `adata recast` can convert one to match. var's columns and varm (a SCope -import's `Regulons`, pycisTopic's `topic_region`) follow `--merge`, compared -on the output's genes: an input lacking some of them, as in an outer join, -counts as not having the entry. Without `--merge`, varm is dropped with a -note. - -`obsp`, `varp` and `raw` are **not** carried over — a pairwise value has no -meaning between cells from different inputs. anndata's own `concat_on_disk` -refuses these too. +and `adata recast` can convert one to match. var's columns, varm (a SCope +import's `Regulons`, pycisTopic's `topic_region`) and varp follow `--merge`, +compared on the output's genes: an input lacking some of them, as in an +outer join, counts as not having the entry. Sparse varp entries are compared +by value, as anndata compares them, whatever their format and dtype. Without +`--merge`, varm and varp are dropped with a note. + +`obsp` and `raw` are **not** carried over — a value between two cells has no +meaning when they come from different inputs. anndata's own +`concat_on_disk` refuses these too. + +Dense matrices and obsm arrays keep the first input's chunking and +compression, as sparse ones do. A dense integer matrix stays integer unless +an outer join fills cells with a `--fill-value` it cannot hold, such as 0.5 +or NaN; then it becomes float64. + +concat opens its inputs as it needs them, at most a quarter of the +process's file descriptor limit (`ulimit -n`) at once, so it takes any +number of inputs. CSC matrices are read once each, a block of entries at a +time, and their columns put together through a scratch file beside the +output, of about 16 bytes plus the value per stored entry of the largest +one. The free-space check counts it, for a Zarr output too. ### Merge strategies diff --git a/src/adata/cli.py b/src/adata/cli.py index d8c4987..f949ec5 100644 --- a/src/adata/cli.py +++ b/src/adata/cli.py @@ -444,7 +444,9 @@ def concat( "(default: drop)", ), fill_value: float = typer.Option( - 0.0, "--fill-value", help="Value for dense cells introduced by an outer join" + 0.0, "--fill-value", + help="Value for dense cells introduced by an outer join; an integer " + "matrix that cannot hold it becomes float64", ), chunk_rows: int = _chunk_option(1024, "Row chunk size while streaming"), zarr_format: Optional[int] = typer.Option( @@ -458,7 +460,8 @@ def concat( Concatenate stores along the obs axis. Streams each input in turn, so memory use is set by --chunk rather than by - the size of the inputs. obsp, varp and raw are not carried over. + the size of the inputs. obsp and raw are not carried over; var's columns, + varm and varp follow --merge. Examples: adata concat a.h5ad b.h5ad -o merged.h5ad diff --git a/src/adata/commands/concat.py b/src/adata/commands/concat.py index b89ffdb..a304771 100644 --- a/src/adata/commands/concat.py +++ b/src/adata/commands/concat.py @@ -9,7 +9,6 @@ from adata.core.concat import MERGE_CHOICES, MERGE_STRATEGIES, concat_on_disk from adata.util.path import check_output -from adata.util.staging import check_free_space, disk_size __all__ = ["MERGE_CHOICES", "MERGE_STRATEGIES", "concat_stores"] @@ -25,9 +24,6 @@ def concat_stores( # concat never replaces an output (it refuses one that exists), but # a new one may still name a path inside a Zarr input. check_output(output, files) - # About the sum of the inputs: an inner join keeps less, an outer join's - # dense fill can add some. - check_free_space( - output, lambda: sum(disk_size(f) for f in files), check=space_check - ) - concat_on_disk(files, output, console, **kwargs) # type: ignore[arg-type] + # The free space is checked by concat_on_disk, once it knows whether a + # CSC matrix needs scratch space beside the output too. + concat_on_disk(files, output, console, space_check=space_check, **kwargs) # type: ignore[arg-type] diff --git a/src/adata/core/concat.py b/src/adata/core/concat.py index 21ec8fa..1638072 100644 --- a/src/adata/core/concat.py +++ b/src/adata/core/concat.py @@ -3,25 +3,33 @@ Mirrors the useful core of `anndata.experimental.concat_on_disk`: inputs are streamed a block of rows at a time and written straight into the output, so peak memory is set by the block size rather than by the total size of the -inputs. +inputs. Inputs are opened as they are needed, a bounded number at a time +(`_Inputs`), so neither file descriptors nor memory grow with how many there +are. Columns of the concatenated axis follow `join`, in obs and in obsm's -dataframes alike; elements of the alternative axis (var's columns, varm) and -of uns follow a `merge` strategy, since there is no single correct way to -reconcile values that differ between inputs. +dataframes alike; elements of the alternative axis (var's columns, varm, +varp) and of uns follow a `merge` strategy, since there is no single correct +way to reconcile values that differ between inputs. """ from __future__ import annotations +import errno import hashlib +import os +import tempfile +from collections import OrderedDict +from contextlib import contextmanager from pathlib import Path -from typing import Any, Dict, List, Optional, Sequence, Tuple +from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Sequence, Tuple +import h5py import numpy as np from rich.console import Console from adata.core.info import axis_len -from adata.core.subset import aligned_rows +from adata.core.subset import _open_output_limit, aligned_rows from adata.core.select import read_names from adata.elements import spec from adata.elements.read import ( @@ -35,6 +43,7 @@ from adata.elements.write import ( ensure_anndata_skeleton, set_column_order, + set_shape_attr, write_categorical, write_dataframe_header, write_dense, @@ -52,15 +61,126 @@ is_zarr_group, Store, is_group, + detect_backend, open_store, + zarr_format_of, ) from adata.util.names import natural_sorted -from adata.util.staging import staged_outputs +from adata.util.staging import check_free_space, disk_size, staged_outputs MERGE_STRATEGIES = ("same", "unique", "first", "only") #: What the CLI accepts. "drop" is the default and maps to no merge at all. MERGE_CHOICES = ("drop",) + MERGE_STRATEGIES +#: How many inputs concat holds open at once: a quarter of the file +#: descriptor limit, as split caps its open outputs. Beyond that, inputs are +#: closed and opened again as they are needed. +MAX_OPEN_INPUTS = _open_output_limit() + + +# --------------------------------------------------------------------------- +# inputs + + +class _Inputs: + """The input stores, opened as they are needed, at most + `MAX_OPEN_INPUTS` of them at once. + + Opening every input up front and holding all of them to the end ran out + of file descriptors at about 253 HDF5 inputs under macOS's default limit + (about 1,020 on Linux), and each open input's caches made memory grow + with the number of inputs whatever --chunk was. An input is opened when + it is next needed, and another closed to make room: the one used most + recently. Every pass over the inputs goes from the first to the last, so + the next pass then finds the first ones still open, where closing the + one used longest ago would close each input just before it is needed. + + So an object read from an input stays valid only until another input is + opened: code that comes back to an input looks its element up again, + through `get`, `each` or a `_Column`, rather than holding on to it. + """ + + def __init__(self, files: Sequence[Path]) -> None: + self.files = list(files) + self._open: "OrderedDict[int, Store]" = OrderedDict() + + def __len__(self) -> int: + return len(self.files) + + def __enter__(self) -> "_Inputs": + return self + + def __exit__(self, exc_type, exc, tb) -> None: + self.close(suppress=exc_type is not None) + + def root(self, i: int) -> Any: + store = self._open.get(i) + if store is not None: + self._open.move_to_end(i) + return store.root + while len(self._open) >= max(1, MAX_OPEN_INPUTS): + self._open.popitem(last=True)[1].close() + store = self._open[i] = _open_input(self.files[i]) + return store.root + + def get(self, i: int, path: str) -> Optional[Any]: + """Input `i`'s element at `path`, or None if it has none.""" + root = self.root(i) + return root[path] if path in root else None + + def each(self, path: str) -> Iterator[Optional[Any]]: + """Each input's element at `path` in turn, None where it has none.""" + for i in range(len(self)): + yield self.get(i, path) + + def keys(self, i: int, path: str) -> List[str]: + """The members of input `i`'s group at `path`; none if it lacks it.""" + group = self.get(i, path) + return list(group.keys()) if group is not None else [] + + def close(self, suppress: bool = False) -> None: + while self._open: + self._open.popitem(last=False)[1].close(suppress=suppress) + + +def _open_input(path: Path) -> Store: + try: + return open_store(path, "r") + except OSError as error: + if error.errno != errno.EMFILE: + raise + raise OSError( + f"Could not open '{path}': too many files are open. concat holds " + f"at most {MAX_OPEN_INPUTS} inputs open at once, but the process " + "has no descriptors left for them; raise its limit with `ulimit -n`." + ) from error + + +class _Column: + """One dataframe column across the inputs, looked up afresh on each pass. + + Iterating gives `(column, frame)` for each input in turn: `column` is + None where the input lacks it, and `frame` is the dataframe it belongs + to, which an anndata 0.7 categorical needs to find its categories. + """ + + def __init__(self, inputs: _Inputs, frame: str, name: str) -> None: + self.inputs, self.frame, self.name = inputs, frame, name + + def __iter__(self) -> Iterator[Tuple[Optional[Any], Any]]: + for i in range(len(self.inputs)): + frame = self.inputs.get(i, self.frame) + present = frame is not None and self.name in frame + yield (frame[self.name] if present else None), frame + + def present(self) -> Iterator[Tuple[Any, Any]]: + """`(column, frame)` for the inputs that have the column.""" + return ((column, frame) for column, frame in self if column is not None) + + +# --------------------------------------------------------------------------- +# indices + def _index_union(per_input: Sequence[List[str]]) -> List[str]: """Union of indices, in order of first appearance.""" @@ -79,6 +199,24 @@ def _index_intersection(per_input: Sequence[List[str]]) -> List[str]: return [name for name in per_input[0] if name in common] +def _var_union(per_input: Sequence[List[str]]) -> List[str]: + """An outer join's var, in the order anndata gives it. + + anndata unions the inputs' indexes with pandas, which keeps the order + when every input has the same var in the same order and sorts the union + otherwise; an input with no variables has no say. First appearance, as + an earlier version used, put the same genes in another order, so code + aligning by position against an anndata-made reference saw another + matrix. + """ + named = [names for names in per_input if names] + if not named: + return [] + if all(names is named[0] or names == named[0] for names in named[1:]): + return list(named[0]) + return sorted(_index_union(named)) + + def _column_map(source: List[str], target: List[str]) -> np.ndarray: """Map each source position to its target position, or -1 if dropped.""" lookup = {name: i for i, name in enumerate(target)} @@ -87,6 +225,15 @@ def _column_map(source: List[str], target: List[str]) -> np.ndarray: ) +def _per_var_list(var_names: List[List[str]], compute: Callable[[List[str]], Any]) -> List[Any]: + """`compute(names)` for each input, once for inputs sharing a var list.""" + done: Dict[int, Any] = {} + for names in var_names: + if id(names) not in done: + done[id(names)] = compute(names) + return [done[id(names)] for names in var_names] + + def _apply_index_unique( names: List[str], key: Optional[str], delimiter: Optional[str] ) -> List[str]: @@ -95,6 +242,30 @@ def _apply_index_unique( return [f"{name}{delimiter}{key}" for name in names] +def _read_axes( + inputs: _Inputs, input_keys: List[str], index_unique: Optional[str] +) -> Tuple[List[List[str]], List[int], List[str], List[Optional[int]]]: + """Each input's var names, obs count and Zarr version, and the output's + obs names, reading the inputs one at a time. + + Inputs with the same var names share one list. An atlas's inputs + usually do, and a copy each made memory grow by a gene list per input. + """ + shared: Dict[Tuple[str, ...], List[str]] = {} + var_names: List[List[str]] = [] + obs_counts: List[int] = [] + formats: List[Optional[int]] = [] + target_obs: List[str] = [] + for i, key in enumerate(input_keys): + root = inputs.root(i) + names = read_names(root, "var") + var_names.append(shared.setdefault(tuple(names), names)) + obs_counts.append(axis_len(root, "obs")) + formats.append(zarr_format_of(root)) + target_obs.extend(_apply_index_unique(read_names(root, "obs"), key, index_unique)) + return var_names, obs_counts, target_obs, formats + + def _resolve_keys( files: Sequence[Path], keys: Optional[Sequence[str]] ) -> List[str]: @@ -136,7 +307,11 @@ def _refuse_repeated_var_names(files: Sequence[Path], var_names: List[List[str]] dense kept the last, and an outer join lost a variable, all with exit 0. anndata.concat refuses the same inputs. """ + checked = set() for path, names in zip(files, var_names): + if id(names) in checked: + continue + checked.add(id(names)) repeated = _duplicates(names) if not repeated: continue @@ -161,43 +336,48 @@ def _duplicates(values: Sequence[str]) -> List[str]: # merge strategies -def _merge_values(values: List[Any], strategy: Optional[str]) -> Tuple[bool, Any]: - """Reconcile one element's value across inputs. +def _merge_choice(values: Iterable[Any], strategy: Optional[str]) -> Optional[int]: + """Which input's value of an element to keep, or None to drop it. + + The strategies match anndata's: "same" keeps a value only when every + input agrees, "unique" when there is exactly one distinct value among + those present, "first" takes the earliest present, and "only" keeps it + when exactly one input has it at all. - Returns ``(keep, value)``. The strategies match anndata's: "same" keeps a - value only when every input agrees, "unique" when there is exactly one - distinct value among those present, "first" takes the earliest present, - and "only" keeps it when exactly one input has it at all. + `values` is read lazily and at most two are held at once, a reference + and the one compared with it, so that comparing an element across many + inputs does not hold every input's copy; and the reading stops as soon + as the answer is known. """ if strategy is None: - return False, None - - present = [v for v in values if v is not _MISSING] - if not present: - return False, None - - if strategy == "first": - return True, present[0] - - if strategy == "only": - return (len(present) == 1), (present[0] if len(present) == 1 else None) - - distinct: List[Any] = [] - for value in present: - if not any(_equal(value, seen) for seen in distinct): - distinct.append(value) - - if strategy == "unique": - return (len(distinct) == 1), (distinct[0] if len(distinct) == 1 else None) + return None + if strategy not in MERGE_STRATEGIES: + raise ValueError( + f"Unknown merge strategy {strategy!r}. " + f"Choose from: {', '.join(MERGE_STRATEGIES)}" + ) + chosen: Optional[int] = None + reference: Any = _MISSING + for i, value in enumerate(values): + if value is _MISSING: + if strategy == "same": + return None + continue + if chosen is None: + if strategy == "first": + return i + chosen, reference = i, value + elif strategy == "only" or not _equal(value, reference): + return None + return chosen - if strategy == "same": - keep = len(distinct) == 1 and len(present) == len(values) - return keep, (distinct[0] if keep else None) - raise ValueError( - f"Unknown merge strategy {strategy!r}. " - f"Choose from: {', '.join(MERGE_STRATEGIES)}" - ) +def _merge_values(values: Iterable[Any], strategy: Optional[str]) -> Tuple[bool, Any]: + """Reconcile one element's value across inputs: `(keep, value)`, as + `_merge_choice` decides it.""" + values = list(values) + chosen = _merge_choice(values, strategy) + return (chosen is not None), (values[chosen] if chosen is not None else None) class _Missing: @@ -312,12 +492,47 @@ def _category_labels(per_input: List[np.ndarray]) -> List[np.ndarray]: ] +def _natural_order(values: List[Any], dtype: np.dtype) -> List[Any]: + """`values` as natsort orders them: numbers by value, text naturally.""" + return sorted(values) if dtype.kind in "biuf" else natural_sorted(values) + + +def _merged_categories(labels: List[np.ndarray], ordered: List[bool]) -> Tuple[np.ndarray, bool]: + """A merged categorical's categories, and whether it stays ordered, as + anndata decides them. + + Unordered, the first input's categories keep their order when they + already hold every input's, as pandas keeps them; otherwise the union is + natsorted, as anndata sorts it. First appearance, as an earlier version + used, gave other codes and another plotting order than anndata.concat. + + Ordered only when every input with categories is ordered and has the + same ones in the same order. Otherwise the first input's order was kept, + silently reversing what another input's meant; anndata gives up the + order (plain values), and so the union is written unordered. + """ + dtype = labels[0].dtype if labels else np.dtype(object) + # Hash lookup, not list membership. `if category not in categories` on a + # list is O(k) per probe and so O(k^2) over the union -- around 2 million + # string comparisons at 1024 categories, and 5 billion at 100k. It reads + # the same number of bytes either way, so only a comparison count sees it. + union: Dict[Any, None] = {} + for source in labels: + for category in source.tolist(): + union.setdefault(category, None) + if any(ordered): + stated = [(c, o) for c, o in zip(labels, ordered) if c.size] + if stated and all( + o and c.size == len(union) and np.array_equal(c, stated[0][0]) for c, o in stated + ): + return stated[0][0], True + elif labels and labels[0].size == len(union): + return labels[0], False + return np.asarray(_natural_order(list(union), dtype), dtype=dtype), False + + def _concat_categorical( - parent: Any, - name: str, - columns: List[Optional[Any]], - frames: List[Any], - lengths: List[int], + parent: Any, name: str, column: _Column, lengths: List[int], console: Console ) -> None: """Concatenate categorical columns, unioning their category sets. @@ -326,25 +541,24 @@ def _concat_categorical( -1, which anndata reads back as a missing value. Categories keep their dtype: integer categories stay integers rather than becoming '1', '2'. """ - present = [(c, f) for c, f in zip(columns, frames) if c is not None] - labels = _category_labels([_categories(c, f) for c, f in present]) - - # Hash lookup, not list membership. `if category not in categories` on a - # list is O(k) per probe and so O(k^2) over the union -- around 2 million - # string comparisons at 1024 categories, and 5 billion at 100k. It reads - # the same number of bytes either way, so only a comparison count sees it. - lookup: Dict[Any, int] = {} - for source in labels: - for category in source.tolist(): - lookup.setdefault(category, len(lookup)) - dtype = labels[0].dtype if labels else np.dtype(object) - categories = np.asarray(list(lookup), dtype=dtype) - ordered = all(_ordered(c, f) for c, f in present) + labels: List[np.ndarray] = [] + ordered: List[bool] = [] + for col, frame in column.present(): + labels.append(_categories(col, frame)) + ordered.append(_ordered(col, frame)) + labels = _category_labels(labels) + categories, keeps_order = _merged_categories(labels, ordered) + if any(ordered) and not keeps_order: + console.print( + f"[yellow]Column {name!r}: the inputs disagree on its category order, " + "so it is written unordered (anndata keeps no order either).[/]" + ) + lookup = {category: i for i, category in enumerate(categories.tolist())} codes = np.full(sum(lengths), -1, dtype=np.int64) offset = 0 source_labels = iter(labels) - for col, length in zip(columns, lengths): + for (col, _), length in zip(column, lengths): if col is not None: source = next(source_labels).tolist() remap = np.array([lookup[c] for c in source] + [-1], dtype=np.int64) @@ -353,11 +567,11 @@ def _concat_categorical( codes[offset : offset + length] = remap[source_codes] offset += length - write_categorical(parent, name, codes, categories, ordered=ordered) + write_categorical(parent, name, codes, categories, ordered=keeps_order) def _concat_numbers( - parent: Any, name: str, columns: List[Optional[Any]], lengths: List[int] + parent: Any, name: str, column: _Column, lengths: List[int] ) -> None: """Concatenate number columns, plain or nullable, as pandas does. @@ -368,18 +582,19 @@ def _concat_numbers( padded with NaN, as pandas pads them. Floats, which have no nullable encoding, hold missing values as NaN. """ - present = [c for c in columns if c is not None] - nullable = any(_column_kind(c) == "masked" for c in present) - dtype = np.result_type( - *[(c["values"] if _column_kind(c) == "masked" else c).dtype for c in present] - ) - gaps = nullable or len(present) < len(columns) + kinds = [ + (_column_kind(c), (c["values"] if _column_kind(c) == "masked" else c).dtype) + for c, _ in column.present() + ] + nullable = any(kind == "masked" for kind, _ in kinds) + dtype = np.result_type(*[d for _, d in kinds]) + gaps = nullable or len(kinds) < len(lengths) total = sum(lengths) values = np.zeros(total, dtype=dtype) missing = np.ones(total, dtype=bool) offset = 0 - for col, length in zip(columns, lengths): + for (col, _), length in zip(column, lengths): if col is not None: if _column_kind(col) == "masked": values[offset : offset + length] = np.asarray(col["values"][...], dtype=dtype) @@ -423,11 +638,7 @@ def _text(col: Any, frame: Any) -> Tuple[np.ndarray, np.ndarray]: def _concat_text( - parent: Any, - name: str, - columns: List[Optional[Any]], - frames: List[Any], - lengths: List[int], + parent: Any, name: str, column: _Column, lengths: List[int] ) -> None: """Concatenate text columns, whatever mix of plain strings, categoricals and nullable strings encodes them, keeping every missing value missing: @@ -440,7 +651,7 @@ def _concat_text( missing = np.ones(total, dtype=bool) nullable = False offset = 0 - for col, frame, length in zip(columns, frames, lengths): + for (col, frame), length in zip(column, lengths): if col is not None: values[offset : offset + length], missing[offset : offset + length] = _text(col, frame) nullable = nullable or ( @@ -484,36 +695,28 @@ def _write_text( def _concat_obs_column( parent: Any, name: str, - columns: List[Optional[Any]], - frames: List[Any], + column: _Column, lengths: List[int], console: Console, ) -> bool: - """Write one concatenated annotation column. Returns whether it was kept. - - `frames` are the dataframes the columns come from, which a 0.7 - categorical needs to find its categories on Zarr. - """ - present = [c for c in columns if c is not None] - families = {_family(c) for c in present} + """Write one concatenated annotation column. Returns whether it was kept.""" + kinds = [(_family(c), _column_kind(c)) for c, _ in column.present()] + families = {family for family, _ in kinds} # Text meeting nullable numbers is kept as text, the numbers rendered and # their mask kept; meeting plain numbers it is dropped below. Both as # before, since anndata cannot write either mix itself. - nullable_numbers = all( - _column_kind(c) == "masked" for c in present if _family(c) == "number" - ) + nullable_numbers = all(kind == "masked" for family, kind in kinds if family == "number") if families == {"categorical"}: - _concat_categorical(parent, name, columns, frames, lengths) + _concat_categorical(parent, name, column, lengths, console) elif families == {"number"}: - _concat_numbers(parent, name, columns, lengths) + _concat_numbers(parent, name, column, lengths) elif families <= {"text", "categorical", "number"} and nullable_numbers: - _concat_text(parent, name, columns, frames, lengths) + _concat_text(parent, name, column, lengths) else: - kinds = sorted({_column_kind(c) for c in columns if c is not None}) console.print( f"[yellow]Dropping column {name!r}: " - f"cannot concatenate encodings {kinds}[/]" + f"cannot concatenate encodings {sorted({kind for _, kind in kinds})}[/]" ) return False return True @@ -544,7 +747,7 @@ def _widen_to_hold(dtype: np.dtype, largest: int) -> np.dtype: def _concat_index_dtypes( - sources: List[Any], n_minor: int + sources: Iterable[Any], n_minor: int ) -> Tuple[np.dtype, np.dtype]: """The `indices` and `indptr` dtypes for a concatenation of `sources`. @@ -555,15 +758,23 @@ def _concat_index_dtypes( 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. + `sources` is read once, so it may be a lazy pass over the inputs. """ - 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) + indices: List[np.dtype] = [] + indptr: List[np.dtype] = [] + nnz = 0 + for source in sources: + indices.append(source["indices"].dtype) + indptr.append(source["indptr"].dtype) + nnz += int(source["data"].shape[0]) + return ( + _widen_to_hold(np.result_type(*indices), n_minor), + _widen_to_hold(np.result_type(*indptr), nnz), + ) def _write_indptr( - group: Any, indptr: List[int], dtype: np.dtype, template: Any + group: Any, indptr: np.ndarray, dtype: np.dtype, template: Any ) -> None: from adata.core.subset import _sparse_dataset @@ -571,10 +782,25 @@ def _write_indptr( _sparse_dataset(group, "indptr", dtype, values.size, template)[:] = values +def _sparse_group(dst_parent: Any, name: str, enc: str, shape: Tuple[int, int]) -> Any: + group = dst_parent.create_group(name) + spec.set_encoding(group, enc) + set_shape_attr(group, shape) + return group + + +def _entries_per_block(nnz: int, n_major: int, chunk_rows: int) -> int: + """How many stored entries make a block: as many as `chunk_rows` rows + hold on average, so a sparse matrix read by entries costs about what one + read by rows does.""" + return max(1, chunk_rows) * max(1, -(-nnz // max(1, n_major))) + + def _concat_sparse( dst_parent: Any, name: str, - sources: List[Any], + inputs: _Inputs, + path: str, col_maps: List[np.ndarray], n_rows: int, n_cols: int, @@ -590,22 +816,19 @@ def _concat_sparse( """ from adata.core.subset import _append, _growable_like - group = dst_parent.create_group(name) - spec.set_encoding(group, spec.CSR_MATRIX) - from adata.elements.write import set_shape_attr - - set_shape_attr(group, (n_rows, n_cols)) - - dtype = np.result_type(*[s["data"].dtype for s in sources]) - 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 - - for source, col_map in zip(sources, col_maps): + group = _sparse_group(dst_parent, name, spec.CSR_MATRIX, (n_rows, n_cols)) + dtype = np.result_type(*[s["data"].dtype for s in inputs.each(path)]) + index_dtype, pointer_dtype = _concat_index_dtypes(inputs.each(path), n_cols) + template = inputs.get(0, path) + out_data = _growable_like(group, "data", dtype, template["data"]) + out_indices = _growable_like(group, "indices", index_dtype, template["indices"]) + # One int64 array for the row pointers. A Python list of int objects + # took 36 to 40 bytes a row against 8: 400 MB at 10M cells. + indptr = np.zeros(n_rows + 1, dtype=np.int64) + row = nnz = 0 + + for i, col_map in enumerate(col_maps): + source = inputs.get(i, path) src_indptr = np.asarray(source["indptr"][...], dtype=np.int64) src_data, src_indices = source["data"], source["indices"] n_src_rows = len(src_indptr) - 1 @@ -637,125 +860,261 @@ def _concat_sparse( mapped, values = mapped[order], values[order] counts = np.bincount(rows, minlength=end - start) - indptr.extend((nnz + np.cumsum(counts)).tolist()) + indptr[row + 1 : row + 1 + end - start] = nnz + np.cumsum(counts) + row += end - start nnz += int(counts.sum()) _append(out_indices, mapped) _append(out_data, values.astype(dtype)) + # Let go of this input's pointers before reading the next one's, + # rather than hold both: 8 bytes a row more at the peak. + del src_indptr + + _write_indptr(group, indptr, pointer_dtype, inputs.get(0, path)["indptr"]) - _write_indptr(group, indptr, pointer_dtype, sources[0]["indptr"]) + +def _dtype_holding(dtype: np.dtype, fill_value: float) -> np.dtype: + """`dtype`, or float64 if it cannot hold `fill_value` exactly.""" + dtype = np.dtype(dtype) + if dtype.kind not in "biu": + return dtype + if not np.isfinite(fill_value) or fill_value != int(fill_value): + return np.dtype(np.float64) + if dtype.kind == "b": + return dtype if fill_value in (0, 1) else np.dtype(np.float64) + info = np.iinfo(dtype) + return dtype if info.min <= fill_value <= info.max else np.dtype(np.float64) + + +def _layout_like(src: Any, dst_parent: Any, shape: Tuple[int, ...], dtype: Any) -> Dict[str, Any]: + """Creation kwargs that give a new array `src`'s chunking and + compression, the chunks fitted to its own shape, as every copied + dataset's are kept.""" + backend = "zarr" if is_zarr_group(dst_parent) else "hdf5" + kw = fit_chunks(dataset_create_kwargs(src, target_backend=backend, dst_parent=dst_parent), + shape, backend=backend) + if np.dtype(dtype) != np.dtype(src.dtype): + # The source's fill value is of its own dtype, not the output's. + kw.pop("fillvalue", None) + kw.pop("fill_value", None) + return kw def _concat_dense( dst_parent: Any, name: str, - sources: List[Any], + inputs: _Inputs, + path: str, col_maps: List[np.ndarray], n_rows: int, n_cols: int, chunk_rows: int, fill_value: float, ) -> None: - """Concatenate dense matrices row-wise, scattering columns into place.""" - dtype = np.result_type(*[s.dtype for s in sources]) - if fill_value != 0 and dtype.kind in ("i", "u", "b"): - dtype = np.dtype("float64") + """Concatenate dense matrices row-wise, scattering columns into place. + + The output keeps the first input's chunking and compression; written + with the defaults, gzip-compressed inputs came out contiguous and + uncompressed, 34 times their size for a matrix of zeros. Cells an input + lacks (an outer join) get `fill_value`, and only then may the dtype + change: to float64, if the inputs' cannot hold it. Promoting whenever a + fill value was given made integer X float64 on an inner join, where + nothing is filled. + """ + dtype = np.result_type(*[s.dtype for s in inputs.each(path)]) + if any(np.count_nonzero(m >= 0) < n_cols for m in col_maps): + dtype = _dtype_holding(dtype, fill_value) - out = create_dataset(dst_parent, name, shape=(n_rows, n_cols), dtype=dtype) + shape = (n_rows, n_cols) + out = create_dataset(dst_parent, name, shape=shape, dtype=dtype, + **_layout_like(inputs.get(0, path), dst_parent, shape, dtype)) spec.set_encoding(out, spec.ARRAY) row_offset = 0 - for source, col_map in zip(sources, col_maps): + for i, col_map in enumerate(col_maps): + source = inputs.get(i, path) keep = col_map >= 0 targets = col_map[keep] + gaps = targets.size < n_cols n_src_rows = source.shape[0] for start in range(0, n_src_rows, chunk_rows): end = min(start + chunk_rows, n_src_rows) - block = np.full((end - start, n_cols), fill_value, dtype=dtype) + block = (np.full((end - start, n_cols), fill_value, dtype=dtype) if gaps + else np.empty((end - start, n_cols), dtype=dtype)) block[:, targets] = np.asarray(source[start:end, :])[:, keep] out[row_offset + start : row_offset + end, :] = block row_offset += n_src_rows -def _inverse_column_map(col_map: np.ndarray, n_cols: int) -> np.ndarray: - """Invert a source->target column map into target->source, -1 where absent.""" - inverse = np.full(n_cols, -1, dtype=np.int64) - present = col_map >= 0 - inverse[col_map[present]] = np.nonzero(present)[0] - return inverse +@contextmanager +def _bucket_scratch(dst_parent: Any) -> Iterator[Any]: + """Where the buckets go: an HDF5 file beside the output, whatever the + output is. `scratch_group` keeps a Zarr output's scratch inside the + output, where every small append to a bucket rewrites that array's + metadata: putting a CSC matrix together there took sixty times as long + as in HDF5.""" + from adata.core.csr_build import scratch_group + + root = getattr(getattr(dst_parent, "store", None), "root", None) + if not is_zarr_group(dst_parent) or root is None: + with scratch_group(dst_parent, "__concat_scratch__") as group: + yield group + return + fd, path = tempfile.mkstemp( + prefix=".adata-", suffix=".concat-scratch.h5", dir=str(Path(root).parent) + ) + os.close(fd) + try: + # No chunk cache, as `scratch_group`'s: each bucket would hold a + # partly written chunk in one. + with h5py.File(path, "w", rdcc_nbytes=0) as handle: + yield handle.create_group("buckets") + finally: + os.unlink(path) + + +def _csc_bounds(indptr: np.ndarray, entries: int) -> np.ndarray: + """Column bounds of the buckets: about `entries` entries each, as a + transpose's are, and a column holding more than that alone in a bucket + of its own, so that it can be streamed rather than read whole.""" + from adata.core.csr_build import MIN_BUCKET_ENTRIES, bucket_bounds + + bounds = bucket_bounds(indptr, chunk=entries) + big = np.flatnonzero(np.diff(indptr) > max(MIN_BUCKET_ENTRIES, entries)) + return np.unique(np.concatenate((bounds, big, big + 1))) + + +def _in_row_order(minors: Any, n: int, entries: int) -> bool: + """Whether a bucket's rows are in order already, read a piece at a time.""" + last = -1 + for lo in range(0, n, entries): + piece = np.asarray(minors[lo : lo + entries], dtype=np.int64) + if piece.size and (piece[0] < last or bool((piece[1:] < piece[:-1]).any())): + return False + last = int(piece[-1]) if piece.size else last + return True + + +def _csc_runs(sorter: Any, entries: int) -> Iterator[Tuple[np.ndarray, np.ndarray]]: + """The buckets' `(rows, values)`, column by column, each in row order. + + A bucket of several columns is read whole and sorted. A bucket of one + column holding more than `entries` is streamed instead: its entries + arrived input by input, each input's in row order, so they are sorted + already -- unless some input stored that column out of order, which a + first pass over the rows finds, and then it is sorted whole after all. + Otherwise a gene expressed in most of a large atlas's cells would put + millions of entries in memory at once. + """ + for b, (majors, minors, values) in enumerate(sorter.buckets): + n = int(majors.shape[0]) + if n == 0: + continue + one_column = int(sorter.bounds[b + 1] - sorter.bounds[b]) == 1 + if one_column and n > entries and _in_row_order(minors, n, entries): + for lo in range(0, n, entries): + yield (np.asarray(minors[lo : lo + entries], dtype=np.int64), + np.asarray(values[lo : lo + entries])) + continue + major = np.asarray(majors[...], dtype=np.int64) + minor = np.asarray(minors[...], dtype=np.int64) + order = np.lexsort((minor, major)) + yield minor[order], np.asarray(values[...])[order] def _concat_csc( dst_parent: Any, name: str, - sources: List[Any], + inputs: _Inputs, + path: str, col_maps: List[np.ndarray], row_counts: List[int], n_rows: int, n_cols: int, + chunk_rows: int, ) -> None: """Concatenate CSC matrices row-wise, keeping the CSC encoding. - A CSC matrix is stored by column, so concatenating along obs means, for - each target column, appending each input's row indices in turn with that - 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. + Each output column holds that gene's entries from every input, so none + is complete until every input has been read. Reading one column of + every input at a time, as an earlier version did, made 2 x n_var x + n_inputs reads, each decompressing a whole chunk (on sharded Zarr, + fetching the shard index too) to take a few hundred entries: 37 times + the input's size read on h5ad, 572 times on Zarr. Instead each input is + read once, in blocks of entries sized by `chunk_rows`, and its entries + scattered into column-range buckets on disk, as a streaming transpose + does (`core/csr_build.py`); the buckets are then read back in turn, + sorted and appended. The output's column pointers are known before any + entry is read, from the inputs' own. """ + from adata.core.csr_build import BucketSorter from adata.core.subset import _append, _growable_like - from adata.elements.write import set_shape_attr - - group = dst_parent.create_group(name) - spec.set_encoding(group, spec.CSC_MATRIX) - set_shape_attr(group, (n_rows, n_cols)) - - dtype = np.result_type(*[s["data"].dtype for s in sources]) - 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 - - inverses = [_inverse_column_map(cm, n_cols) for cm in col_maps] - indptrs = [np.asarray(s["indptr"][...], dtype=np.int64) for s in sources] - row_offsets = np.cumsum([0] + list(row_counts[:-1])) - - for target_col in range(n_cols): - rows: List[np.ndarray] = [] - values: List[np.ndarray] = [] - for source, src_indptr, inverse, offset in zip( - sources, indptrs, inverses, row_offsets - ): - src_col = int(inverse[target_col]) - if src_col < 0: - continue - lo, hi = int(src_indptr[src_col]), int(src_indptr[src_col + 1]) - if hi <= lo: - continue - rows.append(np.asarray(source["indices"][lo:hi], dtype=np.int64) + offset) - values.append(np.asarray(source["data"][lo:hi])) - - if rows: - _append(out_indices, np.concatenate(rows)) - _append(out_data, np.concatenate(values).astype(dtype)) - nnz += int(sum(len(r) for r in rows)) - indptr.append(nnz) - - _write_indptr(group, indptr, pointer_dtype, sources[0]["indptr"]) - -def check_matrix_encodings(roots: List[Any], console: Console) -> None: + group = _sparse_group(dst_parent, name, spec.CSC_MATRIX, (n_rows, n_cols)) + dtype = np.result_type(*[s["data"].dtype for s in inputs.each(path)]) + index_dtype, pointer_dtype = _concat_index_dtypes(inputs.each(path), n_rows) + template = inputs.get(0, path) + out_data = _growable_like(group, "data", dtype, template["data"]) + out_indices = _growable_like(group, "indices", index_dtype, template["indices"]) + + counts = np.zeros(n_cols, dtype=np.int64) + for source, col_map in zip(inputs.each(path), col_maps): + kept = col_map >= 0 + counts[col_map[kept]] += np.diff(np.asarray(source["indptr"][...], dtype=np.int64))[kept] + indptr = np.concatenate(([0], np.cumsum(counts))) + total = int(indptr[-1]) + + if total: + entries = _entries_per_block(total, n_rows, chunk_rows) + with _bucket_scratch(dst_parent) as scratch: + sorter = BucketSorter(scratch, _csc_bounds(indptr, entries), dtype) + offset = 0 + for i, (col_map, n_src_rows) in enumerate(zip(col_maps, row_counts)): + source = inputs.get(i, path) + src_indptr = np.asarray(source["indptr"][...], dtype=np.int64) + src_data, src_indices = source["data"], source["indices"] + stored = int(src_indptr[-1]) if src_indptr.size else 0 + for lo in range(0, stored, entries): + hi = min(lo + entries, stored) + cols = np.searchsorted(src_indptr, np.arange(lo, hi), side="right") - 1 + target = col_map[cols] + keep = target >= 0 + rows = np.asarray(src_indices[lo:hi], dtype=np.int64)[keep] + offset + values = np.asarray(src_data[lo:hi])[keep].astype(dtype, copy=False) + sorter.add(target[keep], rows, values) + offset += n_src_rows + for rows, values in _csc_runs(sorter, entries): + _append(out_indices, rows.astype(index_dtype, copy=False)) + _append(out_data, values) + + _write_indptr(group, indptr, pointer_dtype, inputs.get(0, path)["indptr"]) + + +def check_matrix_encodings( + inputs: _Inputs, console: Console, row_counts: Optional[Sequence[int]] = None +) -> None: """Fail before writing anything if a matrix cannot be concatenated. Checked up front rather than mid-write: discovering this half way through would leave a partial store behind, and skipping the matrix would produce - an output silently missing X. + an output silently missing X. A matrix whose rows are not its input's + obs is refused here too, rather than written out of step with obs. """ - def _check(label: str, sources: List[Any]) -> None: - kinds = {_matrix_kind(s) for s in sources} - if kinds in ({spec.CSR_MATRIX}, {spec.CSC_MATRIX}, {"dense"}): + def _check(label: str, path: str) -> None: + kinds: Dict[str, None] = {} + rows: List[int] = [] + for source in inputs.each(path): + kind = _matrix_kind(source) + kinds.setdefault(kind, None) + rows.append(_shape(source)[0] if kind != "other" else -1) + if set(kinds) in ({spec.CSR_MATRIX}, {spec.CSC_MATRIX}, {"dense"}): + for file, got, expected in zip(inputs.files, rows, row_counts or rows): + if got != expected: + raise ValueError( + f"Cannot concatenate {label!r}: '{file}' holds {got:,} rows " + f"of it but {expected:,} obs." + ) return wanted = sorted(kinds)[0] raise ValueError( @@ -766,31 +1125,33 @@ def _check(label: str, sources: List[Any]) -> None: f"{'dense' if wanted == 'dense' else wanted.removesuffix('_matrix')}`." ) - if all("X" in r for r in roots): - _check("X", [r["X"] for r in roots]) - elif any("X" in r for r in roots): + n = len(inputs) + has_x = [inputs.get(i, "X") is not None for i in range(n)] + if all(has_x): + _check("X", "X") + elif any(has_x): console.print("[yellow]Skipping X: not present in every input[/]") - names = _index_union( - [list(r["layers"].keys()) if "layers" in r else [] for r in roots] - ) - for name in names: - if all("layers" in r and name in r["layers"] for r in roots): - _check(f"layers/{name}", [r["layers"][name] for r in roots]) + for name in _index_union([inputs.keys(i, "layers") for i in range(n)]): + path = f"layers/{name}" + if all(source is not None for source in inputs.each(path)): + _check(path, path) # obsm entries are stacked only when every input stores them alike; a # sparse entry dense elsewhere is refused here, as X would be, rather # than dropped. Entries of different widths are skipped later, whatever # their encodings, and dataframes are not matrices at all. - names = _index_union([list(r["obsm"].keys()) if "obsm" in r else [] for r in roots]) - for name in names: - if not all("obsm" in r and name in r["obsm"] for r in roots): + for name in _index_union([inputs.keys(i, "obsm") for i in range(n)]): + path = f"obsm/{name}" + if not all(source is not None for source in inputs.each(path)): continue - sources = [r["obsm"][name] for r in roots] - kinds = {_matrix_kind(s) for s in sources} - if (kinds & spec.SPARSE_TYPES and kinds <= spec.SPARSE_TYPES | {"dense"} - and len({_trailing_shape(s) for s in sources}) == 1): - _check(f"obsm/{name}", sources) + kinds, widths = set(), set() + for source in inputs.each(path): + kinds.add(_matrix_kind(source)) + if _matrix_kind(source) != "other": + widths.add(_trailing_shape(source)) + if kinds & spec.SPARSE_TYPES and kinds <= spec.SPARSE_TYPES | {"dense"} and len(widths) == 1: + _check(path, path) def _shape(obj: Any) -> Tuple[int, ...]: @@ -808,64 +1169,48 @@ def _trailing_shape(obj: Any) -> Tuple[int, ...]: def _concat_matrix( dst_parent: Any, name: str, - sources: List[Any], + inputs: _Inputs, + path: str, col_maps: List[np.ndarray], row_counts: List[int], n_rows: int, n_cols: int, chunk_rows: int, fill_value: float, - console: Console, -) -> bool: - """Concatenate X or one layer across inputs. Returns whether it was written. +) -> None: + """Concatenate X or one layer across inputs. Encodings are validated by :func:`check_matrix_encodings` before the output store exists, so anything reaching here is concatenable. """ - kinds = {_matrix_kind(s) for s in sources} + kinds = {_matrix_kind(s) for s in inputs.each(path)} if kinds == {spec.CSR_MATRIX}: - _concat_sparse( - dst_parent, name, sources, col_maps, n_rows, n_cols, chunk_rows - ) - return True - - if kinds == {spec.CSC_MATRIX}: - _concat_csc( - dst_parent, name, sources, col_maps, row_counts, n_rows, n_cols - ) - return True - - if kinds == {"dense"}: - _concat_dense( - dst_parent, - name, - sources, - col_maps, - n_rows, - n_cols, - chunk_rows, - fill_value, + _concat_sparse(dst_parent, name, inputs, path, col_maps, n_rows, n_cols, chunk_rows) + elif kinds == {spec.CSC_MATRIX}: + _concat_csc(dst_parent, name, inputs, path, col_maps, row_counts, n_rows, n_cols, chunk_rows) + elif kinds == {"dense"}: + _concat_dense(dst_parent, name, inputs, path, col_maps, n_rows, n_cols, chunk_rows, + fill_value) + else: + # A backstop. `check_matrix_encodings` runs before the output store is + # created and should have raised already; this catches an element it + # does not cover, where failing late still beats writing nonsense. + raise ValueError( + f"Cannot concatenate {name!r}: inputs use {', '.join(sorted(kinds))}." ) - return True - - # A backstop. `check_matrix_encodings` runs before the output store is - # created and should have raised already; this catches an element it - # does not cover, where failing late still beats writing nonsense. - raise ValueError( - f"Cannot concatenate {name!r}: inputs use {', '.join(sorted(kinds))}." - ) def _concat_obsm( dst_parent: Any, name: str, - sources: List[Any], + inputs: _Inputs, + path: str, n_rows: int, chunk_rows: int, console: Console, *, - lengths: Optional[List[int]] = None, + lengths: List[int], join: str = "inner", target_obs: Optional[List[str]] = None, ) -> bool: @@ -874,15 +1219,17 @@ def _concat_obsm( A dataframe in every input (a SCope import's `RegulonsAUC`) is concatenated as obs is, its columns following `join`. """ - if all(spec.encoding_type(s) == spec.DATAFRAME for s in sources) and target_obs is not None: - _concat_frame(dst_parent, name, sources, lengths or [], join, target_obs, console) + if target_obs is not None and all( + spec.encoding_type(s) == spec.DATAFRAME for s in inputs.each(path) + ): + _concat_frame(dst_parent, name, inputs, path, lengths, join, target_obs, console) return True - kinds = {_matrix_kind(s) for s in sources} + kinds = {_matrix_kind(s) for s in inputs.each(path)} if "other" in kinds: console.print(f"[yellow]Skipping obsm/{name}: not an array, or not a dataframe, in every input[/]") return False - widths = {_trailing_shape(s) for s in sources} + widths = {_trailing_shape(s) for s in inputs.each(path)} if len(widths) != 1: console.print( f"[yellow]Skipping obsm/{name}: inputs disagree on shape {sorted(widths)}[/]" @@ -896,21 +1243,22 @@ def _concat_obsm( if kinds != {"dense"}: # Sparse in every input, with one encoding (checked before writing): # stacked as X is, every column kept where it is. - same = [np.arange(trailing[0], dtype=np.int64)] * len(sources) + same = [np.arange(trailing[0], dtype=np.int64)] * len(inputs) if kinds == {spec.CSR_MATRIX}: - _concat_sparse(dst_parent, name, sources, same, n_rows, trailing[0], chunk_rows) + _concat_sparse(dst_parent, name, inputs, path, same, n_rows, trailing[0], chunk_rows) else: - _concat_csc(dst_parent, name, sources, same, [_shape(s)[0] for s in sources], - n_rows, trailing[0]) + _concat_csc(dst_parent, name, inputs, path, same, lengths, n_rows, trailing[0], + chunk_rows) return True - dtype = np.result_type(*[s.dtype for s in sources]) - out = create_dataset( - dst_parent, name, shape=(n_rows,) + trailing, dtype=dtype - ) + + dtype = np.result_type(*[s.dtype for s in inputs.each(path)]) + shape = (n_rows,) + trailing + out = create_dataset(dst_parent, name, shape=shape, dtype=dtype, + **_layout_like(inputs.get(0, path), dst_parent, shape, dtype)) spec.set_encoding(out, spec.ARRAY) offset = 0 - for source in sources: + for source in inputs.each(path): for start in range(0, source.shape[0], chunk_rows): end = min(start + chunk_rows, source.shape[0]) out[offset + start : offset + end, ...] = source[start:end, ...] @@ -921,14 +1269,15 @@ def _concat_obsm( def _concat_frame( dst_parent: Any, name: str, - sources: List[Any], + inputs: _Inputs, + path: str, lengths: List[int], join: str, target_obs: List[str], console: Console, ) -> None: """Dataframes stacked row-wise, as obs is: their columns follow `join`.""" - per_input = [dataframe_columns(s, resolve_index(s, "obs")[1]) for s in sources] + per_input = [dataframe_columns(s, resolve_index(s, "obs")[1]) for s in inputs.each(path)] if join == "inner": shared = set(per_input[0]).intersection(*map(set, per_input[1:])) names = [c for c in per_input[0] if c in shared] @@ -937,14 +1286,13 @@ def _concat_frame( frame = write_dataframe_header(dst_parent, name, target_obs, []) written = [ column for column in names - if _concat_obs_column(frame, column, [s[column] if column in s else None for s in sources], - sources, lengths, console) + if _concat_obs_column(frame, column, _Column(inputs, path, column), lengths, console) ] set_column_order(frame, written) # --------------------------------------------------------------------------- -# varm +# varm and varp def _digest(hasher: Any, values: np.ndarray) -> None: @@ -958,20 +1306,27 @@ def _digest(hasher: Any, values: np.ndarray) -> None: hasher.update(np.ascontiguousarray(values).tobytes()) +def _array_snapshot( + obj: Any, rows: np.ndarray, block_rows: int, columns: Optional[np.ndarray] = None +) -> Any: + """A digest of an array on the target var, read a block of rows at a + time: its rows put in the target's order, and for a pairwise entry its + columns too. A digest, not the values, so that comparing an entry holds + one block in memory rather than every input's whole entry at once.""" + hasher = hashlib.blake2b(digest_size=32) + for start in range(0, len(rows), block_rows): + block = aligned_rows(obj, rows[start:start + block_rows]) + _digest(hasher, block if columns is None else block[:, columns]) + kind = "text" if is_string_dtype(obj.dtype) else str(obj.dtype) + shape = tuple(obj.shape[1:]) if columns is None else (len(columns),) + return ("array", kind, shape, hasher.digest()) + + def _varm_snapshot(obj: Any, where: np.ndarray, block_rows: int) -> Any: """A comparable value of a varm entry, on the target var: a digest of an - array, read a block of rows at a time; of each column, for a dataframe. - - A digest, not the values, so that comparing an entry holds one block - (or one var-length column) in memory rather than every input's whole - entry at once. - """ + array, read a block of rows at a time; of each column, for a dataframe.""" if is_dataset(obj): - hasher = hashlib.blake2b(digest_size=32) - for start in range(0, len(where), block_rows): - _digest(hasher, aligned_rows(obj, where[start:start + block_rows])) - kind = "text" if is_string_dtype(obj.dtype) else str(obj.dtype) - return ("array", kind, tuple(obj.shape[1:]), hasher.digest()) + return _array_snapshot(obj, where, block_rows) if spec.encoding_type(obj) != spec.DATAFRAME: return _Incomparable() columns = [] @@ -1001,21 +1356,134 @@ def _varm_snapshot(obj: Any, where: np.ndarray, block_rows: int) -> Any: def _write_varm_array(parent: Any, name: str, obj: Any, take: np.ndarray, block_rows: int) -> None: """A varm array's rows on the target var, a block at a time, its chunking and compression kept as every other copied dataset's are.""" - backend = "zarr" if is_zarr_group(parent) else "hdf5" shape = (len(take),) + tuple(obj.shape[1:]) - kw = fit_chunks(dataset_create_kwargs(obj, target_backend=backend, dst_parent=parent), - shape, backend=backend) text = is_string_dtype(obj.dtype) - out = create_dataset(parent, name, shape=shape, dtype=obj.dtype, **kw) + out = create_dataset(parent, name, shape=shape, dtype=obj.dtype, + **_layout_like(obj, parent, shape, obj.dtype)) spec.set_encoding(out, spec.STRING_ARRAY if text else spec.ARRAY) for start in range(0, len(take), block_rows): out[start:start + block_rows, ...] = aligned_rows(obj, take[start:start + block_rows]) +def _entries_on(obj: Any, col_map: np.ndarray, entries: int) -> Iterator[Tuple[np.ndarray, np.ndarray, np.ndarray]]: + """A sparse pairwise entry's stored entries on the target var, a block + of `entries` at a time in the order they are stored: `(major, minor, + values)`, both renumbered by `col_map`, and those of genes the target + lacks left out.""" + indptr = np.asarray(obj["indptr"][...], dtype=np.int64) + stored = int(indptr[-1]) if indptr.size else 0 + for lo in range(0, stored, entries): + hi = min(lo + entries, stored) + major = col_map[np.searchsorted(indptr, np.arange(lo, hi), side="right") - 1] + minor = col_map[np.asarray(obj["indices"][lo:hi], dtype=np.int64)] + values = np.asarray(obj["data"][lo:hi]) + keep = (major >= 0) & (minor >= 0) + yield major[keep], minor[keep], values[keep] + + +#: Seeds for the two halves of a sparse entry's 128-bit digest. +_SEEDS = (0x243F6A8885A308D3, 0x13198A2E03707344) + + +def _splitmix(x: np.ndarray) -> np.ndarray: + """splitmix64's finaliser, which spreads every input bit over the output.""" + x = x + np.uint64(0x9E3779B97F4A7C15) + x = (x ^ (x >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9) + x = (x ^ (x >> np.uint64(27))) * np.uint64(0x94D049BB133111EB) + return x ^ (x >> np.uint64(31)) + + +def _sparse_snapshot(obj: Any, col_map: np.ndarray, block_rows: int) -> Any: + """A comparable value of a sparse pairwise entry on the target var. + + anndata compares sparse entries by value: a CSR and a CSC entry holding + the same values are equal, float32 equals float64, and a stored zero is + no different from an absent one. So the digest is of the nonzero + `(row, column, value)` triples, each value as float64, and it is the sum + of each triple's hash: the order the entries are stored in, and the + input's gene order, do not matter, and they are read a block at a time. + """ + csr = spec.encoding_type(obj) == spec.CSR_MATRIX + shape = _shape(obj) + nnz = int(obj["data"].shape[0]) + entries = _entries_per_block(nnz, shape[0], block_rows) + totals = np.zeros(len(_SEEDS), dtype=np.uint64) + count = 0 + for major, minor, values in _entries_on(obj, col_map, entries): + nonzero = values != 0 + rows, cols = (major, minor) if csr else (minor, major) + rows, cols = rows[nonzero].astype(np.uint64), cols[nonzero].astype(np.uint64) + bits = values[nonzero].astype(np.float64) + bits[np.isnan(bits)] = np.nan # one NaN, whatever its payload + bits = bits.view(np.uint64) + for k, seed in enumerate(_SEEDS): + mixed = _splitmix(_splitmix(_splitmix(rows ^ np.uint64(seed)) ^ cols) ^ bits) + totals[k:k + 1] += mixed.sum(dtype=np.uint64) + count += int(rows.size) + return ("sparse", count, tuple(int(t) for t in totals)) + + +def _pairwise_snapshot(obj: Any, take: np.ndarray, col_map: np.ndarray, block_rows: int) -> Any: + """A comparable value of a varp entry on the target var.""" + if is_dataset(obj): + return _array_snapshot(obj, take, block_rows, columns=take) + if spec.encoding_type(obj) in spec.SPARSE_TYPES: + return _sparse_snapshot(obj, col_map, block_rows) + return _Incomparable() + + +def _write_pairwise( + parent: Any, name: str, obj: Any, take: np.ndarray, col_map: np.ndarray, block_rows: int, + console: Console, +) -> None: + """A varp entry on the target var, rows and columns alike, a block at a + time. A sparse entry already on the target's genes is copied as it is; + one in another order is rebuilt through buckets on disk, as a + transpose is, its entries read a block at a time in stored order.""" + from adata.core.csr_build import CsrAppender, write_runs + + n = len(take) + if is_dataset(obj): + out = create_dataset(parent, name, shape=(n, n), dtype=obj.dtype, + **_layout_like(obj, parent, (n, n), obj.dtype)) + spec.set_encoding(out, spec.STRING_ARRAY if is_string_dtype(obj.dtype) else spec.ARRAY) + for start in range(0, n, block_rows): + out[start:start + block_rows, :] = aligned_rows(obj, take[start:start + block_rows])[:, take] + return + enc = spec.encoding_type(obj) + if enc not in spec.SPARSE_TYPES: + console.print(f"[yellow]Skipping varp/{name}: neither an array nor a sparse matrix[/]") + return + if len(take) == _shape(obj)[0] and np.array_equal(take, np.arange(len(take))): + copy_tree(obj, parent, name) + return + + indptr = np.asarray(obj["indptr"][...], dtype=np.int64) + kept = col_map >= 0 + # Entries a major keeps, counted before those of genes the target lacks + # are dropped: an estimate, which only balances the buckets. + counts = np.zeros(n, dtype=np.int64) + counts[col_map[kept]] = np.diff(indptr)[kept] + entries = _entries_per_block(int(indptr[-1]) if indptr.size else 0, n, block_rows) + appender = CsrAppender( + parent, name, shape=(n, n), data_dtype=obj["data"].dtype, + index_dtype=_widen_to_hold(obj["indices"].dtype, n), + templates=(obj["data"], obj["indices"], obj["indptr"]), + ) + write_runs( + _entries_on(obj, col_map, entries), n_rows=n, n_targets=1, + destinations=[[appender]], scratch_parent=parent, grouped=False, canonical=False, + counts=counts, value_dtype=obj["data"].dtype, chunk=entries, sum_duplicates=False, + ) + appender.finish() + # Built by major, which is what CSC's columns are too: n x n either way. + spec.set_encoding(appender.group, enc) + + def _write_varm( dst: Any, - roots: List[Any], - var_names: List[List[str]], + inputs: _Inputs, + positions: List[np.ndarray], target_var: List[str], merge: Optional[str], console: Console, @@ -1028,7 +1496,7 @@ def _write_varm( digests read once per input, `first` and `only` presence alone. Arrays are read and written `chunk_rows` genes at a time. """ - names = _index_union([list(r["varm"].keys()) if "varm" in r else [] for r in roots]) + names = _index_union([inputs.keys(i, "varm") for i in range(len(inputs))]) if not names: return if merge is None: @@ -1037,22 +1505,24 @@ def _write_varm( ) return varm = write_mapping(dst, "varm") - positions = [_column_map(target_var, names_i) for names_i in var_names] + complete = _per_var_list(positions, lambda where: not (where < 0).any()) compares = merge in ("same", "unique") for name in names: - aligned: List[Any] = [] - for root, where in zip(roots, positions): - if "varm" not in root or name not in root["varm"] or (where < 0).any(): - aligned.append(_MISSING) - elif compares: - aligned.append(_varm_snapshot(root["varm"][name], where, chunk_rows)) - else: - aligned.append(_PRESENT) - keep, _ = _merge_values(aligned, merge) - if not keep: + path = f"varm/{name}" + + def aligned() -> Iterator[Any]: + for obj, where, whole in zip(inputs.each(path), positions, complete): + if obj is None or not whole: + yield _MISSING + elif compares: + yield _varm_snapshot(obj, where, chunk_rows) + else: + yield _PRESENT + + source = _merge_choice(aligned(), merge) + if source is None: continue - source = next(i for i, v in enumerate(aligned) if v is not _MISSING) - obj, take = roots[source]["varm"][name], positions[source] + obj, take = inputs.get(source, path), positions[source] if is_dataset(obj): _write_varm_array(varm, name, obj, take, chunk_rows) elif spec.encoding_type(obj) == spec.DATAFRAME: @@ -1065,6 +1535,50 @@ def _write_varm( console.print(f"[yellow]Skipping varm/{name}: neither an array nor a dataframe[/]") +def _write_varp( + dst: Any, + inputs: _Inputs, + positions: List[np.ndarray], + col_maps: List[np.ndarray], + merge: Optional[str], + console: Console, + chunk_rows: int = 1024, +) -> None: + """varp entries the merge strategy keeps, on the target var, as varm's. + + varp relates genes, not cells, so it is merged as anndata merges it, + like varm. It used to be dropped with obsp, whose pairwise values do + have no meaning between cells from different inputs. + """ + names = _index_union([inputs.keys(i, "varp") for i in range(len(inputs))]) + if not names: + return + if merge is None: + console.print( + "[yellow]Dropping varp: pass --merge to keep the entries the inputs agree on.[/]" + ) + return + varp = write_mapping(dst, "varp") + complete = _per_var_list(positions, lambda where: not (where < 0).any()) + compares = merge in ("same", "unique") + for name in names: + path = f"varp/{name}" + + def aligned() -> Iterator[Any]: + for obj, where, col_map, whole in zip(inputs.each(path), positions, col_maps, complete): + if obj is None or not whole: + yield _MISSING + elif compares: + yield _pairwise_snapshot(obj, where, col_map, chunk_rows) + else: + yield _PRESENT + + source = _merge_choice(aligned(), merge) + if source is not None: + _write_pairwise(varp, name, inputs.get(source, path), positions[source], + col_maps[source], chunk_rows, console) + + # --------------------------------------------------------------------------- # orchestration @@ -1072,7 +1586,7 @@ def _write_varm( def _merge_group( dst_parent: Any, name: str, - groups: List[Optional[Any]], + inputs: _Inputs, strategy: Optional[str], console: Console, ) -> None: @@ -1083,12 +1597,12 @@ def _merge_group( target = write_mapping(dst_parent, name) if strategy is None: return - present = [g for g in groups if g is not None] + holders = [i for i in range(len(inputs)) if inputs.get(i, name) is not None] # An input without the mapping holds none of its keys, so `same` keeps # nothing, as anndata's intersection with an empty `uns` does. - if strategy == "same" and len(present) < len(groups): - present = [] - _write_merged(target, _merge_plan(present, strategy)) + if strategy == "same" and len(holders) < len(inputs): + holders = [] + _write_merged(target, _merge_plan(inputs, holders, name, strategy), inputs) def _is_mapping(obj: Any) -> bool: @@ -1100,41 +1614,43 @@ def _is_mapping(obj: Any) -> bool: ) -def _merge_plan(groups: List[Any], strategy: str) -> Dict[str, Any]: - """What to keep of `groups` merged: for each key, the member to copy, or - a nested plan for a mapping. `same` considers the keys every group has, - the other strategies the keys any group has; only `same` and `unique` - compare values, `first` and `only` decide by presence.""" - keys: Dict[str, None] = {} - for group in groups: - for key in group.keys(): - keys.setdefault(key, None) +def _merge_plan(inputs: _Inputs, holders: List[int], path: str, strategy: str) -> Dict[str, Any]: + """What to keep of the mapping at `path` in the inputs `holders`, merged: + for each key, `(input, path)` of the member to copy, or a nested plan for + a mapping. `same` considers the keys every holder has, the other + strategies the keys any holder has; only `same` and `unique` compare + values, `first` and `only` decide by presence.""" + owners: Dict[str, List[int]] = {} + for i in holders: + for key in inputs.keys(i, path): + owners.setdefault(key, []).append(i) if strategy == "same": - keys = {k: None for k in keys if all(k in g for g in groups)} + owners = {k: v for k, v in owners.items() if len(v) == len(holders)} + compares = strategy in ("same", "unique") plan: Dict[str, Any] = {} - for key in keys: - members = [g[key] for g in groups if key in g] - if all(_is_mapping(m) for m in members): - nested = _merge_plan(members, strategy) + for key, members in owners.items(): + member = f"{path}/{key}" + if all(_is_mapping(inputs.get(i, member)) for i in members): + nested = _merge_plan(inputs, members, member, strategy) if nested: plan[key] = nested continue - compares = strategy in ("same", "unique") - keep, _ = _merge_values( - [_readable_value(m) if compares else _PRESENT for m in members], strategy + chosen = _merge_choice( + (_readable_value(inputs.get(i, member)) if compares else _PRESENT for i in members), + strategy, ) - if keep: - plan[key] = members[0] + if chosen is not None: + plan[key] = (members[chosen], member) return plan -def _write_merged(target: Any, plan: Dict[str, Any]) -> None: +def _write_merged(target: Any, plan: Dict[str, Any], inputs: _Inputs) -> None: for key, value in plan.items(): if isinstance(value, dict): - _write_merged(write_mapping(target, key), value) + _write_merged(write_mapping(target, key), value, inputs) else: - copy_tree(value, target, key) + copy_tree(inputs.get(*value), target, key) #: Elements larger than this are not compared value-by-value. @@ -1214,14 +1730,55 @@ def _hashable(value: Any) -> Any: return value -def _inputs_zarr_format(stores: Sequence[Store]) -> Optional[int]: +def _scratch_size(inputs: _Inputs) -> int: + """The most scratch space concat needs at once: for the largest CSC + matrix, each stored entry's row and column as int64 and its value.""" + n = len(inputs) + paths = ["X"] + [f"{group}/{name}" for group in ("layers", "obsm") + for name in _index_union([inputs.keys(i, group) for i in range(n)])] + largest = 0 + for path in paths: + nnz, dtypes = 0, [] + for source in inputs.each(path): + if source is None or spec.encoding_type(source) != spec.CSC_MATRIX: + break + nnz += int(source["data"].shape[0]) + dtypes.append(source["data"].dtype) + else: + largest = max(largest, nnz * (16 + np.result_type(*dtypes).itemsize)) + return largest + + +def _check_space(inputs: _Inputs, output: Path, check: bool) -> None: + """Refuse to start if the disk cannot hold the output and the scratch. + + The output needs about the sum of the inputs (an inner join keeps less, + an outer join's dense fill can add some), and a CSC matrix's buckets go + in an HDF5 file beside it. A full disk part way through an HDF5 write + can abort inside the library with no cleanup, so the scratch is checked + whatever the output is; a Zarr output itself is not, nor its inputs + walked to size it, as `check_free_space` explains. + """ + if not check: + return + if detect_backend(output) == "hdf5": + check_free_space( + output, lambda: _scratch_size(inputs) + sum(disk_size(f) for f in inputs.files) + ) + return + scratch = _scratch_size(inputs) + if scratch: + check_free_space(output.with_name(f".{output.name}.scratch.h5"), scratch) + + +def _inputs_zarr_format(formats: Sequence[Optional[int]]) -> Optional[int]: """The Zarr version the inputs agree on, or None to take the default. A Zarr output follows its inputs, as every other command's follows its source, so v2 inputs are not silently upgraded. h5ad inputs have no version and do not count; Zarr inputs of both versions get the default. """ - versions = {s.zarr_format for s in stores} - {None} + versions = set(formats) - {None} return versions.pop() if len(versions) == 1 else None @@ -1239,17 +1796,18 @@ def concat_on_disk( fill_value: float = 0.0, chunk_rows: int = 1024, zarr_format: Optional[int] = None, + space_check: bool = True, ) -> None: """Concatenate stores along the obs axis, streaming each one in turn. `join` aligns var: "inner" keeps only variables present in every input, "outer" keeps their union. obs columns follow the same rule. Elements of - the var axis and of uns are reconciled by `merge`/`uns_merge`, since - inputs may legitimately disagree about them. + the var axis (var's columns, varm, varp) and of uns are reconciled by + `merge`/`uns_merge`, since inputs may legitimately disagree about them. - obsp and varp are not concatenated -- a pairwise matrix has no meaningful - value between cells that came from different inputs. anndata's own - implementation refuses these too. + obsp is not concatenated -- a pairwise value has no meaningful value + between cells that came from different inputs, and anndata's own + implementation refuses it too. """ if len(files) < 2: raise ValueError("Concatenation needs at least two input stores.") @@ -1259,20 +1817,13 @@ def concat_on_disk( raise FileExistsError(f"'{output}' already exists.") input_keys = _resolve_keys(files, keys) - stores = [open_store(f, "r") for f in files] - try: - roots = [s.root for s in stores] - - var_names = [read_names(r, "var") for r in roots] - obs_names = [read_names(r, "obs") for r in roots] - obs_counts = [axis_len(r, "obs") for r in roots] + with _Inputs(files) as inputs: + var_names, obs_counts, target_obs, formats = _read_axes(inputs, input_keys, index_unique) _refuse_repeated_var_names(files, var_names) target_var = ( - _index_intersection(var_names) - if join == "inner" - else _index_union(var_names) + _index_intersection(var_names) if join == "inner" else _var_union(var_names) ) # Inputs with no variables at all concatenate to n_obs x 0, as in # anndata; the refusal is for inputs that have some but share none. @@ -1282,7 +1833,8 @@ def concat_on_disk( "keep their union instead." ) - col_maps = [_column_map(names, target_var) for names in var_names] + col_maps = _per_var_list(var_names, lambda names: _column_map(names, target_var)) + positions = _per_var_list(var_names, lambda names: _column_map(target_var, names)) n_obs = sum(obs_counts) n_var = len(target_var) @@ -1293,10 +1845,6 @@ def concat_on_disk( for f, key, count in zip(files, input_keys, obs_counts): console.print(f" [dim]{key}[/]: {count} obs from {f.name}") - target_obs: List[str] = [] - for names, key in zip(obs_names, input_keys): - target_obs.extend(_apply_index_unique(names, key, index_unique)) - if len(set(target_obs)) != len(target_obs): console.print( "[yellow]Warning: obs names are not unique across inputs. " @@ -1305,10 +1853,11 @@ def concat_on_disk( # Validated before the output exists, so a mismatch cannot leave a # half-written store behind. - check_matrix_encodings(roots, console) + check_matrix_encodings(inputs, console, obs_counts) + _check_space(inputs, output, space_check) if zarr_format is None: - zarr_format = _inputs_zarr_format(stores) + zarr_format = _inputs_zarr_format(formats) # Written beside the output and renamed into place once complete, so # a failure part way leaves nothing that reads as a finished store. @@ -1317,14 +1866,14 @@ def concat_on_disk( ) as dst_store: dst = dst_store.root _write_obs( - dst, roots, target_obs, obs_counts, join, label, input_keys, console + dst, inputs, target_obs, obs_counts, join, label, input_keys, console ) - _write_var(dst, roots, var_names, target_var, merge, console) + _write_var(dst, inputs, positions, target_var, merge, console) ensure_anndata_skeleton(dst) _write_matrices( dst, - roots, + inputs, col_maps, obs_counts, n_obs, @@ -1333,35 +1882,31 @@ def concat_on_disk( fill_value, console, ) - _write_obsm(dst, roots, n_obs, chunk_rows, console, + _write_obsm(dst, inputs, n_obs, chunk_rows, console, lengths=obs_counts, join=join, target_obs=target_obs) - _write_varm(dst, roots, var_names, target_var, merge, console, chunk_rows) + _write_varm(dst, inputs, positions, target_var, merge, console, chunk_rows) + _write_varp(dst, inputs, positions, col_maps, merge, console, chunk_rows) - uns_groups = [r["uns"] if "uns" in r else None for r in roots] - if any(g is not None for g in uns_groups): + if any(inputs.get(i, "uns") is not None for i in range(len(inputs))): if uns_merge is None: write_mapping(dst, "uns") else: - _merge_group(dst, "uns", uns_groups, uns_merge, console) - - for name in ("obsp", "varp"): - if any(name in r and len(list(r[name].keys())) for r in roots): - console.print( - f"[yellow]Dropping {name}: pairwise values have no " - f"meaning across concatenated inputs.[/]" - ) - if any("raw" in r for r in roots): + _merge_group(dst, "uns", inputs, uns_merge, console) + + if any(inputs.keys(i, "obsp") for i in range(len(inputs))): + console.print( + "[yellow]Dropping obsp: pairwise values have no " + "meaning across concatenated inputs.[/]" + ) + if any(inputs.get(i, "raw") is not None for i in range(len(inputs))): console.print("[yellow]Dropping raw/: not concatenated.[/]") - console.print(f"[green]Wrote[/] {output} ({n_obs} obs x {n_var} var)") - finally: - for store in stores: - store.close() + console.print(f"[green]Wrote[/] {output} ({n_obs} obs x {n_var} var)") def _write_obs( dst: Any, - roots: List[Any], + inputs: _Inputs, target_obs: List[str], obs_counts: List[int], join: str, @@ -1370,9 +1915,8 @@ def _write_obs( console: Console, ) -> None: """Write the concatenated obs frame, plus the batch column if requested.""" - groups = [r["obs"] for r in roots] per_input_cols = [ - dataframe_columns(g, resolve_index(g, "obs")[1]) for g in groups + dataframe_columns(g, resolve_index(g, "obs")[1]) for g in inputs.each("obs") ] if join == "inner": @@ -1385,8 +1929,7 @@ def _write_obs( obs_group = write_dataframe_header(dst, "obs", target_obs, []) for name in names: - columns = [g[name] if name in g else None for g in groups] - if _concat_obs_column(obs_group, name, columns, groups, obs_counts, console): + if _concat_obs_column(obs_group, name, _Column(inputs, "obs", name), obs_counts, console): written.append(name) if label: @@ -1405,8 +1948,8 @@ def _write_obs( def _write_var( dst: Any, - roots: List[Any], - var_names: List[List[str]], + inputs: _Inputs, + positions: List[np.ndarray], target_var: List[str], merge: Optional[str], console: Console, @@ -1416,13 +1959,11 @@ def _write_var( if merge is None: return - groups = [r["var"] for r in roots] per_input_cols = [ - dataframe_columns(g, resolve_index(g, "var")[1]) for g in groups + dataframe_columns(g, resolve_index(g, "var")[1]) for g in inputs.each("var") ] candidates = _index_union(per_input_cols) - - positions = [_column_map(target_var, names) for names in var_names] + complete = _per_var_list(positions, lambda where: not (where < 0).any()) written: List[str] = [] # "first" and "only" decide on presence alone, so the column values are @@ -1430,31 +1971,27 @@ def _write_var( compares = merge in ("same", "unique") for name in candidates: - aligned: List[Any] = [] - for group, where in zip(groups, positions): - if name not in group or (where < 0).any(): - aligned.append(_MISSING) - continue - if not compares: - aligned.append(_PRESENT) - continue - # Read the column once and index the result. Reading it inside the - # generator -- as an earlier version did -- re-read the whole - # column for every target variable, which is quadratic and turns a - # 36k-var merge into hours of pure CPU. - values = read_str_all(group[name], parent_group=group) - aligned.append(tuple(values[i] for i in where)) - - keep, _ = _merge_values(aligned, merge) - if not keep: - continue - source_i = next( - i for i, v in enumerate(aligned) if v is not _MISSING - ) - column = groups[source_i][name] - take = positions[source_i] - _write_var_column(var_group, name, column, take, groups[source_i]) + def aligned() -> Iterator[Any]: + for (column, group), where, whole in zip(_Column(inputs, "var", name), positions, complete): + if column is None or not whole: + yield _MISSING + elif not compares: + yield _PRESENT + else: + # Read the column once and index the result. Reading it + # inside the generator -- as an earlier version did -- + # re-read the whole column for every target variable, + # which is quadratic and turns a 36k-var merge into hours + # of pure CPU. + values = read_str_all(column, parent_group=group) + yield tuple(values[i] for i in where) + + source = _merge_choice(aligned(), merge) + if source is None: + continue + group = inputs.get(source, "var") + _write_var_column(var_group, name, group[name], positions[source], group) written.append(name) set_column_order(var_group, written) @@ -1488,8 +2025,6 @@ def _write_var_column( mask = np.asarray(column["mask"][...], dtype=bool)[take] raw_values = np.asarray(column["values"][...]) if enc == spec.NULLABLE_STRING_ARRAY: - from adata.elements.read import decode_str_array - values = decode_str_array(raw_values)[take].tolist() else: values = raw_values[take] @@ -1503,7 +2038,7 @@ def _write_var_column( def _write_matrices( dst: Any, - roots: List[Any], + inputs: _Inputs, col_maps: List[np.ndarray], row_counts: List[int], n_obs: int, @@ -1513,73 +2048,51 @@ def _write_matrices( console: Console, ) -> None: """Concatenate X and every layer shared by all inputs.""" - if all("X" in r for r in roots): - _concat_matrix( - dst, - "X", - [r["X"] for r in roots], - col_maps, - row_counts, - n_obs, - n_var, - chunk_rows, - fill_value, - console, - ) + if all(x is not None for x in inputs.each("X")): + _concat_matrix(dst, "X", inputs, "X", col_maps, row_counts, n_obs, n_var, + chunk_rows, fill_value) - layer_names = _index_union( - [list(r["layers"].keys()) if "layers" in r else [] for r in roots] - ) + layer_names = _index_union([inputs.keys(i, "layers") for i in range(len(inputs))]) if not layer_names: return layers = write_mapping(dst, "layers") for name in layer_names: - if not all("layers" in r and name in r["layers"] for r in roots): + path = f"layers/{name}" + if not all(layer is not None for layer in inputs.each(path)): console.print( f"[yellow]Skipping layer {name!r}: not present in every input[/]" ) continue - _concat_matrix( - layers, - name, - [r["layers"][name] for r in roots], - col_maps, - row_counts, - n_obs, - n_var, - chunk_rows, - fill_value, - console, - ) + _concat_matrix(layers, name, inputs, path, col_maps, row_counts, n_obs, n_var, + chunk_rows, fill_value) def _write_obsm( dst: Any, - roots: List[Any], + inputs: _Inputs, n_obs: int, chunk_rows: int, console: Console, *, - lengths: Optional[List[int]] = None, + lengths: List[int], join: str = "inner", target_obs: Optional[List[str]] = None, ) -> None: """Concatenate obsm entries shared by all inputs.""" - names = _index_union( - [list(r["obsm"].keys()) if "obsm" in r else [] for r in roots] - ) + names = _index_union([inputs.keys(i, "obsm") for i in range(len(inputs))]) if not names: return obsm = write_mapping(dst, "obsm") for name in names: - if not all("obsm" in r and name in r["obsm"] for r in roots): + path = f"obsm/{name}" + if not all(entry is not None for entry in inputs.each(path)): console.print( f"[yellow]Skipping obsm/{name}: not present in every input[/]" ) continue _concat_obsm( - obsm, name, [r["obsm"][name] for r in roots], n_obs, chunk_rows, console, + obsm, name, inputs, path, n_obs, chunk_rows, console, lengths=lengths, join=join, target_obs=target_obs, ) diff --git a/src/adata/core/recast.py b/src/adata/core/recast.py index 7c227fa..56f1e06 100644 --- a/src/adata/core/recast.py +++ b/src/adata/core/recast.py @@ -24,7 +24,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, Optional, Tuple import numpy as np from rich.console import Console @@ -615,7 +615,9 @@ def sparsify( out_indices = _growable_like( group, "indices", index_dtype, src.obj, compression=compression ) - counts: List[int] = [] + # Per-major counts in one int64 array: a Python list of int objects + # took 36 to 40 bytes a row against 8. + counts = np.zeros(n_major, dtype=np.int64) for lo in range(0, n_major, chunk_rows): hi = min(lo + chunk_rows, n_major) @@ -623,13 +625,13 @@ def sparsify( if not csr: block = block.T # iterate majors as rows either way nonzero_major, nonzero_minor = np.nonzero(block) - counts.extend(np.bincount(nonzero_major, minlength=hi - lo).tolist()) + counts[lo:hi] = np.bincount(nonzero_major, minlength=hi - lo) _append(out_indices, nonzero_minor.astype(index_dtype, copy=False)) _append(out_data, block[nonzero_major, nonzero_minor].astype( data_dtype, copy=False )) - indptr = np.concatenate(([0], np.cumsum(counts, dtype=np.int64))) + indptr = np.concatenate(([0], np.cumsum(counts))) kw = {} if compression is not None: kw = with_compression( diff --git a/tests/test_concat_axis_elements.py b/tests/test_concat_axis_elements.py index a17654c..d11ff71 100644 --- a/tests/test_concat_axis_elements.py +++ b/tests/test_concat_axis_elements.py @@ -17,6 +17,7 @@ ad = pytest.importorskip("anndata") pd = pytest.importorskip("pandas") +sparse = pytest.importorskip("scipy.sparse") QUIET = Console(quiet=True, record=True) @@ -192,3 +193,149 @@ def test_varm_differing_by_a_fraction_is_not_the_same(tmp_path): a = _varm_store(tmp_path / "a.h5ad", genes, {"x": values}) b = _varm_store(tmp_path / "b.h5ad", genes, {"x": values + 0.5}) assert "x" not in _concat(tmp_path, [a, b], merge="same").varm + + +# --------------------------------------------------------------------------- +# F118: varp + + +def _varp_store(path: Path, genes, entries, *, n_cells: int = 2): + obj = ad.AnnData( + X=np.zeros((n_cells, len(genes)), dtype=np.float32), + obs=pd.DataFrame(index=[f"{path.stem}{i}" for i in range(n_cells)]), + var=pd.DataFrame(index=list(genes)), + ) + for key, value in entries.items(): + obj.varp[key] = value + if path.suffix == ".zarr": + obj.write_zarr(path) + else: + obj.write_h5ad(path) + return path + + +def _pairwise(genes, seed: int = 0, density: float = 0.4) -> "pd.DataFrame": + """A gene x gene matrix, labelled so it can be put in any gene order.""" + rng = np.random.default_rng(seed) + values = rng.integers(1, 9, (len(genes), len(genes))) * (rng.random((len(genes),) * 2) < density) + return pd.DataFrame(values.astype(np.float64), index=list(genes), columns=list(genes)) + + +def _as(layout: str, frame: "pd.DataFrame", genes=None): + values = (frame.loc[list(genes), list(genes)] if genes is not None else frame).to_numpy() + return {"dense": np.asarray, "csr": sparse.csr_matrix, "csc": sparse.csc_matrix}[layout](values) + + +def _dense(value): + return value.toarray() if sparse.issparse(value) else np.asarray(value) + + + +@pytest.mark.parametrize("layout", ["dense", "csr", "csc"]) +@pytest.mark.parametrize("merge", ["same", "unique", "first", "only"]) +def test_varp_follows_merge_as_anndata_does(tmp_path, merge, layout): + """F118: varp was dropped whatever --merge said, with a reason ('pairwise + values have no meaning between cells from different inputs') that holds + for obsp only: varp relates genes, and anndata merges it as it merges + varm. b's genes are in another order, so its entries are compared and + written on a's.""" + genes = [f"g{i}" for i in range(7)] + shuffled = list(np.random.default_rng(1).permutation(genes)) + shared, own_a, own_b = _pairwise(genes, 0), _pairwise(genes, 1), _pairwise(genes, 2) + a = _varp_store(tmp_path / "a.h5ad", genes, {"corr": _as(layout, shared), "own": _as(layout, own_a)}) + b = _varp_store(tmp_path / "b.h5ad", shuffled, { + "corr": _as(layout, shared, shuffled), "own": _as(layout, own_b, shuffled), + "only_b": _as(layout, own_b, shuffled), + }) + got = _concat(tmp_path, [a, b], merge=merge, chunk_rows=2) + ref = ad.concat([ad.read_h5ad(a), ad.read_h5ad(b)], merge=merge) + assert sorted(got.varp) == sorted(ref.varp) + assert got.varp, "every strategy keeps something here" + for key in ref.varp: + assert np.array_equal(_dense(got.varp[key]), _dense(ref.varp[key])), key + + +@pytest.mark.parametrize("ext", [".h5ad", ".zarr"]) +def test_varp_in_another_gene_order_is_written_on_the_output_genes(tmp_path, ext): + """F118: the entry kept comes from an input whose genes are in another + order, so a sparse one is rebuilt on the output's genes through buckets + on disk, and a dense one a block of rows at a time.""" + genes = [f"g{i}" for i in range(30)] + shuffled = list(np.random.default_rng(3).permutation(genes)) + frame = _pairwise(genes, 4) + entries = {layout: _as(layout, frame, shuffled) for layout in ("dense", "csr", "csc")} + bare = _varp_store(tmp_path / f"bare{ext}", genes, {}) + b = _varp_store(tmp_path / f"b{ext}", shuffled, entries) + out = tmp_path / f"out{ext}" + concat_on_disk([bare, b], out, QUIET, merge="first", chunk_rows=3) + got = ad.read_zarr(out) if ext == ".zarr" else ad.read_h5ad(out) + ref = ad.concat([ad.read_h5ad(bare) if ext == ".h5ad" else ad.read_zarr(bare), + ad.read_h5ad(b) if ext == ".h5ad" else ad.read_zarr(b)], merge="first") + assert list(got.var_names) == genes + for layout in ("dense", "csr", "csc"): + mine = got.varp[layout] + assert np.array_equal(_dense(mine), _dense(ref.varp[layout])), layout + assert np.array_equal(_dense(mine), frame.to_numpy()), layout + if layout != "dense": + assert mine.format == layout and mine.has_sorted_indices + + +def test_sparse_varp_is_compared_by_value_as_anndata_compares_it(tmp_path): + """F118: anndata's `same` finds a CSR and a CSC entry with the same + values equal, and float32 equal to float64, and ignores stored zeros; + a dense entry must match its dtype as well, as varm's does.""" + genes = [f"g{i}" for i in range(5)] + frame = _pairwise(genes, 5) + with_zero = sparse.csr_matrix(frame.to_numpy()) + with_zero.data[0] = 0.0 + without = frame.to_numpy().copy() + without[np.unravel_index(np.flatnonzero(without)[0], without.shape)] = 0.0 + a = _varp_store(tmp_path / "a.h5ad", genes, { + "formats": sparse.csr_matrix(frame.to_numpy()), "dtypes": sparse.csr_matrix(frame.to_numpy()), + "zeros": with_zero, "dense_dtypes": frame.to_numpy(), + }) + b = _varp_store(tmp_path / "b.h5ad", genes, { + "formats": sparse.csc_matrix(frame.to_numpy()), + "dtypes": sparse.csr_matrix(frame.to_numpy().astype(np.float32)), + "zeros": sparse.csr_matrix(without), "dense_dtypes": frame.to_numpy().astype(np.float32), + }) + got = _concat(tmp_path, [a, b], merge="same") + ref = ad.concat([ad.read_h5ad(a), ad.read_h5ad(b)], merge="same") + assert sorted(got.varp) == sorted(ref.varp) == ["dtypes", "formats", "zeros"] + + +def test_an_outer_join_drops_varp_an_input_cannot_cover(tmp_path): + """F118, as varm: an input lacking some of the output's genes has no + entry for them, so it counts as not having the entry.""" + a = _varp_store(tmp_path / "a.h5ad", [f"g{i}" for i in range(6)], + {"corr": _pairwise([f"g{i}" for i in range(6)]).to_numpy()}) + b = _varp_store(tmp_path / "b.h5ad", [f"g{i}" for i in range(3, 8)], + {"corr": _pairwise([f"g{i}" for i in range(3, 8)]).to_numpy()}) + assert len(_concat(tmp_path, [a, b], join="outer", merge="same").varp) == 0 + assert len(_concat(tmp_path, [a, b], join="outer", merge="first").varp) == 0 + + +def test_varp_is_dropped_without_merge_with_a_note_and_obsp_always(tmp_path): + """F118: without --merge varp is dropped as varm is, with a note that + says how to keep it; obsp is dropped whatever --merge says.""" + genes = [f"g{i}" for i in range(4)] + paths = [] + for tag in "ab": + obj = ad.AnnData(X=np.ones((3, 4), dtype=np.float32), + obs=pd.DataFrame(index=[f"{tag}{i}" for i in range(3)]), + var=pd.DataFrame(index=genes)) + obj.varp["corr"] = np.eye(4) + obj.obsp["distances"] = sparse.csr_matrix(np.eye(3)) + obj.write_h5ad(tmp_path / f"{tag}.h5ad") + paths.append(tmp_path / f"{tag}.h5ad") + console = Console(record=True, width=200) + concat_on_disk(paths, tmp_path / "plain.h5ad", console) + text = console.export_text() + assert "Dropping varp: pass --merge" in text and "Dropping obsp" in text + plain = ad.read_h5ad(tmp_path / "plain.h5ad") + assert len(plain.varp) == 0 and len(plain.obsp) == 0 + console = Console(record=True, width=200) + concat_on_disk(paths, tmp_path / "merged.h5ad", console, merge="same") + merged = ad.read_h5ad(tmp_path / "merged.h5ad") + assert list(merged.varp) == ["corr"] and len(merged.obsp) == 0 + assert "Dropping obsp" in console.export_text() diff --git a/tests/test_concat_correctness.py b/tests/test_concat_correctness.py index 09dfafe..2db137c 100644 --- a/tests/test_concat_correctness.py +++ b/tests/test_concat_correctness.py @@ -56,10 +56,11 @@ def _store( obs: Optional[Dict[str, Writer]] = None, var: Optional[Dict[str, Writer]] = None, obsm: Optional[Dict[str, Writer]] = None, + layers: Optional[Dict[str, Writer]] = None, zarr_format: Optional[int] = None, ) -> Path: - """A store with these names, whose columns and obsm entries are each - written by the function given for them.""" + """A store with these names, whose columns, layers and obsm entries are + each written by the function given for them; X too, if given one.""" with open_store(path, "w", zarr_format=zarr_format) as handle: root = handle.root for axis, names, columns in (("obs", cells, obs), ("var", genes, var)): @@ -69,9 +70,10 @@ def _store( ew.ensure_anndata_skeleton(root) if X is None: X = np.arange(1, len(cells) * len(genes) + 1, dtype=np.float32).reshape(len(cells), len(genes)) - _matrix(root, "X", X) - for name, write in (obsm or {}).items(): - write(root["obsm"], name) + X(root, "X") if callable(X) else _matrix(root, "X", X) + for group, entries in (("obsm", obsm), ("layers", layers)): + for name, write in (entries or {}).items(): + write(root[group], name) return path @@ -109,6 +111,27 @@ def matrix(value: Any) -> Writer: return lambda g, n: _matrix(g, n, value) +def chunked(value: Any, chunks: tuple) -> Writer: + """A dense array stored in `chunks` and compressed with gzip, which no + backend's defaults give: a copy that keeps them kept them on purpose.""" + def write(parent: Any, name: str) -> None: + value_ = np.asarray(value) + if isinstance(parent, h5py.Group): + out = parent.create_dataset(name, data=value_, chunks=chunks, + compression="gzip", compression_opts=6) + elif parent.metadata.zarr_format == 3: + from zarr.codecs import GzipCodec + + out = parent.create_array(name, data=value_, chunks=chunks, + compressors=[GzipCodec(level=6)]) + else: + from numcodecs import GZip + + out = parent.create_array(name, data=value_, chunks=chunks, compressors=GZip(level=6)) + spec.set_encoding(out, spec.ARRAY) + return write + + # --------------------------------------------------------------------------- # reading results @@ -720,3 +743,226 @@ def test_inputs_sharing_no_variables_are_still_refused_under_an_inner_join(tmp_p b = _store(tmp_path / "b.h5ad", ["b0"], ["g2"]) with pytest.raises(ValueError, match="--join outer"): concat_on_disk([a, b], tmp_path / "out.h5ad", QUIET) + + +# --------------------------------------------------------------------------- +# F075: dense layout + + +def _layout(path: Path, element: str) -> tuple: + """A dense array's chunks and compression, as its backend spells them.""" + with open_store(path, "r") as handle: + array = handle.root[element] + if isinstance(array, h5py.Dataset): + return array.chunks, array.compression + return tuple(array.chunks), [type(c).__name__ for c in array.compressors] + + +@pytest.mark.parametrize("join", ["inner", "outer"]) +def test_dense_matrices_and_obsm_arrays_keep_the_sources_chunking_and_compression(tmp_path, backend, join): + """F075: dense X, layers and obsm arrays were written contiguous and + uncompressed whatever the inputs were, so 109 KB of gzip-compressed + input came out as 3.7 MB. The sparse path and varm kept theirs.""" + rng = np.random.default_rng(0) + + def make(stem, cells, genes): + x = rng.poisson(1, (len(cells), len(genes))).astype(np.float32) + return _on(backend, tmp_path, stem, cells, genes, X=chunked(x, (3, 2)), + layers={"counts": chunked(x * 2, (3, 2))}, + obsm={"X_pca": chunked(rng.normal(size=(len(cells), 4)), (3, 4))}) + + a = make("a", [f"a{i}" for i in range(7)], ["g0", "g1", "g2", "g3"]) + b = make("b", [f"b{i}" for i in range(5)], ["g0", "g1", "g2", "g4"]) + out = backend_path(tmp_path, backend, "out") + got = _concat([a, b], out, join=join, chunk_rows=4) + for element in ("X", "layers/counts", "obsm/X_pca"): + assert _layout(out, element) == _layout(a, element), element + ref = _reference([a, b], join=join, fill_value=0) + assert np.array_equal(got.X, ref.X) + assert np.array_equal(got.layers["counts"], ref.layers["counts"]) + assert np.array_equal(got.obsm["X_pca"], ref.obsm["X_pca"]) + + +# --------------------------------------------------------------------------- +# F113: an outer join's var order + + +@pytest.mark.parametrize("layout", ["csr", "csc", "dense"]) +@pytest.mark.parametrize( + "genes_b", + [["g2", "g0", "g1"], ["g0", "g1", "g2"], ["g3", "g0"], ["g2", "g0"]], + ids=["same-order", "reordered", "new-genes", "a-subset"], +) +def test_an_outer_join_orders_var_as_anndata_does(tmp_path, backend, layout, genes_b): + """F113: an outer join's var came out in order of first appearance, where + anndata.concat sorts the union unless every input has the same var in + the same order. Values matched once aligned, but code that aligns by + position against an anndata-made reference saw another matrix.""" + genes_a = ["g2", "g0", "g1"] + as_layout = {"csr": sparse.csr_matrix, "csc": sparse.csc_matrix, "dense": np.asarray}[layout] + + def X(n_cells, n_genes, start): + return as_layout(np.arange(start, start + n_cells * n_genes, dtype=np.float32).reshape(n_cells, n_genes)) + + a = _on(backend, tmp_path, "a", ["a0", "a1"], genes_a, X=X(2, 3, 1)) + b = _on(backend, tmp_path, "b", ["b0", "b1", "b2"], genes_b, X=X(3, len(genes_b), 10)) + got = _concat([a, b], backend_path(tmp_path, backend, "out"), join="outer") + ref = _reference([a, b], join="outer", fill_value=0) + assert list(got.var_names) == list(ref.var_names) + dense = lambda m: m.toarray() if sparse.issparse(m) else np.asarray(m) + assert np.array_equal(dense(got.X), dense(ref.X)) + + +def test_an_outer_join_keeps_a_shared_var_order_beside_inputs_without_variables(tmp_path): + """F113: an input with no variables has no say in the order, as in + anndata, where a union with an empty index is the other index as is.""" + genes = ["g2", "g0", "g1"] + a = _store(tmp_path / "a.h5ad", ["a0"], genes) + empty = _store(tmp_path / "e.h5ad", ["e0"], [], X=np.zeros((1, 0), dtype=np.float32)) + b = _store(tmp_path / "b.h5ad", ["b0"], genes) + got = _concat([a, empty, b], tmp_path / "out.h5ad", join="outer") + ref = _reference([a, empty, b], join="outer", fill_value=0) + assert list(got.var_names) == list(ref.var_names) == genes + assert np.array_equal(got.X, ref.X) + + +# --------------------------------------------------------------------------- +# F114: the order of merged categories + + +@pytest.mark.parametrize( + "first, second", + [ + (["t2", "t10", "b"], ["t1", "b"]), + (["b", "a"], ["a", "b"]), + (["c", "b", "a"], ["a"]), + (["a"], ["c", "b", "a"]), + (["b", "a"], ["b", "a"]), + ([10, 2], [3, 10]), + ([2.5, -1.0], [0.5]), + ], + ids=["new-categories", "same-set-reordered", "first-covers-the-union", + "second-covers-the-union", "identical", "integers", "floats"], +) +def test_merged_categories_are_ordered_as_anndata_orders_them(tmp_path, backend, first, second): + """F114: merged categories came out in order of first appearance, where + anndata natsorts the union -- unless the first input already holds every + category, whose order pandas then keeps. `.cat.codes` and plotting order + differed from anndata.concat's.""" + a = _on(backend, tmp_path, "a", [f"a{i}" for i in range(len(first))], ["g"], + obs={"c": categorical(first, list(range(len(first))))}) + b = _on(backend, tmp_path, "b", [f"b{i}" for i in range(len(second))], ["g"], + obs={"c": categorical(second, list(range(len(second))))}) + got = _concat([a, b], backend_path(tmp_path, backend, "out")) + ref = _reference([a, b]) + assert list(got.obs["c"].cat.categories) == list(ref.obs["c"].cat.categories) + assert list(got.obs["c"]) == list(ref.obs["c"]) + assert not got.obs["c"].cat.ordered + + +@pytest.mark.parametrize("join", ["inner", "outer"]) +def test_merged_categories_follow_the_first_input_that_has_the_column(tmp_path, join): + """F114 under an outer join: the first input without the column has no + say, as in anndata.""" + a = _store(tmp_path / "a.h5ad", ["a0"], ["g"], obs={"other": strings("x")}) + b = _store(tmp_path / "b.h5ad", ["b0", "b1", "b2"], ["g"], + obs={"other": strings("y", "y", "y"), "c": categorical(["c", "b", "a"], [0, 1, 2])}) + c = _store(tmp_path / "c.h5ad", ["c0"], ["g"], + obs={"other": strings("z"), "c": categorical(["a"], [0])}) + got = _concat([a, b, c], tmp_path / "out.h5ad", join=join) + ref = _reference([a, b, c], join=join) + if join == "outer": + assert list(got.obs["c"].cat.categories) == list(ref.obs["c"].cat.categories) + assert _cells(got.obs["c"]) == _cells(ref.obs["c"]) + else: + assert "c" not in got.obs and "c" not in ref.obs + + +@pytest.mark.parametrize( + "second, keeps_order", + [ + ((["lo", "mid", "hi"], True), True), + ((["hi", "mid", "lo"], True), False), + ((["lo", "mid", "hi", "max"], True), False), + ((["lo", "mid", "hi"], False), False), + (([], True), True), + ], + ids=["same-order", "reversed", "another-set", "second-unordered", "second-all-missing"], +) +def test_ordered_categoricals_stay_ordered_only_when_every_input_agrees(tmp_path, backend, second, keeps_order): + """F114: 'ordered' categoricals whose orders conflict kept the first + input's order, silently reversing what the second input's meant. + anndata gives up on the order (plain values); concat now writes them + unordered, with natsorted categories and a warning.""" + first = ["lo", "mid", "hi"] + labels, ordered = second + a = _on(backend, tmp_path, "a", ["a0", "a1", "a2"], ["g"], + obs={"c": categorical(first, [0, 1, 2], ordered=True)}) + codes = [len(labels) - 1, 0] if labels else [-1, -1] + b = _on(backend, tmp_path, "b", ["b0", "b1"], ["g"], + obs={"c": categorical(labels, codes, ordered=ordered)}) + console = Console(record=True, width=200) + out = backend_path(tmp_path, backend, "out") + concat_on_disk([a, b], out, console) + got, ref = _read(out), _reference([a, b]) + assert _cells(got.obs["c"]) == _cells(ref.obs["c"]) + warned = "written unordered" in " ".join(console.export_text().split()) + if keeps_order: + assert ref.obs["c"].cat.ordered and got.obs["c"].cat.ordered + assert list(got.obs["c"].cat.categories) == list(ref.obs["c"].cat.categories) == first + assert not warned + else: + assert ref.obs["c"].dtype == object + assert not got.obs["c"].cat.ordered + union = sorted(set(first) | set(labels)) + assert list(got.obs["c"].cat.categories) == union + assert warned + + +# --------------------------------------------------------------------------- +# F116: --fill-value and integer matrices + + +@pytest.mark.parametrize("layout", ["dense", "csr"]) +def test_a_fill_value_leaves_integer_matrices_alone_when_nothing_is_filled(tmp_path, backend, layout): + """F116: --fill-value made integer X and layers float64 on an inner join, + where no cell is filled, doubling X; anndata.concat keeps them int32.""" + as_layout = {"csr": sparse.csr_matrix, "dense": np.asarray}[layout] + x = np.arange(6, dtype=np.int32).reshape(2, 3) + a = _on(backend, tmp_path, "a", ["a0", "a1"], ["g0", "g1", "g2"], X=as_layout(x), + layers={"counts": matrix(as_layout(x * 2))}) + b = _on(backend, tmp_path, "b", ["b0", "b1"], ["g2", "g1", "g3"], X=as_layout(x + 10), + layers={"counts": matrix(as_layout(x * 3))}) + got = _concat([a, b], backend_path(tmp_path, backend, "out"), fill_value=-1) + ref = _reference([a, b], fill_value=-1) + dense = lambda m: m.toarray() if sparse.issparse(m) else np.asarray(m) + for mine, theirs in ((got.X, ref.X), (got.layers["counts"], ref.layers["counts"])): + assert mine.dtype == theirs.dtype == np.int32 + assert np.array_equal(dense(mine), dense(theirs)) + + +@pytest.mark.parametrize( + "dtype, fill, expected", + [ + (np.int32, -1, np.int32), + (np.int32, 7, np.int32), + (np.int32, 0.5, np.float64), + (np.int32, np.nan, np.float64), + (np.uint8, -1, np.float64), + (np.uint8, 300, np.float64), + (np.bool_, 1, np.bool_), + (np.bool_, 2, np.float64), + (np.float32, 0.5, np.float32), + ], +) +def test_an_outer_join_widens_a_matrix_only_for_a_fill_value_it_cannot_hold(tmp_path, dtype, fill, expected): + """F116: when an outer join does fill cells, an integer matrix that can + hold the fill value keeps its dtype; one that cannot becomes float64. + The values are anndata.concat's either way.""" + x = np.array([[1, 0], [1, 1]], dtype=dtype) + a = _store(tmp_path / "a.h5ad", ["a0", "a1"], ["g0", "g1"], X=x) + b = _store(tmp_path / "b.h5ad", ["b0", "b1"], ["g0", "g2"], X=x) + got = _concat([a, b], tmp_path / "out.h5ad", join="outer", fill_value=fill) + ref = _reference([a, b], join="outer", fill_value=fill) + assert got.X.dtype == expected + assert np.array_equal(got.X.astype(np.float64), ref.X.astype(np.float64), equal_nan=True) diff --git a/tests/test_concat_scale.py b/tests/test_concat_scale.py new file mode 100644 index 0000000..4f47c05 --- /dev/null +++ b/tests/test_concat_scale.py @@ -0,0 +1,361 @@ +"""`concat` at scale: CSC inputs, and more inputs than can be open at once. + +What these check is the result, against `anndata.concat`, on the paths that +scale: CSC matrices rebuilt through buckets on disk (F011), and inputs opened +as they are needed rather than all at once (F076). How much they cost is +`test_performance.py`'s business. +""" + +from __future__ import annotations + +import subprocess +import sys +import textwrap +from pathlib import Path + +import h5py +import numpy as np +import pytest +from rich.console import Console + +from adata.core import concat as concat_mod +from adata.core.concat import concat_on_disk +from adata.storage import open_store + +from tests.conftest import backend_path, backend_zarr_format + +ad = pytest.importorskip("anndata") +pd = pytest.importorskip("pandas") +sparse = pytest.importorskip("scipy.sparse") + +QUIET = Console(quiet=True) + + +def _read(path: Path): + return ad.read_zarr(path) if path.suffix == ".zarr" else ad.read_h5ad(path) + + +def _write(obj, path: Path, backend: str) -> Path: + if path.suffix == ".zarr": + with ad.settings.override(zarr_write_format=backend_zarr_format(backend)): + obj.write_zarr(path) + else: + obj.write_h5ad(path) + return path + + +def _csc_input(path: Path, backend: str, tag: str, n_cells: int, genes, *, seed: int) -> Path: + rng = np.random.default_rng(seed) + values = rng.poisson(1.2, (n_cells, len(genes))).astype(np.float32) + values[:, 0] = 0 # an empty column + obj = ad.AnnData( + X=sparse.csc_matrix(values), + obs=pd.DataFrame(index=[f"{tag}{i}" for i in range(n_cells)]), + var=pd.DataFrame(index=list(genes)), + ) + obj.layers["counts"] = sparse.csc_matrix(values * 2) + return _write(obj, path, backend) + + +def _stored_csc(path: Path, element: str): + with open_store(path, "r") as handle: + group = handle.root[element] + assert group.attrs["encoding-type"] == "csc_matrix" + return sparse.csc_matrix( + (group["data"][...], group["indices"][...], group["indptr"][...]), + shape=tuple(int(n) for n in group.attrs["shape"]), + ) + + +# --------------------------------------------------------------------------- +# F011: CSC inputs + + +@pytest.mark.parametrize("join", ["inner", "outer"]) +def test_csc_inputs_concatenate_as_anndata_concatenates_them(tmp_path, backend, join): + """F011: CSC concat read each column of each input separately, 2 x n_var x + n_inputs reads that each decompressed a whole chunk: 37 times the input + on h5ad, 572 times on Zarr. Each input is now read once, in blocks, and + the output's columns put together through buckets on disk. Tiny blocks + and enough entries for several buckets exercise every seam; one input + orders its genes differently and lacks some, so columns are remapped.""" + genes = [f"g{i:02d}" for i in range(60)] + shuffled = list(np.random.default_rng(9).permutation(genes[:50])) + ["x1", "x2"] + a = _csc_input(backend_path(tmp_path, backend, "a"), backend, "a", 150, genes, seed=0) + b = _csc_input(backend_path(tmp_path, backend, "b"), backend, "b", 120, shuffled, seed=1) + c = _csc_input(backend_path(tmp_path, backend, "c"), backend, "c", 90, genes, seed=2) + out = backend_path(tmp_path, backend, "out") + concat_on_disk([a, b, c], out, QUIET, join=join, chunk_rows=2) + + got = _read(out) + ref = ad.concat([_read(a), _read(b), _read(c)], join=join) + assert list(got.var_names) == list(ref.var_names) + for element, mine, theirs in (("X", got.X, ref.X), ("layers/counts", got.layers["counts"], ref.layers["counts"])): + stored = _stored_csc(out, element) + assert stored.has_sorted_indices, element + assert stored.nnz >= 8192, "enough entries for more than one bucket" + assert np.array_equal(mine.toarray(), theirs.toarray()), element + + +def test_csc_concat_keeps_the_index_dtypes_and_layout(tmp_path): + """F011: the rewrite keeps what the per-column version kept: int32 + indices and pointers from int32 inputs, and the source's compression.""" + genes = [f"g{i}" for i in range(20)] + paths = [] + for tag, seed in (("a", 0), ("b", 1)): + obj = ad.AnnData( + X=sparse.csc_matrix(np.random.default_rng(seed).poisson(1, (30, 20)).astype(np.float32)), + obs=pd.DataFrame(index=[f"{tag}{i}" for i in range(30)]), + var=pd.DataFrame(index=genes), + ) + obj.write_h5ad(tmp_path / f"{tag}.h5ad", compression="gzip") + paths.append(tmp_path / f"{tag}.h5ad") + out = tmp_path / "out.h5ad" + concat_on_disk(paths, out, QUIET, chunk_rows=4) + with h5py.File(out) as f, h5py.File(paths[0]) as src: + for name in ("data", "indices", "indptr"): + assert f[f"X/{name}"].dtype == src[f"X/{name}"].dtype, name + assert f["X/data"].compression == src["X/data"].compression == "gzip" + ref = ad.concat([ad.read_h5ad(p) for p in paths]) + assert np.array_equal(ad.read_h5ad(out).X.toarray(), ref.X.toarray()) + + +def test_csc_inputs_with_nothing_stored_concatenate(tmp_path, backend): + """F011: no entries at all needs no buckets.""" + genes = ["g0", "g1"] + paths = [] + for tag in "ab": + obj = ad.AnnData(X=sparse.csc_matrix((3, 2), dtype=np.float32), + obs=pd.DataFrame(index=[f"{tag}{i}" for i in range(3)]), + var=pd.DataFrame(index=genes)) + paths.append(_write(obj, backend_path(tmp_path, backend, tag), backend)) + out = backend_path(tmp_path, backend, "out") + concat_on_disk(paths, out, QUIET) + got = _read(out) + assert got.shape == (6, 2) and got.X.nnz == 0 and got.X.format == "csc" + + +# --------------------------------------------------------------------------- +# F076: inputs opened as they are needed + + +def _rich_input(path: Path, tag: str, n: int, genes, *, seed: int) -> Path: + """An h5ad holding something of every kind concat reads per input.""" + rng = np.random.default_rng(seed) + cells = [f"{tag}{i}" for i in range(n)] + obs = pd.DataFrame(index=cells) + obs["cell_type"] = pd.Categorical(rng.choice(["T", "B", f"NK{tag}"], n)) + obs["n_counts"] = rng.integers(0, 100, n) + obs["score"] = pd.array([None if i % 3 == 0 else i for i in range(n)], dtype="Int64") + obs["note"] = [f"{tag}-{i}" for i in range(n)] + var = pd.DataFrame({"gene_ids": [f"E{g}" for g in genes]}, index=list(genes)) + obj = ad.AnnData(X=sparse.csc_matrix(rng.poisson(1, (n, len(genes))).astype(np.float32)), + obs=obs, var=var) + obj.layers["dense"] = rng.normal(size=(n, len(genes))) + obj.layers["csr"] = sparse.csr_matrix(rng.poisson(1, (n, len(genes))).astype(np.float32)) + obj.obsm["X_pca"] = rng.normal(size=(n, 3)) + obj.obsm["frame"] = pd.DataFrame({"r0": rng.random(n), "kind": pd.Categorical(["x"] * n)}, index=cells) + obj.varm["loadings"] = np.arange(len(genes) * 2, dtype=np.float64).reshape(len(genes), 2) + obj.varp["corr"] = sparse.csr_matrix(np.eye(len(genes))) + obj.uns["title"] = "atlas" + obj.uns["nested"] = {"a": np.arange(3), "b": "x"} + obj.write_h5ad(path) + return path + + +@pytest.mark.parametrize("join", ["inner", "outer"]) +def test_concat_with_one_input_open_at_a_time_matches_anndata(tmp_path, monkeypatch, join): + """F076: concat opened every input up front and held them all, so it ran + out of file descriptors at about 253 HDF5 inputs. Inputs are now opened + as they are needed, at most `MAX_OPEN_INPUTS` at once. With the limit at + one, an object read from an input closed since -- which h5py invalidates -- + would fail here, so every element has to be looked up again.""" + genes = [f"g{i}" for i in range(8)] + paths = [ + _rich_input(tmp_path / "a.h5ad", "a", 6, genes, seed=0), + _rich_input(tmp_path / "b.h5ad", "b", 5, genes[2:] + ["x"], seed=1), + _rich_input(tmp_path / "c.h5ad", "c", 4, genes[::-1], seed=2), + ] + kwargs = dict(join=join, merge="first", uns_merge="same", label="batch", keys=["a", "b", "c"], + chunk_rows=2) + monkeypatch.setattr(concat_mod, "MAX_OPEN_INPUTS", 1) + concat_on_disk(paths, tmp_path / "one.h5ad", QUIET, **kwargs) + monkeypatch.setattr(concat_mod, "MAX_OPEN_INPUTS", 64) + concat_on_disk(paths, tmp_path / "all.h5ad", QUIET, **kwargs) + + one, every = ad.read_h5ad(tmp_path / "one.h5ad"), ad.read_h5ad(tmp_path / "all.h5ad") + ref = ad.concat([ad.read_h5ad(p) for p in paths], join=join, merge="first", + uns_merge="same", label="batch", keys=["a", "b", "c"], fill_value=0) + for got in (one, every): + assert list(got.var_names) == list(ref.var_names) + assert np.array_equal(got.X.toarray(), ref.X.toarray()) + assert np.allclose(got.layers["dense"], ref.layers["dense"]) + assert np.array_equal(got.layers["csr"].toarray(), ref.layers["csr"].toarray()) + assert np.allclose(got.obsm["X_pca"], ref.obsm["X_pca"]) + assert list(got.obs.columns) == list(ref.obs.columns) + for column in ref.obs.columns: + assert [str(v) for v in got.obs[column]] == [str(v) for v in ref.obs[column]], column + assert got.uns["title"] == ref.uns["title"] == "atlas" + if join == "inner": + # Under an outer join anndata pads an input's var columns, varm + # and varp for the genes it lacks; concat counts that input as + # lacking them, so `first` finds none here. + assert sorted(got.varm) == sorted(ref.varm) == ["loadings"] + assert sorted(got.varp) == sorted(ref.varp) == ["corr"] + assert list(got.var["gene_ids"]) == list(ref.var["gene_ids"]) + pd.testing.assert_frame_equal(one.obs, every.obs) + assert np.array_equal(one.X.toarray(), every.X.toarray()) + + +@pytest.mark.skipif(sys.platform == "win32", reason="POSIX file descriptor limits") +def test_more_inputs_than_the_file_descriptor_limit_concatenate(tmp_path): + """F076: with every input held open, concat stopped with 'Too many open + files' at about 253 HDF5 inputs under macOS's default limit of 256. Here + the limit is 64 and there are 100 inputs.""" + paths = [] + for i in range(100): + obj = ad.AnnData(X=np.full((2, 3), i, dtype=np.float32), + obs=pd.DataFrame(index=[f"c{i}-0", f"c{i}-1"]), + var=pd.DataFrame(index=["g0", "g1", "g2"])) + obj.write_h5ad(tmp_path / f"in{i:03d}.h5ad") + paths.append(str(tmp_path / f"in{i:03d}.h5ad")) + out = tmp_path / "out.h5ad" + script = textwrap.dedent( + """ + import resource, sys + soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) + resource.setrlimit(resource.RLIMIT_NOFILE, (64, hard)) + from adata.cli import app + app(sys.argv[1:]) + """ + ) + result = subprocess.run( + [sys.executable, "-c", script, "concat", *paths, "-o", str(out)], + capture_output=True, text=True, + ) + assert result.returncode == 0, result.stderr + result.stdout + got = ad.read_h5ad(out) + assert got.shape == (200, 3) + assert np.array_equal(got.X[::2, 0], np.arange(100, dtype=np.float32)) + + +def _tiny_disk(monkeypatch, free: int = 1) -> None: + import shutil + + from adata.util import staging + + usage = shutil.disk_usage(".") + monkeypatch.setattr(staging.shutil, "disk_usage", + lambda path: type(usage)(usage.total, usage.total - free, free)) + + +def test_the_csc_scratch_is_checked_for_space_even_for_a_zarr_output(tmp_path, monkeypatch): + """F011: a CSC matrix's buckets go in an HDF5 file beside the output, and + a full disk part way through an HDF5 write can abort with no cleanup. + A Zarr output is not itself measured, but its scratch is, without + walking the inputs to size them.""" + from adata.util import staging + + genes = [f"g{i}" for i in range(10)] + a = _csc_input(tmp_path / "a.zarr", "zarr3", "a", 20, genes, seed=0) + b = _csc_input(tmp_path / "b.zarr", "zarr3", "b", 20, genes, seed=1) + _tiny_disk(monkeypatch) + monkeypatch.setattr(staging, "disk_size", lambda path: pytest.fail("walked an input")) + out = tmp_path / "out.zarr" + with pytest.raises(OSError, match="Not enough free space"): + concat_on_disk([a, b], out, QUIET) + assert not out.exists() and not list(tmp_path.glob(".*scratch*")) + concat_on_disk([a, b], out, QUIET, space_check=False) + assert _read(out).shape == (40, 10) + + +def test_an_h5ad_outputs_estimate_adds_the_csc_scratch(tmp_path, monkeypatch): + """F011: the scratch is there while the output is written, so an HDF5 + output needs room for both: the inputs' size, and for the largest CSC + matrix 16 bytes of row and column plus the value per stored entry.""" + from adata.util.staging import disk_size + + genes = [f"g{i}" for i in range(10)] + a = _csc_input(tmp_path / "a.h5ad", "h5ad", "a", 20, genes, seed=0) + b = _csc_input(tmp_path / "b.h5ad", "h5ad", "b", 20, genes, seed=1) + seen = [] + monkeypatch.setattr(concat_mod, "check_free_space", + lambda out, needed, check=True: seen.append(needed() if callable(needed) else needed)) + concat_on_disk([a, b], tmp_path / "out.h5ad", QUIET) + nnz = sum(int(_read(p).layers["counts"].nnz) for p in (a, b)) + assert seen == [disk_size(a) + disk_size(b) + nnz * (16 + 4)] + + +def _csc_store(path: Path, cells, genes, data, indices, indptr) -> Path: + """A CSC X stored exactly as given, unsorted rows within a column too.""" + from adata.elements import spec + from adata.elements import write as ew + + with open_store(path, "w") as handle: + root = handle.root + ew.write_dataframe_header(root, "obs", list(cells), []) + ew.write_dataframe_header(root, "var", list(genes), []) + ew.ensure_anndata_skeleton(root) + ew.write_sparse(root, "X", np.asarray(data, dtype=np.float32), np.asarray(indices), + np.asarray(indptr), (len(cells), len(genes)), spec.CSC_MATRIX) + return path + + +@pytest.mark.parametrize("unsorted", [False, True], ids=["sorted", "an-input-unsorted"]) +def test_a_column_bigger_than_a_block_is_streamed_and_comes_out_sorted(tmp_path, unsorted): + """F011: a gene stored in most cells fills a bucket on its own, which is + streamed a block at a time rather than read whole -- its entries arrive + input by input, each input's in row order. If an input stored that + column out of order, the bucket is sorted after all.""" + n = 2600 + paths = [] + for tag, seed in (("a", 0), ("b", 1)): + rng = np.random.default_rng(seed) + rows = np.arange(n) + if unsorted and tag == "b": + rows = rng.permutation(rows) + sparse_rows = np.flatnonzero(rng.random(n) < 0.1) + data = np.concatenate((rng.integers(1, 9, n), rng.integers(1, 9, sparse_rows.size))) + indices = np.concatenate((rows, sparse_rows)) + indptr = [0, n, n + sparse_rows.size] + paths.append(_csc_store(tmp_path / f"{tag}.h5ad", [f"{tag}{i}" for i in range(n)], + ["dense", "sparse"], data, indices, indptr)) + out = tmp_path / "out.h5ad" + concat_on_disk(paths, out, QUIET, chunk_rows=64) + stored = _stored_csc(out, "X") + assert stored.has_sorted_indices + ref = ad.concat([ad.read_h5ad(p) for p in paths]) + assert np.array_equal(stored.toarray(), ref.X.toarray()) + + +def test_a_matrix_out_of_step_with_obs_is_refused_before_writing(tmp_path): + """The CSR path now fills one int64 array of row pointers sized by obs, + so a matrix whose rows are not its input's obs is refused up front, + naming the input, rather than failing half way through the write.""" + genes = ["g0", "g1"] + good = _csc_store(tmp_path / "a.h5ad", ["a0", "a1"], genes, [1, 2], [0, 1], [0, 1, 2]) + bad = _csc_store(tmp_path / "b.h5ad", ["b0", "b1", "b2"], genes, [1], [0], [0, 1, 1]) + with h5py.File(bad, "a") as f: + f["X"].attrs["shape"] = np.array([2, 2]) + out = tmp_path / "out.h5ad" + with pytest.raises(ValueError, match=r"b\.h5ad.* 2 rows of it but 3 obs"): + concat_on_disk([good, bad], out, QUIET) + assert not out.exists() + + +def test_running_out_of_file_descriptors_says_how_to_raise_the_limit(tmp_path, monkeypatch): + """F076: if the process has no descriptors left even for the inputs + concat holds open, the error names the limit and `ulimit -n`, not just + 'Too many open files'.""" + import errno + + genes = ["g0"] + paths = [_csc_store(tmp_path / f"{t}.h5ad", [f"{t}0"], genes, [1], [0], [0, 1]) for t in "ab"] + + def exhausted(path, mode, **kwargs): + raise OSError(errno.EMFILE, "Too many open files") + + monkeypatch.setattr(concat_mod, "open_store", exhausted) + with pytest.raises(OSError, match=r"ulimit -n") as raised: + concat_on_disk(paths, tmp_path / "out.h5ad", QUIET) + assert "a.h5ad" in str(raised.value) diff --git a/tests/test_concat_versions.py b/tests/test_concat_versions.py index f386843..1c4251b 100644 --- a/tests/test_concat_versions.py +++ b/tests/test_concat_versions.py @@ -68,3 +68,33 @@ def test_an_outer_join_missing_a_text_column_opens_in_the_release_that_wrote_the ["a", "bb", "ccc", "d", "ee", "f"] ) assert read_back(out, "str(a.obs['cell_type'].dtype)") == "category" + + +def test_a_csc_concat_with_varp_opens_in_the_release_that_wrote_the_inputs( + store, tmp_path, fmt, read_back +): + """F011, F075, F118, F114: CSC X put together through buckets, dense + layers written with the input's layout, varp merged under --merge and an + ordered categorical kept ordered -- all opened by the release that wrote + the inputs, 0.8 to 0.13.""" + csc = tmp_path / f"csc.{fmt}" + result = runner.invoke(app, [ + "recast", str(store), "X", "layers/counts", "-o", str(csc), "--layout", "csc", + ]) + assert result.exit_code == 0, _out(result) + other = tmp_path / f"other.{fmt}" + shutil.copytree(csc, other) if csc.is_dir() else shutil.copy2(csc, other) + + out = tmp_path / f"merged.{fmt}" + result = runner.invoke(app, [ + "concat", str(csc), str(other), "-o", str(out), "--merge", "same", + "--index-unique", "-", "--chunk", "2", + ]) + assert result.exit_code == 0, _out(result) + + assert read_back(out, "a.shape") == f"({2 * N_OBS}, 4)" + assert read_back(out, "a.X.format") == "csc" + assert read_back(out, "bool((a.X.toarray()[:6] == a.X.toarray()[6:]).all())") == "True" + assert read_back(out, "bool((a.layers['counts'].toarray() == a.layers['dense']).all())") == "True" + assert read_back(out, "sorted(a.varp)") == "['corr']" + assert read_back(out, "a.obs['cell_type'].cat.ordered") == "True" diff --git a/tests/test_performance.py b/tests/test_performance.py index 4871722..9d041cb 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -33,8 +33,9 @@ from __future__ import annotations import pathlib +from contextlib import contextmanager from pathlib import Path -from typing import List, Optional +from typing import Iterator, List, Optional import numpy as np import pytest @@ -433,6 +434,219 @@ def measure(n: int) -> int: ) +# --------------------------------------------------------------------------- +# concat: CSC inputs, open inputs and per-row counts + + +def _csc_inputs(tmp_path: Path, tag: str, n_var: int, ext: str, n_obs: int = 256) -> List[Path]: + """Two inputs whose CSC X holds one entry a row, so nnz stays n_obs + whatever n_var is: only the number of columns changes.""" + directory = tmp_path / tag + directory.mkdir(parents=True, exist_ok=True) + paths = [] + for name in "ab": + rows = np.arange(n_obs) + obj = ad.AnnData( + X=sparse.csc_matrix((np.ones(n_obs, dtype="float32"), (rows, rows % n_var)), + shape=(n_obs, n_var)), + obs=pd.DataFrame(index=[f"{name}c{i}" for i in range(n_obs)]), + var=pd.DataFrame(index=[f"g{i}" for i in range(n_var)]), + ) + path = directory / f"{name}{ext}" + obj.write_zarr(path) if ext == ".zarr" else obj.write_h5ad(path) + paths.append(path) + return paths + + +def _count_x_reads(monkeypatch) -> List[int]: + """Read calls on X's own arrays, on either backend: each one fetches and + decompresses at least a chunk, however few entries it takes.""" + import h5py + import zarr + + calls = [0] + h5_get, zarr_get = h5py.Dataset.__getitem__, zarr.Array.__getitem__ + + def h5(self, key): + calls[0] += self.name.startswith("/X/") + return h5_get(self, key) + + def zr(self, key): + calls[0] += (self.path or "").startswith("X/") + return zarr_get(self, key) + + monkeypatch.setattr(h5py.Dataset, "__getitem__", h5) + monkeypatch.setattr(zarr.Array, "__getitem__", zr) + return calls + + +@pytest.mark.parametrize("ext", [".h5ad", ".zarr"]) +def test_csc_concat_reads_each_input_in_blocks_not_a_column_at_a_time(tmp_path, monkeypatch, ext): + """F011: CSC inputs were read one column at a time, 2 x n_var x n_inputs + reads that each decompressed a chunk (and on sharded Zarr fetched the + shard index) to take a few entries: 37 times the input read on h5ad and + 572 times on Zarr, so atlas-sized CSC inputs never finished. At a fixed + number of entries, reads must not grow with the number of columns.""" + calls = _count_x_reads(monkeypatch) + + def measure(n_var: int) -> int: + files = _csc_inputs(tmp_path, f"csc{ext}{n_var}", n_var, ext) + calls[0] = 0 + concat_on_disk(files, tmp_path / f"ocsc{n_var}{ext}", QUIET) + assert calls[0] >= 6, "indptr, indices and data of each input" + return calls[0] + + assert_independent_of(measure, what=f"CSC concat reads of X ({ext})", axis="n_var") + + +@contextmanager +def _count_open_files() -> Iterator[List[int]]: + """The most HDF5 files open at once while the block runs, beyond those + open before it, as `[n]`: sampled each time one is opened.""" + import h5py + + def open_now() -> int: + return h5py.h5f.get_obj_count(h5py.h5f.OBJ_ALL, h5py.h5f.OBJ_FILE) + + peak, before, init = [0], open_now(), h5py.File.__init__ + + def counting(self, *args, **kwargs): + init(self, *args, **kwargs) + peak[0] = max(peak[0], open_now() - before) + + h5py.File.__init__ = counting + try: + yield peak + finally: + h5py.File.__init__ = init + + +def test_concat_holds_a_bounded_number_of_inputs_open(tmp_path, monkeypatch): + """F076: every input was opened up front and held to the end, so concat + ran out of file descriptors at about 253 HDF5 inputs under macOS's + default limit, and each input's caches added 0.8 to 2.7 MB of memory + whatever --chunk was. At most `MAX_OPEN_INPUTS` are open at once now, + however many there are, through every step: obs, var, matrices, obsm, + varm and uns.""" + from adata.core import concat + + monkeypatch.setattr(concat, "MAX_OPEN_INPUTS", 4) + + def measure(n: int) -> int: + files = _inputs(tmp_path, f"open{n}", count=n, n_obs=2, n_var=4, n_categories=2, + frames=True) + with _count_open_files() as peak: + concat_on_disk(files, tmp_path / f"oopen{n}.h5ad", QUIET, merge="same", + uns_merge="same") + assert peak[0] >= 2, "an input and the output, at least" + return peak[0] + + assert_independent_of(measure, what="HDF5 files open at once", axis="n_inputs", + sizes=(8, 64)) + + +def test_one_input_over_the_limit_costs_about_one_reopen_a_pass(tmp_path, monkeypatch): + """F076: concat passes over its inputs, first to last, once per column + and element. Closing the input used longest ago to make room would close + each one just before the pass reached it again, reopening every input on + every pass; closing the one used last keeps all but about one open.""" + from adata.core import concat + + n = 8 + files = _inputs(tmp_path, "reopen", count=n, n_obs=2, n_var=4, n_categories=2, frames=True) + opens = [0] + original = concat._open_input + + def counting(path): + opens[0] += 1 + return original(path) + + monkeypatch.setattr(concat, "_open_input", counting) + + def measure(limit: int) -> int: + monkeypatch.setattr(concat, "MAX_OPEN_INPUTS", limit) + opens[0] = 0 + concat_on_disk(files, tmp_path / f"oreopen{limit}.h5ad", QUIET, merge="same", + uns_merge="same") + return opens[0] + + assert measure(n) == n, "each input opened once when all fit" + passes = measure(1) / n + assert passes >= 10, "enough passes over the inputs to tell the policies apart" + reopened = measure(n - 1) - n + assert reopened <= 2 * passes, ( + f"{reopened} reopens over about {passes:.0f} passes with one input over " + f"the limit; closing the least recently used would reopen all {n} each pass." + ) + + +def _matrix_step_allocation(monkeypatch) -> List[int]: + """Peak Python allocation of the matrix step alone, past what was already + held when it began (obs names, which concat keeps by design).""" + import tracemalloc + + from adata.core import concat + + peaks: List[int] = [] + original = concat._write_matrices + + def measured(*args, **kwargs): + tracemalloc.reset_peak() + base = tracemalloc.get_traced_memory()[0] + original(*args, **kwargs) + peaks.append(tracemalloc.get_traced_memory()[1] - base) + + monkeypatch.setattr(concat, "_write_matrices", measured) + return peaks + + +def test_csr_concat_keeps_row_pointers_in_numpy_arrays(tmp_path, monkeypatch): + """F132: each output row's pointer went into a Python list of int + objects, 36 to 40 bytes a row against numpy's 8, before being written: + about 400 MB for 10M cells, beside the obs names. What the matrix step + holds per row is its int64 pointers and one input's, and no more.""" + peaks = _matrix_step_allocation(monkeypatch) + + def measure(n: int) -> int: + files = _inputs(tmp_path, f"ptr{n}", n_obs=n, n_var=4) + with count_allocations(): + concat_on_disk(files, tmp_path / f"optr{n}.h5ad", QUIET) + return peaks[-1] + + small, large = 4096, 16384 + per_row = (measure(large) - measure(small)) / (2 * (large - small)) + assert 8 <= per_row <= 20, ( + f"{per_row:.1f} bytes a row of output held by the matrix step; an int64 " + f"indptr is 8, plus up to 8 for the input being read. A Python list is 36+." + ) + + +def test_sparsify_keeps_per_row_counts_in_numpy_arrays(tmp_path): + """F132, in `recast --layout csr`: sparsify extended a Python list with + every row's count of nonzeros, 36 bytes a row once a count passes 256, + before turning it into the indptr. Here each row holds 300.""" + import h5py + + from adata.core.recast import Matrix, sparsify + from adata.elements import spec + + def measure(n: int) -> int: + with h5py.File(tmp_path / f"sp{n}.h5", "w") as f: + src = f.create_dataset("X", data=np.ones((n, 300), dtype="int8")) + with count_allocations() as peak: + sparsify(Matrix(src, "dense", (n, 300), src.dtype), f, "out", + enc=spec.CSR_MATRIX, data_dtype=np.dtype("int8"), + index_dtype=np.dtype("int32"), chunk_rows=256) + return peak[0] + + small, large = 2048, 8192 + per_row = (measure(large) - measure(small)) / (large - small) + assert per_row <= 30, ( + f"{per_row:.1f} bytes a row held by sparsify: the counts, their sum " + f"and the indptr are 8 each as int64. A Python list of counts is 36 more." + ) + + # --------------------------------------------------------------------------- # subset, split and conversion # From 33441fa6df1d6d22d615816b0febaf0b3dc502d2 Mon Sep 17 00:00:00 2001 From: Aljes Date: Fri, 30 Oct 2026 17:55:38 +0000 Subject: [PATCH 2/2] Review: exact CSC scratch, a capped input pool, anndata's category order MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit From the independent review of #116: - CSC buckets get fixed regions of three contiguous datasets, since every output column's size is known up front. BucketSorter's three growable datasets per bucket held about 57 KB of memory and a 1.3 MB chunk of disk each, and buckets multiply as --chunk shrinks, so the scratch came to 16 times the free-space estimate at --chunk 1 and memory grew with n_obs / --chunk. The scratch is now exactly the estimate. - At most 64 inputs open at once (fewer under a low descriptor limit): each costs about a megabyte, and clusters set the limit in thousands. - natural_key compares text in NFD form, as natsort does ("Müller" after "Muscle"); integer categories meeting float ones lose the first input's order, as in anndata. - Dense obsm out of step with obs is refused, as sparse obsm was. - A float32 matrix becomes float64 for a fill beyond its range. - Docs: a sparse varp entry never equals a dense one; the pool's cap. Co-Authored-By: Claude Opus 5.5 --- docs/reference/commands.md | 10 +- src/adata/core/concat.py | 198 ++++++++++++++++++----------- src/adata/util/names.py | 6 +- tests/test_concat_axis_elements.py | 7 +- tests/test_concat_correctness.py | 45 +++++-- tests/test_concat_scale.py | 32 +++++ tests/test_performance.py | 40 ++++++ 7 files changed, 250 insertions(+), 88 deletions(-) diff --git a/docs/reference/commands.md b/docs/reference/commands.md index 1704828..2d6dbfe 100644 --- a/docs/reference/commands.md +++ b/docs/reference/commands.md @@ -338,7 +338,9 @@ and `adata recast` can convert one to match. var's columns, varm (a SCope import's `Regulons`, pycisTopic's `topic_region`) and varp follow `--merge`, compared on the output's genes: an input lacking some of them, as in an outer join, counts as not having the entry. Sparse varp entries are compared -by value, as anndata compares them, whatever their format and dtype. Without +by value, whatever their format (CSR or CSC) and dtype, and stored zeros +count as absent, as in anndata; a sparse entry never equals a dense one, +where anndata's answer depends on which input comes first. Without `--merge`, varm and varp are dropped with a note. `obsp` and `raw` are **not** carried over — a value between two cells has no @@ -350,9 +352,9 @@ compression, as sparse ones do. A dense integer matrix stays integer unless an outer join fills cells with a `--fill-value` it cannot hold, such as 0.5 or NaN; then it becomes float64. -concat opens its inputs as it needs them, at most a quarter of the -process's file descriptor limit (`ulimit -n`) at once, so it takes any -number of inputs. CSC matrices are read once each, a block of entries at a +concat opens its inputs as it needs them, at most 64 at once (or a quarter +of the process's file descriptor limit, `ulimit -n`, if that is fewer), so +it takes any number of inputs. CSC matrices are read once each, a block of entries at a time, and their columns put together through a scratch file beside the output, of about 16 bytes plus the value per stored entry of the largest one. The free-space check counts it, for a Zarr output too. diff --git a/src/adata/core/concat.py b/src/adata/core/concat.py index 1638072..f38d8a4 100644 --- a/src/adata/core/concat.py +++ b/src/adata/core/concat.py @@ -72,10 +72,12 @@ #: What the CLI accepts. "drop" is the default and maps to no merge at all. MERGE_CHOICES = ("drop",) + MERGE_STRATEGIES -#: How many inputs concat holds open at once: a quarter of the file -#: descriptor limit, as split caps its open outputs. Beyond that, inputs are -#: closed and opened again as they are needed. -MAX_OPEN_INPUTS = _open_output_limit() +#: How many inputs concat holds open at once: 64, or a quarter of the file +#: descriptor limit if that is fewer, as split caps its open outputs. Beyond +#: that, inputs are closed and opened again as they are needed. Each open +#: HDF5 input costs about a megabyte, so the limit is not left to grow with +#: the descriptor limit, which a cluster may set in the thousands. +MAX_OPEN_INPUTS = min(64, _open_output_limit()) # --------------------------------------------------------------------------- @@ -497,13 +499,16 @@ def _natural_order(values: List[Any], dtype: np.dtype) -> List[Any]: return sorted(values) if dtype.kind in "biuf" else natural_sorted(values) -def _merged_categories(labels: List[np.ndarray], ordered: List[bool]) -> Tuple[np.ndarray, bool]: +def _merged_categories( + labels: List[np.ndarray], ordered: List[bool], first_kept: bool = True +) -> Tuple[np.ndarray, bool]: """A merged categorical's categories, and whether it stays ordered, as anndata decides them. Unordered, the first input's categories keep their order when they - already hold every input's, as pandas keeps them; otherwise the union is - natsorted, as anndata sorts it. First appearance, as an earlier version + already hold every input's and needed no cast (`first_kept`: integer + categories meeting float ones are recast, and pandas then reorders + them); otherwise the union is natsorted, as anndata sorts it. First appearance, as an earlier version used, gave other codes and another plotting order than anndata.concat. Ordered only when every input with categories is ordered and has the @@ -526,7 +531,7 @@ def _merged_categories(labels: List[np.ndarray], ordered: List[bool]) -> Tuple[n o and c.size == len(union) and np.array_equal(c, stated[0][0]) for c, o in stated ): return stated[0][0], True - elif labels and labels[0].size == len(union): + elif first_kept and labels and labels[0].size == len(union): return labels[0], False return np.asarray(_natural_order(list(union), dtype), dtype=dtype), False @@ -546,8 +551,11 @@ def _concat_categorical( for col, frame in column.present(): labels.append(_categories(col, frame)) ordered.append(_ordered(col, frame)) + stored = [c.dtype for c in labels] labels = _category_labels(labels) - categories, keeps_order = _merged_categories(labels, ordered) + categories, keeps_order = _merged_categories( + labels, ordered, first_kept=bool(labels) and stored[0] == labels[0].dtype + ) if any(ordered) and not keeps_order: console.print( f"[yellow]Column {name!r}: the inputs disagree on its category order, " @@ -873,8 +881,12 @@ def _concat_sparse( def _dtype_holding(dtype: np.dtype, fill_value: float) -> np.dtype: - """`dtype`, or float64 if it cannot hold `fill_value` exactly.""" + """`dtype`, or float64 if it cannot hold `fill_value`: exactly, for + integers and booleans; within its range, for floats.""" dtype = np.dtype(dtype) + if dtype.kind == "f": + fits = not np.isfinite(fill_value) or abs(fill_value) <= float(np.finfo(dtype).max) + return dtype if fits else np.dtype(np.float64) if dtype.kind not in "biu": return dtype if not np.isfinite(fill_value) or fill_value != int(fill_value): @@ -949,26 +961,26 @@ def _concat_dense( @contextmanager def _bucket_scratch(dst_parent: Any) -> Iterator[Any]: - """Where the buckets go: an HDF5 file beside the output, whatever the - output is. `scratch_group` keeps a Zarr output's scratch inside the - output, where every small append to a bucket rewrites that array's - metadata: putting a CSC matrix together there took sixty times as long - as in HDF5.""" - from adata.core.csr_build import scratch_group - - root = getattr(getattr(dst_parent, "store", None), "root", None) - if not is_zarr_group(dst_parent) or root is None: - with scratch_group(dst_parent, "__concat_scratch__") as group: - yield group - return + """A group for the buckets, in a temporary HDF5 file beside the output, + whatever the output is, and gone afterwards. + + Beside the output rather than in the system temp directory, which on a + cluster node is often far too small. HDF5 for a Zarr output too: inside + a Zarr store every write to a bucket rewrites that array's metadata, and + putting a CSC matrix together there took sixty times as long. + """ + if is_zarr_group(dst_parent): + root = getattr(dst_parent.store, "root", None) + where = Path(root).parent if root is not None else None + else: + where = Path(str(dst_parent.file.filename)).parent fd, path = tempfile.mkstemp( - prefix=".adata-", suffix=".concat-scratch.h5", dir=str(Path(root).parent) + prefix=".adata-", suffix=".concat-scratch.h5", + dir=str(where) if where is not None and where.is_dir() else None, ) os.close(fd) try: - # No chunk cache, as `scratch_group`'s: each bucket would hold a - # partly written chunk in one. - with h5py.File(path, "w", rdcc_nbytes=0) as handle: + with h5py.File(path, "w") as handle: yield handle.create_group("buckets") finally: os.unlink(path) @@ -985,42 +997,83 @@ def _csc_bounds(indptr: np.ndarray, entries: int) -> np.ndarray: return np.unique(np.concatenate((bounds, big, big + 1))) -def _in_row_order(minors: Any, n: int, entries: int) -> bool: - """Whether a bucket's rows are in order already, read a piece at a time.""" - last = -1 - for lo in range(0, n, entries): - piece = np.asarray(minors[lo : lo + entries], dtype=np.int64) - if piece.size and (piece[0] < last or bool((piece[1:] < piece[:-1]).any())): - return False - last = int(piece[-1]) if piece.size else last - return True - +class _ColumnBuckets: + """Entries scattered into column-range buckets on disk, and read back a + bucket at a time, each in column order and then row order. -def _csc_runs(sorter: Any, entries: int) -> Iterator[Tuple[np.ndarray, np.ndarray]]: - """The buckets' `(rows, values)`, column by column, each in row order. - - A bucket of several columns is read whole and sorted. A bucket of one - column holding more than `entries` is streamed instead: its entries - arrived input by input, each input's in row order, so they are sorted - already -- unless some input stored that column out of order, which a - first pass over the rows finds, and then it is sorted whole after all. - Otherwise a gene expressed in most of a large atlas's cells would put - millions of entries in memory at once. + Every output column's size is known before an entry is read, so each + bucket has a fixed region of three contiguous datasets (column, row and + value) and a cursor into it. `csr_build.BucketSorter` grows three + chunked datasets per bucket instead, and the buckets multiply as --chunk + shrinks: each held about 57 KB of memory open and a whole chunk of each + dataset, 1.3 MB, on disk however few entries it had. A CSC concat's + memory and scratch grew with n_obs / --chunk that way. Here three + datasets are open, and the scratch is exactly 16 bytes plus the value + for each entry, what the free-space check counts. """ - for b, (majors, minors, values) in enumerate(sorter.buckets): - n = int(majors.shape[0]) - if n == 0: - continue - one_column = int(sorter.bounds[b + 1] - sorter.bounds[b]) == 1 - if one_column and n > entries and _in_row_order(minors, n, entries): - for lo in range(0, n, entries): - yield (np.asarray(minors[lo : lo + entries], dtype=np.int64), - np.asarray(values[lo : lo + entries])) - continue - major = np.asarray(majors[...], dtype=np.int64) - minor = np.asarray(minors[...], dtype=np.int64) - order = np.lexsort((minor, major)) - yield minor[order], np.asarray(values[...])[order] + + def __init__(self, scratch: Any, indptr: np.ndarray, entries: int, dtype: np.dtype) -> None: + self.bounds = _csc_bounds(indptr, entries) + self.starts = indptr[self.bounds] + self.cursor = self.starts[:-1].copy() + total = int(indptr[-1]) + self.columns = scratch.create_dataset("columns", shape=(total,), dtype=np.int64) + self.rows = scratch.create_dataset("rows", shape=(total,), dtype=np.int64) + self.values = scratch.create_dataset("values", shape=(total,), dtype=dtype) + + def add(self, columns: np.ndarray, rows: np.ndarray, values: np.ndarray) -> None: + """Entries in any order: one write per bucket they fall in.""" + if columns.size == 0: + return + which = np.searchsorted(self.bounds, columns, side="right") - 1 + order = np.argsort(which, kind="stable") + which = which[order] + cuts = np.concatenate(([0], np.flatnonzero(np.diff(which)) + 1, [which.size])) + for lo, hi in zip(cuts[:-1], cuts[1:]): + bucket, pick = int(which[lo]), order[lo:hi] + at = int(self.cursor[bucket]) + if at + pick.size > self.starts[bucket + 1]: + raise ValueError("A CSC matrix holds more entries than its pointers say.") + self.columns[at : at + pick.size] = columns[pick] + self.rows[at : at + pick.size] = rows[pick] + self.values[at : at + pick.size] = values[pick] + self.cursor[bucket] = at + pick.size + + def runs(self, entries: int) -> Iterator[Tuple[np.ndarray, np.ndarray]]: + """The buckets' `(rows, values)`, column by column, each in row order. + + A bucket of several columns is read whole and sorted. A bucket of one + column holding more than `entries` is streamed instead: its entries + arrived input by input, each input's in row order, so they are sorted + already -- unless some input stored that column out of order, which a + first pass over the rows finds, and then it is sorted whole after + all. Otherwise a gene expressed in most of a large atlas's cells + would put millions of entries in memory at once. + """ + if not np.array_equal(self.cursor, self.starts[1:]): + raise ValueError("A CSC matrix holds fewer entries than its pointers say.") + for b in range(len(self.bounds) - 1): + lo, hi = int(self.starts[b]), int(self.starts[b + 1]) + if hi == lo: + continue + one_column = int(self.bounds[b + 1] - self.bounds[b]) == 1 + if one_column and hi - lo > entries and self._in_row_order(lo, hi, entries): + for start in range(lo, hi, entries): + stop = min(start + entries, hi) + yield np.asarray(self.rows[start:stop]), np.asarray(self.values[start:stop]) + continue + columns, rows = np.asarray(self.columns[lo:hi]), np.asarray(self.rows[lo:hi]) + order = np.lexsort((rows, columns)) + yield rows[order], np.asarray(self.values[lo:hi])[order] + + def _in_row_order(self, lo: int, hi: int, entries: int) -> bool: + last = -1 + for start in range(lo, hi, entries): + piece = np.asarray(self.rows[start : min(start + entries, hi)]) + if piece[0] < last or bool((piece[1:] < piece[:-1]).any()): + return False + last = int(piece[-1]) + return True def _concat_csc( @@ -1043,12 +1096,11 @@ def _concat_csc( fetching the shard index too) to take a few hundred entries: 37 times the input's size read on h5ad, 572 times on Zarr. Instead each input is read once, in blocks of entries sized by `chunk_rows`, and its entries - scattered into column-range buckets on disk, as a streaming transpose - does (`core/csr_build.py`); the buckets are then read back in turn, + scattered into column-range buckets on disk (`_ColumnBuckets`), as a + streaming transpose's are; the buckets are then read back in turn, sorted and appended. The output's column pointers are known before any entry is read, from the inputs' own. """ - from adata.core.csr_build import BucketSorter from adata.core.subset import _append, _growable_like group = _sparse_group(dst_parent, name, spec.CSC_MATRIX, (n_rows, n_cols)) @@ -1068,23 +1120,24 @@ def _concat_csc( if total: entries = _entries_per_block(total, n_rows, chunk_rows) with _bucket_scratch(dst_parent) as scratch: - sorter = BucketSorter(scratch, _csc_bounds(indptr, entries), dtype) + buckets = _ColumnBuckets(scratch, indptr, entries, dtype) offset = 0 for i, (col_map, n_src_rows) in enumerate(zip(col_maps, row_counts)): source = inputs.get(i, path) src_indptr = np.asarray(source["indptr"][...], dtype=np.int64) src_data, src_indices = source["data"], source["indices"] + first = int(src_indptr[0]) if src_indptr.size else 0 stored = int(src_indptr[-1]) if src_indptr.size else 0 - for lo in range(0, stored, entries): + for lo in range(first, stored, entries): hi = min(lo + entries, stored) cols = np.searchsorted(src_indptr, np.arange(lo, hi), side="right") - 1 target = col_map[cols] keep = target >= 0 rows = np.asarray(src_indices[lo:hi], dtype=np.int64)[keep] + offset values = np.asarray(src_data[lo:hi])[keep].astype(dtype, copy=False) - sorter.add(target[keep], rows, values) + buckets.add(target[keep], rows, values) offset += n_src_rows - for rows, values in _csc_runs(sorter, entries): + for rows, values in buckets.runs(entries): _append(out_indices, rows.astype(index_dtype, copy=False)) _append(out_data, values) @@ -1139,8 +1192,10 @@ def _check(label: str, path: str) -> None: # obsm entries are stacked only when every input stores them alike; a # sparse entry dense elsewhere is refused here, as X would be, rather - # than dropped. Entries of different widths are skipped later, whatever - # their encodings, and dataframes are not matrices at all. + # than dropped, and so is one whose rows are not its input's obs, which + # was stacked out of line with the cells. Entries of different widths + # are skipped later, whatever their encodings, and dataframes are not + # matrices at all. for name in _index_union([inputs.keys(i, "obsm") for i in range(n)]): path = f"obsm/{name}" if not all(source is not None for source in inputs.each(path)): @@ -1150,7 +1205,7 @@ def _check(label: str, path: str) -> None: kinds.add(_matrix_kind(source)) if _matrix_kind(source) != "other": widths.add(_trailing_shape(source)) - if kinds & spec.SPARSE_TYPES and kinds <= spec.SPARSE_TYPES | {"dense"} and len(widths) == 1: + if "other" not in kinds and len(widths) == 1: _check(path, path) @@ -1371,8 +1426,9 @@ def _entries_on(obj: Any, col_map: np.ndarray, entries: int) -> Iterator[Tuple[n values)`, both renumbered by `col_map`, and those of genes the target lacks left out.""" indptr = np.asarray(obj["indptr"][...], dtype=np.int64) + first = int(indptr[0]) if indptr.size else 0 stored = int(indptr[-1]) if indptr.size else 0 - for lo in range(0, stored, entries): + for lo in range(first, stored, entries): hi = min(lo + entries, stored) major = col_map[np.searchsorted(indptr, np.arange(lo, hi), side="right") - 1] minor = col_map[np.asarray(obj["indices"][lo:hi], dtype=np.int64)] diff --git a/src/adata/util/names.py b/src/adata/util/names.py index c8687fd..050e44a 100644 --- a/src/adata/util/names.py +++ b/src/adata/util/names.py @@ -3,6 +3,7 @@ from __future__ import annotations import re +import unicodedata from collections import Counter from typing import Dict, Iterable, List, Sequence, Set @@ -93,8 +94,9 @@ def make_index_unique(names: Sequence[str], join: str = "-") -> List[str]: def natural_key(text: str) -> tuple: - """Sort key putting "A2" before "A10", as natsort's default does.""" - parts = _DIGITS.split(text) + """Sort key putting "A2" before "A10", as natsort's default does, on the + text in NFD form as natsort compares it: "Müller" after "Muscle".""" + parts = _DIGITS.split(unicodedata.normalize("NFD", text)) return tuple(int(part) if i % 2 else part for i, part in enumerate(parts)) diff --git a/tests/test_concat_axis_elements.py b/tests/test_concat_axis_elements.py index d11ff71..3967516 100644 --- a/tests/test_concat_axis_elements.py +++ b/tests/test_concat_axis_elements.py @@ -283,7 +283,9 @@ def test_varp_in_another_gene_order_is_written_on_the_output_genes(tmp_path, ext def test_sparse_varp_is_compared_by_value_as_anndata_compares_it(tmp_path): """F118: anndata's `same` finds a CSR and a CSC entry with the same values equal, and float32 equal to float64, and ignores stored zeros; - a dense entry must match its dtype as well, as varm's does.""" + a dense entry must match its dtype as well, as varm's does. A dense + entry never equals a sparse one here; anndata agrees when the dense one + comes first, and finds them equal the other way round.""" genes = [f"g{i}" for i in range(5)] frame = _pairwise(genes, 5) with_zero = sparse.csr_matrix(frame.to_numpy()) @@ -292,12 +294,13 @@ def test_sparse_varp_is_compared_by_value_as_anndata_compares_it(tmp_path): without[np.unravel_index(np.flatnonzero(without)[0], without.shape)] = 0.0 a = _varp_store(tmp_path / "a.h5ad", genes, { "formats": sparse.csr_matrix(frame.to_numpy()), "dtypes": sparse.csr_matrix(frame.to_numpy()), - "zeros": with_zero, "dense_dtypes": frame.to_numpy(), + "zeros": with_zero, "dense_dtypes": frame.to_numpy(), "dense_first": frame.to_numpy(), }) b = _varp_store(tmp_path / "b.h5ad", genes, { "formats": sparse.csc_matrix(frame.to_numpy()), "dtypes": sparse.csr_matrix(frame.to_numpy().astype(np.float32)), "zeros": sparse.csr_matrix(without), "dense_dtypes": frame.to_numpy().astype(np.float32), + "dense_first": sparse.csr_matrix(frame.to_numpy()), }) got = _concat(tmp_path, [a, b], merge="same") ref = ad.concat([ad.read_h5ad(a), ad.read_h5ad(b)], merge="same") diff --git a/tests/test_concat_correctness.py b/tests/test_concat_correctness.py index 2db137c..36fd581 100644 --- a/tests/test_concat_correctness.py +++ b/tests/test_concat_correctness.py @@ -840,21 +840,32 @@ def test_an_outer_join_keeps_a_shared_var_order_beside_inputs_without_variables( (["b", "a"], ["b", "a"]), ([10, 2], [3, 10]), ([2.5, -1.0], [0.5]), + ([10, 2], [2.0]), + (["Mz", "Müller glia"], ["Muscle"]), ], ids=["new-categories", "same-set-reordered", "first-covers-the-union", - "second-covers-the-union", "identical", "integers", "floats"], + "second-covers-the-union", "identical", "integers", "floats", + "integers-meeting-floats", "accented"], ) def test_merged_categories_are_ordered_as_anndata_orders_them(tmp_path, backend, first, second): """F114: merged categories came out in order of first appearance, where anndata natsorts the union -- unless the first input already holds every category, whose order pandas then keeps. `.cat.codes` and plotting order - differed from anndata.concat's.""" + differed from anndata.concat's. Integer categories meeting float ones + are recast, which loses the first input's order in anndata too; and + text is compared in NFD form, as natsort compares it.""" + def stored(categories): + # Numbers as an array, which keeps their dtype; a list is text. + values = categories if isinstance(categories[0], str) else np.asarray(categories) + return lambda g, name: ew.write_categorical(g, name, list(range(len(categories))), values) + a = _on(backend, tmp_path, "a", [f"a{i}" for i in range(len(first))], ["g"], - obs={"c": categorical(first, list(range(len(first))))}) + obs={"c": stored(first)}) b = _on(backend, tmp_path, "b", [f"b{i}" for i in range(len(second))], ["g"], - obs={"c": categorical(second, list(range(len(second))))}) + obs={"c": stored(second)}) got = _concat([a, b], backend_path(tmp_path, backend, "out")) ref = _reference([a, b]) + assert got.obs["c"].cat.categories.dtype == ref.obs["c"].cat.categories.dtype assert list(got.obs["c"].cat.categories) == list(ref.obs["c"].cat.categories) assert list(got.obs["c"]) == list(ref.obs["c"]) assert not got.obs["c"].cat.ordered @@ -923,8 +934,9 @@ def test_ordered_categoricals_stay_ordered_only_when_every_input_agrees(tmp_path # F116: --fill-value and integer matrices +@pytest.mark.parametrize("fill", [-1, np.nan]) @pytest.mark.parametrize("layout", ["dense", "csr"]) -def test_a_fill_value_leaves_integer_matrices_alone_when_nothing_is_filled(tmp_path, backend, layout): +def test_a_fill_value_leaves_integer_matrices_alone_when_nothing_is_filled(tmp_path, backend, layout, fill): """F116: --fill-value made integer X and layers float64 on an inner join, where no cell is filled, doubling X; anndata.concat keeps them int32.""" as_layout = {"csr": sparse.csr_matrix, "dense": np.asarray}[layout] @@ -933,8 +945,8 @@ def test_a_fill_value_leaves_integer_matrices_alone_when_nothing_is_filled(tmp_p layers={"counts": matrix(as_layout(x * 2))}) b = _on(backend, tmp_path, "b", ["b0", "b1"], ["g2", "g1", "g3"], X=as_layout(x + 10), layers={"counts": matrix(as_layout(x * 3))}) - got = _concat([a, b], backend_path(tmp_path, backend, "out"), fill_value=-1) - ref = _reference([a, b], fill_value=-1) + got = _concat([a, b], backend_path(tmp_path, backend, "out"), fill_value=fill) + ref = _reference([a, b], fill_value=fill) dense = lambda m: m.toarray() if sparse.issparse(m) else np.asarray(m) for mine, theirs in ((got.X, ref.X), (got.layers["counts"], ref.layers["counts"])): assert mine.dtype == theirs.dtype == np.int32 @@ -953,12 +965,14 @@ def test_a_fill_value_leaves_integer_matrices_alone_when_nothing_is_filled(tmp_p (np.bool_, 1, np.bool_), (np.bool_, 2, np.float64), (np.float32, 0.5, np.float32), + (np.float32, 1e40, np.float64), ], ) def test_an_outer_join_widens_a_matrix_only_for_a_fill_value_it_cannot_hold(tmp_path, dtype, fill, expected): """F116: when an outer join does fill cells, an integer matrix that can - hold the fill value keeps its dtype; one that cannot becomes float64. - The values are anndata.concat's either way.""" + hold the fill value keeps its dtype; one that cannot becomes float64, as + does a float32 one for a value beyond its range, which it would have + held as infinity. The values are anndata.concat's either way.""" x = np.array([[1, 0], [1, 1]], dtype=dtype) a = _store(tmp_path / "a.h5ad", ["a0", "a1"], ["g0", "g1"], X=x) b = _store(tmp_path / "b.h5ad", ["b0", "b1"], ["g0", "g2"], X=x) @@ -966,3 +980,16 @@ def test_an_outer_join_widens_a_matrix_only_for_a_fill_value_it_cannot_hold(tmp_ ref = _reference([a, b], join="outer", fill_value=fill) assert got.X.dtype == expected assert np.array_equal(got.X.astype(np.float64), ref.X.astype(np.float64), equal_nan=True) + + +def test_dense_obsm_out_of_step_with_obs_is_refused_before_writing(tmp_path): + """An obsm array with more or fewer rows than its input's obs was + stacked anyway, every later input's rows shifted against their cells, + with exit 0. Sparse ones were refused; dense ones are now too.""" + a = _store(tmp_path / "a.h5ad", ["a0", "a1", "a2", "a3"], ["g"], + obsm={"X_pca": dense(*[[1.0, 1.0]] * 3)}) + b = _store(tmp_path / "b.h5ad", ["b0", "b1"], ["g"], obsm={"X_pca": dense(*[[2.0, 2.0]] * 2)}) + out = tmp_path / "out.h5ad" + with pytest.raises(ValueError, match=r"obsm/X_pca.*a\.h5ad.* 3 rows of it but 4 obs"): + concat_on_disk([a, b], out, QUIET) + _no_output(out) diff --git a/tests/test_concat_scale.py b/tests/test_concat_scale.py index 4f47c05..9104821 100644 --- a/tests/test_concat_scale.py +++ b/tests/test_concat_scale.py @@ -359,3 +359,35 @@ def exhausted(path, mode, **kwargs): with pytest.raises(OSError, match=r"ulimit -n") as raised: concat_on_disk(paths, tmp_path / "out.h5ad", QUIET) assert "a.h5ad" in str(raised.value) + + +@pytest.mark.parametrize("chunk", [1, 1024]) +def test_the_csc_scratch_on_disk_is_what_the_estimate_says(tmp_path, monkeypatch, chunk): + """F011, from the review: the buckets were three growing datasets each, + and HDF5 gives each a whole 65,536-entry chunk however little is in it, + so the scratch came to 1.6 times the estimate at the default --chunk and + 16 times at --chunk 1, where there are hundreds of buckets. Each bucket + now has a fixed region of three datasets, and the scratch on disk is + the estimate: 16 bytes plus the value for each entry.""" + import os + + genes = [f"g{i}" for i in range(200)] + paths = [_csc_input(tmp_path / f"{t}.h5ad", "h5ad", t, 300, genes, seed=s) + for t, s in (("a", 0), ("b", 1))] + sizes = [] + original = concat_mod._ColumnBuckets.runs + + def measured(self, entries): + self.columns.file.flush() + sizes.append(os.path.getsize(self.columns.file.filename)) + return original(self, entries) + + monkeypatch.setattr(concat_mod._ColumnBuckets, "runs", measured) + out = tmp_path / "out.h5ad" + concat_on_disk(paths, out, QUIET, chunk_rows=chunk) + nnz = sum(int(_read(p).X.nnz) for p in paths) + estimate = nnz * (16 + 4) + assert len(sizes) == 2, "X and layers/counts" + assert all(estimate <= size <= estimate + (64 << 10) for size in sizes), (sizes, estimate) + assert np.array_equal(_read(out).X.toarray(), + ad.concat([_read(p) for p in paths]).X.toarray()) diff --git a/tests/test_performance.py b/tests/test_performance.py index 9d041cb..a05d8d3 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -530,6 +530,9 @@ def test_concat_holds_a_bounded_number_of_inputs_open(tmp_path, monkeypatch): varm and uns.""" from adata.core import concat + # About a megabyte for each open HDF5 input, so never more than 64, + # however high a cluster sets the descriptor limit. + assert concat.MAX_OPEN_INPUTS <= 64 monkeypatch.setattr(concat, "MAX_OPEN_INPUTS", 4) def measure(n: int) -> int: @@ -580,6 +583,43 @@ def measure(limit: int) -> int: ) +def test_csc_concat_holds_its_scratch_open_in_a_fixed_number_of_datasets(tmp_path): + """F011, from the review: the buckets were three growing datasets each, + all held open for the whole matrix, about 57 KB of memory a bucket; and + the buckets multiply with n_obs / --chunk, so memory grew with the + input after all (1.8 GB at 30,000 buckets). However many buckets there + are, the number of HDF5 datasets open at once must not grow.""" + import h5py + + def measure(n_obs: int) -> int: + directory = tmp_path / f"buckets{n_obs}" + directory.mkdir() + paths = [] + for name, seed in (("a", 0), ("b", 1)): + values = np.random.default_rng(seed).poisson(1.0, (n_obs, 40)).astype("float32") + obj = ad.AnnData(X=sparse.csc_matrix(values), + obs=pd.DataFrame(index=[f"{name}{i}" for i in range(n_obs)]), + var=pd.DataFrame(index=[f"g{i}" for i in range(40)])) + obj.write_h5ad(directory / f"{name}.h5ad") + paths.append(directory / f"{name}.h5ad") + peak, write = [0], h5py.Dataset.__setitem__ + + def counting(self, key, value): + peak[0] = max(peak[0], h5py.h5f.get_obj_count(h5py.h5f.OBJ_ALL, h5py.h5f.OBJ_DATASET)) + return write(self, key, value) + + h5py.Dataset.__setitem__ = counting + try: + concat_on_disk(paths, directory / "out.h5ad", QUIET, chunk_rows=4) + finally: + h5py.Dataset.__setitem__ = write + assert peak[0] >= 3, "the scratch's datasets, at least" + return peak[0] + + assert_independent_of(measure, what="HDF5 datasets open during a CSC concat", + axis="n_obs (and so buckets)", sizes=(256, 4096)) + + def _matrix_step_allocation(monkeypatch) -> List[int]: """Peak Python allocation of the matrix step alone, past what was already held when it began (obs names, which concat keeps by design)."""