From 13c8c8fb4cf9b3fbdc0ca55d44a97a61679c68c5 Mon Sep 17 00:00:00 2001 From: Aljes Date: Thu, 15 Oct 2026 16:09:07 +0100 Subject: [PATCH] Groundwork for MTX folders: shared CSR builder and feature table No change in behaviour; this is what #24's MTX import builds on. - core/csr_build.py: the bucket stage of the streaming transpose (count, nnz-balanced bounds, scratch group, bucket sort), taken out of recast so an MTX import can feed it entries in any order. On the way, two costs that grew with the matrix rather than the block: pass 1 ran bincount over every output row for each block, and pass 2 compared each block against every bucket. - conversion/features.py: the feature-table half of the 10x HDF5 converter (targets and --split-feature-types, var names, obs/var writing, var -> id/name/type on export), shared with the MTX one. - staging: only a Zarr store's temporary gets a .zarr suffix; an MTX output folder keeps its name. is_zarr_path moves to util.path so util stays below storage. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 7 +- src/adata/conversion/features.py | 380 +++++++++++++++++++++++++++++++ src/adata/conversion/tenx.py | 316 ++++--------------------- src/adata/core/csr_build.py | 178 +++++++++++++++ src/adata/core/recast.py | 121 ++-------- src/adata/storage/__init__.py | 12 +- src/adata/util/path.py | 12 + src/adata/util/staging.py | 7 +- tests/test_csr_build.py | 96 ++++++++ tests/test_staging.py | 15 ++ 10 files changed, 754 insertions(+), 390 deletions(-) create mode 100644 src/adata/conversion/features.py create mode 100644 src/adata/core/csr_build.py create mode 100644 tests/test_csr_build.py diff --git a/AGENTS.md b/AGENTS.md index c45acfe..fa84d22 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -61,6 +61,9 @@ Each layer calls only the layers below it: - `concat.py` concatenates on disk; - `recast.py` handles dtype, layout and density changes, including an out-of-core transpose; + - `csr_build.py` builds a CSR matrix from entries in any order, through + row-range buckets on disk. The transpose uses it, and so does anything + that reads entries out of order; - `query.py` and `select.py` hold the expression language behind `--obs-query`. - **`elements/`**: the AnnData on-disk spec. @@ -82,7 +85,9 @@ Each layer calls only the layers below it: other formats. To add a format: - add one `Format` entry in `conversion/__init__.py::FORMATS` with a content sniffer; - - write a converter module (`tenx.py` is the model); + - write a converter module (`tenx.py` is the model). The feature table that + every 10x-style format shares (var names, feature-type splits, var to + id/name/type on export) is in `features.py`; reuse it; - add Cell Ranger-style h5py fixtures in `tests/conftest.py`; - add a scanpy comparison in `tests/test_external_readers.py`. diff --git a/src/adata/conversion/features.py b/src/adata/conversion/features.py new file mode 100644 index 0000000..d7c33b6 --- /dev/null +++ b/src/adata/conversion/features.py @@ -0,0 +1,380 @@ +"""The feature table shared by every 10x-style format. + +A 10x HDF5 matrix and a 10x MTX folder record the same things about each +feature -- an id, a name, a type -- and import and export them the same way. +What is common lives here, so `tenx` and `mtx` differ only in how they read +and write their files: + +* import: which features each output keeps (`plan_targets`), how they are + named (`annotate`), and how obs and var are written (`write_obs_index`, + `write_var`); +* export: which matrix to write (`plan_matrix_source`) and each feature's id, + name and type, worked out from var (`export_features`). + +The AnnData produced matches what scanpy's `read_10x_h5` and `read_10x_mtx` +give with `gex_only=False`. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple + +import numpy as np + +from adata.conversion.base import Options +from adata.core.recast import describe +from adata.elements import spec +from adata.elements.read import ( + dataframe_columns, + element_len, + read_str_all, + resolve_index, +) +from adata.elements.strings import string_dtype_for +from adata.elements.write import ( + write_categorical, + write_dataframe_header, + write_dense, + write_string_array, +) +from adata.storage import create_dataset, is_group, is_zarr_group, open_store +from adata.util.names import count_duplicates, make_index_unique, unique_names + +VAR_NAMES = ("gene_symbols", "gene_ids") + +#: Rows of a string column handled per read when streaming one. +STRING_STEP = 1 << 18 + + +# --------------------------------------------------------------------------- +# import + + +@dataclass +class Target: + """One AnnData to write: which features it keeps, and their annotation.""" + + label: Optional[str] + keep: Optional[np.ndarray] + var_names: List[str] + columns: Dict[str, np.ndarray] + + +def type_counts(types: np.ndarray) -> List[Tuple[str, int]]: + """Feature types in order of first appearance, with how many of each.""" + labels, first, counts = np.unique(types, return_index=True, return_counts=True) + order = np.argsort(first) + return [(str(labels[i]), int(counts[i])) for i in order] + + +def check_import_options(options: Options) -> None: + if options.var_names not in VAR_NAMES: + raise ValueError(f"--var-names must be one of: {', '.join(VAR_NAMES)}.") + if options.feature_type is not None and options.split_feature_types: + raise ValueError( + "--feature-type keeps one type and --split-feature-types writes " + "every type separately; pass one or the other." + ) + + +def annotate( + names: np.ndarray, + columns: Dict[str, np.ndarray], + keep: Optional[np.ndarray], + options: Options, + label: Optional[str], +) -> Target: + """var for the features in `keep`: their names, and their columns.""" + + def pick(values: np.ndarray) -> np.ndarray: + return values if keep is None else values[keep] + + picked = {key: pick(values) for key, values in columns.items()} + if options.var_names == "gene_ids": + # As scanpy's read_10x_mtx does it: ids become the names, and the + # symbols move into a column where the ids were. + index = list(picked["gene_ids"]) + picked = { + ("gene_symbols" if key == "gene_ids" else key): ( + pick(names) if key == "gene_ids" else values + ) + for key, values in picked.items() + } + else: + index = list(pick(names)) + if options.make_unique: + index = make_index_unique(index) + return Target(label=label, keep=keep, var_names=index, columns=picked) + + +def plan_targets( + output: Path, + names: np.ndarray, + columns: Dict[str, np.ndarray], + types: Optional[np.ndarray], + keep: np.ndarray, + options: Options, + *, + noun: str, +) -> Tuple[List[Path], List[Target], List[str]]: + """The outputs to write, the features each keeps, and any warnings. + + One output, unless `--split-feature-types` asks for one per type. `keep` + is what survives filters the caller applied first, such as `--genome`; + `noun` names the input in messages ("file", "folder"). + """ + groups: List[Tuple[Optional[str], np.ndarray]] = [] + if options.feature_type is not None: + available = [t for t, _ in type_counts(types)] + if options.feature_type not in available: + raise ValueError( + f"No features of type {options.feature_type!r}; this {noun} " + f"has: {', '.join(available)}." + ) + groups.append((options.feature_type, keep & (types == options.feature_type))) + elif options.split_feature_types: + for label, _ in type_counts(types[keep]): + groups.append((label, keep & (types == label))) + else: + groups.append((None, keep)) + + stems = unique_names([label for label, _ in groups if label is not None]) + outputs: List[Path] = [] + targets: List[Target] = [] + warnings: List[str] = [] + for label, mask in groups: + path = output + if options.split_feature_types: + path = output.with_name(f"{output.stem}.{stems[label]}{output.suffix}") + selection = None if mask.all() else np.flatnonzero(mask) + target = annotate(names, columns, selection, options, label) + duplicated = count_duplicates(target.var_names) + if duplicated: + where = f" in {path.name}" if options.split_feature_types else "" + warnings.append( + f"{duplicated:,} var names{where} repeat an earlier one. " + "Pass --make-unique to suffix them as anndata's " + "var_names_make_unique does." + ) + outputs.append(path) + targets.append(target) + return outputs, targets, warnings + + +def describe_import( + n_barcodes: int, + n_features: int, + types: Optional[np.ndarray], + targets: Sequence[Target], +) -> str: + """`1,222 barcodes x 33,538 features (Gene Expression 33,538)`, and so on.""" + described = f"{n_barcodes:,} barcodes x {n_features:,} features" + if types is not None: + described += " (" + ", ".join(f"{t} {n:,}" for t, n in type_counts(types)) + ")" + if len(targets) > 1: + described += f", into {len(targets)} outputs by feature type" + elif targets[0].keep is not None: + described += f", keeping {len(targets[0].keep):,} features" + return described + + +def write_obs_index(root: Any, n: int, blocks: Iterable[Sequence[str]]) -> Any: + """obs, holding only its index: the barcodes, written a block at a time. + + Returns the group, for the caller to add columns to. A source that gives + more or fewer than `n` barcodes is refused. + """ + obs = root.create_group("obs") + spec.set_encoding(obs, spec.DATAFRAME) + obs.attrs["_index"] = "_index" + + backend = "zarr" if is_zarr_group(root) else "hdf5" + index = create_dataset(obs, "_index", shape=(n,), dtype=string_dtype_for(backend)) + written = 0 + for block in blocks: + if written + len(block) > n: + raise ValueError(f"There are more than {n:,} barcodes.") + index[written : written + len(block)] = np.asarray(block, dtype=object) + written += len(block) + if written != n: + raise ValueError(f"There are {written:,} barcodes, not {n:,}.") + spec.set_encoding(index, spec.STRING_ARRAY) + return obs + + +def write_var(root: Any, target: Target) -> None: + """var, with string columns encoded as anndata's writer would encode them. + + anndata turns any string column with a repeated value into a categorical + when it writes (`strings_to_categoricals`), so an h5ad scanpy wrote from + the same file holds `feature_types` and `genome` as categoricals. Doing + the same keeps the two interchangeable, and those columns a few bytes + instead of a string per feature. + """ + var = write_dataframe_header(root, "var", target.var_names, list(target.columns)) + for key, values in target.columns.items(): + if values.dtype.kind == "b": + write_dense(var, key, values) + continue + categories, codes = np.unique(values, return_inverse=True) + if len(categories) < len(values): + write_categorical(var, key, codes, list(categories)) + else: + write_string_array(var, key, values) + + +# --------------------------------------------------------------------------- +# export + + +def matrix_path(layer: Optional[str]) -> str: + """`--layer` as a path: X, raw/X, a layer's name, or layers/NAME.""" + if layer is None or layer in ("X", "raw/X") or layer.startswith("layers/"): + return layer or "X" + return f"layers/{layer}" + + +def resolve(root: Any, path: str) -> Any: + node = root + for part in path.split("/"): + if not is_group(node) or part not in node: + raise KeyError(f"{path!r} not found in the store.") + node = node[part] + return node + + +@dataclass +class ExportFeatures: + ids: List[str] + names: List[str] + types: List[str] + genomes: List[str] + tags: Dict[str, List[str]] + + +@dataclass +class MatrixSource: + """The matrix an export writes, and what is known about it before writing.""" + + path: str + n_obs: int + n_var: int + dtype: np.dtype + features: ExportFeatures + + +def plan_matrix_source( + source: Path, options: Options, warnings: List[str] +) -> MatrixSource: + """Check the matrix `--layer` names against obs and var; read var's features. + + Warns when the matrix holds floats, which 10x formats are not expected to. + """ + path = matrix_path(options.layer) + with open_store(source, "r", require_anndata=False) as store: + root = store.root + matrix = describe(resolve(root, path)) + var_path = "raw/var" if path.startswith("raw/") else "var" + var, obs = resolve(root, var_path), resolve(root, "obs") + obs_index, _ = resolve_index(obs, "obs") + var_index, _ = resolve_index(var, "var") + n_obs, n_var = element_len(obs_index), element_len(var_index) + if matrix.shape != (n_obs, n_var): + raise ValueError( + f"{path} is {matrix.shape[0]} x {matrix.shape[1]}, but obs has " + f"{n_obs} and {var_path} has {n_var} entries." + ) + features = export_features(root, var, var_index, options, warnings) + + if matrix.dtype.kind == "f": + warnings.append( + f"{path} holds {matrix.dtype} values. Cell Ranger writes integer " + "counts, and tools that read its output may expect them: export " + "a counts layer with --layer, or recast to int32 first." + ) + return MatrixSource( + path=path, n_obs=n_obs, n_var=n_var, dtype=matrix.dtype, features=features + ) + + +def export_features( + root: Any, + var: Any, + var_index: Any, + options: Options, + warnings: List[str], +) -> ExportFeatures: + """Work out each feature's id, name, type and genome from var.""" + columns = dataframe_columns(var) + n = element_len(var_index) + + def column(name: str) -> List[str]: + return read_str_all(var[name], parent_group=var) + + def requested(flag: str, name: Optional[str]) -> Optional[str]: + if name is not None and name not in columns: + raise ValueError( + f"{flag} {name}: no such var column. It has: " + f"{', '.join(columns) or 'none'}." + ) + return name + + ids_col = requested("--gene-ids-column", options.gene_ids_column) + names_col = requested("--gene-names-column", options.gene_names_column) + types_col = requested("--feature-types-column", options.feature_types_column) + index = read_str_all(var_index) + + missing = [] + if ids_col is None and names_col is None and "gene_ids" not in columns and "gene_symbols" in columns: + # var_names are ids already, as `--var-names gene_ids` leaves them. + ids, names = index, column("gene_symbols") + else: + names = column(names_col) if names_col else index + if ids_col is None and "gene_ids" in columns: + ids_col = "gene_ids" + if ids_col is not None: + ids = column(ids_col) + else: + missing.append("gene ids (a gene_ids column, or --gene-ids-column)") + ids = names + if types_col is None and "feature_types" in columns: + types_col = "feature_types" + if types_col is not None: + types = column(types_col) + else: + missing.append("feature types (a feature_types column, or --feature-types-column)") + types = ["Gene Expression"] * n + + if missing: + if not options.force: + raise ValueError( + f"var has no {' and no '.join(missing)}, which a 10x matrix " + "records for every feature. Pass --force to write placeholders: " + "the names as ids, and 'Gene Expression' as the type." + ) + warnings.append(f"Writing placeholders for {' and '.join(missing)}.") + + genomes = column("genome") if "genome" in columns else [""] * n + tags: Dict[str, List[str]] = {} + for key in recorded_tag_keys(root): + if key != "genome" and key in columns: + tags[key] = column(key) + return ExportFeatures(ids=ids, names=names, types=types, genomes=genomes, tags=tags) + + +def cellranger_meta(root: Any) -> Optional[Any]: + """`uns/cellranger`, as an import recorded it, if there is one.""" + if "uns" not in root or "cellranger" not in root["uns"]: + return None + meta = root["uns"]["cellranger"] + return meta if is_group(meta) else None + + +def recorded_tag_keys(root: Any) -> List[str]: + meta = cellranger_meta(root) + if meta is None or "feature_tag_keys" not in meta: + return [] + return read_str_all(meta["feature_tag_keys"]) + diff --git a/src/adata/conversion/tenx.py b/src/adata/conversion/tenx.py index 40c5dc2..edd0f1d 100644 --- a/src/adata/conversion/tenx.py +++ b/src/adata/conversion/tenx.py @@ -32,23 +32,32 @@ from rich.console import Console from adata.conversion.base import ConversionPlan, Converter, Options +from adata.conversion.features import ( + STRING_STEP, + ExportFeatures, + Target, + cellranger_meta, + check_import_options, + describe_import, + plan_matrix_source, + plan_targets, + resolve, + write_obs_index, + write_var, +) from adata.core.recast import describe, sparsify, transpose_sparse_streaming from adata.core.subset import fan_out_sparse_matrix from adata.elements import spec from adata.elements.read import ( - dataframe_columns, decode_str_array, element_len, read_str_all, read_str_chunk, resolve_index, ) -from adata.elements.strings import string_dtype_for from adata.elements.write import ( ensure_anndata_skeleton, set_column_order, - write_categorical, - write_dataframe_header, write_dense, write_mapping, write_scalar, @@ -57,14 +66,10 @@ from adata.storage import ( GZIP_LEVEL, copy_dataset, - create_dataset, has_valid_anndata_root_attrs, is_dataset, - is_group, - is_zarr_group, open_store, ) -from adata.util.names import count_duplicates, make_index_unique, unique_names #: What a group must hold to be a matrix, in each layout. _V3_KEYS = frozenset({"barcodes", "data", "indices", "indptr", "shape", "features"}) @@ -87,10 +92,6 @@ "version", ) -#: Rows of a string column handled per read when streaming one. -_STRING_STEP = 1 << 18 - -VAR_NAMES = ("gene_symbols", "gene_ids") # --------------------------------------------------------------------------- @@ -260,23 +261,6 @@ def __getitem__(self, key: str) -> Any: # import -@dataclass -class _Output: - """One AnnData to write: which features it keeps, and their annotation.""" - - label: Optional[str] - keep: Optional[np.ndarray] - var_names: List[str] - columns: Dict[str, np.ndarray] - - -def _type_counts(types: np.ndarray) -> List[Tuple[str, int]]: - """Feature types in order of first appearance, with how many of each.""" - labels, first, counts = np.unique(types, return_index=True, return_counts=True) - order = np.argsort(first) - return [(str(labels[i]), int(counts[i])) for i in order] - - class TenxImport(Converter): options = frozenset( {"--genome", "--var-names", "--make-unique", "--feature-type", @@ -291,13 +275,7 @@ def plan( dst_format: str, options: Options, ) -> ConversionPlan: - if options.var_names not in VAR_NAMES: - raise ValueError(f"--var-names must be one of: {', '.join(VAR_NAMES)}.") - if options.feature_type is not None and options.split_feature_types: - raise ValueError( - "--feature-type keeps one type and --split-feature-types writes " - "every type separately; pass one or the other." - ) + check_import_options(options) with h5py.File(source, "r") as handle: matrix = _read_matrix(handle, options.genome) @@ -324,49 +302,12 @@ def plan( keep &= genomes == options.genome types = matrix.feature_types - groups: List[Tuple[Optional[str], np.ndarray]] = [] - if options.feature_type is not None: - available = [t for t, _ in _type_counts(types)] - if options.feature_type not in available: - raise ValueError( - f"No features of type {options.feature_type!r}; this file " - f"has: {', '.join(available)}." - ) - groups.append((options.feature_type, keep & (types == options.feature_type))) - elif options.split_feature_types: - for label, _ in _type_counts(types[keep]): - groups.append((label, keep & (types == label))) - else: - groups.append((None, keep)) - - stems = unique_names([label for label, _ in groups if label is not None]) - outputs: List[Path] = [] - targets: List[_Output] = [] - warnings: List[str] = [] - for label, mask in groups: - path = output - if options.split_feature_types: - path = output.with_name(f"{output.stem}.{stems[label]}{output.suffix}") - selection = None if mask.all() else np.flatnonzero(mask) - target = self._annotate(matrix, selection, options, label) - duplicated = count_duplicates(target.var_names) - if duplicated: - where = f" in {path.name}" if options.split_feature_types else "" - warnings.append( - f"{duplicated:,} var names{where} repeat an earlier one. " - "Pass --make-unique to suffix them as anndata's " - "var_names_make_unique does." - ) - outputs.append(path) - targets.append(target) - - described = f"{matrix.n_barcodes:,} barcodes x {matrix.n_features:,} features" - if types is not None: - described += " (" + ", ".join(f"{t} {n:,}" for t, n in _type_counts(types)) + ")" - if len(outputs) > 1: - described += f", into {len(outputs)} outputs by feature type" - elif targets[0].keep is not None: - described += f", keeping {len(targets[0].keep):,} features" + outputs, targets, warnings = plan_targets( + output, matrix.names, matrix.columns, types, keep, options, noun="file" + ) + described = describe_import( + matrix.n_barcodes, matrix.n_features, types, targets + ) return ConversionPlan( source=source, @@ -383,39 +324,12 @@ def plan( }, ) - @staticmethod - def _annotate( - matrix: _Matrix, - keep: Optional[np.ndarray], - options: Options, - label: Optional[str], - ) -> _Output: - def pick(values: np.ndarray) -> np.ndarray: - return values if keep is None else values[keep] - - columns = {key: pick(values) for key, values in matrix.columns.items()} - if options.var_names == "gene_ids": - # As scanpy's read_10x_mtx does it: ids become the names, and the - # symbols move into a column where the ids were. - index = list(columns["gene_ids"]) - columns = { - ("gene_symbols" if key == "gene_ids" else key): ( - pick(matrix.names) if key == "gene_ids" else values - ) - for key, values in columns.items() - } - else: - index = list(pick(matrix.names)) - if options.make_unique: - index = make_index_unique(index) - return _Output(label=label, keep=keep, var_names=index, columns=columns) - def write( self, plan: ConversionPlan, staged: Sequence[Path], console: Console ) -> None: detail = plan.detail matrix: _Matrix = detail["matrix"] - targets: List[_Output] = detail["targets"] + targets: List[Target] = detail["targets"] with h5py.File(plan.source, "r") as handle, ExitStack() as stack: group = handle[matrix.group] @@ -428,7 +342,7 @@ def write( for root, target in zip(roots, targets): with console.status("[cyan]Writing obs and var...[/]"): _write_obs(group, root, detail["filtered_barcodes"]) - _write_var(root, target) + write_var(root, target) ensure_anndata_skeleton(root) _write_uns(root, matrix) @@ -443,18 +357,16 @@ def write( def _write_obs(group: h5py.Group, root: Any, filtered_barcodes: bool) -> None: """obs: the barcodes, streamed, and Space Ranger's `filtered_barcodes`.""" - obs = root.create_group("obs") - spec.set_encoding(obs, spec.DATAFRAME) - obs.attrs["_index"] = "_index" - barcodes = group["barcodes"] n = barcodes.shape[0] - backend = "zarr" if is_zarr_group(root) else "hdf5" - index = create_dataset(obs, "_index", shape=(n,), dtype=string_dtype_for(backend)) - for lo in range(0, n, _STRING_STEP): - hi = min(lo + _STRING_STEP, n) - index[lo:hi] = np.asarray(_strings_slice(barcodes, lo, hi), dtype=object) - spec.set_encoding(index, spec.STRING_ARRAY) + obs = write_obs_index( + root, + n, + ( + _strings_slice(barcodes, lo, min(lo + STRING_STEP, n)) + for lo in range(0, n, STRING_STEP) + ), + ) columns = [] if filtered_barcodes: @@ -467,27 +379,6 @@ def _strings_slice(dataset: Any, lo: int, hi: int) -> List[str]: return [str(v) for v in decode_str_array(np.asarray(dataset[lo:hi]))] -def _write_var(root: Any, target: _Output) -> None: - """var, with string columns encoded as anndata's writer would encode them. - - anndata turns any string column with a repeated value into a categorical - when it writes (`strings_to_categoricals`), so an h5ad scanpy wrote from - the same file holds `feature_types` and `genome` as categoricals. Doing - the same keeps the two interchangeable, and those columns a few bytes - instead of a string per feature. - """ - var = write_dataframe_header(root, "var", target.var_names, list(target.columns)) - for key, values in target.columns.items(): - if values.dtype.kind == "b": - write_dense(var, key, values) - continue - categories, codes = np.unique(values, return_inverse=True) - if len(categories) < len(values): - write_categorical(var, key, codes, list(categories)) - else: - write_string_array(var, key, values) - - def _write_uns(root: Any, matrix: _Matrix) -> None: """The run's metadata, which scanpy drops and an export puts back.""" meta = write_mapping(root["uns"], "cellranger") @@ -521,22 +412,6 @@ def _write_value(parent: Any, key: str, value: Any) -> None: # export -def _matrix_path(layer: Optional[str]) -> str: - """`--layer` as a path: X, raw/X, a layer's name, or layers/NAME.""" - if layer is None or layer in ("X", "raw/X") or layer.startswith("layers/"): - return layer or "X" - return f"layers/{layer}" - - -def _resolve(root: Any, path: str) -> Any: - node = root - for part in path.split("/"): - if not is_group(node) or part not in node: - raise KeyError(f"{path!r} not found in the store.") - node = node[part] - return node - - def _read_value(obj: Any) -> Any: """An AnnData element read back as the attribute value it came from.""" enc = spec.encoding_type(obj) @@ -552,15 +427,6 @@ def _read_value(obj: Any) -> Any: return value -@dataclass -class _Features: - ids: List[str] - names: List[str] - types: List[str] - genomes: List[str] - tags: Dict[str, List[str]] - - class TenxExport(Converter): options = frozenset( {"--layer", "--gene-ids-column", "--gene-names-column", "--feature-types-column"} @@ -574,31 +440,11 @@ def plan( dst_format: str, options: Options, ) -> ConversionPlan: - path = _matrix_path(options.layer) warnings: List[str] = [] + matrix = plan_matrix_source(source, options, warnings) with open_store(source, "r", require_anndata=False) as store: - root = store.root - matrix = describe(_resolve(root, path)) - var_path = "raw/var" if path.startswith("raw/") else "var" - var, obs = _resolve(root, var_path), _resolve(root, "obs") - obs_index, _ = resolve_index(obs, "obs") - var_index, index_name = resolve_index(var, "var") - n_obs, n_var = element_len(obs_index), element_len(var_index) - if matrix.shape != (n_obs, n_var): - raise ValueError( - f"{path} is {matrix.shape[0]} x {matrix.shape[1]}, but obs has " - f"{n_obs} and {var_path} has {n_var} entries." - ) - - features = _features(root, var, var_index, options, warnings) - attrs = _run_attrs(root) - - if matrix.dtype.kind == "f": - warnings.append( - f"{path} holds {matrix.dtype} values. Cell Ranger writes integer " - "counts, and tools that read its output may expect them: export " - "a counts layer with --layer, or recast to int32 first." - ) + attrs = _run_attrs(store.root) + path, n_obs, n_var = matrix.path, matrix.n_obs, matrix.n_var return ConversionPlan( source=source, @@ -610,19 +456,19 @@ def plan( f"from {path}" ), warnings=warnings, - detail={"path": path, "features": features, "attrs": attrs}, + detail={"path": path, "features": matrix.features, "attrs": attrs}, ) def write( self, plan: ConversionPlan, staged: Sequence[Path], console: Console ) -> None: detail = plan.detail - features: _Features = detail["features"] + features: ExportFeatures = detail["features"] with open_store(plan.source, "r", require_anndata=False) as store, h5py.File( staged[0], "w" ) as out: root = store.root - matrix = describe(_resolve(root, detail["path"])) + matrix = describe(resolve(root, detail["path"])) with console.status(f"[cyan]Copying {detail['path']}...[/]"): group = _write_matrix(matrix, out) n_obs, n_var = matrix.shape @@ -647,90 +493,10 @@ def write( out.attrs[key] = value -def _features( - root: Any, - var: Any, - var_index: Any, - options: Options, - warnings: List[str], -) -> _Features: - """Work out each feature's id, name, type and genome from var.""" - columns = dataframe_columns(var) - n = element_len(var_index) - - def column(name: str) -> List[str]: - return read_str_all(var[name], parent_group=var) - - def requested(flag: str, name: Optional[str]) -> Optional[str]: - if name is not None and name not in columns: - raise ValueError( - f"{flag} {name}: no such var column. It has: " - f"{', '.join(columns) or 'none'}." - ) - return name - - ids_col = requested("--gene-ids-column", options.gene_ids_column) - names_col = requested("--gene-names-column", options.gene_names_column) - types_col = requested("--feature-types-column", options.feature_types_column) - index = read_str_all(var_index) - - missing = [] - if ids_col is None and names_col is None and "gene_ids" not in columns and "gene_symbols" in columns: - # var_names are ids already, as `--var-names gene_ids` leaves them. - ids, names = index, column("gene_symbols") - else: - names = column(names_col) if names_col else index - if ids_col is None and "gene_ids" in columns: - ids_col = "gene_ids" - if ids_col is not None: - ids = column(ids_col) - else: - missing.append("gene ids (a gene_ids column, or --gene-ids-column)") - ids = names - if types_col is None and "feature_types" in columns: - types_col = "feature_types" - if types_col is not None: - types = column(types_col) - else: - missing.append("feature types (a feature_types column, or --feature-types-column)") - types = ["Gene Expression"] * n - - if missing: - if not options.force: - raise ValueError( - f"var has no {' and no '.join(missing)}, which a 10x matrix " - "records for every feature. Pass --force to write placeholders: " - "the names as ids, and 'Gene Expression' as the type." - ) - warnings.append(f"Writing placeholders for {' and '.join(missing)}.") - - genomes = column("genome") if "genome" in columns else [""] * n - tags: Dict[str, List[str]] = {} - for key in _recorded_tag_keys(root): - if key != "genome" and key in columns: - tags[key] = column(key) - return _Features(ids=ids, names=names, types=types, genomes=genomes, tags=tags) - - -def _cellranger_meta(root: Any) -> Optional[Any]: - """`uns/cellranger`, as an import recorded it, if there is one.""" - if "uns" not in root or "cellranger" not in root["uns"]: - return None - meta = root["uns"]["cellranger"] - return meta if is_group(meta) else None - - -def _recorded_tag_keys(root: Any) -> List[str]: - meta = _cellranger_meta(root) - if meta is None or "feature_tag_keys" not in meta: - return [] - return read_str_all(meta["feature_tag_keys"]) - - def _run_attrs(root: Any) -> Dict[str, Any]: """Root attributes: the recorded run's, with Cell Ranger's format markers.""" attrs: Dict[str, Any] = {} - meta = _cellranger_meta(root) + meta = cellranger_meta(root) if meta is not None and "attrs" in meta: for key in _RUN_ATTRS: if key in meta["attrs"] and is_dataset(meta["attrs"][key]): @@ -802,11 +568,11 @@ def _write_fixed_strings_streamed( n = element_len(column) cache: Dict[str, np.ndarray] = {} width = 1 - for lo in range(0, n, _STRING_STEP): - block = read_str_chunk(column, lo, min(lo + _STRING_STEP, n), cache, frame) + for lo in range(0, n, STRING_STEP): + block = read_str_chunk(column, lo, min(lo + STRING_STEP, n), cache, frame) width = max(width, _fixed_width(block)) dataset = _create_fixed(parent, name, n, width) - for lo in range(0, n, _STRING_STEP): - hi = min(lo + _STRING_STEP, n) + for lo in range(0, n, STRING_STEP): + hi = min(lo + STRING_STEP, n) block = read_str_chunk(column, lo, hi, cache, frame) dataset[lo:hi] = np.array([v.encode("utf-8") for v in block], dtype=f"S{width}") diff --git a/src/adata/core/csr_build.py b/src/adata/core/csr_build.py new file mode 100644 index 0000000..3c38cb7 --- /dev/null +++ b/src/adata/core/csr_build.py @@ -0,0 +1,178 @@ +"""Building a compressed sparse matrix from entries that arrive out of order. + +A streaming transpose and an MTX import have the same problem: entries come +in an order unrelated to the rows they belong to, and the output cannot be +written until each row is complete. Both solve it the same way, with memory +set by the bucket size rather than by nnz: + +1. count the entries per output row (`count_keys`), which gives `indptr`; +2. scatter the entries into buckets, each covering a slice of the output rows + and holding about as many entries as one read (`bucket_bounds`, + `BucketSorter.add`), kept on disk in a scratch group (`scratch_group`); +3. read one bucket at a time, sort it, and append it to the output + (`BucketSorter.runs`). +""" + +from __future__ import annotations + +import os +import tempfile +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Iterator, List, Optional, Tuple + +import h5py +import numpy as np + +from adata.storage import is_hdf5_group + +#: Smallest bucket worth making. Below this the per-bucket overhead +#: dominates and the extra passes buy nothing. +MIN_BUCKET_ENTRIES = 1 << 12 + + +def count_keys(counts: np.ndarray, keys: np.ndarray) -> None: + """Add one to `counts[k]` for every `k` in `keys`. + + `np.bincount(keys, minlength=len(counts))` costs the length of `counts` + on every call, which for a block of a few thousand entries against + millions of rows is almost all waste -- and quadratic over a matrix read + in many blocks. Count over the span the block actually covers instead, + or by sorting when that span is sparse. + """ + if keys.size == 0: + return + lo, hi = int(keys.min()), int(keys.max()) + span = hi - lo + 1 + if span <= 4 * keys.size: + counts[lo : hi + 1] += np.bincount(keys - lo, minlength=span) + else: + unique, seen = np.unique(keys, return_counts=True) + counts[unique] += seen + + +def bucket_bounds( + indptr: np.ndarray, *, chunk: int, bucket_entries: Optional[int] = None +) -> np.ndarray: + """Row bounds splitting the output into buckets of about equal nnz. + + One bucket per slice of the output rows, holding about as many nonzeros + as one read. Tying this to `chunk` rather than a constant is what makes + the peak follow the setting the caller chose: with a fixed bucket size, + anything below it went into a single bucket and the "streaming" path + quietly held the whole matrix. An explicit `bucket_entries` is honoured + as given; the floor applies only to the value derived from `chunk`, + where it stops a tiny chunk producing thousands of buckets whose + overhead outweighs the saving. + + Split by nonzero count, not by coordinate. Equal-width bounds put nearly + everything in one bucket whenever the matrix is skewed -- and single-cell + matrices are: a handful of genes carry most of the counts. Measured on + one such matrix, the largest of three equal-width buckets held 88% of + the entries, so the bucket, not the chunk, set the peak. + """ + n_rows = int(indptr.size - 1) + nnz = int(indptr[-1]) if indptr.size else 0 + per_bucket = ( + max(1, int(bucket_entries)) + if bucket_entries + else max(MIN_BUCKET_ENTRIES, int(chunk)) + ) + n_buckets = max(1, int(np.ceil(nnz / per_bucket))) if nnz else 1 + n_buckets = min(n_buckets, n_rows) or 1 + + targets = np.linspace(0, nnz, n_buckets + 1)[1:-1] + bounds = np.concatenate(( + [0], + np.searchsorted(indptr, targets, side="left").astype(np.int64), + [n_rows], + )) + return np.unique(bounds) + + +@contextmanager +def scratch_group(dst_parent: Any, name: str) -> Iterator[Any]: + """A group for intermediate buckets, gone afterwards. + + HDF5 never returns the space of a deleted dataset to the file, so scratch + written into an HDF5 output stayed in it as dead space: a streamed + transpose came out three times the size of an in-memory one. On HDF5 the + scratch is a separate file beside the output instead. Deleting from a + Zarr store frees the space, so there it stays in place. + """ + if not is_hdf5_group(dst_parent): + group = dst_parent.create_group(name) + try: + yield group + finally: + del dst_parent[name] + return + + # Beside the output rather than in the system temp directory, which on + # a cluster node is often far too small for a matrix's worth of entries. + where = Path(str(dst_parent.file.filename)).parent + fd, path = tempfile.mkstemp( + prefix=".adata-", suffix=".transpose-scratch.h5", + dir=str(where) if where.is_dir() else None, + ) + os.close(fd) + try: + with h5py.File(path, "w") as handle: + yield handle.create_group(name) + finally: + os.unlink(path) + + +class BucketSorter: + """Scatter (row, column, value) entries into row-range buckets on disk. + + `add` takes entries in any order; `runs` gives them back one bucket at a + time, sorted by row and then column, buckets in row order. Concatenated, + the runs are the matrix in CSR order. + """ + + def __init__(self, scratch: Any, bounds: np.ndarray, value_dtype: Any) -> None: + # Looked up through the module rather than imported by name, so a + # test that watches `_append` sees every bucket write. + from adata.core import subset + + self._subset = subset + self.bounds = np.asarray(bounds, dtype=np.int64) + self.n_buckets = len(self.bounds) - 1 + self.buckets: List[Tuple[Any, Any, Any]] = [ + ( + subset._growable(scratch, f"major{b}", np.int64), + subset._growable(scratch, f"minor{b}", np.int64), + subset._growable(scratch, f"value{b}", value_dtype), + ) + for b in range(self.n_buckets) + ] + + def add(self, major: np.ndarray, minor: np.ndarray, values: np.ndarray) -> None: + if major.size == 0: + return + which = np.clip( + np.searchsorted(self.bounds, major, side="right") - 1, + 0, + self.n_buckets - 1, + ) + # One stable sort by bucket, then a slice per bucket the block + # touches, rather than a full-block comparison for every bucket. + order = np.argsort(which, kind="stable") + starts = np.searchsorted(which[order], np.arange(self.n_buckets + 1)) + for b in np.flatnonzero(np.diff(starts)): + pick = order[starts[b] : starts[b + 1]] + majors, minors, vals = self.buckets[b] + self._subset._append(majors, major[pick]) + self._subset._append(minors, minor[pick]) + self._subset._append(vals, values[pick]) + + def runs(self) -> Iterator[Tuple[np.ndarray, np.ndarray, np.ndarray]]: + for majors, minors, vals in self.buckets: + if majors.shape[0] == 0: + continue + major = np.asarray(majors[...], dtype=np.int64) + minor = np.asarray(minors[...], dtype=np.int64) + values = np.asarray(vals[...]) + order = np.lexsort((minor, major)) + yield major[order], minor[order], values[order] diff --git a/src/adata/core/recast.py b/src/adata/core/recast.py index 850dab8..7c227fa 100644 --- a/src/adata/core/recast.py +++ b/src/adata/core/recast.py @@ -23,19 +23,20 @@ from __future__ import annotations -import os -import tempfile -from contextlib import contextmanager from dataclasses import dataclass -from pathlib import Path -from typing import Any, Dict, Iterator, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple -import h5py import numpy as np from rich.console import Console from adata.elements import spec from adata.elements.write import set_shape_attr +from adata.core.csr_build import ( + BucketSorter, + bucket_bounds, + count_keys, + scratch_group, +) from adata.core.subset import _growable_like, _sparse_dataset from adata.storage import ( copy_attrs, @@ -44,7 +45,6 @@ fit_chunks, is_dataset, is_group, - is_hdf5_group, is_zarr_group, with_compression, zarr_format_of, @@ -74,9 +74,6 @@ #: Values per read while streaming `data`/`indices`. DEFAULT_CHUNK = 1 << 20 -#: Smallest bucket worth making during a streaming transpose. Below this -#: the per-bucket overhead dominates and the extra passes buy nothing. -MIN_BUCKET_ENTRIES = 1 << 12 def parse_dtype(text: str, *, allowed: Tuple[str, ...] = DTYPES) -> np.dtype: @@ -414,37 +411,8 @@ def transpose_sparse_in_memory( ) -@contextmanager -def _scratch_group(dst_parent: Any, name: str) -> Iterator[Any]: - """A group for the transpose's intermediate buckets, gone afterwards. - - HDF5 never returns the space of a deleted dataset to the file, so scratch - written into an HDF5 output stayed in it as dead space: a streamed - transpose came out three times the size of an in-memory one. On HDF5 the - scratch is a separate file beside the output instead. Deleting from a - Zarr store frees the space, so there it stays in place. - """ - if not is_hdf5_group(dst_parent): - group = dst_parent.create_group(name) - try: - yield group - finally: - del dst_parent[name] - return - - # Beside the output rather than in the system temp directory, which on - # a cluster node is often far too small for a matrix's worth of entries. - where = Path(str(dst_parent.file.filename)).parent - fd, path = tempfile.mkstemp( - prefix=".adata-", suffix=".transpose-scratch.h5", - dir=str(where) if where.is_dir() else None, - ) - os.close(fd) - try: - with h5py.File(path, "w") as handle: - yield handle.create_group(name) - finally: - os.unlink(path) +#: Kept under its old name: the scratch group moved to `core.csr_build`. +_scratch_group = scratch_group def transpose_sparse_streaming( @@ -476,8 +444,6 @@ def transpose_sparse_streaming( Costs about two extra passes over nnz, which is the price of not holding the matrix. `--in-memory` is there for when you would rather pay in RAM. """ - from adata.core.subset import _append, _growable - target = ( spec.CSC_MATRIX if src.kind == spec.CSR_MATRIX else spec.CSR_MATRIX ) @@ -492,39 +458,10 @@ def transpose_sparse_streaming( counts = np.zeros(n_major_out, dtype=np.int64) for start in range(0, nnz, chunk): block = np.asarray(source_indices[start : min(start + chunk, nnz)], dtype=np.int64) - counts += np.bincount(block, minlength=n_major_out) + count_keys(counts, block) out_indptr = np.concatenate(([0], np.cumsum(counts))) - # One bucket per slice of the output major axis, holding about as many - # nonzeros as one read. Tying this to `chunk` rather than a constant is - # what makes the peak follow the setting the caller chose: with a fixed - # bucket size, anything below it went into a single bucket and the - # "streaming" path quietly held the whole matrix. - # An explicit request is honoured as given; the floor applies only to - # the value derived from `chunk`, where it stops a tiny chunk producing - # thousands of buckets whose overhead outweighs the saving. - per_bucket = ( - max(1, int(bucket_entries)) - if bucket_entries - else max(MIN_BUCKET_ENTRIES, int(chunk)) - ) - n_buckets = max(1, int(np.ceil(nnz / per_bucket))) if nnz else 1 - n_buckets = min(n_buckets, n_major_out) or 1 - - # Split by nonzero count, not by coordinate. Equal-width bounds put - # nearly everything in one bucket whenever the matrix is skewed -- and - # single-cell matrices are: a handful of genes carry most of the - # counts. Measured on one such matrix, the largest of three equal-width - # buckets held 88% of the entries, so the bucket, not the chunk, set - # the peak. `counts` is already to hand from pass 1. - cumulative = out_indptr - targets = np.linspace(0, nnz, n_buckets + 1)[1:-1] - bounds = np.concatenate(( - [0], - np.searchsorted(cumulative, targets, side="left").astype(np.int64), - [n_major_out], - )) - bounds = np.unique(bounds) + bounds = bucket_bounds(out_indptr, chunk=chunk, bucket_entries=bucket_entries) n_buckets = len(bounds) - 1 if console is not None and n_buckets > 1: console.print( @@ -542,15 +479,8 @@ def transpose_sparse_streaming( ) written = 0 - with _scratch_group(dst_parent, f"__{name}_transpose_scratch__") as scratch: - buckets = [ - ( - _growable(scratch, f"major{b}", np.int64), - _growable(scratch, f"minor{b}", np.int64), - _growable(scratch, f"value{b}", src.dtype), - ) - for b in range(n_buckets) - ] + with scratch_group(dst_parent, f"__{name}_transpose_scratch__") as scratch: + sorter = BucketSorter(scratch, bounds, src.dtype) # Pass 2: scatter the input into buckets, one input block at a time. # Step over the major axis in blocks holding about `chunk` nonzeros, @@ -568,29 +498,16 @@ def transpose_sparse_streaming( major = np.repeat( np.arange(lo, hi, dtype=np.int64), np.diff(indptr[lo : hi + 1]) ) - which = np.clip(np.searchsorted(bounds, minor, side="right") - 1, 0, n_buckets - 1) - for b in range(n_buckets): - pick = which == b - if not pick.any(): - continue - _append(buckets[b][0], minor[pick]) - _append(buckets[b][1], major[pick]) - _append(buckets[b][2], values[pick]) + # The input's minor axis is the output's major one. + sorter.add(minor, major, values) # Pass 3: each bucket in turn, sorted into output order. - for b in range(n_buckets): - new_major, new_minor, value = buckets[b] - if new_major.shape[0] == 0: - continue - majors = np.asarray(new_major[...], dtype=np.int64) - minors = np.asarray(new_minor[...], dtype=np.int64) - values = np.asarray(value[...]) - order = np.lexsort((minors, majors)) - count = order.size - out_indices[written : written + count] = minors[order].astype( + for _, new_minor, values in sorter.runs(): + count = new_minor.size + out_indices[written : written + count] = new_minor.astype( index_dtype, copy=False ) - out_data[written : written + count] = values[order].astype( + out_data[written : written + count] = values.astype( data_dtype, copy=False ) written += count diff --git a/src/adata/storage/__init__.py b/src/adata/storage/__init__.py index 148bd9f..c2c0ee8 100644 --- a/src/adata/storage/__init__.py +++ b/src/adata/storage/__init__.py @@ -15,6 +15,8 @@ import numpy as np +from adata.util.path import is_zarr_path # noqa: F401 - re-exported + ROOT_ENCODING_TYPE = "anndata" ROOT_ENCODING_VERSION = "0.1.0" @@ -92,16 +94,6 @@ def is_dataset(obj: Any) -> bool: return is_hdf5_dataset(obj) or is_zarr_array(obj) -def is_zarr_path(path: Path) -> bool: - if not path.exists() or not path.is_dir(): - return False - if (path / "zarr.json").exists(): - return True - if (path / ".zgroup").exists() or (path / ".zattrs").exists(): - return True - return False - - def detect_backend(path: Path) -> str: if path.exists(): if path.is_dir(): diff --git a/src/adata/util/path.py b/src/adata/util/path.py index 24d3832..7cfb103 100644 --- a/src/adata/util/path.py +++ b/src/adata/util/path.py @@ -17,3 +17,15 @@ def same_path(left: Path, right: Path) -> bool: return Path(left).resolve() == Path(right).resolve() except OSError: # pragma: no cover - unresolvable path return Path(left).absolute() == Path(right).absolute() + + +def is_zarr_path(path: Path) -> bool: + """Is this an existing Zarr store (v2 or v3), judged by its metadata files?""" + path = Path(path) + if not path.exists() or not path.is_dir(): + return False + if (path / "zarr.json").exists(): + return True + if (path / ".zgroup").exists() or (path / ".zattrs").exists(): + return True + return False diff --git a/src/adata/util/staging.py b/src/adata/util/staging.py index df2f7c7..779c3ac 100644 --- a/src/adata/util/staging.py +++ b/src/adata/util/staging.py @@ -19,6 +19,8 @@ from pathlib import Path from typing import Iterator, List, Sequence, Tuple +from adata.util.path import is_zarr_path + def _sibling(path: Path, tag: str) -> Path: """`out.zarr` -> `out..zarr`, keeping the suffix that picks the backend.""" @@ -27,9 +29,10 @@ def _sibling(path: Path, tag: str) -> Path: def _temp_for(path: Path) -> Path: temp = _sibling(path, "adata-tmp") - if path.is_dir() and path.suffix != ".zarr": + if path.suffix != ".zarr" and is_zarr_path(path): # A Zarr store named without the suffix. A temporary named the same - # way would not exist yet, and so be opened as HDF5. + # way would not exist yet, and so be opened as HDF5. Only a Zarr + # store: any other directory, such as an MTX folder, keeps its name. temp = temp.with_name(f"{temp.name}.zarr") return temp diff --git a/tests/test_csr_build.py b/tests/test_csr_build.py new file mode 100644 index 0000000..47b3731 --- /dev/null +++ b/tests/test_csr_build.py @@ -0,0 +1,96 @@ +"""Tests for `adata.core.csr_build`, the bucket sort behind streaming transposes. + +The streaming transpose is tested end to end in test_recast.py; these pin the +pieces directly, because the MTX import builds on them with inputs a +transpose never produces: blocks whose keys cover a narrow span, and entries +in no particular order. +""" + +from __future__ import annotations + +import h5py +import numpy as np +import pytest + +from adata.core.csr_build import BucketSorter, bucket_bounds, count_keys, scratch_group + + +@pytest.mark.parametrize("span", [10, 900_000], ids=["dense span", "sparse span"]) +def test_count_keys_matches_bincount(span): + rng = np.random.default_rng(0) + counts = np.zeros(1_000_000, dtype=np.int64) + expected = np.zeros_like(counts) + for _ in range(5): + keys = rng.integers(500, 500 + span, size=1000) + count_keys(counts, keys) + expected += np.bincount(keys, minlength=counts.size) + assert np.array_equal(counts, expected) + + +def test_count_keys_ignores_an_empty_block(): + counts = np.zeros(4, dtype=np.int64) + count_keys(counts, np.array([], dtype=np.int64)) + assert not counts.any() + + +def test_bucket_bounds_cover_every_row_and_balance_nonzeros(): + # One heavy row among many light ones: bounds follow the counts. + counts = np.ones(1000, dtype=np.int64) + counts[10] = 5000 + indptr = np.concatenate(([0], np.cumsum(counts))) + + bounds = bucket_bounds(indptr, chunk=1000, bucket_entries=1000) + + assert bounds[0] == 0 and bounds[-1] == 1000 + assert np.all(np.diff(bounds) > 0) + per_bucket = np.diff(indptr[bounds]) + assert per_bucket.max() <= 5000 + 1000 + + +@pytest.mark.parametrize("order", ["sorted", "reversed", "shuffled"]) +def test_the_runs_are_the_entries_in_csr_order_whatever_order_they_came_in( + tmp_path, order +): + rng = np.random.default_rng(1) + n_rows, n_cols, nnz = 300, 50, 4000 + flat = rng.choice(n_rows * n_cols, size=nnz, replace=False) + rows, cols = flat // n_cols, flat % n_cols + values = rng.integers(1, 100, size=nnz).astype(np.int32) + + arrival = { + "sorted": np.lexsort((cols, rows)), + "reversed": np.lexsort((cols, rows))[::-1], + "shuffled": rng.permutation(nnz), + }[order] + + counts = np.bincount(rows, minlength=n_rows) + indptr = np.concatenate(([0], np.cumsum(counts))) + bounds = bucket_bounds(indptr, chunk=500, bucket_entries=500) + + with h5py.File(tmp_path / "out.h5", "w") as handle: + with scratch_group(handle, "scratch") as scratch: + sorter = BucketSorter(scratch, bounds, np.int32) + for block in np.array_split(arrival, 7): + sorter.add(rows[block], cols[block], values[block]) + runs = list(sorter.runs()) + + assert len(runs) > 1, "expected several buckets" + got_rows = np.concatenate([r for r, _, _ in runs]) + got_cols = np.concatenate([c for _, c, _ in runs]) + got_values = np.concatenate([v for _, _, v in runs]) + + expected = np.lexsort((cols, rows)) + assert np.array_equal(got_rows, rows[expected]) + assert np.array_equal(got_cols, cols[expected]) + assert np.array_equal(got_values, values[expected]) + + +def test_hdf5_scratch_is_a_separate_file_removed_afterwards(tmp_path): + with h5py.File(tmp_path / "out.h5", "w") as handle: + with scratch_group(handle, "scratch") as scratch: + assert scratch.file.filename != handle.filename + during = sorted(p.name for p in tmp_path.iterdir()) + after = sorted(p.name for p in tmp_path.iterdir()) + + assert any("scratch" in name for name in during) + assert after == ["out.h5"] diff --git a/tests/test_staging.py b/tests/test_staging.py index a781822..c45d863 100644 --- a/tests/test_staging.py +++ b/tests/test_staging.py @@ -111,3 +111,18 @@ def test_a_zarr_store_named_without_the_suffix_stays_zarr(tmp_path): assert (out / "zarr.json").read_text() == '{"new": true}' assert _names(tmp_path) == ["store"] + + +def test_a_directory_that_is_not_zarr_keeps_its_name(tmp_path): + """An MTX output folder is a directory too, and must not be staged as Zarr.""" + out = tmp_path / "filtered_feature_bc_matrix" + out.mkdir() + (out / "matrix.mtx.gz").write_bytes(b"old") + + with staged_outputs([out]) as (temp,): + assert temp.name == "filtered_feature_bc_matrix.adata-tmp" + temp.mkdir() + (temp / "matrix.mtx.gz").write_bytes(b"new") + + assert (out / "matrix.mtx.gz").read_bytes() == b"new" + assert _names(tmp_path) == ["filtered_feature_bc_matrix"]