From 1a671d14f135d496de5206c88c509a70f4461dde Mon Sep 17 00:00:00 2001 From: Aljes Date: Wed, 21 Oct 2026 20:08:04 +0100 Subject: [PATCH 1/3] concat: stack obsm dataframes, merge varm SCope and pycisTopic imports leave dataframes in obsm (RegulonsAUC) and varm (Regulons, topic_region). concat skipped the obsm ones ("not a dense array in every input") and never wrote varm at all, though its docstring said the alternative axis follows --merge, as anndata's concat does. - obsm: a dataframe in every input is stacked as obs is, its columns following --join, through the obs column writers. - varm: merged as var's columns are, on the output's genes. same/unique compare a snapshot read once per input (exact bytes, so NaN equals NaN); first/only go by presence. An input lacking some genes (an outer join) counts as missing. Without --merge it is dropped, with a note. Checked against anndata.concat for each join and strategy, and on two h5ads converted from the Zenodo pySCENIC loom. Co-Authored-By: Claude Opus 5.5 --- docs/COMMANDS.md | 7 ++ docs/TESTING.md | 1 + src/adata/core/concat.py | 149 +++++++++++++++++++++++++++-- tests/test_concat_axis_elements.py | 102 ++++++++++++++++++++ tests/test_performance.py | 34 +++++++ 5 files changed, 285 insertions(+), 8 deletions(-) create mode 100644 tests/test_concat_axis_elements.py 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..d5ffef2 100644 --- a/src/adata/core/concat.py +++ b/src/adata/core/concat.py @@ -5,9 +5,10 @@ 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 @@ -689,10 +690,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 +730,116 @@ 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 _varm_snapshot(obj: Any, where: np.ndarray) -> Any: + """A comparable value of a varm entry, on the target var: its bytes for + an array; each column's, for a dataframe. Exact, so NaN equals NaN.""" + if is_dataset(obj): + values = np.asarray(obj[...])[where] + return ("array", str(values.dtype), values.shape, values.tobytes()) + 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) + if kind == "numeric": + values = np.asarray(column[...])[where] + columns.append((name, str(values.dtype), values.tobytes())) + elif kind == "categorical": + categories = [str(c) for c in read_categories(column)] + codes = np.asarray(column["codes"][...], dtype=np.int64)[where] + columns.append((name, tuple(categories[c] if c >= 0 else None for c in codes))) + elif kind == "masked": + mask = np.asarray(column["mask"][...], dtype=bool)[where] + values = np.asarray(column["values"][...])[where] + columns.append((name, values.tobytes() if values.dtype.kind != "O" else tuple(values.tolist()), + mask.tobytes())) + else: + text = read_str_all(column) + columns.append((name, tuple(text[i] for i in where))) + return ("dataframe", tuple(columns)) + + +def _write_varm( + dst: Any, + roots: List[Any], + var_names: List[List[str]], + target_var: List[str], + merge: Optional[str], + console: Console, +) -> 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 + values read once per input, `first` and `only` presence alone. + """ + 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)) + 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_dense(varm, name, np.asarray(obj[...])[take]) + 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 +1067,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) 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 +1289,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 +1314,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..2154746 --- /dev/null +++ b/tests/test_concat_axis_elements.py @@ -0,0 +1,102 @@ +"""`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" 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: From 65060c69019ebc4ccb9a78479c5b21457ae0fdc7 Mon Sep 17 00:00:00 2001 From: Aljes Date: Thu, 22 Oct 2026 14:36:03 +0100 Subject: [PATCH 2/3] concat: stream varm a block of genes at a time, keep its layout Comparing varm under same/unique read every input's whole entry and kept its bytes; now each input is reduced to a digest a block of rows at a time, and the kept entry is copied the same way, with the source's chunking and compression. Text arrays are compared by their text, not by the object pointers their bytes held. Co-Authored-By: Claude Opus 5.5 --- src/adata/core/concat.py | 105 +++++++++++++++++++++++------ tests/test_concat_axis_elements.py | 70 +++++++++++++++++++ 2 files changed, 155 insertions(+), 20 deletions(-) diff --git a/src/adata/core/concat.py b/src/adata/core/concat.py index d5ffef2..0a49e7a 100644 --- a/src/adata/core/concat.py +++ b/src/adata/core/concat.py @@ -13,6 +13,7 @@ from __future__ import annotations +import hashlib from pathlib import Path from typing import Any, Dict, List, Optional, Sequence, Tuple @@ -24,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, @@ -40,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, @@ -212,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" @@ -759,36 +764,94 @@ def _concat_frame( # varm -def _varm_snapshot(obj: Any, where: np.ndarray) -> Any: - """A comparable value of a varm entry, on the target var: its bytes for - an array; each column's, for a dataframe. Exact, so NaN equals NaN.""" +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): - values = np.asarray(obj[...])[where] - return ("array", str(values.dtype), values.shape, values.tobytes()) + 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] - columns.append((name, str(values.dtype), values.tobytes())) + _digest(hasher, values) + columns.append((name, str(values.dtype), hasher.digest())) elif kind == "categorical": - categories = [str(c) for c in read_categories(column)] + categories = np.asarray([str(c) for c in read_categories(column)] + [""], dtype=object) codes = np.asarray(column["codes"][...], dtype=np.int64)[where] - columns.append((name, tuple(categories[c] if c >= 0 else None for c in codes))) + _digest(hasher, codes < 0) + _digest(hasher, categories[codes]) + columns.append((name, "categorical", hasher.digest())) elif kind == "masked": - mask = np.asarray(column["mask"][...], dtype=bool)[where] - values = np.asarray(column["values"][...])[where] - columns.append((name, values.tobytes() if values.dtype.kind != "O" else tuple(values.tolist()), - mask.tobytes())) + _digest(hasher, np.asarray(column["mask"][...], dtype=bool)[where]) + _digest(hasher, np.asarray(column["values"][...])[where]) + columns.append((name, "masked", hasher.digest())) else: - text = read_str_all(column) - columns.append((name, tuple(text[i] for i in where))) + _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], @@ -796,12 +859,14 @@ def _write_varm( 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 - values read once per input, `first` and `only` presence alone. + 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: @@ -820,7 +885,7 @@ def _write_varm( 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)) + aligned.append(_varm_snapshot(root["varm"][name], where, chunk_rows)) else: aligned.append(_PRESENT) keep, _ = _merge_values(aligned, merge) @@ -829,7 +894,7 @@ def _write_varm( 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_dense(varm, name, np.asarray(obj[...])[take]) + _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]) @@ -1069,7 +1134,7 @@ def concat_on_disk( ) _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) + _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): diff --git a/tests/test_concat_axis_elements.py b/tests/test_concat_axis_elements.py index 2154746..d203c67 100644 --- a/tests/test_concat_axis_elements.py +++ b/tests/test_concat_axis_elements.py @@ -100,3 +100,73 @@ def test_an_outer_join_drops_varm_an_input_cannot_cover(tmp_path): 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) From 800c409dfc222589ffb664dd61cc3e624b62d979 Mon Sep 17 00:00:00 2001 From: Aljes Date: Thu, 22 Oct 2026 16:01:16 +0100 Subject: [PATCH 3/3] tests: varm in another dtype merges as anndata merges it Co-Authored-By: Claude Opus 5.5 --- tests/test_concat_axis_elements.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/test_concat_axis_elements.py b/tests/test_concat_axis_elements.py index d203c67..a17654c 100644 --- a/tests/test_concat_axis_elements.py +++ b/tests/test_concat_axis_elements.py @@ -170,3 +170,25 @@ def test_a_kept_varm_array_keeps_its_compression(tmp_path): 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