diff --git a/docs/COMMANDS.md b/docs/COMMANDS.md index d2ad5dc..0b3310e 100644 --- a/docs/COMMANDS.md +++ b/docs/COMMANDS.md @@ -678,6 +678,13 @@ obs columns keep their dtypes: categoricals union their category sets, nullable columns keep their masks, and a column missing from one input is padded rather than dropped. +obsm entries present in every input are stacked: arrays as they are, and +dataframes (a SCope import's `RegulonsAUC`) as obs is, their columns +following `--join`. 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. diff --git a/docs/TESTING.md b/docs/TESTING.md index a22bb8a..ed01ef4 100644 --- a/docs/TESTING.md +++ b/docs/TESTING.md @@ -35,6 +35,7 @@ separate jobs. | `test_hdf5.py` | One dataset of a plain HDF5 file: the listing without `--key`, `read_hdf`'s `row_names` / `col_names` and other names, `--transpose` CSR and dense, slabs of chunks, `--from hdf5` on `.h5` | | `test_cistopic.py` | pycisTopic pickles read through the allowlist: a CistopicObject's matrices, frames, model and projections, an imputed-features pickle, a pickle naming `os.system` refused unrun, the missing-extra hint | | `test_pandas_frame.py` | Each pandas dtype written as anndata writes it, on every backend | +| `test_concat_axis_elements.py` | `concat` on obsm dataframes and varm, against `anndata.concat` for each join and merge strategy | | `test_scope.py` | SCope looms as pySCENIC's `export2loom` writes them: regulons as dataframes, embeddings and clusterings named from `MetaData`, plain or compressed `MetaData`, region looms, R's logicals and factors, renamed fields, small blocks | | `test_table.py` | The streaming CSV reader against pandas: column types, quoting, row subsets, blocks, refusals | | `test_mtx_codec.py`, `test_csr_build.py` | The streaming Matrix Market reader and writer against scipy, and the bucket sort they and `recast` share | diff --git a/src/adata/core/concat.py b/src/adata/core/concat.py index 5d43ec2..0a49e7a 100644 --- a/src/adata/core/concat.py +++ b/src/adata/core/concat.py @@ -5,13 +5,15 @@ peak memory is set by the block size rather than by the total size of the inputs. -Columns of the concatenated axis follow `join`; elements of the alternative -axis and of uns follow a `merge` strategy, since there is no single correct -way to reconcile values that differ between inputs. +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. """ from __future__ import annotations +import hashlib from pathlib import Path from typing import Any, Dict, List, Optional, Sequence, Tuple @@ -23,6 +25,7 @@ from adata.elements import spec from adata.elements.read import ( dataframe_columns, + decode_str_array, element_len, read_str_all, read_categories, @@ -39,10 +42,15 @@ write_mapping, write_string_array, ) +from adata.elements.strings import is_string_dtype from adata.storage import ( copy_tree, create_dataset, + dataset_create_kwargs, + fit_chunks, is_dataset, + is_zarr_array, + is_zarr_group, Store, is_group, open_store, @@ -211,8 +219,6 @@ def _column_kind(obj: Any) -> str: if enc in spec.MASKED_TYPES or (is_group(obj) and "values" in obj): return "masked" if is_dataset(obj): - from adata.elements.strings import is_string_dtype - return "string" if is_string_dtype(obj.dtype) else "numeric" return "other" @@ -689,10 +695,21 @@ def _concat_obsm( n_rows: int, chunk_rows: int, console: Console, + *, + lengths: Optional[List[int]] = None, + join: str = "inner", + target_obs: Optional[List[str]] = None, ) -> bool: - """Concatenate an obsm entry row-wise; its columns are not an axis.""" + """Concatenate an obsm entry row-wise; its columns are not an axis. + + 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) + return True if not all(is_dataset(s) for s in sources): - console.print(f"[yellow]Skipping obsm/{name}: not a dense array in every input[/]") + console.print(f"[yellow]Skipping obsm/{name}: not an array, or not a dataframe, in every input[/]") return False widths = {tuple(s.shape[1:]) for s in sources} @@ -718,6 +735,176 @@ def _concat_obsm( return True +def _concat_frame( + dst_parent: Any, + name: str, + sources: List[Any], + 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] + 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] + else: + names = _index_union(per_input) + 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], + lengths, console) + ] + set_column_order(frame, written) + + +# --------------------------------------------------------------------------- +# varm + + +def _aligned_rows(obj: Any, rows: np.ndarray) -> np.ndarray: + """`obj[rows]`, for target rows in any order, one block at a time. + + Rows in order (an inner join keeps the first input's) are read as the + one slab they span; rows out of order by index, which both backends + take only sorted, then put back in the target's order. + """ + order = np.argsort(rows, kind="stable") + wanted = rows[order] + lo, hi = int(wanted[0]), int(wanted[-1]) + 1 + if hi - lo <= 2 * len(rows): + block = np.asarray(obj[lo:hi, ...])[wanted - lo] + elif is_zarr_array(obj): + block = np.asarray(obj.oindex[(wanted,) + (slice(None),) * (obj.ndim - 1)]) + else: + block = np.asarray(obj[wanted, ...]) + out = np.empty_like(block) + out[order] = block + # h5py writes text only from str objects, not numpy's ` None: + """Feed values to `hasher` exactly: bytes for numbers, so NaN equals NaN, + and the text itself for strings, never object pointers.""" + if values.dtype.kind in "OSUT": + for text in decode_str_array(values).reshape(-1).tolist(): + data = text.encode("utf-8", "surrogatepass") + hasher.update(len(data).to_bytes(8, "little") + data) + else: + hasher.update(np.ascontiguousarray(values).tobytes()) + + +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. + """ + 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()) + if spec.encoding_type(obj) != spec.DATAFRAME: + return _Incomparable() + columns = [] + for name in dataframe_columns(obj, resolve_index(obj, "var")[1]): + column = obj[name] + kind = _column_kind(column) + hasher = hashlib.blake2b(digest_size=32) + if kind == "numeric": + values = np.asarray(column[...])[where] + _digest(hasher, values) + columns.append((name, str(values.dtype), hasher.digest())) + elif kind == "categorical": + categories = np.asarray([str(c) for c in read_categories(column)] + [""], dtype=object) + codes = np.asarray(column["codes"][...], dtype=np.int64)[where] + _digest(hasher, codes < 0) + _digest(hasher, categories[codes]) + columns.append((name, "categorical", hasher.digest())) + elif kind == "masked": + _digest(hasher, np.asarray(column["mask"][...], dtype=bool)[where]) + _digest(hasher, np.asarray(column["values"][...])[where]) + columns.append((name, "masked", hasher.digest())) + else: + _digest(hasher, np.asarray(read_str_all(column), dtype=object)[where]) + columns.append((name, "text", hasher.digest())) + return ("dataframe", tuple(columns)) + + +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) + 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 _write_varm( + dst: Any, + roots: List[Any], + var_names: List[List[str]], + target_var: List[str], + merge: Optional[str], + console: Console, + chunk_rows: int = 1024, +) -> None: + """varm entries the merge strategy keeps, on the target var. + + As var's columns: an input lacking the entry, or some of the target's + genes (an outer join), counts as missing; `same` and `unique` compare + 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]) + if not names: + return + if merge is None: + console.print( + "[yellow]Dropping varm: pass --merge to keep the entries the inputs agree on.[/]" + ) + return + varm = write_mapping(dst, "varm") + positions = [_column_map(target_var, names_i) for names_i in var_names] + 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: + continue + source = next(i for i, v in enumerate(aligned) if v is not _MISSING) + obj, take = roots[source]["varm"][name], positions[source] + if is_dataset(obj): + _write_varm_array(varm, name, obj, take, chunk_rows) + elif spec.encoding_type(obj) == spec.DATAFRAME: + frame = write_dataframe_header(varm, name, target_var, []) + columns = dataframe_columns(obj, resolve_index(obj, "var")[1]) + for column in columns: + _write_var_column(frame, column, obj[column], take) + set_column_order(frame, columns) + else: + console.print(f"[yellow]Skipping varm/{name}: neither an array nor a dataframe[/]") + + # --------------------------------------------------------------------------- # orchestration @@ -945,7 +1132,9 @@ def concat_on_disk( fill_value, console, ) - _write_obsm(dst, roots, n_obs, chunk_rows, console) + _write_obsm(dst, roots, 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) 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): @@ -1165,7 +1354,15 @@ def _write_matrices( def _write_obsm( - dst: Any, roots: List[Any], n_obs: int, chunk_rows: int, console: Console + dst: Any, + roots: List[Any], + n_obs: int, + chunk_rows: int, + console: Console, + *, + lengths: Optional[List[int]] = None, + join: str = "inner", + target_obs: Optional[List[str]] = None, ) -> None: """Concatenate obsm entries shared by all inputs.""" names = _index_union( @@ -1182,5 +1379,6 @@ def _write_obsm( ) continue _concat_obsm( - obsm, name, [r["obsm"][name] for r in roots], n_obs, chunk_rows, console + obsm, name, [r["obsm"][name] for r in roots], n_obs, chunk_rows, console, + lengths=lengths, join=join, target_obs=target_obs, ) diff --git a/tests/test_concat_axis_elements.py b/tests/test_concat_axis_elements.py new file mode 100644 index 0000000..a17654c --- /dev/null +++ b/tests/test_concat_axis_elements.py @@ -0,0 +1,194 @@ +"""`concat` on what a SCope or pycisTopic import leaves in obsm and varm: +dataframes in obsm (`RegulonsAUC`), and varm (`Regulons`, `topic_region`). + +anndata's own `concat` is the reference: our output, read back, must equal +`anndata.concat(..., join=..., merge=...)` on the same inputs. +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +from rich.console import Console + +from adata.core.concat import concat_on_disk + +ad = pytest.importorskip("anndata") +pd = pytest.importorskip("pandas") + +QUIET = Console(quiet=True, record=True) + + +def _store(tmp_path: Path, tag: str, n: int, genes, *, seed: int, regulons=("ATF3_(+)", "SOX2_(+)"), + shift: float = 0.0): + rng = np.random.default_rng(seed) + cells = [f"{tag}{i}" for i in range(n)] + obj = ad.AnnData( + X=rng.poisson(1, (n, len(genes))).astype(np.float32), + obs=pd.DataFrame(index=cells), + var=pd.DataFrame(index=list(genes)), + ) + obj.obsm["RegulonsAUC"] = pd.DataFrame(rng.random((n, len(regulons))), index=cells, columns=list(regulons)) + obj.obsm["X_umap"] = rng.normal(size=(n, 2)) + membership = (np.arange(len(genes))[:, None] + np.arange(len(regulons))) % 2 + obj.varm["Regulons"] = pd.DataFrame(membership, index=list(genes), columns=list(regulons)) + obj.varm["topic_region"] = np.arange(len(genes) * 3, dtype=np.float64).reshape(len(genes), 3) + shift + path = tmp_path / f"{tag}.h5ad" + obj.write_h5ad(path) + return path + + +def _concat(tmp_path, paths, **kw): + out = tmp_path / f"out{len(list(tmp_path.iterdir()))}.h5ad" + concat_on_disk(paths, out, QUIET, **kw) + return ad.read_h5ad(out) + + +@pytest.mark.parametrize("join", ["inner", "outer"]) +def test_obsm_dataframes_concatenate_as_anndata_concatenates_them(tmp_path, join): + genes = [f"g{i}" for i in range(6)] + a = _store(tmp_path, "a", 5, genes, seed=0) + b = _store(tmp_path, "b", 4, genes, seed=1, regulons=("ATF3_(+)", "PAX6_(-)")) + got = _concat(tmp_path, [a, b], join=join) + ref = ad.concat([ad.read_h5ad(a), ad.read_h5ad(b)], join=join) + frame, expected = got.obsm["RegulonsAUC"], ref.obsm["RegulonsAUC"] + assert isinstance(frame, pd.DataFrame) + assert list(frame.index) == list(got.obs_names) + assert sorted(frame.columns) == sorted(expected.columns) + pd.testing.assert_frame_equal(frame[expected.columns], expected, check_dtype=False) + assert np.array_equal(got.obsm["X_umap"], ref.obsm["X_umap"]) + + +@pytest.mark.parametrize("merge", ["same", "unique", "first", "only"]) +def test_varm_follows_merge_as_anndata_does(tmp_path, merge): + """varm was never written: concat dropped it whatever --merge said.""" + genes = [f"g{i}" for i in range(6)] + a = _store(tmp_path, "a", 5, genes, seed=0) + b = _store(tmp_path, "b", 4, genes[::-1], seed=1, shift=1.0) # topic_region differs + got = _concat(tmp_path, [a, b], merge=merge) + ref = ad.concat([ad.read_h5ad(a), ad.read_h5ad(b)], merge=merge) + assert sorted(got.varm.keys()) == sorted(ref.varm.keys()) + if merge == "first": + assert sorted(got.varm.keys()) == ["Regulons", "topic_region"] + for key in ref.varm: + mine, theirs = got.varm[key], ref.varm[key] + if isinstance(theirs, pd.DataFrame): + pd.testing.assert_frame_equal(mine, theirs, check_dtype=False) + else: + assert np.array_equal(mine, theirs) + + +def test_varm_is_dropped_without_merge_with_a_note(tmp_path): + genes = [f"g{i}" for i in range(6)] + a, b = _store(tmp_path, "a", 5, genes, seed=0), _store(tmp_path, "b", 4, genes, seed=1) + console = Console(record=True, width=200) + out = tmp_path / "plain.h5ad" + concat_on_disk([a, b], out, console) + assert "Dropping varm: pass --merge" in console.export_text() + assert len(ad.read_h5ad(out).varm) == 0 + + +def test_an_outer_join_drops_varm_an_input_cannot_cover(tmp_path): + """With genes one input lacks, that input's varm has no rows for them: + it counts as missing, so `same` drops the entry and `first` keeps the + input that covers every gene, if any.""" + a = _store(tmp_path, "a", 5, [f"g{i}" for i in range(6)], seed=0) + b = _store(tmp_path, "b", 4, [f"g{i}" for i in range(3, 8)], seed=1) + got = _concat(tmp_path, [a, b], join="outer", merge="same") + assert len(got.varm) == 0 + got = _concat(tmp_path, [a, b], join="outer", merge="first") + assert len(got.varm) == 0, "neither input has every gene of the union" + + +def _varm_store(path: Path, genes, entries, *, compress: bool = False): + obj = ad.AnnData( + X=np.zeros((2, len(genes)), dtype=np.float32), + obs=pd.DataFrame(index=[f"{path.stem}{i}" for i in range(2)]), + var=pd.DataFrame(index=list(genes)), + ) + for key, value in entries.items(): + obj.varm[key] = value + if path.suffix == ".zarr": + obj.write_zarr(path) + else: + obj.write_h5ad(path, compression="gzip" if compress else None) + return path + + +@pytest.mark.parametrize("ext", [".h5ad", ".zarr"]) +def test_varm_on_genes_in_another_order_is_compared_and_written_a_block_at_a_time(tmp_path, ext): + """varm was read whole, every input's at once, to compare it: memory grew + with the entry and the number of inputs rather than with --chunk.""" + genes = [f"g{i}" for i in range(40)] + shuffled = list(np.random.default_rng(0).permutation(genes)) + values = np.arange(40 * 3, dtype=np.float64).reshape(40, 3) + by_gene = dict(zip(genes, values)) + a = _varm_store(tmp_path / f"a{ext}", genes, {"topic_region": values}) + b = _varm_store(tmp_path / f"b{ext}", shuffled, + {"topic_region": np.stack([by_gene[g] for g in shuffled])}) + c = _varm_store(tmp_path / f"c{ext}", shuffled, + {"topic_region": np.stack([by_gene[g] for g in shuffled]) + 1}) + for inputs, kept in (([a, b], True), ([a, c], False)): + out = tmp_path / f"out{len(list(tmp_path.iterdir()))}{ext}" + concat_on_disk(inputs, out, QUIET, merge="same", chunk_rows=4) + got = ad.read_zarr(out) if ext == ".zarr" else ad.read_h5ad(out) + assert ("topic_region" in got.varm) is kept + if kept: + assert np.array_equal(got.varm["topic_region"], values) + # The first input has no varm, so `first` takes c's, onto the first's genes. + bare = _varm_store(tmp_path / f"bare{ext}", genes, {}) + out = tmp_path / f"first{ext}" + concat_on_disk([bare, c], out, QUIET, merge="first", chunk_rows=4) + got = ad.read_zarr(out) if ext == ".zarr" else ad.read_h5ad(out) + assert np.array_equal(got.varm["topic_region"], values + 1) + + +def test_varm_text_arrays_are_compared_by_their_text(tmp_path): + """Text arrays were compared by their bytes, which for strings read from + HDF5 are pointers to Python objects: equal text compared unequal.""" + genes = [f"g{i}" for i in range(5)] + labels = np.asarray([[f"t{i}", f"u{i}"] for i in range(5)], dtype=object) + a = _varm_store(tmp_path / "a.h5ad", genes, {"labels": labels}) + b = _varm_store(tmp_path / "b.h5ad", genes, {"labels": labels.copy()}) + got = _concat(tmp_path, [a, b], merge="same") + assert [list(row) for row in got.varm["labels"]] == labels.tolist() + + +def test_a_kept_varm_array_keeps_its_compression(tmp_path): + """varm was written with the writer's defaults, uncompressed whatever the + input was.""" + import h5py + + genes = [f"g{i}" for i in range(50)] + values = np.tile(np.arange(50, dtype=np.float64)[:, None], (1, 4)) + a = _varm_store(tmp_path / "a.h5ad", genes, {"topic_region": values}, compress=True) + b = _varm_store(tmp_path / "b.h5ad", genes, {"topic_region": values}, compress=True) + out = tmp_path / "out.h5ad" + concat_on_disk([a, b], out, QUIET, merge="same", chunk_rows=8) + with h5py.File(out) as f, h5py.File(a) as src: + assert f["varm/topic_region"].compression == src["varm/topic_region"].compression == "gzip" + assert np.array_equal(ad.read_h5ad(out).varm["topic_region"], values) + + +@pytest.mark.parametrize("other", [np.int32, np.int64, np.float32, np.float64]) +def test_varm_in_another_dtype_merges_as_anndata_merges_it(tmp_path, other): + """anndata's `same` keeps an array only when the dtype matches as well as + the values (0.13: int64 against int32, float32 or float64 is dropped), so + the dtype stays part of what is compared.""" + genes = [f"g{i}" for i in range(6)] + values = np.arange(12, dtype=np.int64).reshape(6, 2) + a = _varm_store(tmp_path / "a.h5ad", genes, {"counts": values}) + b = _varm_store(tmp_path / "b.h5ad", genes, {"counts": values.astype(other)}) + got = _concat(tmp_path, [a, b], merge="same") + ref = ad.concat([ad.read_h5ad(a), ad.read_h5ad(b)], merge="same") + assert sorted(got.varm) == sorted(ref.varm) + + +def test_varm_differing_by_a_fraction_is_not_the_same(tmp_path): + genes = [f"g{i}" for i in range(6)] + values = np.arange(12, dtype=np.float64).reshape(6, 2) + 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 diff --git a/tests/test_performance.py b/tests/test_performance.py index 31f2b90..ca20ea5 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -90,6 +90,7 @@ def _store( n_categories: int = 0, obs_kind: Optional[str] = None, shared_var: bool = True, + frames: bool = False, ) -> Path: """One input store. Every argument is an axis a guard can scale.""" obs = pd.DataFrame(index=[f"{name}c{i}" for i in range(n_obs)]) @@ -124,6 +125,15 @@ def _store( var=var, ) obj.obsm["X_pca"] = np.zeros((n_obs, 3), dtype="float32") + if frames: + # What a SCope import leaves: dataframes in obsm and varm, and an array in varm. + obj.obsm["RegulonsAUC"] = pd.DataFrame( + {"r0": np.arange(n_obs, dtype="float64"), "r1": np.zeros(n_obs)}, index=obs.index + ) + obj.varm["Regulons"] = pd.DataFrame( + {"r0": np.arange(n_var) % 2, "r1": np.ones(n_var, dtype="int64")}, index=var.index + ) + obj.varm["topic_region"] = np.ones((n_var, 3)) obj.write_h5ad(path) return path @@ -207,6 +217,30 @@ def measure(n: int) -> int: assert_grows_linearly(measure, what=f"concat --merge {merge}", axis="n_var") +@pytest.mark.parametrize("merge", ["same", "unique", "first", "only"]) +def test_concat_varm_reads_grow_linearly_in_var_count(tmp_path, merge): + """varm is merged as var's columns are: read once per input, not per gene.""" + + def measure(n: int) -> int: + files = _inputs(tmp_path, f"varm{merge}{n}", n_var=n, frames=True) + with count_io() as io: + concat_on_disk(files, tmp_path / f"om{merge}{n}.h5ad", QUIET, merge=merge) + return io.reads + + assert_grows_linearly(measure, what=f"concat varm --merge {merge}", axis="n_var") + + +@pytest.mark.parametrize("join", ["inner", "outer"]) +def test_concat_obsm_frames_grow_linearly_in_obs_count(tmp_path, join): + def measure(n: int) -> int: + files = _inputs(tmp_path, f"obsmf{join}{n}", n_obs=n, n_var=4, frames=True) + with count_io() as io: + concat_on_disk(files, tmp_path / f"of{join}{n}.h5ad", QUIET, join=join) + return io.reads + + assert_grows_linearly(measure, what=f"concat obsm dataframes ({join})", axis="n_obs") + + @pytest.mark.parametrize("join", ["inner", "outer"]) def test_concat_reads_grow_linearly_in_obs_count(tmp_path, join): def measure(n: int) -> int: