diff --git a/docs/reference/commands.md b/docs/reference/commands.md index 6402055..c468bb8 100644 --- a/docs/reference/commands.md +++ b/docs/reference/commands.md @@ -305,7 +305,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) | @@ -317,11 +317,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 @@ -342,15 +352,30 @@ 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, 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 +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 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. ### Merge strategies diff --git a/src/adata/cli.py b/src/adata/cli.py index a9103a6..93281ef 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 b5b6a07..2cb539a 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,128 @@ 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: 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()) + + +# --------------------------------------------------------------------------- +# 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 +201,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 +227,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 +244,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 +309,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 +338,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 +494,50 @@ 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], 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 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 + 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 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 + + 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 +546,27 @@ 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)) + stored = [c.dtype for c in labels] + labels = _category_labels(labels) + 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, " + "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 +575,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 +590,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 +646,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 +659,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 +703,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 +755,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 +766,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 +790,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 +824,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 +868,306 @@ 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, sources[0]["indptr"]) + _write_indptr(group, indptr, pointer_dtype, inputs.get(0, path)["indptr"]) + + +def _dtype_holding(dtype: np.dtype, fill_value: float) -> np.dtype: + """`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): + 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]: + """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(where) if where is not None and where.is_dir() else None, + ) + os.close(fd) + try: + with h5py.File(path, "w") 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))) + + +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. + + 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. + """ + + 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( 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 (`_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.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: + 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(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) + buckets.add(target[keep], rows, values) + offset += n_src_rows + for rows, values in buckets.runs(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 +1178,35 @@ 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): + # 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)): 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 "other" not in kinds and len(widths) == 1: + _check(path, path) def _shape(obj: Any) -> Tuple[int, ...]: @@ -808,64 +1224,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 +1274,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 +1298,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 +1324,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 +1341,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 +1361,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 +1411,135 @@ 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) + first = int(indptr[0]) if indptr.size else 0 + stored = int(indptr[-1]) if indptr.size else 0 + 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)] + 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 +1552,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 +1561,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 +1591,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 +1642,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 +1653,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 +1670,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. @@ -1218,14 +1790,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 @@ -1243,17 +1856,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.") @@ -1263,20 +1877,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. @@ -1286,7 +1893,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) @@ -1297,10 +1905,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. " @@ -1309,10 +1913,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. @@ -1321,14 +1926,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, @@ -1337,35 +1942,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, @@ -1374,9 +1975,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": @@ -1389,8 +1989,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: @@ -1409,8 +2008,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, @@ -1420,13 +2019,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 @@ -1434,31 +2031,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) @@ -1492,8 +2085,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] @@ -1507,7 +2098,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, @@ -1517,73 +2108,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/src/adata/util/names.py b/src/adata/util/names.py index 010a973..b233990 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, Optional, Sequence, Set @@ -120,8 +121,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 a17654c..3967516 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,152 @@ 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. 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()) + 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(), "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") + 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..36fd581 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,253 @@ 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]), + ([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", + "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. 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": stored(first)}) + b = _on(backend, tmp_path, "b", [f"b{i}" for i in range(len(second))], ["g"], + 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 + + +@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("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, 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] + 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=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 + 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), + (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, 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) + 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) + + +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 new file mode 100644 index 0000000..9104821 --- /dev/null +++ b/tests/test_concat_scale.py @@ -0,0 +1,393 @@ +"""`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) + + +@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_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..a05d8d3 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,259 @@ 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 + + # 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: + 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 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).""" + 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 #