diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml
new file mode 100644
index 0000000..938aba5
--- /dev/null
+++ b/.github/workflows/benchmark.yml
@@ -0,0 +1,137 @@
+# Comparative benchmark: adata-cli against anndata and scanpy.
+#
+# Report-only, by design. Wall time on a shared runner is too noisy to gate a
+# release on, and a benchmark that can block a publish stops being run. What
+# *does* gate merges is tests/test_performance.py, which counts operations
+# rather than seconds and lives in the ordinary test job.
+#
+# Its own file rather than a job in publish.yml: at the `ci` tier this takes
+# the better part of an hour, and hanging that off the release graph would
+# either delay the PyPI publish or paint the release run red for a report.
+name: Benchmark
+
+on:
+ push:
+ tags: ["*"]
+ workflow_dispatch:
+ inputs:
+ tier:
+ description: Input size
+ type: choice
+ options: [smoke, ci, large]
+ default: ci
+ publish:
+ description: Commit the results to docs/ on main
+ type: boolean
+ default: false
+
+concurrency:
+ group: benchmark-${{ github.ref }}
+ cancel-in-progress: true
+
+jobs:
+ benchmark:
+ name: adata-cli vs anndata/scanpy
+ runs-on: ubuntu-latest
+ timeout-minutes: 90
+ # Job level, not step level: a report must never mark a release red.
+ continue-on-error: true
+ permissions:
+ contents: write
+ steps:
+ - uses: actions/checkout@v4
+ with:
+ fetch-depth: 0
+
+ - uses: astral-sh/setup-uv@v3
+ with:
+ enable-cache: true
+
+ # The checked-out source, not the published wheel. Benchmarking the
+ # wheel would mean waiting on publish-pypi and then on index
+ # propagation, for a measurement that would come out the same.
+ - name: Install adata-cli
+ run: uv sync --extra dev --frozen
+
+ - name: Show free space
+ run: df -h /
+
+ - name: Run
+ run: |
+ uv run python -m benchmarks.run \
+ --tier "${{ inputs.tier || 'ci' }}" \
+ --out results.json \
+ --work "$RUNNER_TEMP/bench"
+
+ - name: Render
+ run: |
+ uv run python -m benchmarks.report results.json --out summary.md
+ cat summary.md >> "$GITHUB_STEP_SUMMARY"
+
+ - uses: actions/upload-artifact@v4
+ if: always()
+ with:
+ name: benchmark-${{ github.ref_name }}
+ path: |
+ results.json
+ summary.md
+ retention-days: 90
+
+ # docs/ is the durable home: artifacts expire, and one absolute number
+ # with nothing to compare it against says very little. The page is
+ # already served by Pages from docs/, so this needs no extra machinery.
+ - name: Publish to docs
+ if: github.ref_type == 'tag' || inputs.publish
+ run: |
+ git config user.name "github-actions[bot]"
+ git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
+ # A tag build is on a detached HEAD; the page belongs on main.
+ git fetch origin main
+ git checkout -B main origin/main
+ uv run python -m benchmarks.report results.json --docs docs --publish
+ git add docs/BENCHMARKS.md docs/benchmarks
+ if git diff --cached --quiet; then
+ echo "nothing to publish"; exit 0
+ fi
+ git commit -m "Benchmark results for ${{ github.ref_name }} [skip ci]"
+ # Another job may have landed on main in the meantime.
+ for attempt in 1 2 3; do
+ if git push origin main; then exit 0; fi
+ git pull --rebase origin main
+ done
+ echo "could not push benchmark results"; exit 1
+
+ # Best effort, and last: releases here are cut by hand, so this job may
+ # well run before one exists. The docs page is authoritative either way.
+ - name: Append to the release notes
+ if: github.ref_type == 'tag'
+ continue-on-error: true
+ env:
+ GH_TOKEN: ${{ github.token }}
+ run: |
+ tag="${{ github.ref_name }}"
+ gh release view "$tag" >/dev/null 2>&1 || {
+ echo "no release for $tag yet; skipping"; exit 0
+ }
+ gh release view "$tag" --json body -q .body > body.md
+ # Idempotent: replace any block this job wrote before.
+ python - <<'PY'
+ import pathlib, re
+ body = pathlib.Path("body.md").read_text()
+ summary = pathlib.Path("summary.md").read_text()
+ block = (
+ "\n\n"
+ "Benchmark vs anndata/scanpy
\n\n"
+ f"{summary}\n[.json` keeps the series, and
+`docs/BENCHMARKS.md` renders the latest run plus a history table. Both are in
+git and both are already served by GitHub Pages, so no extra machinery.
+
+ uv run python -m benchmarks.report results.json
+ uv run python -m benchmarks.report results.json --docs docs --publish
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import shutil
+from pathlib import Path
+from typing import Dict, List, Optional
+
+from benchmarks.cases import by_name
+
+
+def _bytes(n: int) -> str:
+ if not n:
+ return "-"
+ for unit, scale in (("GB", 1e9), ("MB", 1e6), ("kB", 1e3)):
+ if n >= scale:
+ return f"{n / scale:.1f} {unit}"
+ return f"{n} B"
+
+
+def _cell(record: Dict) -> tuple:
+ """(wall, rss, output) as they should read in the table."""
+ status = record.get("status")
+ if status == "n/a":
+ return "n/a", "n/a", "n/a"
+ if status == "timeout":
+ # This is how the 0.5.1 hang presented: still running, output stuck at
+ # the size it reached in the first second. Say so, rather than
+ # reporting a large time as though the run had completed.
+ return (
+ f"**timeout** (>{record['wall_s']:.0f} s)",
+ _bytes(record["maxrss_bytes"]),
+ f"{_bytes(record['output_bytes'])} and not growing",
+ )
+ if status == "oom":
+ return "**out of memory**", "hit the ceiling", "-"
+ if status == "failed":
+ return "**failed**", _bytes(record["maxrss_bytes"]), "-"
+
+ output = _bytes(record["output_bytes"])
+ if record.get("output_files", 0) > 1:
+ # Zarr: a directory of small chunks looks far larger under `du` than
+ # the bytes it holds, so the file count goes alongside.
+ output += f" in {record['output_files']} files"
+ return f"{record['wall_s']:.2f} s", _bytes(record["maxrss_bytes"]), output
+
+
+def render(payload: Dict, previous: Optional[Dict] = None) -> str:
+ cases = by_name()
+ before = (
+ {(r["case"], r["contender"]): r for r in previous["results"]}
+ if previous
+ else {}
+ )
+
+ lines: List[str] = []
+ # H2, because this is embedded under the page's own H1 by `build_page`.
+ # The framing lives in `page_template.md`; repeating it here would print
+ # it twice on the published page.
+ lines.append(f"## Results — `{payload['ref']}`")
+ lines.append("")
+
+ n_obs, n_var, density = payload["tier_shape"]
+ vh_obs, vh_var = payload["var_heavy_shape"]
+ cache = (
+ "dropped between runs"
+ if payload.get("page_cache_dropped")
+ else "**warm** (could not be dropped; reads are served from RAM, "
+ "which understates the streaming advantage)"
+ )
+ lines.append("| | |")
+ lines.append("|---|---|")
+ lines.append(f"| Tier | `{payload['tier']}`, {n_obs:,} obs x {n_var:,} var, "
+ f"{density:.0%} dense CSR |")
+ lines.append(f"| Var-heavy shape | {vh_obs:,} obs x {vh_var:,} var |")
+ lines.append(f"| Page cache | {cache} |")
+ lines.append(f"| Address-space ceiling | "
+ f"{_bytes(payload['memory_limit_bytes'])} per process |")
+ lines.append(f"| Repeats | {payload['repeats']} (fastest shown) |")
+ lines.append(f"| Platform | {payload['platform']}, Python "
+ f"{payload['python']} |")
+ versions = ", ".join(f"{k} {v}" for k, v in payload["versions"].items())
+ lines.append(f"| Versions | {versions} |")
+ lines.append(f"| Generated | {payload['generated']} |")
+ lines.append("")
+
+ grouped: Dict[str, List[Dict]] = {}
+ for record in payload["results"]:
+ grouped.setdefault(record["case"], []).append(record)
+
+ for name, records in grouped.items():
+ case = cases.get(name)
+ lines.append(f"### `{name}`")
+ lines.append("")
+ if case:
+ lines.append(f"{case.question}")
+ lines.append("")
+ lines.append("| Contender | Wall time | Peak RSS | Output | vs previous |")
+ lines.append("|---|---|---|---|---|")
+ for record in records:
+ wall, rss, output = _cell(record)
+ lines.append(
+ f"| {record['contender']} | {wall} | {rss} | {output} | "
+ f"{_delta(record, before.get((name, record['contender'])))} |"
+ )
+ note = record.get("note") or (
+ record.get("stderr_tail") if record.get("status") == "n/a" else None
+ )
+ if note:
+ lines.append(f"| | *{note}* | | | |")
+ lines.append("")
+
+ lines.append("---")
+ lines.append("")
+ lines.append(
+ "Rows marked `n/a` are operations the baseline does not offer; that "
+ "is a result, not a gap in the measurement. Where a baseline is "
+ "faster, the row stands as it is -- the trade this tool makes is "
+ "memory for time, and hiding the cost would make the table useless."
+ )
+ lines.append("")
+ lines.append(
+ "**Output size** is the file as the filesystem reports it, which for "
+ "HDF5 includes allocation slack. On a small tier that slack can be "
+ "several times the stored bytes and the column says more about the "
+ "writer's allocation strategy than about the data; at `ci` and above "
+ "it is noise. Zarr stores report a file count alongside, because a "
+ "directory of small chunks measures much larger than it holds."
+ )
+ return "\n".join(lines) + "\n"
+
+
+def _delta(record: Dict, previous: Optional[Dict]) -> str:
+ """Wall time and peak RSS against the last published run."""
+ if not previous or record.get("status") != "ok" or previous.get("status") != "ok":
+ return "-"
+ parts = []
+ for key, label in (("wall_s", "time"), ("maxrss_bytes", "RSS")):
+ old, new = previous.get(key), record.get(key)
+ if not old:
+ continue
+ change = (new - old) / old
+ if abs(change) < 0.10: # within run-to-run noise on a shared runner
+ continue
+ parts.append(f"{label} {change:+.0%}")
+ return ", ".join(parts) if parts else "no change"
+
+
+def history_table(docs: Path) -> str:
+ """One row per published run, newest first."""
+ runs = []
+ for path in sorted((docs / "benchmarks").glob("*.json")):
+ try:
+ payload = json.loads(path.read_text())
+ except Exception: # pragma: no cover
+ continue
+ rows = {(r["case"], r["contender"]): r for r in payload["results"]}
+ key = ("concat-inner", "adata-cli")
+ headline = rows.get(key, {})
+ runs.append(
+ (
+ payload.get("generated", ""),
+ payload.get("ref", path.stem),
+ payload.get("tier", "?"),
+ headline.get("wall_s"),
+ headline.get("maxrss_bytes"),
+ path.name,
+ )
+ )
+ if not runs:
+ return ""
+
+ lines = [
+ "### History",
+ "",
+ "`concat-inner` on adata-cli, run by run. Full results for each are "
+ "in [`docs/benchmarks/`](benchmarks/).",
+ "",
+ "| Run | Ref | Tier | Wall time | Peak RSS | Raw |",
+ "|---|---|---|---|---|---|",
+ ]
+ for generated, ref, tier, wall, rss, filename in sorted(runs, reverse=True):
+ lines.append(
+ f"| {generated} | `{ref}` | {tier} | "
+ f"{f'{wall:.2f} s' if wall else '-'} | "
+ f"{_bytes(rss) if rss else '-'} | [json](benchmarks/{filename}) |"
+ )
+ return "\n".join(lines) + "\n"
+
+
+#: The standing prose of the docs page, with `` marking where
+#: the tables go. Kept as a file rather than inline so the words that explain
+#: the numbers live in one place and survive every republish -- an earlier
+#: version overwrote the whole page with bare tables, which would have thrown
+#: away the framing on the first tag.
+PAGE_TEMPLATE = Path(__file__).with_name("page_template.md")
+
+#: What the template shows before any run has happened.
+NOT_YET_RUN = (
+ "## Results\n\nNo run has been published yet. The next tag fills this in; "
+ "until then, produce one locally with the commands below.\n"
+)
+
+
+def build_page(body: str) -> str:
+ """Wrap rendered results in the page's standing explanation."""
+ template = PAGE_TEMPLATE.read_text()
+ if "" not in template: # pragma: no cover - template edited
+ return template.rstrip() + "\n\n" + body
+ return template.replace("", body.strip())
+
+
+def publish(payload: Dict, results: Path, docs: Path) -> Path:
+ """Copy the raw results in and rewrite the docs page.
+
+ Only the results section is replaced; everything explaining what the
+ numbers mean comes from `page_template.md` and is reinstated every time.
+ """
+ store = docs / "benchmarks"
+ store.mkdir(parents=True, exist_ok=True)
+ ref = payload["ref"].replace("/", "-")
+ shutil.copyfile(results, store / f"{ref}.json")
+
+ previous = _previous(store, skip=f"{ref}.json")
+ page = docs / "BENCHMARKS.md"
+ page.write_text(
+ build_page(render(payload, previous) + "\n" + history_table(docs))
+ )
+ return page
+
+
+def _previous(store: Path, *, skip: str) -> Optional[Dict]:
+ candidates = sorted(
+ (p for p in store.glob("*.json") if p.name != skip),
+ key=lambda p: p.stat().st_mtime,
+ )
+ if not candidates:
+ return None
+ try:
+ return json.loads(candidates[-1].read_text())
+ except Exception: # pragma: no cover
+ return None
+
+
+def main() -> int: # pragma: no cover - CLI entry
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("results", type=Path)
+ parser.add_argument("--docs", type=Path, default=Path("docs"))
+ parser.add_argument(
+ "--publish",
+ action="store_true",
+ help="write docs/benchmarks/][.json and rewrite docs/BENCHMARKS.md",
+ )
+ parser.add_argument("--out", type=Path, help="also write the markdown here")
+ args = parser.parse_args()
+
+ payload = json.loads(args.results.read_text())
+ if args.publish:
+ page = publish(payload, args.results, args.docs)
+ print(f"wrote {page}")
+ text = page.read_text()
+ else:
+ text = render(payload, _previous(args.docs / "benchmarks", skip=""))
+ print(text)
+ if args.out:
+ args.out.write_text(text)
+ return 0
+
+
+if __name__ == "__main__": # pragma: no cover
+ raise SystemExit(main())
diff --git a/benchmarks/run.py b/benchmarks/run.py
new file mode 100644
index 0000000..290667b
--- /dev/null
+++ b/benchmarks/run.py
@@ -0,0 +1,342 @@
+"""Run the benchmark and write results.json.
+
+Not a pytest module. Each contender has to run in its own process for peak
+RSS to mean anything, and the whole thing takes tens of minutes at the `ci`
+tier -- neither of which belongs in a suite that gates merges. The guards
+that *do* gate merges are in `tests/test_performance.py`, and they measure
+operation counts rather than time.
+
+ uv run python -m benchmarks.run --tier smoke
+ uv run python -m benchmarks.run --tier ci --out results.json
+
+Baselines run from environments built once, up front, rather than through
+`uv run --with`. `reference_stores.py` uses the latter to build fixtures,
+where cost is irrelevant; here the first invocation would resolve and download
+several hundred megabytes of wheels straight into the measured wall time, and
+uv's own memory into the measured peak.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import platform
+import shutil
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Dict, List, Optional
+
+from benchmarks import datasets
+from benchmarks._measure import (
+ DEFAULT_MEMORY_LIMIT,
+ Measurement,
+ drop_page_cache,
+ measure,
+)
+from benchmarks.cases import CASES, Case, Contender
+
+#: Packages each baseline environment needs.
+#: `dask` is there for `concat_on_disk`, which imports it to concatenate a
+#: dense element such as obsm and raises ModuleNotFoundError without it.
+#: Installing it is the fair thing to do -- the rule is to give the baseline
+#: the best idiom it has -- and the dependency is itself worth knowing about.
+ENVIRONMENTS: Dict[str, List[str]] = {
+ "anndata": ["anndata", "scipy", "pandas", "h5py", "zarr", "dask"],
+ "scanpy": ["scanpy", "anndata", "scipy", "pandas", "h5py", "zarr", "dask"],
+}
+
+
+def build_environments(root: Path, wanted: List[str]) -> Dict[str, Path]:
+ """Create one venv per baseline and return its interpreter.
+
+ The package list is recorded beside the venv and rebuilt when it changes.
+ Checking only that the interpreter exists means a reused `--work`
+ directory keeps whatever was installed the first time: adding `dask` to
+ ENVIRONMENTS had no effect on an existing tree, and `concat_on_disk`
+ went on failing with ModuleNotFoundError as though that were a finding
+ about anndata.
+ """
+ interpreters: Dict[str, Path] = {}
+ for name in wanted:
+ venv = root / name
+ python = venv / "bin" / "python"
+ stamp = venv / ".packages"
+ wanted_packages = "\n".join(sorted(ENVIRONMENTS[name]))
+ current = stamp.read_text() if stamp.exists() else None
+
+ if not python.exists() or current != wanted_packages:
+ if current is not None and current != wanted_packages:
+ print(f"[env] {name}: package list changed, rebuilding", flush=True)
+ else:
+ print(f"[env] building {name}", flush=True)
+ # --clear because a rebuild runs over an existing tree; without
+ # it `uv venv` refuses and the run dies before any measurement.
+ done = subprocess.run(
+ ["uv", "venv", "--clear", str(venv)],
+ capture_output=True,
+ text=True,
+ )
+ if done.returncode != 0:
+ raise RuntimeError(
+ f"could not create the {name} environment at {venv}:\n"
+ + (done.stderr or done.stdout)
+ )
+ subprocess.run(
+ ["uv", "pip", "install", "--python", str(python), *ENVIRONMENTS[name]],
+ check=True,
+ )
+ stamp.write_text(wanted_packages)
+ interpreters[name] = python
+ return interpreters
+
+
+def _versions(interpreters: Dict[str, Path]) -> Dict[str, str]:
+ """Record exactly what was compared, so a table is interpretable later."""
+ found: Dict[str, str] = {}
+ try:
+ found["adata-cli"] = subprocess.run(
+ ["adata", "--version"], capture_output=True, text=True, timeout=120
+ ).stdout.strip()
+ except Exception: # pragma: no cover
+ found["adata-cli"] = "unknown"
+ for name, python in interpreters.items():
+ code = (
+ "import importlib.metadata as m;"
+ f"print(m.version({name!r}))"
+ )
+ try:
+ found[name] = subprocess.run(
+ [str(python), "-c", code], capture_output=True, text=True, timeout=120
+ ).stdout.strip()
+ except Exception: # pragma: no cover
+ found[name] = "unknown"
+ return found
+
+
+def _run_contender(
+ case: Case,
+ contender: Contender,
+ inputs: List[Path],
+ workdir: Path,
+ interpreters: Dict[str, Path],
+ *,
+ timeout_s: float,
+ scripts_dir: Path,
+ shape: tuple,
+) -> Measurement:
+ workdir.mkdir(parents=True, exist_ok=True)
+ output = workdir / f"{case.name}{case.output_suffix}"
+ outdir = workdir / f"{case.name}-out"
+
+ # The harness never overwrites, so clear any previous run first.
+ for stale in (output, outdir):
+ if stale.is_dir():
+ shutil.rmtree(stale)
+ elif stale.exists():
+ stale.unlink()
+
+ sidecar = _build_sidecar(case, inputs, workdir)
+ # The tier's shape, so a case that builds rather than reads -- `create`
+ # has no input to take its size from -- scales with the tier instead of
+ # making the smoke run as heavy as the ci one.
+ substitutions = {
+ "output": str(output),
+ "outdir": str(outdir),
+ "sidecar_csv": str(sidecar) if sidecar else "",
+ "n_obs": str(shape[0]),
+ "n_var": str(shape[1]),
+ **{f"input{i}": str(p) for i, p in enumerate(inputs)},
+ }
+ env = {"BENCH_N_OBS": str(shape[0]), "BENCH_N_VAR": str(shape[1])}
+ if sidecar:
+ env["BENCH_SIDECAR_CSV"] = str(sidecar)
+
+ if contender.argv is not None:
+ command = [part.format(**substitutions) for part in contender.argv]
+ target = outdir if "{outdir}" in " ".join(contender.argv) else output
+ else:
+ safe = f"{case.name}-{contender.label}".replace(" ", "_").replace("(", "").replace(")", "")
+ script = scripts_dir / f"{safe}.py"
+ script.write_text(contender.script or "")
+ # A script that makes a directory is handed the directory; everything
+ # else is handed the single output path.
+ target = outdir if "makedirs" in (contender.script or "") else output
+ command = [
+ str(interpreters[contender.env]), str(script), *map(str, inputs), str(target)
+ ]
+
+ # Read-only cases produce nothing to size.
+ watched = None if (case.output_suffix == "" and target == output) else target
+ # A contender whose binary is not installed -- `h5ls` ships with the HDF5
+ # tools and is often absent -- comes back as `n/a` from `measure`.
+ return measure(command, output=watched, timeout_s=timeout_s, env=env)
+
+
+def _build_sidecar(case: Case, inputs: List[Path], workdir: Path) -> Optional[Path]:
+ """Build the auxiliary input a case declares, once per case.
+
+ Derived from the real store rather than invented, so an import writes
+ back something the file could plausibly have held.
+ """
+ if case.sidecar != "csv":
+ return None
+ path = workdir / f"{case.name}-sidecar.csv"
+ if path.exists():
+ return path
+ import h5py
+ from anndata.io import read_elem
+
+ with h5py.File(inputs[0], "r") as handle:
+ read_elem(handle["obs"]).to_csv(path)
+ return path
+
+
+def run(
+ tier_name: str,
+ *,
+ out: Path,
+ work: Path,
+ only: Optional[List[str]] = None,
+ repeats: int = 3,
+ timeout_s: float = 900.0,
+ drop_cache: bool = True,
+) -> Dict:
+ tier = datasets.TIERS[tier_name]
+ work.mkdir(parents=True, exist_ok=True)
+ fixtures = work / "fixtures"
+ scripts = work / "scripts"
+ scripts.mkdir(parents=True, exist_ok=True)
+
+ selected = [c for c in CASES if only is None or c.name in only]
+ needed = sorted(
+ {c.env for case in selected for c in case.contenders if c.script}
+ )
+ interpreters = build_environments(work / "envs", needed)
+
+ print(f"[data] building {tier.name} fixtures ({tier.n_obs} x {tier.n_var})",
+ flush=True)
+ main_inputs = datasets.build(fixtures, tier, count=2)
+ var_heavy_inputs = datasets.build(fixtures, datasets.VAR_HEAVY, count=2)
+
+ cache_dropped = drop_page_cache() if drop_cache else False
+
+ results: List[Dict] = []
+ for case in selected:
+ pool = var_heavy_inputs if case.var_heavy else main_inputs
+ inputs = pool[: case.inputs]
+ print(f"[case] {case.name}", flush=True)
+ for contender in case.contenders:
+ if contender.unsupported:
+ results.append(
+ {
+ "case": case.name,
+ "contender": contender.label,
+ "status": "n/a",
+ "note": contender.unsupported,
+ }
+ )
+ print(f" {contender.label}: n/a ({contender.unsupported})")
+ continue
+
+ runs: List[Measurement] = []
+ for _ in range(repeats):
+ if drop_cache:
+ drop_page_cache()
+ runs.append(
+ _run_contender(
+ case,
+ contender,
+ inputs,
+ work / "out",
+ interpreters,
+ timeout_s=timeout_s,
+ scripts_dir=scripts,
+ shape=(
+ (datasets.VAR_HEAVY.n_obs, datasets.VAR_HEAVY.n_var)
+ if case.var_heavy
+ else (tier.n_obs, tier.n_var)
+ ),
+ )
+ )
+ if runs[-1].status != "ok":
+ break # a failure repeats identically; do not pay for it twice
+
+ best = min(runs, key=lambda m: m.wall_s)
+ record = {
+ "case": case.name,
+ "contender": contender.label,
+ # The fastest run, because a slower one differs only by what
+ # else the machine was doing.
+ **best.as_dict(),
+ "runs": len(runs),
+ "wall_s_all": [m.wall_s for m in runs],
+ }
+ results.append(record)
+ print(
+ f" {contender.label}: {best.status} "
+ f"{best.wall_s:.2f}s {best.maxrss_bytes / 1e6:.0f} MB",
+ flush=True,
+ )
+
+ payload = {
+ "tier": tier.name,
+ "tier_shape": [tier.n_obs, tier.n_var, tier.density],
+ "var_heavy_shape": [datasets.VAR_HEAVY.n_obs, datasets.VAR_HEAVY.n_var],
+ "generated": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+ "ref": _git_ref(),
+ "platform": f"{platform.system()} {platform.machine()}",
+ "python": sys.version.split()[0],
+ "memory_limit_bytes": DEFAULT_MEMORY_LIMIT,
+ "page_cache_dropped": cache_dropped,
+ "repeats": repeats,
+ "versions": _versions(interpreters),
+ "results": results,
+ }
+ out.parent.mkdir(parents=True, exist_ok=True)
+ out.write_text(json.dumps(payload, indent=2) + "\n")
+ print(f"\nwrote {out}")
+ return payload
+
+
+def _git_ref() -> str:
+ for args in (["git", "describe", "--tags", "--exact-match"],
+ ["git", "rev-parse", "--short", "HEAD"]):
+ try:
+ done = subprocess.run(args, capture_output=True, text=True, timeout=30)
+ if done.returncode == 0 and done.stdout.strip():
+ return done.stdout.strip()
+ except Exception: # pragma: no cover
+ pass
+ return "unknown"
+
+
+def main() -> int: # pragma: no cover - CLI entry
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--tier", default="smoke", choices=sorted(datasets.TIERS))
+ parser.add_argument("--out", type=Path, default=Path("results.json"))
+ parser.add_argument("--work", type=Path, default=Path(".bench"))
+ parser.add_argument("--case", action="append", dest="only")
+ parser.add_argument("--repeats", type=int, default=3)
+ parser.add_argument("--timeout", type=float, default=900.0)
+ parser.add_argument(
+ "--keep-cache",
+ action="store_true",
+ help="do not try to drop the page cache between runs",
+ )
+ args = parser.parse_args()
+ run(
+ args.tier,
+ out=args.out,
+ work=args.work,
+ only=args.only,
+ repeats=args.repeats,
+ timeout_s=args.timeout,
+ drop_cache=not args.keep_cache,
+ )
+ return 0
+
+
+if __name__ == "__main__": # pragma: no cover
+ raise SystemExit(main())
diff --git a/docs/BENCHMARKS.md b/docs/BENCHMARKS.md
new file mode 100644
index 0000000..de52293
--- /dev/null
+++ b/docs/BENCHMARKS.md
@@ -0,0 +1,78 @@
+# Benchmarks
+
+adata-cli measured against anndata, and against scanpy wherever scanpy has a real
+equivalent. The [`Benchmark`](../.github/workflows/benchmark.yml) workflow rewrites
+this page on every tag and keeps each run's raw numbers in `docs/benchmarks/`.
+
+**Peak RSS is the headline, not wall time.** This tool exists so that memory is set
+by `--chunk` rather than by the size of the input. For data that fits in RAM, loading
+the whole thing is often faster — so the tables below carry the rows where adata-cli
+is the slower of the two, because those are the same measurement. The trade is memory
+for time, and a table that hid the cost would not be worth publishing.
+
+## What is measured
+
+Fifteen cases on a GitHub-hosted runner: concat (inner, outer, `--merge same`),
+subset by query, h5ad→zarr, `view`, `ls`, `create`, `export` dataframe / array /
+sparse, `import dataframe`, `split`, and a peak-RSS-against-input-size sweep. Inputs
+are 50,000 × 20,000 at 5% density, plus the 2,000 × 36,601 shape that hung in 0.5.1.
+
+Every run also reports a startup floor (`--version` for each contender), because the
+CLI costs 0.3–1 s to import and `import scanpy` costs 3–8 s; on small inputs that is
+the entire measurement.
+
+`export image`, `export dict`, `import image` and `import dict` are not here. No
+library offers an equivalent, so the rows would only ever read `n/a`. They are covered
+by the complexity guards instead.
+
+## Results
+
+No tagged run has been published yet — the next tag fills this in. These
+numbers are from a `ci` run during development, as an indication of what the
+tables will say.
+
+| Case | adata-cli | best baseline |
+|---|---|---|
+| `concat-inner` | **202 MB**, 2.35 s | 439 MB, 16.44 s (`concat_on_disk`) |
+| `concat-outer` | **202 MB**, 2.60 s | 436 MB, **2.24 s** (`concat_on_disk`) |
+| `create` | **78 MB**, 0.28 s | 2,225 MB, 2.85 s |
+| `inspect` | **63 MB**, 0.31 s | 138 MB, 0.79 s (`read_elem`) |
+| `import-dataframe` | **107 MB**, 0.41 s | 571 MB, 1.84 s |
+| `export-sparse` | **68 MB**, 10.79 s | 759 MB, **1.90 s** (full load) |
+| `ls` | 63 MB, 0.36 s | **7 MB, 0.03 s** (`h5ls -r`) |
+
+The last two rows are why this is not a leaderboard. `export sparse` streams in a
+tenth of the memory and takes five times as long. `h5ls -r` walks the file in a
+hundredth of our time and a ninth of our memory, being a C program rather than a
+Python process that must import typer, rich, h5py and zarr before it starts.
+
+## How peak memory is measured
+
+Each measured command is forked from a small shim process, not from the benchmark
+runner itself. On Linux a forked child inherits its parent's resident pages and
+`execve` folds that into the `maxrss` the kernel reports, so a child of a fat parent
+cannot appear small: with the runner holding 330 MB, a process allocating nothing
+measured 326 MB. The shim brings that floor down to about 8 MB, uniform across every
+contender and visible in the `startup` row.
+
+## A caveat on the streaming claim
+
+Peak memory is not flat in input size. Over a 256× span at a fixed chunk, peak
+allocation grows 1.6× for `export array`, 2.2× for `export dataframe`, 6.5× for
+`export sparse` and 46× for `subset` — far below the input curve, but not constant.
+obs columns are read whole and a dense block is `--chunk` × n_var. The benchmark
+exists to keep that curve visible rather than to assert a guarantee the code does
+not yet meet.
+
+## Running it yourself
+
+```bash
+uv run python -m benchmarks.run --tier smoke --out results.json
+uv run python -m benchmarks.report results.json
+```
+
+Tiers are `smoke` (seconds), `ci` (50,000 × 20,000) and `large` (dispatch only, where
+the in-memory baseline is expected to hit the address-space ceiling).
+
+See [TESTING.md](TESTING.md#performance) for how the baselines are kept honest, and
+how this differs from the complexity guards that gate every pull request.
diff --git a/docs/COMMANDS.md b/docs/COMMANDS.md
index a5d27c6..4cac9b8 100644
--- a/docs/COMMANDS.md
+++ b/docs/COMMANDS.md
@@ -81,6 +81,52 @@ adata subset data.h5ad --inplace --obs barcodes.txt
`raw/` is carried over and matched against its **own** var axis, which usually
holds more genes than the main object.
+## `convert`
+
+Change a matrix's dtype, layout or density. Counts held as float64 cost twice
+the disk and twice the read for no information; a tool that wants CSC cannot
+use a CSR store; and `concat` refuses inputs whose encodings disagree.
+
+```bash
+adata convert data.h5ad X -o out.h5ad --dtype float32
+adata convert data.h5ad X -o out.h5ad --layout csc
+adata convert data.h5ad X --inplace --dtype float32 --indices-dtype int32
+adata convert data.h5ad --all -o out.h5ad --dtype float32
+adata convert data.h5ad layers/counts -o out.h5ad --layout dense --force
+```
+
+| Flag | Meaning |
+|---|---|
+| `--output`, `-o` | Output path. Required unless `--inplace` |
+| `--inplace` | Replace the source (written to a temporary path first) |
+| `--all` | Convert `X`, every layer and `raw/X` |
+| `--dtype` | New dtype for the values, e.g. `float32` |
+| `--indices-dtype` | New dtype for sparse indices: `int32` or `int64` |
+| `--layout` | `csr`, `csc`, `dense` or `sparse` |
+| `--force` | Convert despite a lossy cast or a large size increase |
+| `--in-memory` | Transpose in memory rather than streaming |
+| `--chunk`, `-C` | Row chunk size for dense matrices |
+| `--zarr-format` | Zarr version to write; defaults to the source's |
+
+Two things are refused before anything is written. A cast that would not
+round-trip -- checked by casting every value and casting it back, not by
+comparing dtypes -- and a densification that would inflate the store beyond
+four times its size. `--force` overrides either, and the message says how
+many values would change or how large the result would be.
+
+The index dtype is **preserved** unless `--indices-dtype` asks otherwise, so
+narrowing the values does not silently widen the indices and leave the file
+bigger than it started. `indices` and `indptr` keep their own widths, which
+can differ: a narrow matrix with more than 2^31 nonzeros needs int64 offsets
+over int32 coordinates, and both are range-checked before writing.
+
+The output path may not name the input; use `--inplace`, which writes to a
+temporary file and swaps it in only once the conversion has finished.
+
+Transposing streams by default and works on matrices too large to load, at
+the cost of two extra passes over the nonzeros. `--in-memory` is faster when
+the matrix fits.
+
## `split`
One store per distinct value of a column.
@@ -118,7 +164,7 @@ adata concat a.h5ad b.h5ad -o m.h5ad --keys a,b --index-unique - --uns-merge sam
| `--label` | Add an obs column recording each cell's source |
| `--keys` | Names for the inputs; defaults to their filenames |
| `--index-unique` | Delimiter for suffixing obs names with their key |
-| `--merge` / `--uns-merge` | `same`, `unique`, `first`, `only`; default drops |
+| `--merge` / `--uns-merge` | `drop` (default), `same`, `unique`, `first`, `only` |
| `--fill-value` | Value for dense cells introduced by an outer join |
obs columns keep their dtypes: categoricals union their category sets, nullable
diff --git a/docs/TESTING.md b/docs/TESTING.md
index 7159a14..d363cb1 100644
--- a/docs/TESTING.md
+++ b/docs/TESTING.md
@@ -4,6 +4,7 @@
uv sync --extra dev
uv run pytest # everything
uv run pytest -m "not integration" # fast: no environment building, ~40s
+uv run pytest -m perf # just the tracing complexity guards
uv run pytest -m integration # compatibility across anndata releases
```
@@ -22,6 +23,8 @@ and the compatibility suite as a separate job.
| `test_anndata_roundtrip.py` | anndata writes the fixtures, reads back our output |
| `test_anndata_versions.py` | Compatibility with six real anndata releases (see below) |
| `test_commands_phase2.py`, `test_commands_coverage.py`, `test_cli.py` | Command surfaces and error paths |
+| `test_performance.py` | Complexity guards -- what an operation costs, not how long it takes (see below) |
+| `test_benchmark_harness.py` | That the benchmark's measurement is trustworthy |
| `test_subset.py`, `test_export.py`, `test_import.py`, `test_info_read.py`, `test_zarr.py`, `test_query.py` | Per-feature unit tests |
## Writing a test
@@ -44,6 +47,301 @@ For CLI surfaces, use the module-level `CliRunner` and assert on
ANSI before matching message text — Rich also wraps long lines, so collapse
whitespace.
+## Performance
+
+Two mechanisms, and confusing them is the main way this gets misused.
+
+| | `tests/test_performance.py` | `benchmarks/` |
+|---|---|---|
+| Measures | operation **counts** | wall time, peak RSS, output size |
+| Runs | every CI job, both interpreters | tags, and `workflow_dispatch` |
+| Gates | **yes** -- it fails the build | never |
+| Data | 64-4096 elements | 50,000 x 20,000 |
+
+### Why the guards do not use a clock
+
+A test that can fail because a runner was busy does not belong in a merge
+gate. Every number in `test_performance.py` is deterministic: the same input
+gives the same count on every machine. That is what lets it block a merge.
+
+This is not new -- `test_commands_phase2.py` already counts `read_str_all`
+calls for exactly this reason. The performance file generalises the idea; it
+does not replace those tests, and where an exact count is derivable an exact
+count is still better than a ratio, because it says what the number *should*
+be rather than only that it did not grow.
+
+### What is counted
+
+`tests/perf_counters.py` patches four seams, all verified against the pinned
+h5py 3.15.1 and zarr 3.1.5:
+
+| Hook | Sees |
+|---|---|
+| `h5py.Dataset.__getitem__` | HDF5 reads |
+| `h5py.Dataset.__setitem__` | HDF5 writes |
+| `h5py.Group.create_dataset` | HDF5 writes made at creation |
+| `zarr.Array.__getitem__` | Zarr reads, v2 and v3 alike |
+| `zarr.storage.LocalStore.get` | chunk and metadata fetches |
+| `LocalStore.set` / `.delete` | write storms |
+
+Both write seams are needed. `create_dataset(name, data=...)` writes its payload
+during creation and never touches `__setitem__`, so hooking only the latter leaves
+every `import` and `create` guard measuring zero — which is how they were first
+written, and what the lower bound below caught.
+
+Use `io.reads` for a read path, `io.work` (reads plus writes) for `import` and
+`create`.
+
+The libraries are patched rather than anything in `src/adata/`. An in-repo
+seam would only see the call sites that remembered to use it, which is the
+wrong property for a guard meant to catch the read nobody thought of -- and
+the matrix paths slice the backend objects directly anyway.
+
+Elements are counted, not just calls: the `--merge` bug was n calls each
+reading n elements, and counting calls alone would miss a vectorised variant
+that reads n elements n times in one call.
+
+**Known bypasses.** `Dataset.read_direct`, `np.asarray(dataset)` (it goes via
+`__array__`), `dataset.asstr()[...]` and `.fields()` reach the file without
+passing any hook. None are used today. If you add one, the counters will
+quietly report less work than happened -- which is why every guard also
+asserts a lower bound, and why there are two canary tests with known absolute
+counts. If a canary fails, fix the hooks before trusting anything else here.
+
+### The invariant
+
+Three sizes at 4x spacing, comparing successive **increments**:
+
+```python
+d1 = c(4n) - c(n)
+d2 = c(16n) - c(4n)
+assert d2 <= 6 * d1 # grows no faster than linearly
+assert d1 >= 4n # and the counters actually saw something
+```
+
+The increment form is the point. Comparing raw counts needs an additive
+constant to absorb fixed setup cost, and there is no principled value for
+one: too small and it is flaky, too large and a quadratic with a small
+coefficient hides underneath at n=64. Any constant appears in both
+differences and cancels exactly.
+
+Why 6: at 4x spacing the increment ratio is 4.0 for linear work, about 4.4
+for n log n, 8 for n^1.5 and 16 for quadratic. 6 sits in the gap with room on
+both sides. `test_growth_limit_separates_linear_from_super_linear` asserts
+that calibration rather than leaving it as a comment, and it fails if anyone
+changes `SIZES` without recomputing the limit.
+
+**Scale one axis per test, and name it in the test id.** Scaling two at once
+makes legitimate work look quadratic -- an outer-join concat really does
+produce n_obs x n_var_union cells. The axes worth separate coverage are
+`n_var`, `n_obs`, `n_inputs`, `n_obs_columns`, `n_var_columns`,
+`n_categories` and `n_groups`.
+
+### Three invariants, not one
+
+`assert_grows_linearly` only catches super-linear growth. Most commands need exactly
+that, but two claims this tool makes are stronger, and a linear guard would happily
+accept a 64x increase in a command that is supposed to read nothing.
+
+| Helper | Claim | Used for |
+|---|---|---|
+| `assert_grows_linearly` | no worse than linear in one axis | most commands |
+| `assert_grows_slower_than_input` | grows at least N times slower than the input | streaming at a fixed chunk |
+| `assert_independent_of` | does not grow at all | inspection, and grouping work per row |
+
+**Inspection is free, and that is an exact number.** `view` and `ls` reach only
+`.shape`, `.dtype` and `.attrs`: `axis_len` goes through `element_len`, which reads a
+shape, and `_array_details` and `_infer_untagged` never touch a value. So the guard
+asserts **zero** data elements read rather than a ratio — 0 against 0 is not a
+meaningful ratio, and the moment inspection reads one column the answer stops being
+zero however the store scales. A companion test exports the same fixture to prove
+there was data there to read, so the zero cannot pass because the fixture was empty.
+
+**Streaming is bounded well below the input, which is weaker than it sounds and is
+what the measurements support.** Over a 256x span at a fixed chunk:
+
+| | peak allocation growth |
+|---|---|
+| `export array` | 1.6x |
+| `export dataframe` | 2.2x |
+| `export sparse` | 6.5x |
+| `subset` | 46x |
+
+Only `export array` is close to flat, so only it is held to a near-flat bound. None is
+asserted as flat outright. `subset` is the weakest because obs columns are
+materialised one at a time — a known gap that `benchmarks/` reports rather than these
+guards conceal. Streamed `export sparse` is separately asserted to stay under a
+quarter of what `--in-memory` costs, both measured in the same run so the factor holds
+on any machine.
+
+### Where counting is not enough
+
+Two real defects in `concat.py` were invisible to all of the above, and both
+needed their own instrument:
+
+- **Category merging** used `if category not in categories` on a list. The
+ category lists are read once either way, so no read counter sees it, and
+ `x not in lst` is a single bytecode, so the line tracer does not either --
+ the quadratic is inside C-level list membership. Counting string
+ comparisons through a `str` subclass is what made it visible: 2,096,128
+ comparisons at k=1024.
+- **A per-element loop over obs rows.** Linear, just with a fat constant, so
+ no ratio catches it. The guard compares executed Python lines per row
+ against the numeric column path measured in the same run -- self-calibrating,
+ so it needs no hand-tuned budget and holds across interpreters.
+
+A third case needed a third instrument. `split --by` grouped rows with
+`np.nonzero(values == label)` inside a loop over distinct labels, rescanning each
+chunk once per label: O(n_rows x n_groups), which at a million cells and a thousand
+samples is 10^9 comparisons. No read counter moves, because the chunk is already in
+memory; nothing lasting is allocated; and the Python line count per label is constant.
+`count_scanned_elements` counts what is handed to numpy's scanning primitives and
+makes it visible — 16,384 elements at 4 groups against 1,048,576 at 256.
+
+It is a **floor, not a measurement**: an operator such as `values == label` dispatches
+to the ufunc in C and never passes the patched `np.equal`, and `arr.argsort()` is
+invisible for the same reason. That is the right property for a guard — it can only
+under-report — but it means a guard using it must also assert a lower bound, so that
+under-reporting to nothing fails instead of passing.
+
+`count_allocations` (tracemalloc), `count_lines` (`sys.settrace`) and
+`count_scanned_elements` are the tools for all of this. The line tracer costs a 10-50x
+slowdown, so its tests carry the `perf` marker and stay small.
+
+### What the guards deliberately do not claim
+
+Obs columns are read whole, so peak allocation for a concat is O(n_obs), not
+O(`--chunk`). The streaming guarantee holds for X, not for obs annotation.
+The guards only stop that getting worse than linear; `benchmarks/` reports
+the actual curve.
+
+### What is covered
+
+Every subcommand has a guard. When you add a command, add one: pick the axis its cost
+should scale with, scale only that, and hold everything else fixed.
+
+| Command | Axis scaled | Invariant |
+|---|---|---|
+| `view`, `view --types`, `ls`, `ls --long`, `ls --plain` | n_obs, n_var | **zero** data reads |
+| `create` | n_obs | linear in writes, both generated names and a name file |
+| `concat` | n_var, n_obs, n_inputs, n_obs_columns, n_var_columns | linear; every `--merge` strategy and both joins |
+| `concat --label`, `--index-unique` | n_obs | linear |
+| `concat` category union | n_categories | bounded string comparisons |
+| `subset` by name, by query | n_obs, n_var | linear |
+| `split --by`, `--axis var` | n_groups | linear reads, **flat** scan work |
+| `export dataframe` | n_obs, n_obs_columns | linear |
+| `export array`, `sparse` (both paths), `dict`, `image` | elements, nnz, keys, pixels | linear |
+| `import dataframe`, `array`, `sparse`, `dict`, `image` | rows, elements, nnz, keys, pixels | linear in writes |
+| h5ad to zarr | n_obs | linear |
+
+The `slow` marker is on the three guards that build a 65,536-row store; they still gate
+merges, and `-m "not slow"` skips them locally. `perf` is on the tracing guards.
+
+### Reading a failure
+
+It says cost grew super-linearly on the named axis. The assertion message
+carries the whole measured series and the computed ratio. Usually the code is
+wrong. Occasionally the expectation is -- an operation legitimately gained
+work -- and then the new number needs a comment saying why, in the same style
+as the exact counts in `test_commands_phase2.py`.
+
+## The comparative benchmark
+
+```bash
+uv run python -m benchmarks.run --tier smoke --out results.json
+uv run python -m benchmarks.report results.json
+```
+
+Tiers are `smoke` (1,000 x 2,000, seconds), `ci` (50,000 x 20,000) and
+`large` (500,000 x 20,000, dispatch only, where the in-memory baseline is
+expected to hit the ceiling). Every tier also builds the 2,000 x 36,601
+var-heavy shape, which is what hung in 0.5.1.
+
+`ci` is cheaper than it looks: 580 MB of fixtures and three cases took 43
+seconds end to end on a laptop, so the full set with three repeats is minutes
+rather than the hour the workflow allows. The 90-minute timeout is headroom
+for `large`, not an estimate.
+
+A `ci` run measured while writing this, for a sense of what the tables say — and of
+what they are for. Peak RSS first, wall time second:
+
+| Case | adata-cli | best baseline |
+|---|---|---|
+| `concat-inner` | **202 MB**, 2.35 s | 439 MB, 16.44 s (`concat_on_disk`) |
+| `concat-outer` | **202 MB**, 2.60 s | 436 MB, **2.24 s** (`concat_on_disk`) |
+| `inspect` | **63 MB**, 0.31 s | 138 MB, 0.79 s (`read_elem`) |
+| `create` | **78 MB**, 0.28 s | 2,225 MB, 2.85 s |
+| `import-dataframe` | **107 MB**, 0.41 s | 571 MB, 1.84 s |
+| `export-sparse` | **68 MB**, 10.79 s | 759 MB, **1.90 s** (full load) |
+| `ls` | 63 MB, 0.25 s | **7 MB, 0.01 s** (`h5ls -r`) |
+
+The last two rows are the reason the report is not a leaderboard. `export sparse`
+streams in a tenth of the memory and takes five times as long; `h5ls` walks the file
+in a hundredth of our time and a ninth of our memory, because it is C and does not
+start a Python interpreter. Both belong in the table. A benchmark that only published
+the rows we win would not be worth running.
+
+Note also that 202 MB is not flat in input size — obs columns are read whole, and a
+dense block is `--chunk` x n_var. The benchmark exists to keep that curve visible
+rather than to assert a claim the code does not yet meet.
+
+Results are published to `docs/BENCHMARKS.md` and `docs/benchmarks/.json`
+on every tag, and appended to the GitHub release notes if a release exists.
+Artifacts expire; the docs page is the durable series.
+
+### How it measures
+
+- **Every command is forked from a small shim, not from the runner.** On Linux a
+ forked child inherits its parent's resident pages and `execve` folds that into the
+ reported `maxrss`, so a child of a process that has imported anndata cannot measure
+ below roughly 200 MB — which would have flattened every row in the table. Measured
+ with a 330 MB parent: 326 MB for a no-op child directly, 8 MB through the shim.
+ `posix_spawn` does not help (329 MB), so dropping `preexec_fn` is not a fix.
+- **`os.wait4`, not `resource.getrusage`.** `RUSAGE_CHILDREN` is a running
+ maximum over every child ever reaped, so one large case would poison every
+ later row. Output goes to temporary files rather than pipes, because
+ `communicate()` reaps the child and loses its rusage.
+- **`RLIMIT_AS` at 12 GiB on every child.** Cases where the in-memory
+ baseline cannot cope are the whole point, but an uncontained OOM kills the
+ runner agent and the job ends with no report. With a ceiling it is a row
+ that says `out of memory` at a limit we can state. macOS refuses
+ `RLIMIT_AS`, so a local run is unbounded.
+- **Baselines run from venvs built up front**, not `uv run --with`. The
+ latter is right for building fixtures, as `reference_stores.py` does, and
+ wrong here: the first invocation would put several hundred megabytes of
+ wheel downloads into the measured wall time and uv's own memory into the
+ measured peak. scanpy therefore never enters `uv.lock` or the image.
+
+### Rules that keep the comparison honest
+
+These live in `benchmarks/cases.py` as well, because this is the part that
+decays fastest.
+
+- **Use the best idiom the baseline has.** Comparing `export dataframe`
+ against a full `read_h5ad()` is a strawman -- anndata reads just `obs` via
+ `read_elem`. The good idiom is the primary row; the naive load may appear
+ only as a clearly labelled second row.
+- **Pin compression on both sides.** adata-cli forwards the source's
+ settings; `write_h5ad` defaults to none.
+- **`n/a` is a result.** Where scanpy has no equivalent, or
+ `concat_on_disk` refuses, say which and why. An omitted row reads as an
+ oversight.
+- **Include the startup floor.** The CLI costs 0.3-1 s to import, and
+ `import scanpy` 3-8 s; on a small tier that is the entire measurement.
+- **Do not hide the rows where the baseline wins.** `_concat_csr` loops per
+ row in Python and scipy's C `vstack` will often beat it on time at many
+ times the memory. That trade is the argument for the tool.
+- **Give the baseline everything it needs.** `concat_on_disk` imports `dask` to
+ concatenate a dense element and fails outright without it, so the baseline
+ environments install it. Measuring a library crippled by a missing optional
+ dependency would be measuring our own setup.
+- **Four commands are deliberately not benchmarked.** `export image`, `export dict`,
+ `import image` and `import dict` have no library equivalent, so their rows would only
+ ever read `n/a` while adding runtime to every tag. The complexity guards cover them.
+- **Say whether the page cache was dropped.** A fixture written seconds ago
+ is entirely in RAM, which understates streaming. The runner can drop it; a
+ laptop usually cannot, and the report states which happened.
+
## Compatibility testing against real anndata releases
`test_anndata_versions.py` does not trust this repo's idea of the format. For
diff --git a/docs/index.md b/docs/index.md
index ed04b98..b90e217 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -34,8 +34,11 @@ docker run --rm -it -v /path/to/data:/data \
out in `.h5ad`, and what this tool does with it.
- **[Element spec: Zarr](ELEMENTS_zarr.md)** — the same for `.zarr`, including
the v2/v3 differences.
-- **[Testing](TESTING.md)** — how the suite is organised, and how compatibility
- is verified against six real anndata releases.
+- **[Testing](TESTING.md)** — how the suite is organised, how compatibility is
+ verified against six real anndata releases, and the complexity guards that
+ keep cost regressions out.
+- **[Benchmarks](BENCHMARKS.md)** — peak memory and wall time against anndata
+ and scanpy, remeasured and republished on every release.
## At a glance
@@ -48,6 +51,8 @@ adata subset data.h5ad -o cortex.h5ad --obs-query "cluster == Cortex_2"
adata split data.h5ad --by sample -o per_sample/
adata concat per_sample/*.h5ad -o merged.h5ad --join outer --label sample
+adata convert data.h5ad X -o small.h5ad --dtype float32
+
adata create new.h5ad --n-obs 5000 --n-var 2000
adata import sparse new.h5ad X counts.mtx --inplace
```
diff --git a/pyproject.toml b/pyproject.toml
index 7767eca..c4de0fa 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -3,7 +3,7 @@
# so the distribution is published as `pyadata-cli`. The import package
# and the command are both still `adata`.
name = "pyadata-cli"
-version = "0.5.1"
+version = "0.6.0"
description = "Streaming CLI for exploring and editing large AnnData .h5ad and .zarr stores"
readme = "README.md"
requires-python = ">=3.12"
diff --git a/pytest.ini b/pytest.ini
index e7b2969..a116cbf 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -6,4 +6,5 @@ python_functions = test_*
addopts = -v --strict-markers --tb=short --timeout=300 --timeout-method=thread
markers =
slow: marks tests as slow (deselect with '-m "not slow"')
+ perf: complexity guards that trace Python execution (10-50x slower)
integration: builds environments with uv and needs the network on first run
diff --git a/src/adata/cli.py b/src/adata/cli.py
index ed7dc8d..cd496a5 100644
--- a/src/adata/cli.py
+++ b/src/adata/cli.py
@@ -6,9 +6,11 @@
from rich.console import Console
import typer
+from adata.core.convert import LAYOUTS
from adata.commands import (
- MERGE_STRATEGIES,
+ MERGE_CHOICES,
concat_stores,
+ convert_store,
create_store,
split_store,
list_store,
@@ -413,12 +415,14 @@ def concat(
merge: Optional[str] = typer.Option(
None,
"--merge",
- help="How to reconcile var columns: same, unique, first, only (default: drop)",
+ help="How to reconcile var columns: drop, same, unique, first, only "
+ "(default: drop)",
),
uns_merge: Optional[str] = typer.Option(
None,
"--uns-merge",
- help="How to reconcile uns: same, unique, first, only (default: drop)",
+ help="How to reconcile uns: drop, same, unique, first, only "
+ "(default: drop)",
),
fill_value: float = typer.Option(
0.0, "--fill-value", help="Value for dense cells introduced by an outer join"
@@ -442,13 +446,18 @@ def concat(
adata concat a.h5ad b.h5ad -o m.h5ad --keys a,b --index-unique - --uns-merge same
"""
for name, value in (("--merge", merge), ("--uns-merge", uns_merge)):
- if value is not None and value not in MERGE_STRATEGIES:
+ if value is not None and value not in MERGE_CHOICES:
console.print(
f"[bold red]Error:[/] {name} must be one of: "
- f"{', '.join(MERGE_STRATEGIES)}"
+ f"{', '.join(MERGE_CHOICES)}"
)
raise typer.Exit(code=1)
+ # "drop" is the default, and naming it explicitly has to be allowed: a
+ # config that spells out the default should not be rejected.
+ merge = None if merge == "drop" else merge
+ uns_merge = None if uns_merge == "drop" else uns_merge
+
try:
concat_stores(
files,
@@ -469,6 +478,134 @@ def concat(
raise typer.Exit(code=1)
+# ============================================================================
+# CONVERT command
+# ============================================================================
+@app.command("convert")
+def convert(
+ file: Path = typer.Argument(
+ ...,
+ help="Input .h5ad/.zarr",
+ exists=True,
+ readable=True,
+ dir_okay=True,
+ file_okay=True,
+ ),
+ entries: Optional[List[str]] = typer.Argument(
+ None,
+ help="Matrix paths to convert, e.g. 'X', 'layers/counts', 'raw/X'",
+ ),
+ output: Optional[Path] = typer.Option(
+ None,
+ "--output",
+ "-o",
+ help="Output .h5ad/.zarr path. Required unless --inplace.",
+ dir_okay=True,
+ file_okay=True,
+ ),
+ inplace: bool = typer.Option(
+ False,
+ "--inplace",
+ help="Modify source file directly.",
+ ),
+ convert_all: bool = typer.Option(
+ False,
+ "--all",
+ help="Convert X, every layer, and raw/X",
+ ),
+ dtype: Optional[str] = typer.Option(
+ None,
+ "--dtype",
+ help="New dtype for the values, e.g. float32",
+ ),
+ indices_dtype: Optional[str] = typer.Option(
+ None,
+ "--indices-dtype",
+ help="New dtype for sparse indices: int32 or int64",
+ ),
+ layout: Optional[str] = typer.Option(
+ None,
+ "--layout",
+ help="Target layout: csr, csc, dense or sparse",
+ ),
+ force: bool = typer.Option(
+ False,
+ "--force",
+ help="Convert despite a lossy cast or a large size increase",
+ ),
+ in_memory: bool = typer.Option(
+ False,
+ "--in-memory",
+ help="Transpose in memory instead of streaming (faster if it fits)",
+ ),
+ chunk_rows: int = typer.Option(
+ 1024,
+ "--chunk",
+ "-C",
+ help="Row chunk size for dense matrices",
+ ),
+ zarr_format: Optional[int] = typer.Option(
+ None,
+ "--zarr-format",
+ help="Zarr spec version to write (defaults to the source store's)",
+ ),
+) -> None:
+ """
+ Change a matrix's dtype, layout or density.
+
+ Counts held as float64 cost twice the disk and twice the read for no
+ information; a tool that wants CSC cannot use a CSR store; and concat
+ refuses inputs whose encodings disagree. All three are this command.
+
+ A cast that would not round-trip is refused before anything is written,
+ as is a densification that would inflate the store; --force overrides
+ both. Transposing streams by default, so it works on matrices too large
+ to load.
+
+ Examples:
+ adata convert data.h5ad X -o out.h5ad --dtype float32
+ adata convert data.h5ad X -o out.h5ad --layout csc
+ adata convert data.h5ad X --inplace --dtype float32 --indices-dtype int32
+ adata convert data.h5ad --all -o out.h5ad --dtype float32
+ """
+ if not inplace and output is None:
+ console.print(
+ "[bold red]Error:[/] Output file is required. "
+ "Use --output/-o or --inplace.",
+ )
+ raise typer.Exit(code=1)
+
+ if layout is not None and layout not in LAYOUTS:
+ console.print(
+ f"[bold red]Error:[/] --layout must be one of: {', '.join(LAYOUTS)}."
+ )
+ raise typer.Exit(code=1)
+
+ if zarr_format is not None and zarr_format not in (2, 3):
+ console.print("[bold red]Error:[/] --zarr-format must be 2 or 3.")
+ raise typer.Exit(code=1)
+
+ try:
+ convert_store(
+ file=file,
+ entries=list(entries) if entries else None,
+ output=output,
+ console=console,
+ dtype=dtype,
+ indices_dtype=indices_dtype,
+ layout=layout,
+ convert_all=convert_all,
+ inplace=inplace,
+ chunk_rows=chunk_rows,
+ in_memory=in_memory,
+ force=force,
+ zarr_format=zarr_format,
+ )
+ except Exception as e:
+ console.print(f"[bold red]Error:[/] {e}")
+ raise typer.Exit(code=1)
+
+
# ============================================================================
# SPLIT command
# ============================================================================
diff --git a/src/adata/commands/__init__.py b/src/adata/commands/__init__.py
index 70e0690..5ed2e4c 100644
--- a/src/adata/commands/__init__.py
+++ b/src/adata/commands/__init__.py
@@ -5,4 +5,5 @@
from adata.commands.ls import list_store
from adata.commands.create import create_store
from adata.commands.split import split_store
-from adata.commands.concat import MERGE_STRATEGIES, concat_stores
+from adata.commands.concat import MERGE_CHOICES, MERGE_STRATEGIES, concat_stores
+from adata.commands.convert import convert_store
diff --git a/src/adata/commands/concat.py b/src/adata/commands/concat.py
index 650b854..befda91 100644
--- a/src/adata/commands/concat.py
+++ b/src/adata/commands/concat.py
@@ -7,9 +7,9 @@
from rich.console import Console
-from adata.core.concat import MERGE_STRATEGIES, concat_on_disk
+from adata.core.concat import MERGE_CHOICES, MERGE_STRATEGIES, concat_on_disk
-__all__ = ["MERGE_STRATEGIES", "concat_stores"]
+__all__ = ["MERGE_CHOICES", "MERGE_STRATEGIES", "concat_stores"]
def concat_stores(
diff --git a/src/adata/commands/convert.py b/src/adata/commands/convert.py
new file mode 100644
index 0000000..54844ff
--- /dev/null
+++ b/src/adata/commands/convert.py
@@ -0,0 +1,218 @@
+"""The `convert` command: rewrite matrices, copy everything else.
+
+Writes a whole new store rather than editing one in place, for the same
+reason `subset` does: a conversion that fails half way leaves the original
+untouched, and `--inplace` becomes an atomic swap of a finished file rather
+than a partial edit of a live one.
+"""
+
+from __future__ import annotations
+
+import shutil
+from pathlib import Path
+from typing import Any, List, Optional
+
+from rich.console import Console
+
+from adata.core.convert import (
+ DEFAULT_CHUNK,
+ LAYOUTS,
+ Plan,
+ convert_matrix,
+ parse_dtype,
+ plan_conversion,
+)
+from adata.elements import spec
+from adata.elements.write import ensure_anndata_skeleton
+from adata.storage import copy_tree, detect_backend, is_group, open_store
+
+#: What `--all` means. The matrices an AnnData object is built from, not
+#: every 2-D array in the file: obsm/varm hold embeddings, whose dtype is
+#: rarely what anyone is trying to shrink. Those still convert by path.
+ALL_PREFIXES = ("X", "layers", "raw/X")
+
+
+def discover_matrices(root: Any) -> List[str]:
+ """Entry paths `--all` resolves to, in a stable order."""
+ found: List[str] = []
+ if "X" in root:
+ found.append("X")
+ if "layers" in root and is_group(root["layers"]):
+ found.extend(f"layers/{key}" for key in sorted(root["layers"].keys()))
+ if "raw" in root and is_group(root["raw"]) and "X" in root["raw"]:
+ found.append("raw/X")
+ return found
+
+
+def _same_path(left: Path, right: Path) -> bool:
+ """Do these name the same store, through symlinks and `..` alike?"""
+ try:
+ return left.resolve() == right.resolve()
+ except OSError: # pragma: no cover - unresolvable path
+ return left.absolute() == right.absolute()
+
+
+def _resolve(root: Any, path: str) -> Any:
+ obj = root
+ for part in path.split("/"):
+ if part not in obj:
+ raise KeyError(f"{path!r} not found in the store.")
+ obj = obj[part]
+ return obj
+
+
+def convert_store(
+ file: Path,
+ entries: Optional[List[str]],
+ output: Optional[Path],
+ console: Console,
+ *,
+ dtype: Optional[str] = None,
+ indices_dtype: Optional[str] = None,
+ layout: Optional[str] = None,
+ convert_all: bool = False,
+ inplace: bool = False,
+ chunk: int = DEFAULT_CHUNK,
+ chunk_rows: int = 1024,
+ in_memory: bool = False,
+ force: bool = False,
+ zarr_format: Optional[int] = None,
+) -> None:
+ """Write a copy of `file` with the named matrices converted."""
+ if not inplace and output is None:
+ raise ValueError("Output file is required unless --inplace is specified.")
+ if dtype is None and indices_dtype is None and layout is None:
+ raise ValueError(
+ "Nothing to do: pass at least one of --dtype, --indices-dtype "
+ "or --layout."
+ )
+ if layout is not None and layout not in LAYOUTS:
+ raise ValueError(f"--layout must be one of: {', '.join(LAYOUTS)}.")
+
+ data_dtype = parse_dtype(dtype) if dtype else None
+ index_dtype = (
+ parse_dtype(indices_dtype, allowed=("int32", "int64"))
+ if indices_dtype
+ else None
+ )
+
+ if output is not None and not inplace and _same_path(file, output):
+ # Opening the destination "w" clears it while the source is still
+ # being read from it. HDF5 refuses; Zarr does not, and quietly
+ # produced an empty store where the data used to be.
+ raise ValueError(
+ f"Output path is the input: {output}. Use --inplace to replace "
+ "it, which writes to a temporary file first."
+ )
+
+ if inplace:
+ backend = detect_backend(file)
+ if backend == "zarr":
+ base = file.stem if file.suffix else file.name
+ dst_path = file.with_name(f"{base}.convert-tmp.zarr")
+ else:
+ dst_path = file.with_name(f"{file.name}.convert-tmp")
+ if dst_path.exists():
+ raise FileExistsError(f"Temporary path already exists: {dst_path}")
+ else:
+ dst_path = output
+
+ if zarr_format is None and detect_backend(file) == "zarr":
+ with open_store(file, "r") as probe:
+ zarr_format = probe.zarr_format
+
+ with open_store(file, "r") as src_store:
+ src = src_store.root
+ targets = discover_matrices(src) if convert_all else list(entries or [])
+ if not targets:
+ raise ValueError(
+ "No matrices selected. Name one (e.g. `X`) or pass --all."
+ )
+ # Every check first, before the destination exists. A refusal must
+ # not leave behind a store holding a copy of obs and var and nothing
+ # else -- the caller cannot tell that from a finished conversion.
+ plans = {
+ path: plan_conversion(
+ _resolve(src, path),
+ path,
+ dtype=data_dtype,
+ index_dtype=index_dtype,
+ layout=layout,
+ chunk=chunk,
+ force=force,
+ console=console,
+ )
+ for path in targets
+ }
+
+ console.print(
+ f"[cyan]Converting {len(targets)} matrix/matrices:[/] "
+ + ", ".join(targets)
+ )
+
+ with open_store(dst_path, "w", zarr_format=zarr_format) as dst_store:
+ dst = dst_store.root
+ _write(
+ src, dst, plans,
+ chunk=chunk, chunk_rows=chunk_rows, in_memory=in_memory,
+ console=console,
+ )
+ ensure_anndata_skeleton(dst)
+
+ if inplace:
+ if file.is_dir():
+ shutil.rmtree(file)
+ elif file.exists():
+ file.unlink()
+ if dst_path.is_dir():
+ shutil.move(str(dst_path), str(file))
+ else:
+ dst_path.replace(file)
+ console.print(f"[green]Converted[/] {file}")
+ else:
+ console.print(f"[green]Wrote[/] {dst_path}")
+
+
+def _write(src: Any, dst: Any, plans: dict, **options: Any) -> None:
+ """Copy the store across, converting the targeted entries as they pass."""
+ by_parent: dict = {}
+ for path in plans:
+ parent, _, leaf = path.rpartition("/")
+ by_parent.setdefault(parent, {})[leaf] = plans[path]
+
+ for key in src.keys():
+ if key in by_parent.get("", {}):
+ convert_matrix(
+ src[key], dst, key, plan=by_parent[""][key], **options
+ )
+ elif key in by_parent:
+ # Any parent, not just layers and raw. Restricting it to those
+ # two meant an explicitly named `obsm/X_pca` was copied
+ # unconverted and the command still reported success.
+ _write_group(src[key], dst, key, by_parent, **options)
+ else:
+ copy_tree(src[key], dst, key)
+
+
+def _write_group(
+ group: Any, dst: Any, name: str, by_parent: dict, **options: Any
+) -> None:
+ """Recreate one container, converting the members that were named."""
+ from adata.storage import copy_attrs, is_zarr_group
+
+ out = dst.create_group(name)
+ copy_attrs(
+ group.attrs,
+ out.attrs,
+ target_backend="zarr" if is_zarr_group(dst) else "hdf5",
+ )
+ if not spec.encoding_type(out):
+ spec.set_encoding(out, spec.RAW if name == "raw" else spec.DICT)
+
+
+ wanted = by_parent.get(name, {})
+ for key in group.keys():
+ if key in wanted:
+ convert_matrix(group[key], out, key, plan=wanted[key], **options)
+ else:
+ copy_tree(group[key], out, key)
diff --git a/src/adata/core/concat.py b/src/adata/core/concat.py
index 4982916..39aa7ea 100644
--- a/src/adata/core/concat.py
+++ b/src/adata/core/concat.py
@@ -48,6 +48,8 @@
)
MERGE_STRATEGIES = ("same", "unique", "first", "only")
+#: What the CLI accepts. "drop" is the default and maps to no merge at all.
+MERGE_CHOICES = ("drop",) + MERGE_STRATEGIES
def _index_union(per_input: Sequence[List[str]]) -> List[str]:
@@ -174,6 +176,16 @@ def __repr__(self) -> str: # pragma: no cover - debugging aid
_MISSING = _Missing()
+class _Present:
+ """Stands for a column whose value was not read, only its presence."""
+
+ def __repr__(self) -> str: # pragma: no cover - debugging aid
+ return ""
+
+
+_PRESENT = _Present()
+
+
def _equal(a: Any, b: Any) -> bool:
if isinstance(a, _Incomparable) or isinstance(b, _Incomparable):
return False
@@ -212,15 +224,20 @@ def _concat_categorical(
concatenated directly. A row from an input lacking the column gets code
-1, which anndata reads back as a missing value.
"""
- categories: List[str] = []
+ # Hash lookup, not list membership. `if category not in categories` on a
+ # list is O(k) per probe and so O(k^2) over the union -- around 2 million
+ # string comparisons at 1024 categories, and 5 billion at 100k. It reads
+ # the same number of bytes either way, so only a comparison count sees it.
+ lookup: Dict[str, int] = {}
for col in columns:
if col is None:
continue
for category in read_categories(col):
- if category not in categories:
- categories.append(str(category))
+ text = str(category)
+ if text not in lookup:
+ lookup[text] = len(lookup)
- lookup = {c: i for i, c in enumerate(categories)}
+ categories = list(lookup)
ordered = all(is_ordered(c) for c in columns if c is not None)
codes = np.full(sum(lengths), -1, dtype=np.int64)
@@ -277,33 +294,35 @@ def _concat_masked(
total = sum(lengths)
mask = np.ones(total, dtype=bool)
- values: List[Any] = [None] * total
+ present = [c for c in columns if c is not None]
+
+ # Fill a typed buffer by slice, rather than a Python list of length
+ # n_obs one element at a time. The list cost an object per row and three
+ # full passes over it, which is what made obs concatenation the most
+ # allocation-hungry part of a streamed concat.
+ if enc == spec.NULLABLE_STRING_ARRAY:
+ from adata.elements.read import decode_str_array
+
+ filled = np.empty(total, dtype=object)
+ filled[:] = ""
+ else:
+ decode_str_array = None
+ dtype = np.result_type(*[c["values"].dtype for c in present])
+ filled = np.zeros(total, dtype=dtype)
offset = 0
for col, length in zip(columns, lengths):
if col is not None:
chunk = np.asarray(col["values"][...])
- chunk_mask = np.asarray(col["mask"][...], dtype=bool)
- for i in range(length):
- values[offset + i] = chunk[i]
- mask[offset : offset + length] = chunk_mask
+ if decode_str_array is not None:
+ chunk = decode_str_array(chunk)
+ filled[offset : offset + length] = chunk
+ mask[offset : offset + length] = np.asarray(
+ col["mask"][...], dtype=bool
+ )
offset += length
- if enc == spec.NULLABLE_STRING_ARRAY:
- from adata.elements.read import decode_str_array
-
- filled = [
- "" if v is None else decode_str_array(np.asarray([v]))[0] for v in values
- ]
- write_masked(parent, name, filled, mask, enc)
- return
-
- present = [c for c in columns if c is not None]
- dtype = np.result_type(*[c["values"].dtype for c in present])
- filled_num = np.array(
- [0 if v is None else v for v in values], dtype=dtype
- )
- write_masked(parent, name, filled_num, mask, enc)
+ write_masked(parent, name, filled, mask, enc)
def _concat_string(
@@ -542,10 +561,13 @@ def _check(label: str, sources: List[Any]) -> None:
kinds = {_matrix_kind(s) for s in sources}
if kinds in ({spec.CSR_MATRIX}, {spec.CSC_MATRIX}, {"dense"}):
return
+ wanted = sorted(kinds)[0]
raise ValueError(
f"Cannot concatenate {label!r}: inputs use "
f"{', '.join(sorted(kinds))}. Every input must use the same "
- "encoding -- convert them to match first."
+ f"encoding. Convert them to match first, e.g. "
+ f"`adata convert INPUT {label} -o converted.h5ad --layout "
+ f"{'dense' if wanted == 'dense' else wanted.removesuffix('_matrix')}`."
)
if all("X" in r for r in roots):
@@ -605,6 +627,9 @@ def _concat_matrix(
)
return True
+ # A backstop. `check_matrix_encodings` runs before the output store is
+ # created and should have raised already; this catches an element it
+ # does not cover, where failing late still beats writing nonsense.
raise ValueError(
f"Cannot concatenate {name!r}: inputs use {', '.join(sorted(kinds))}."
)
@@ -945,13 +970,25 @@ def _write_var(
positions = [_column_map(target_var, names) for names in var_names]
written: List[str] = []
+ # "first" and "only" decide on presence alone, so the column values are
+ # never read for them.
+ compares = merge in ("same", "unique")
+
for name in candidates:
aligned: List[Any] = []
for group, where in zip(groups, positions):
if name not in group or (where < 0).any():
aligned.append(_MISSING)
continue
- aligned.append(tuple(read_str_all(group[name])[i] for i in where))
+ if not compares:
+ aligned.append(_PRESENT)
+ continue
+ # Read the column once and index the result. Reading it inside the
+ # generator -- as an earlier version did -- re-read the whole
+ # column for every target variable, which is quadratic and turns a
+ # 36k-var merge into hours of pure CPU.
+ values = read_str_all(group[name])
+ aligned.append(tuple(values[i] for i in where))
keep, _ = _merge_values(aligned, merge)
if not keep:
diff --git a/src/adata/core/convert.py b/src/adata/core/convert.py
new file mode 100644
index 0000000..fb3cfd3
--- /dev/null
+++ b/src/adata/core/convert.py
@@ -0,0 +1,906 @@
+"""Rewriting a matrix's dtype, layout or density, on disk.
+
+Three conversions, one command. Counts stored as float64 cost twice the disk
+and twice the read for no information (issue #13); a CSR store handed to a
+tool that wants CSC has to be transposed somewhere; and `concat` refuses
+inputs whose encodings disagree, which until now left nothing to do about it.
+
+Everything here streams. The matrix these conversions matter for is the one
+too large to load, so a converter that loads it would only work on the files
+that did not need converting.
+
+Two safety rules, both of which fail before anything is written -- the same
+principle `check_matrix_encodings` states for concat, because a half-written
+store is worse than a refusal:
+
+* a cast that does not round-trip is refused unless forced, so `int32`
+ indices that would overflow, or a float downcast that would lose real
+ precision, are reported rather than silently written;
+* densifying is refused when it would inflate the store beyond
+ `MAX_GROWTH_FACTOR`, because a 5% dense matrix becomes ten times its size
+ and that should not be a surprise.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+from rich.console import Console
+
+from adata.elements import spec
+from adata.elements.write import set_shape_attr
+from adata.storage import (
+ copy_attrs,
+ create_dataset,
+ dataset_create_kwargs,
+ is_dataset,
+ is_group,
+ is_zarr_group,
+)
+
+#: Layouts the user can ask for. "sparse" means "whichever sparse encoding
+#: keeps the major axis it already has", so densify/sparsify round-trips.
+LAYOUTS = ("csr", "csc", "dense", "sparse")
+
+#: What `--dtype` accepts. An allowlist rather than `np.dtype(text)`: that
+#: would cheerfully accept "S10" or "datetime64[ns]" and produce a store no
+#: reader expects from a matrix.
+DTYPES = (
+ "float16", "float32", "float64",
+ "int8", "int16", "int32", "int64",
+ "uint8", "uint16", "uint32", "uint64",
+ "bool",
+)
+
+#: Indices and indptr must be integers, and signed: anndata and scipy both
+#: expect a signed index type.
+INDEX_DTYPES = ("int32", "int64")
+
+#: How much larger a densified store may get before it needs --force.
+MAX_GROWTH_FACTOR = 4.0
+
+#: 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:
+ """Resolve a user-supplied dtype name, or say what is accepted."""
+ name = text.strip().lower()
+ if name not in allowed:
+ raise ValueError(
+ f"Unknown dtype {text!r}. Choose from: {', '.join(allowed)}."
+ )
+ return np.dtype(name)
+
+
+# ---------------------------------------------------------------------------
+# safety
+
+
+@dataclass
+class CastReport:
+ """What a cast would do to the values, measured rather than assumed."""
+
+ total: int = 0
+ changed: int = 0
+ worst_absolute: float = 0.0
+ overflowed: bool = False
+
+ @property
+ def lossless(self) -> bool:
+ return self.changed == 0 and not self.overflowed
+
+ def describe(self, source: Any, target: Any) -> str:
+ if self.lossless:
+ return f"{source} -> {target}: every value round-trips"
+ share = self.changed / self.total if self.total else 0.0
+ detail = (
+ "values exceed its range"
+ if self.overflowed
+ else f"largest change {self.worst_absolute:g}"
+ )
+ return (
+ f"{source} -> {target} is lossy: {self.changed:,} of "
+ f"{self.total:,} values ({share:.2%}) do not round-trip, {detail}"
+ )
+
+
+def check_cast(dataset: Any, target: np.dtype, *, chunk: int = DEFAULT_CHUNK) -> CastReport:
+ """Would casting `dataset` to `target` lose anything?
+
+ Streams the values, casts each block and casts it back. An exact
+ round-trip is the only honest test: whether float64 counts survive
+ float32 depends on the counts, not on the dtypes, and that is precisely
+ the question issue #13 is asking.
+ """
+ report = CastReport()
+ source = np.dtype(getattr(dataset, "dtype", "float64"))
+ n = int(dataset.shape[0]) if getattr(dataset, "shape", None) else 0
+
+ info = np.finfo(target) if target.kind == "f" else (
+ np.iinfo(target) if target.kind in "iu" else None
+ )
+
+ for start in range(0, n, chunk):
+ block = np.asarray(dataset[start : min(start + chunk, n)])
+ report.total += block.size
+ if block.size == 0:
+ continue
+
+ with np.errstate(invalid="ignore", over="ignore"):
+ cast = block.astype(target)
+ back = cast.astype(source)
+
+ if info is not None:
+ finite = block[np.isfinite(block)] if source.kind == "f" else block
+ if finite.size and (
+ float(finite.max()) > float(info.max)
+ or float(finite.min()) < float(info.min)
+ ):
+ report.overflowed = True
+
+ # NaN never equals itself, so compare those separately rather than
+ # counting every missing value as a loss.
+ differs = back != block
+ if source.kind == "f":
+ both_nan = np.isnan(block) & np.isnan(back)
+ differs &= ~both_nan
+ count = int(differs.sum())
+ if count:
+ report.changed += count
+ delta = np.abs(
+ block[differs].astype("float64") - back[differs].astype("float64")
+ )
+ finite_delta = delta[np.isfinite(delta)]
+ if finite_delta.size:
+ report.worst_absolute = max(
+ report.worst_absolute, float(finite_delta.max())
+ )
+
+ return report
+
+
+def check_index_dtype(
+ group: Any,
+ index_dtype: np.dtype,
+ pointer_dtype: np.dtype,
+ shape: Tuple[int, int],
+) -> None:
+ """Refuse index dtypes that cannot address this matrix.
+
+ Cheaper than `check_cast`: the largest value each array must hold is
+ bounded by the dimensions and the nonzero count, so no pass over the
+ data is needed. `indices` holds coordinates, bounded by the larger
+ dimension; `indptr` holds offsets, bounded by nnz. They are checked
+ separately because they can legitimately need different widths.
+ """
+ nnz = int(group["indices"].shape[0])
+ for name, dtype, largest, what in (
+ ("indices", index_dtype, max(int(shape[0]), int(shape[1])), "coordinates"),
+ ("indptr", pointer_dtype, nnz, "offsets"),
+ ):
+ limit = int(np.iinfo(dtype).max)
+ if largest > limit:
+ raise ValueError(
+ f"{dtype} cannot hold this matrix's {what}: {name} must reach "
+ f"{largest:,} for a {shape[0]:,} x {shape[1]:,} matrix with "
+ f"{nnz:,} nonzeros, and {dtype} tops out at {limit:,}. "
+ "Use int64."
+ )
+
+
+def _stored_bytes(obj: Any) -> int:
+ """Bytes this element occupies, as best the backend will say."""
+ total = 0
+ targets = [obj] if is_dataset(obj) else [
+ obj[k] for k in ("data", "indices", "indptr") if k in obj
+ ]
+ for item in targets:
+ try:
+ total += int(item.nbytes)
+ except Exception: # pragma: no cover - backend without nbytes
+ shape = getattr(item, "shape", ()) or ()
+ size = int(np.prod(shape)) if shape else 0
+ total += size * int(getattr(item.dtype, "itemsize", 8) or 8)
+ return total
+
+
+def check_growth(
+ current: int, projected: int, *, force: bool, what: str
+) -> None:
+ """Refuse a conversion that inflates the store, unless asked twice."""
+ if force or current <= 0 or projected <= current * MAX_GROWTH_FACTOR:
+ return
+ raise ValueError(
+ f"{what} would grow from {current / 1e6:,.0f} MB to "
+ f"{projected / 1e6:,.0f} MB ({projected / current:.1f}x). That is "
+ f"above the {MAX_GROWTH_FACTOR:g}x limit; pass --force if it is what "
+ "you want."
+ )
+
+
+# ---------------------------------------------------------------------------
+# reading a source
+
+
+@dataclass
+class Matrix:
+ """A matrix on disk, described enough to convert it."""
+
+ obj: Any
+ kind: str # "csr_matrix" | "csc_matrix" | "dense"
+ shape: Tuple[int, int]
+ dtype: np.dtype
+
+ @property
+ def sparse(self) -> bool:
+ return self.kind in spec.SPARSE_TYPES
+
+ @property
+ def nnz(self) -> int:
+ return int(self.obj["indices"].shape[0]) if self.sparse else 0
+
+
+def describe(obj: Any) -> Matrix:
+ """Classify a matrix element, or say why it is not one."""
+ enc = spec.encoding_type(obj)
+ if is_group(obj) and enc in spec.SPARSE_TYPES:
+ shape = obj.attrs.get("shape", None)
+ if shape is None:
+ raise ValueError("Sparse matrix group is missing its 'shape' attribute.")
+ return Matrix(obj, enc, (int(shape[0]), int(shape[1])), obj["data"].dtype)
+
+ if is_dataset(obj):
+ if getattr(obj, "ndim", 0) != 2:
+ raise ValueError(
+ f"Only 2-D matrices can be converted; this one has "
+ f"{getattr(obj, 'ndim', '?')} dimension(s)."
+ )
+ return Matrix(obj, "dense", (int(obj.shape[0]), int(obj.shape[1])), obj.dtype)
+
+ raise ValueError(
+ f"Not a matrix: encoding {enc!r}. `convert` handles X, layers, raw/X "
+ "and any 2-D dense array."
+ )
+
+
+def resolve_layout(source: Matrix, requested: Optional[str]) -> str:
+ """The concrete target encoding for a possibly-vague request."""
+ if requested is None:
+ return source.kind
+ if requested == "dense":
+ return "dense"
+ if requested == "sparse":
+ # Keep the major axis it already had, so sparse -> dense -> sparse is
+ # the identity rather than a silent transpose.
+ return source.kind if source.sparse else spec.CSR_MATRIX
+ return spec.CSR_MATRIX if requested == "csr" else spec.CSC_MATRIX
+
+
+# ---------------------------------------------------------------------------
+# writers
+
+
+def _new_sparse_group(
+ src: Matrix, dst_parent: Any, name: str, enc: str, shape: Tuple[int, int]
+) -> Any:
+ group = dst_parent.create_group(name)
+ backend = "zarr" if is_zarr_group(dst_parent) else "hdf5"
+ copy_attrs(src.obj.attrs, group.attrs, target_backend=backend)
+ spec.set_encoding(group, enc)
+ set_shape_attr(group, shape)
+ return group
+
+
+def _sparse_dataset(
+ group: Any, name: str, dtype: np.dtype, n: int, template: Any
+) -> Any:
+ """A 1-D dataset of `n` elements laid out like `template`.
+
+ Forwarding compression and chunking matters more here than anywhere
+ else: the point of a dtype change is usually to make the file smaller,
+ and creating the output with a fixed 65,536-element chunk made a
+ 2,400-nonzero matrix allocate 786 KB of mostly empty chunk -- eight
+ times the source, from a conversion asked for to halve it.
+ """
+ from adata.core.subset import _clamp_chunks
+
+ backend = "zarr" if is_zarr_group(group) else "hdf5"
+ kw = dataset_create_kwargs(template, target_backend=backend, dst_parent=group)
+ kw = _clamp_chunks(kw, max(1, n))
+ if "chunks" not in kw and n:
+ kw["chunks"] = (min(n, 1 << 16),)
+ return create_dataset(group, name, shape=(n,), dtype=dtype, **kw)
+
+
+def _growable_like(group: Any, name: str, dtype: np.dtype, template: Any) -> Any:
+ """Like `_sparse_dataset`, but extensible for a size not yet known."""
+ backend = "zarr" if is_zarr_group(group) else "hdf5"
+ kw = dataset_create_kwargs(template, target_backend=backend, dst_parent=group)
+ kw.pop("shards", None)
+ chunks = kw.pop("chunks", None)
+ step = int(chunks[0]) if chunks else 1 << 16
+ if is_zarr_group(group):
+ return group.create_array(name, shape=(0,), dtype=dtype, chunks=(step,), **kw)
+ return group.create_dataset(
+ name, shape=(0,), maxshape=(None,), dtype=dtype, chunks=(step,), **kw
+ )
+
+
+def _write_sparse_arrays(
+ group: Any,
+ data: np.ndarray,
+ indices: np.ndarray,
+ indptr: np.ndarray,
+ *,
+ data_dtype: np.dtype,
+ index_dtype: np.dtype,
+ pointer_dtype: np.dtype,
+ source: Any,
+) -> None:
+ for name, values, dtype, template in (
+ ("data", data, data_dtype, source["data"]),
+ ("indices", indices, index_dtype, source["indices"]),
+ ("indptr", indptr, pointer_dtype, source["indptr"]),
+ ):
+ cast = values.astype(dtype, copy=False)
+ dataset = _sparse_dataset(group, name, dtype, cast.size, template)
+ if cast.size:
+ dataset[:] = cast
+
+
+def cast_sparse(
+ src: Matrix,
+ dst_parent: Any,
+ name: str,
+ *,
+ data_dtype: np.dtype,
+ index_dtype: np.dtype,
+ pointer_dtype: np.dtype,
+ chunk: int = DEFAULT_CHUNK,
+) -> None:
+ """Rewrite a sparse matrix with new dtypes, keeping its layout.
+
+ The cheap case, and the one issue #13 asks for: the structure is
+ untouched, so this is a straight streamed copy of `data` and `indices`
+ into differently typed datasets. Sized up front, because nnz is known.
+ """
+ group = _new_sparse_group(src, dst_parent, name, src.kind, src.shape)
+ source_data, source_indices = src.obj["data"], src.obj["indices"]
+ nnz = int(source_data.shape[0])
+
+ out_data = _sparse_dataset(group, "data", data_dtype, nnz, source_data)
+ out_indices = _sparse_dataset(group, "indices", index_dtype, nnz, source_indices)
+
+ for start in range(0, nnz, chunk):
+ end = min(start + chunk, nnz)
+ out_data[start:end] = np.asarray(source_data[start:end]).astype(
+ data_dtype, copy=False
+ )
+ out_indices[start:end] = np.asarray(source_indices[start:end]).astype(
+ index_dtype, copy=False
+ )
+
+ indptr = np.asarray(src.obj["indptr"][...]).astype(pointer_dtype, copy=False)
+ out_indptr = _sparse_dataset(
+ group, "indptr", pointer_dtype, indptr.size, src.obj["indptr"]
+ )
+ out_indptr[:] = indptr
+
+
+def transpose_sparse_in_memory(
+ src: Matrix,
+ dst_parent: Any,
+ name: str,
+ *,
+ data_dtype: np.dtype,
+ index_dtype: np.dtype,
+ pointer_dtype: np.dtype = np.dtype("int64"),
+) -> None:
+ """Swap CSR<->CSC by loading the matrix and sorting it once.
+
+ Faster than streaming whenever the matrix fits, and the whole matrix is
+ what it needs -- so it is opt-in, never the default.
+ """
+ target = (
+ spec.CSC_MATRIX if src.kind == spec.CSR_MATRIX else spec.CSR_MATRIX
+ )
+ n_major_in = src.shape[0] if src.kind == spec.CSR_MATRIX else src.shape[1]
+ n_major_out = src.shape[1] if src.kind == spec.CSR_MATRIX else src.shape[0]
+
+ indptr = np.asarray(src.obj["indptr"][...], dtype=np.int64)
+ minor = np.asarray(src.obj["indices"][...], dtype=np.int64)
+ values = np.asarray(src.obj["data"][...])
+ major = np.repeat(np.arange(n_major_in, dtype=np.int64), np.diff(indptr))
+
+ # Sort by the new major axis, then the new minor, which is what both
+ # encodings require of `indices` within a row.
+ order = np.lexsort((major, minor))
+ out_indptr = np.concatenate(
+ ([0], np.cumsum(np.bincount(minor, minlength=n_major_out)))
+ )
+
+ group = _new_sparse_group(src, dst_parent, name, target, src.shape)
+ _write_sparse_arrays(
+ group,
+ values[order],
+ major[order],
+ out_indptr,
+ data_dtype=data_dtype,
+ index_dtype=index_dtype,
+ pointer_dtype=pointer_dtype,
+ source=src.obj,
+ )
+
+
+def transpose_sparse_streaming(
+ src: Matrix,
+ dst_parent: Any,
+ name: str,
+ *,
+ data_dtype: np.dtype,
+ index_dtype: np.dtype,
+ pointer_dtype: np.dtype = np.dtype("int64"),
+ chunk: int = DEFAULT_CHUNK,
+ bucket_entries: Optional[int] = None,
+ console: Optional[Console] = None,
+) -> None:
+ """Swap CSR<->CSC without loading the matrix.
+
+ A transpose cannot be done in one pass: the first entry of the output
+ may come from the last row of the input. Three passes instead, with
+ memory set by `bucket_entries` rather than by nnz:
+
+ 1. count nonzeros per output major, by streaming `indices` alone. That
+ gives the output `indptr` by cumulative sum.
+ 2. stream the input again, splitting each block's entries into buckets by
+ which slice of the output they land in, and append each bucket to
+ scratch datasets.
+ 3. read one bucket at a time, sort it, and append to the output in order.
+
+ 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
+ )
+ n_major_in = src.shape[0] if src.kind == spec.CSR_MATRIX else src.shape[1]
+ n_major_out = src.shape[1] if src.kind == spec.CSR_MATRIX else src.shape[0]
+
+ indptr = np.asarray(src.obj["indptr"][...], dtype=np.int64)
+ source_indices, source_data = src.obj["indices"], src.obj["data"]
+ nnz = int(source_indices.shape[0])
+
+ # Pass 1: how many entries land in each output major.
+ 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)
+ 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)
+ n_buckets = len(bounds) - 1
+ if console is not None and n_buckets > 1:
+ console.print(
+ f"[dim]Transposing {nnz:,} nonzeros through {n_buckets} buckets[/]"
+ )
+
+ # nnz is invariant under a transpose, so the output can be sized now and
+ # written by slice rather than grown block by block.
+ group = _new_sparse_group(src, dst_parent, name, target, src.shape)
+ out_data = _sparse_dataset(group, "data", data_dtype, nnz, source_data)
+ out_indices = _sparse_dataset(group, "indices", index_dtype, nnz, source_indices)
+ written = 0
+
+ scratch_name = f"__{name}_transpose_scratch__"
+ scratch = dst_parent.create_group(scratch_name)
+ try:
+ 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)
+ ]
+
+ # Pass 2: scatter the input into buckets, one input block at a time.
+ # Step over the major axis in blocks holding about `chunk` nonzeros,
+ # so the read size is set by the data rather than by how many rows
+ # happen to be empty.
+ average = max(1, nnz // max(1, n_major_in))
+ major_step = max(1, chunk // average)
+ for lo in range(0, n_major_in, major_step):
+ hi = min(lo + major_step, n_major_in)
+ start, end = int(indptr[lo]), int(indptr[hi])
+ if end <= start:
+ continue
+ minor = np.asarray(source_indices[start:end], dtype=np.int64)
+ values = np.asarray(source_data[start:end])
+ 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])
+
+ # 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(
+ index_dtype, copy=False
+ )
+ out_data[written : written + count] = values[order].astype(
+ data_dtype, copy=False
+ )
+ written += count
+ finally:
+ del dst_parent[scratch_name]
+
+ cast_indptr = out_indptr.astype(pointer_dtype, copy=False)
+ dataset = _sparse_dataset(
+ group, "indptr", pointer_dtype, cast_indptr.size, src.obj["indptr"]
+ )
+ dataset[:] = cast_indptr
+
+
+def densify(
+ src: Matrix,
+ dst_parent: Any,
+ name: str,
+ *,
+ data_dtype: np.dtype,
+ chunk_rows: int = 1024,
+) -> None:
+ """Write a sparse matrix out as a dense array, a block of rows at a time.
+
+ The zeros are the point: nothing here materialises the whole grid, so a
+ matrix too large to densify in memory still converts -- it just produces
+ a file that is honestly much larger, which `check_growth` warns about
+ before any of it is written.
+ """
+ n_rows, n_cols = src.shape
+ backend = "zarr" if is_zarr_group(dst_parent) else "hdf5"
+ dst = create_dataset(
+ dst_parent, name, shape=(n_rows, n_cols), dtype=data_dtype
+ )
+ copy_attrs(src.obj.attrs, dst.attrs, target_backend=backend)
+ spec.set_encoding(dst, spec.ARRAY)
+ # `shape` belongs to the sparse encoding and would contradict the array's
+ # own shape if it were carried over.
+ if "shape" in dst.attrs:
+ del dst.attrs["shape"]
+
+ indptr = np.asarray(src.obj["indptr"][...], dtype=np.int64)
+ indices, values = src.obj["indices"], src.obj["data"]
+ csr = src.kind == spec.CSR_MATRIX
+ n_major = n_rows if csr else n_cols
+
+ for lo in range(0, n_major, chunk_rows):
+ hi = min(lo + chunk_rows, n_major)
+ start, end = int(indptr[lo]), int(indptr[hi])
+ block = np.zeros(
+ (hi - lo, n_cols) if csr else (n_rows, hi - lo), dtype=data_dtype
+ )
+ if end > start:
+ minor = np.asarray(indices[start:end], dtype=np.int64)
+ data = np.asarray(values[start:end])
+ major = np.repeat(
+ np.arange(hi - lo, dtype=np.int64), np.diff(indptr[lo : hi + 1])
+ )
+ # Repeated coordinates are legal in a CSR/CSC store and mean
+ # their sum, which is what scipy's own `toarray` produces.
+ # Plain assignment keeps whichever came last, so a
+ # non-canonical input silently changed value on densifying.
+ if csr:
+ np.add.at(block, (major, minor), data)
+ else:
+ np.add.at(block, (minor, major), data)
+ if csr:
+ dst[lo:hi, :] = block
+ else:
+ dst[:, lo:hi] = block
+
+
+def sparsify(
+ src: Matrix,
+ dst_parent: Any,
+ name: str,
+ *,
+ enc: str,
+ data_dtype: np.dtype,
+ index_dtype: np.dtype,
+ pointer_dtype: np.dtype = np.dtype("int64"),
+ chunk_rows: int = 1024,
+ console: Optional[Console] = None,
+) -> Tuple[int, float]:
+ """Write a dense array out as CSR or CSC, a block at a time.
+
+ Returns `(nnz, density)` so the caller can say whether it was worth it --
+ a matrix that is half nonzero gets bigger, not smaller, and the user
+ should hear that from us rather than from `du`.
+ """
+ from adata.core.subset import _append, _growable
+
+ n_rows, n_cols = src.shape
+ csr = enc == spec.CSR_MATRIX
+ n_major = n_rows if csr else n_cols
+
+ group = dst_parent.create_group(name)
+ backend = "zarr" if is_zarr_group(dst_parent) else "hdf5"
+ copy_attrs(src.obj.attrs, group.attrs, target_backend=backend)
+ spec.set_encoding(group, enc)
+ set_shape_attr(group, src.shape)
+
+ out_data = _growable_like(group, "data", data_dtype, src.obj)
+ out_indices = _growable_like(group, "indices", index_dtype, src.obj)
+ counts: List[int] = []
+
+ for lo in range(0, n_major, chunk_rows):
+ hi = min(lo + chunk_rows, n_major)
+ block = np.asarray(src.obj[lo:hi, :] if csr else src.obj[:, lo:hi])
+ if not csr:
+ block = block.T # iterate majors as rows either way
+ nonzero_major, nonzero_minor = np.nonzero(block)
+ counts.extend(np.bincount(nonzero_major, minlength=hi - lo).tolist())
+ _append(out_indices, nonzero_minor.astype(index_dtype, copy=False))
+ _append(out_data, block[nonzero_major, nonzero_minor].astype(
+ data_dtype, copy=False
+ ))
+
+ indptr = np.concatenate(([0], np.cumsum(counts, dtype=np.int64)))
+ create_dataset(group, "indptr", data=indptr.astype(pointer_dtype, copy=False))
+
+
+ nnz = int(indptr[-1])
+ density = nnz / max(1, n_rows * n_cols)
+ if console is not None and density > 0.5:
+ console.print(
+ f"[yellow]{name} is {density:.0%} nonzero; the sparse form is "
+ f"larger than the dense one below about 33%.[/]"
+ )
+ return nnz, density
+
+
+def cast_dense(
+ src: Matrix,
+ dst_parent: Any,
+ name: str,
+ *,
+ data_dtype: np.dtype,
+ chunk_rows: int = 1024,
+) -> None:
+ """Rewrite a dense matrix with a new dtype, a block of rows at a time."""
+ n_rows, n_cols = src.shape
+ backend = "zarr" if is_zarr_group(dst_parent) else "hdf5"
+ kw = dataset_create_kwargs(
+ src.obj, target_backend=backend, dst_parent=dst_parent
+ )
+ from adata.core.subset import _clamp_chunks
+
+ dst = create_dataset(
+ dst_parent,
+ name,
+ shape=(n_rows, n_cols),
+ dtype=data_dtype,
+ **_clamp_chunks(kw, n_rows, n_cols),
+ )
+ copy_attrs(src.obj.attrs, dst.attrs, target_backend=backend)
+
+ for lo in range(0, n_rows, chunk_rows):
+ hi = min(lo + chunk_rows, n_rows)
+ dst[lo:hi, :] = np.asarray(src.obj[lo:hi, :]).astype(
+ data_dtype, copy=False
+ )
+
+
+# ---------------------------------------------------------------------------
+# dispatch
+
+
+@dataclass
+class Plan:
+ """A checked conversion, ready to run.
+
+ Separating the decision from the writing is what makes "fails before
+ anything is written" true rather than aspirational: the caller resolves
+ every plan first, and only then creates the destination store. An
+ earlier version checked inside the writer, so a refusal still left an
+ output file holding a copy of obs and var.
+ """
+
+ source: Matrix
+ layout: str
+ data_dtype: np.dtype
+ index_dtype: np.dtype
+ #: `indptr` is tracked apart from `indices` because the two can
+ #: legitimately differ: a narrow matrix with more than 2^31 nonzeros
+ #: needs int64 offsets over int32 column indices. Inferring one from
+ #: the other silently overflowed the offsets and corrupted the matrix.
+ pointer_dtype: np.dtype = np.dtype("int64")
+ report: Optional[CastReport] = None
+
+
+def plan_conversion(
+ obj: Any,
+ name: str,
+ *,
+ dtype: Optional[np.dtype] = None,
+ index_dtype: Optional[np.dtype] = None,
+ layout: Optional[str] = None,
+ chunk: int = DEFAULT_CHUNK,
+ force: bool = False,
+ console: Optional[Console] = None,
+) -> Plan:
+ """Decide what to do, and refuse here if it should not be done."""
+ src = describe(obj)
+ target_layout = resolve_layout(src, layout)
+ data_dtype = np.dtype(dtype) if dtype is not None else src.dtype
+
+ if index_dtype is not None:
+ idx_dtype = ptr_dtype = np.dtype(index_dtype)
+ elif src.sparse:
+ # Keep what the source used, each independently. Defaulting to
+ # int64 doubled the index arrays of every int32 store; inferring
+ # indptr from indices narrowed the offsets of every store that
+ # needed them wider.
+ idx_dtype = np.dtype(src.obj["indices"].dtype)
+ ptr_dtype = np.dtype(src.obj["indptr"].dtype)
+ else:
+ idx_dtype = ptr_dtype = np.dtype("int64")
+
+ report: Optional[CastReport] = None
+ if dtype is not None and data_dtype != src.dtype:
+ values = src.obj["data"] if src.sparse else src.obj
+ report = check_cast(values, data_dtype, chunk=chunk)
+ message = report.describe(src.dtype, data_dtype)
+ if not report.lossless and not force:
+ raise ValueError(f"{name}: {message}. Pass --force to convert anyway.")
+ if console is not None:
+ colour = "dim" if report.lossless else "yellow"
+ console.print(f"[{colour}]{name}: {message}[/]")
+
+ # Always, not only when asked: an inferred dtype can be too narrow too,
+ # and a silently overflowed offset is indistinguishable from corruption.
+ if src.sparse and not force:
+ check_index_dtype(src.obj, idx_dtype, ptr_dtype, src.shape)
+
+ if target_layout == "dense" and src.sparse:
+ projected = src.shape[0] * src.shape[1] * data_dtype.itemsize
+ check_growth(
+ _stored_bytes(src.obj), projected, force=force, what=f"{name} as dense"
+ )
+
+ return Plan(src, target_layout, data_dtype, idx_dtype, ptr_dtype, report)
+
+
+def convert_matrix(
+ obj: Any,
+ dst_parent: Any,
+ name: str,
+ *,
+ plan: Optional[Plan] = None,
+ dtype: Optional[np.dtype] = None,
+ index_dtype: Optional[np.dtype] = None,
+ layout: Optional[str] = None,
+ chunk: int = DEFAULT_CHUNK,
+ chunk_rows: int = 1024,
+ in_memory: bool = False,
+ force: bool = False,
+ console: Optional[Console] = None,
+) -> None:
+ """Write `obj` into `dst_parent` under `name`, converted.
+
+ Pass a `plan` from `plan_conversion` to have the checks already done;
+ without one they run here, which is convenient for a direct caller but
+ means the destination already exists by the time a refusal is raised.
+ """
+ if plan is None:
+ plan = plan_conversion(
+ obj, name, dtype=dtype, index_dtype=index_dtype, layout=layout,
+ chunk=chunk, force=force, console=console,
+ )
+ src = plan.source
+ target_layout = plan.layout
+ data_dtype = plan.data_dtype
+ idx_dtype = plan.index_dtype
+ ptr_dtype = plan.pointer_dtype
+
+ # --- then write -------------------------------------------------------
+ if target_layout == "dense":
+ if src.sparse:
+ densify(src, dst_parent, name, data_dtype=data_dtype, chunk_rows=chunk_rows)
+ else:
+ cast_dense(src, dst_parent, name, data_dtype=data_dtype, chunk_rows=chunk_rows)
+ return
+
+ if not src.sparse:
+ sparsify(
+ src,
+ dst_parent,
+ name,
+ enc=target_layout,
+ data_dtype=data_dtype,
+ index_dtype=idx_dtype,
+ pointer_dtype=ptr_dtype,
+ chunk_rows=chunk_rows,
+ console=console,
+ )
+ return
+
+ if target_layout == src.kind:
+ cast_sparse(
+ src,
+ dst_parent,
+ name,
+ data_dtype=data_dtype,
+ index_dtype=idx_dtype,
+ pointer_dtype=ptr_dtype,
+ chunk=chunk,
+ )
+ return
+
+ transpose = (
+ transpose_sparse_in_memory if in_memory else transpose_sparse_streaming
+ )
+ extra: Dict[str, Any] = (
+ {} if in_memory else {"chunk": chunk, "console": console}
+ )
+ transpose(
+ src,
+ dst_parent,
+ name,
+ data_dtype=data_dtype,
+ index_dtype=idx_dtype,
+ pointer_dtype=ptr_dtype,
+ **extra,
+ )
diff --git a/src/adata/core/select.py b/src/adata/core/select.py
index f1049bc..49edb43 100644
--- a/src/adata/core/select.py
+++ b/src/adata/core/select.py
@@ -100,11 +100,31 @@ def group_indices(
values = np.asarray(
col_chunk_as_strings(group, column, start, end, cache), dtype=str
)
- for label in dict.fromkeys(values.tolist()):
+
+ # One pass over the chunk, whatever the number of distinct labels.
+ # Taking `np.nonzero(values == label)` per label -- as an earlier
+ # version did -- rescanned the whole chunk once per label, so the
+ # cost was O(n_rows * n_groups): a million cells split by a thousand
+ # samples came to 10^9 comparisons, and `split` looked like a hang.
+ uniques, first_seen, codes = np.unique(
+ values, return_index=True, return_inverse=True
+ )
+ codes = codes.ravel()
+
+ # Sorting the codes puts each label's positions in one contiguous
+ # run, so every group is a slice rather than a search.
+ by_code = np.argsort(codes, kind="stable")
+ run_starts = np.searchsorted(codes[by_code], np.arange(len(uniques)), "left")
+ run_ends = np.searchsorted(codes[by_code], np.arange(len(uniques)), "right")
+
+ # `np.unique` sorts; `order` has to stay in order of first appearance,
+ # because it names the output files.
+ for code in np.argsort(first_seen, kind="stable"):
+ label = str(uniques[code])
if label not in buckets:
buckets[label] = []
order.append(label)
- buckets[label].append(np.nonzero(values == label)[0] + start)
+ buckets[label].append(by_code[run_starts[code] : run_ends[code]] + start)
return (
{k: np.concatenate(v).astype(np.int64) for k, v in buckets.items()},
diff --git a/src/adata/storage/__init__.py b/src/adata/storage/__init__.py
index cfea148..025bce5 100644
--- a/src/adata/storage/__init__.py
+++ b/src/adata/storage/__init__.py
@@ -450,30 +450,86 @@ def _is_string_src(src: Any) -> bool:
#: comfortably above a typical 1 MiB Lustre stripe while bounding peak memory.
TARGET_READ_BYTES = 32 * 1024 * 1024
-#: Assumed width of a variable-length string element, for read sizing only.
-#: h5py reports `itemsize` 8 for vlen strings because the value is a pointer,
-#: which would overestimate the row count by an order of magnitude and break
-#: the memory bound. Real cell and gene names sit well under this.
+#: Fallback width of a variable-length string element, used only when the
+#: real width cannot be sampled. h5py reports `itemsize` 8 for vlen strings
+#: because the value is a pointer, which would overestimate the row count by
+#: an order of magnitude. Real cell and gene names sit well under this.
VLEN_ELEMENT_BYTES = 64
+#: Elements read when sampling a variable-length array's real element width.
+#: One small read, against a copy that is about to stream the whole array.
+VLEN_SAMPLE_ROWS = 256
-def _row_bytes(shape: Sequence[int], dtype: Any) -> int:
- """In-memory size of one row along the first axis."""
+
+def _is_vlen(dtype: Any) -> bool:
+ """Does this dtype hide its real element width behind a pointer?
+
+ 'O' is how h5py spells vlen str, 'T' is numpy StringDType as reported by
+ zarr-python 3. Neither `itemsize` reflects the bytes actually stored.
+ """
+ itemsize = int(getattr(dtype, "itemsize", 0) or 0)
+ return itemsize <= 0 or getattr(dtype, "kind", None) in ("O", "T")
+
+
+def _sample_element_bytes(src: Any, n_rows: int) -> int:
+ """Mean stored width of a variable-length element, by reading a few.
+
+ An assumed width is not a bound. Estimating 64 bytes and reading
+ `TARGET_READ_BYTES // 64` elements means the read is 32 MiB only if the
+ guess holds: at 4 KiB elements it is 2 GiB, and for an array shorter than
+ the computed step the whole thing is read at once. Measured before this:
+ copying 200,000 strings of 4 KiB peaked at 827 MB against a stated 32 MiB
+ budget.
+
+ One small read fixes that, and it is negligible beside the copy it is
+ about to size -- `uns` can hold arbitrary text, so the width is not
+ something this layer can assume.
+ """
+ try:
+ sample = np.asarray(src[: min(VLEN_SAMPLE_ROWS, max(1, n_rows))])
+ except Exception: # pragma: no cover - unreadable source
+ return VLEN_ELEMENT_BYTES
+
+ total = 0
+ count = 0
+ for value in sample.reshape(-1)[:VLEN_SAMPLE_ROWS]:
+ try:
+ total += len(value)
+ except TypeError: # pragma: no cover - non-sized element
+ total += VLEN_ELEMENT_BYTES
+ count += 1
+
+ if not count:
+ return VLEN_ELEMENT_BYTES
+ # The floor keeps a column of empty strings from producing an unbounded
+ # step; the object header dominates at that size anyway.
+ return max(VLEN_ELEMENT_BYTES, total // count)
+
+
+def _row_bytes(
+ shape: Sequence[int], dtype: Any, element_bytes: int = VLEN_ELEMENT_BYTES
+) -> int:
+ """In-memory size of one row along the first axis.
+
+ `element_bytes` is the measured width for a variable-length dtype; the
+ default is the fallback for callers that have no sample to offer.
+ """
width = 1
for dim in shape[1:]:
width *= max(1, int(dim))
itemsize = int(getattr(dtype, "itemsize", 0) or 0)
- # 'O' is how h5py spells vlen str, 'T' is numpy StringDType as reported by
- # zarr-python 3. Neither itemsize reflects the bytes actually stored.
- if itemsize <= 0 or getattr(dtype, "kind", None) in ("O", "T"):
- itemsize = VLEN_ELEMENT_BYTES
+ if _is_vlen(dtype):
+ itemsize = element_bytes
return max(1, width * itemsize)
def _chunk_step(
- shape: Sequence[int], chunks: Optional[Sequence[int]], dtype: Any
+ shape: Sequence[int],
+ chunks: Optional[Sequence[int]],
+ dtype: Any,
+ element_bytes: int = VLEN_ELEMENT_BYTES,
) -> int:
"""Rows to copy per read, sized for the filesystem rather than the source.
@@ -492,7 +548,7 @@ def _chunk_step(
if chunks is not None and len(chunks) > 0 and chunks[0]:
chunk_rows = max(1, int(chunks[0]))
- step = max(1, TARGET_READ_BYTES // _row_bytes(shape, dtype))
+ step = max(1, TARGET_READ_BYTES // _row_bytes(shape, dtype, element_bytes))
if chunk_rows:
# Never go below a single chunk: a partial read still decompresses the
# whole thing, so a smaller step costs the same I/O for less data. That
@@ -532,7 +588,11 @@ def copy_dataset(src: Any, dst_group: Any, name: str) -> Any:
ds[()] = src[()]
return ds
- step = _chunk_step(shape, getattr(src, "chunks", None), src.dtype)
+ element_bytes = (
+ _sample_element_bytes(src, shape[0]) if _is_vlen(src.dtype)
+ else VLEN_ELEMENT_BYTES
+ )
+ step = _chunk_step(shape, getattr(src, "chunks", None), src.dtype, element_bytes)
for start in range(0, shape[0], step):
end = min(start + step, shape[0])
if len(shape) == 1:
diff --git a/tests/perf_counters.py b/tests/perf_counters.py
new file mode 100644
index 0000000..cd7e257
--- /dev/null
+++ b/tests/perf_counters.py
@@ -0,0 +1,529 @@
+"""Deterministic cost counters for the complexity guards.
+
+These measure *how much work* an operation does, never how long it takes. A
+test that can fail because a CI runner was busy does not belong in a merge
+gate, so nothing here touches the clock.
+
+Why patch the libraries rather than `adata`
+-------------------------------------------
+The hooks wrap h5py and zarr themselves. A seam inside `src/adata/` would only
+measure the call sites that remembered to use it, which is precisely the wrong
+property for a guard whose job is to catch the read you did not think of. The
+matrix paths in particular slice the backend objects directly -- see
+`core/subset.py` `_copy_dense_rows` and `core/concat.py` `_concat_csr` -- and
+no in-repo helper sees them at all.
+
+What is counted
+---------------
+Elements and bytes, not just calls. The `--merge` hang was n calls each reading
+n elements; counting calls alone would also miss a vectorised-but-quadratic
+variant that reads n elements n times in one call.
+
+Known bypasses
+--------------
+These reach the file without passing through any hook below:
+
+* ``h5py.Dataset.read_direct``
+* ``np.asarray(dataset)`` and ``dataset[...]`` via ``__array__`` (verified:
+ ``np.asarray`` adds nothing to the element count)
+* ``dataset.asstr()[...]``, which goes through ``AsStrWrapper.__getitem__``
+* ``dataset.fields(...)``
+
+None are used in `src/adata/` today. If you add one, the counters will silently
+report less work than actually happened -- which is why every guard also
+asserts a *lower* bound, and why `test_performance.py` keeps a canary test with
+a known absolute count.
+"""
+
+from __future__ import annotations
+
+import sys
+import tracemalloc
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, Iterator, List
+
+import numpy as np
+
+
+@dataclass
+class IOCounts:
+ """Live tally of backend reads and writes.
+
+ The object is yielded before the work runs and mutated in place, so read
+ the fields after the `with` block closes.
+ """
+
+ h5_calls: int = 0
+ h5_elements: int = 0
+ h5_write_calls: int = 0
+ h5_written_elements: int = 0
+ zarr_calls: int = 0
+ zarr_elements: int = 0
+ store_get: int = 0
+ store_get_bytes: int = 0
+ store_set: int = 0
+ store_delete: int = 0
+
+ #: Per-dataset element counts, for diagnosing which column blew up.
+ by_name: Dict[str, int] = field(default_factory=dict)
+
+ @property
+ def elements(self) -> int:
+ """Elements read through either backend's array interface."""
+ return self.h5_elements + self.zarr_elements
+
+ @property
+ def calls(self) -> int:
+ return self.h5_calls + self.zarr_calls
+
+ @property
+ def writes(self) -> int:
+ """Elements written through either backend.
+
+ Zarr is counted in store keys rather than elements -- a chunk write
+ is one `set` -- so the two are not the same unit. For a guard that
+ only compares a command against itself at two sizes, that is fine.
+ """
+ return self.h5_written_elements + self.store_set
+
+ @property
+ def work(self) -> int:
+ """Reads plus writes, for commands whose job is mostly writing.
+
+ `import` and `create` read almost nothing; measuring only reads made
+ their guards vacuous, which the lower bound in `assert_grows_linearly`
+ caught.
+ """
+ return self.reads + self.writes
+
+ @property
+ def reads(self) -> int:
+ """The headline number: array elements plus store-level fetches.
+
+ Covers both the "read the same column n times" shape and the
+ "re-fetch the same chunk n times" shape.
+ """
+ return self.elements + self.store_get
+
+ def __str__(self) -> str: # pragma: no cover - diagnostic only
+ return (
+ f"read={self.elements} (h5={self.h5_elements} "
+ f"zarr={self.zarr_elements}) wrote={self.h5_written_elements} "
+ f"calls={self.calls} store: get={self.store_get} "
+ f"set={self.store_set} delete={self.store_delete}"
+ )
+
+
+def _size(value: Any) -> int:
+ try:
+ return int(np.asarray(value).size)
+ except Exception: # pragma: no cover - exotic dtypes
+ return 1
+
+
+@contextmanager
+def count_io() -> Iterator[IOCounts]:
+ """Count backend reads and writes for the duration of the block.
+
+ Patches four seams, all verified against h5py 3.15.1 and zarr 3.1.5:
+
+ ====================================== =================================
+ ``h5py.Dataset.__getitem__`` HDF5 reads
+ ``h5py.Dataset.__setitem__`` HDF5 writes
+ ``h5py.Group.create_dataset`` HDF5 writes made at creation
+ ``zarr.Array.__getitem__`` Zarr reads, v2 and v3 alike
+ ``zarr.storage.LocalStore.get`` chunk and metadata fetches
+ ``LocalStore.set`` / ``.delete`` write storms
+ ====================================== =================================
+
+ Both write seams are needed: `create_dataset(name, data=...)` writes its
+ payload during creation and never touches `__setitem__`, so hooking only
+ the latter leaves `import` and `create` measuring zero.
+
+ The store hooks are the ones that see a metadata write storm: attribute
+ writes never touch `Array.__getitem__`, so an O(k^2) attribute rewrite is
+ invisible at the array level and obvious at the store level.
+
+ The `LocalStore` methods are coroutines and are wrapped as coroutines.
+ """
+ import h5py
+ import zarr
+ from zarr.storage import LocalStore
+
+ counts = IOCounts()
+
+ h5_get = h5py.Dataset.__getitem__
+ h5_set = h5py.Dataset.__setitem__
+ h5_create = h5py.Group.create_dataset
+ zarr_get = zarr.Array.__getitem__
+ store_get = LocalStore.get
+ store_set = LocalStore.set
+ store_delete = LocalStore.delete
+
+ def _record(name: str, n: int) -> None:
+ counts.by_name[name] = counts.by_name.get(name, 0) + n
+
+ def h5_wrapper(self: Any, key: Any) -> Any:
+ result = h5_get(self, key)
+ n = _size(result)
+ counts.h5_calls += 1
+ counts.h5_elements += n
+ _record(getattr(self, "name", "?"), n)
+ return result
+
+ def zarr_wrapper(self: Any, key: Any) -> Any:
+ result = zarr_get(self, key)
+ n = _size(result)
+ counts.zarr_calls += 1
+ counts.zarr_elements += n
+ _record(getattr(self, "path", "?") or "/", n)
+ return result
+
+ def h5_set_wrapper(self: Any, key: Any, value: Any) -> Any:
+ counts.h5_write_calls += 1
+ counts.h5_written_elements += _size(value)
+ return h5_set(self, key, value)
+
+ def h5_create_wrapper(self: Any, name: Any, *args: Any, **kwargs: Any) -> Any:
+ counts.h5_write_calls += 1
+ data = kwargs.get("data")
+ if data is None and args:
+ # positional signature is (name, shape, dtype, data, ...)
+ data = args[2] if len(args) > 2 else None
+ if data is not None:
+ counts.h5_written_elements += _size(data)
+ return h5_create(self, name, *args, **kwargs)
+
+ async def get_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+ result = await store_get(self, *args, **kwargs)
+ counts.store_get += 1
+ if result is not None:
+ try:
+ counts.store_get_bytes += len(result)
+ except TypeError: # pragma: no cover - prototype buffers
+ pass
+ return result
+
+ async def set_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+ counts.store_set += 1
+ return await store_set(self, *args, **kwargs)
+
+ async def delete_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+ counts.store_delete += 1
+ return await store_delete(self, *args, **kwargs)
+
+ h5py.Dataset.__getitem__ = h5_wrapper
+ h5py.Dataset.__setitem__ = h5_set_wrapper
+ h5py.Group.create_dataset = h5_create_wrapper
+ zarr.Array.__getitem__ = zarr_wrapper
+ LocalStore.get = get_wrapper
+ LocalStore.set = set_wrapper
+ LocalStore.delete = delete_wrapper
+ try:
+ yield counts
+ finally:
+ h5py.Dataset.__getitem__ = h5_get
+ h5py.Dataset.__setitem__ = h5_set
+ h5py.Group.create_dataset = h5_create
+ zarr.Array.__getitem__ = zarr_get
+ LocalStore.get = store_get
+ LocalStore.set = store_set
+ LocalStore.delete = store_delete
+
+
+@contextmanager
+def count_allocations() -> Iterator[List[int]]:
+ """Peak Python allocation during the block, in bytes, as `[peak]`.
+
+ `tracemalloc` sees Python-level allocation, which includes the numpy array
+ *objects* and every Python list and string built along the way, but not
+ memory a C extension takes outside the allocator. That is the right scope
+ here: the defects this catches are whole-column Python lists, and h5py's
+ own buffers are not what we are policing.
+
+ Deterministic for deterministic input, which is why it can gate a merge.
+ """
+ peak = [0]
+ started = not tracemalloc.is_tracing()
+ if started:
+ tracemalloc.start()
+ else: # pragma: no cover - nested use
+ tracemalloc.reset_peak()
+ try:
+ yield peak
+ finally:
+ peak[0] = tracemalloc.get_traced_memory()[1]
+ if started:
+ tracemalloc.stop()
+
+
+@contextmanager
+def count_lines(prefix: str) -> Iterator[List[int]]:
+ """Count executed Python line events in files under `prefix`, as `[n]`.
+
+ A deterministic stand-in for CPU time. It is what catches a quadratic that
+ performs no extra I/O and allocates nothing extra -- `if x not in a_list`
+ inside a loop, say. Costs a 10-50x slowdown, so keep it to the few paths
+ with known Python-level inner loops and to small inputs.
+
+ Blind to work done inside numpy's C code, so it complements the allocation
+ and I/O counters rather than replacing them.
+ """
+ total = [0]
+
+ def local_trace(frame: Any, event: str, arg: Any) -> Any:
+ if event == "line":
+ total[0] += 1
+ return local_trace
+
+ def trace(frame: Any, event: str, arg: Any) -> Any:
+ if event == "call" and frame.f_code.co_filename.startswith(prefix):
+ return local_trace
+ return None
+
+ previous = sys.gettrace()
+ sys.settrace(trace)
+ try:
+ yield total
+ finally:
+ sys.settrace(previous)
+
+
+#: numpy entry points that scan their operand, and the argument that is
+#: scanned. Patched by name, so only calls written as `np.f(...)` are seen.
+_SCANNING = {
+ "nonzero": 0,
+ "flatnonzero": 0,
+ "unique": 0,
+ "argsort": 0,
+ "sort": 0,
+ "searchsorted": 0,
+ "bincount": 0,
+ "equal": 0,
+ "isin": 0,
+}
+
+
+@contextmanager
+def count_scanned_elements() -> Iterator[List[int]]:
+ """Elements handed to numpy's scanning primitives, as `[n]`.
+
+ The third instrument, and the one the other two cannot replace. A loop of
+ the shape::
+
+ for label in distinct_labels:
+ np.nonzero(values == label)
+
+ reads nothing extra (the chunk is already in memory), allocates nothing
+ that lasts, and executes a constant number of Python lines per label. It
+ is invisible to `count_io`, to `count_allocations` and to `count_lines`
+ alike, and it is O(len(values) * len(distinct_labels)).
+
+ Counting what is passed into numpy makes it visible, in the same spirit
+ as the `_CountingStr` guard in `test_performance.py`: when the cost lives
+ inside C, count what is handed to C.
+
+ What this does and does not see
+ -------------------------------
+ Only calls that go through the `numpy` module namespace by name. An
+ operator -- `values == label` -- dispatches to `ndarray.__eq__` and then
+ to the ufunc in C, so it never passes the patched `np.equal`; a scan
+ written purely as `(a == b).sum()` is invisible. Method form,
+ `arr.argsort()`, is invisible for the same reason.
+
+ So this is a floor on the real work, not a measurement of it, which is
+ exactly what a guard needs: it can only under-report, and the guards
+ using it assert a lower bound so that under-reporting to nothing fails
+ rather than passes.
+ """
+ total = [0]
+ originals = {name: getattr(np, name) for name in _SCANNING}
+
+ def wrap(name: str, position: int):
+ original = originals[name]
+
+ def wrapper(*args: Any, **kwargs: Any) -> Any:
+ if len(args) > position:
+ try:
+ total[0] += int(np.size(args[position]))
+ except Exception: # pragma: no cover - exotic operands
+ pass
+ return original(*args, **kwargs)
+
+ return wrapper
+
+ for name, position in _SCANNING.items():
+ setattr(np, name, wrap(name, position))
+ try:
+ yield total
+ finally:
+ for name, original in originals.items():
+ setattr(np, name, original)
+
+
+# ---------------------------------------------------------------------------
+# the growth assertion
+
+
+#: Ratio of successive increments that counts as super-linear.
+#:
+#: Measured at 4x spacing, the increment ratio is 4.0 for linear work, about
+#: 4.4 for n log n, 8 for n**1.5 and 16 for quadratic. 6 sits in the gap with
+#: room on both sides, so it needs no per-test tuning.
+GROWTH_LIMIT = 6.0
+
+#: Sizes every growth guard runs at. 4x spacing is what makes GROWTH_LIMIT
+#: mean what it says; changing one without the other invalidates the constant.
+SIZES = (64, 256, 1024)
+
+
+def assert_grows_linearly(
+ measure: Callable[[int], int],
+ *,
+ what: str,
+ axis: str,
+ sizes: tuple = SIZES,
+ limit: float = GROWTH_LIMIT,
+) -> None:
+ """Assert `measure` grows no faster than linearly in one axis.
+
+ Compares successive *increments* rather than raw counts::
+
+ d1 = c(4n) - c(n)
+ d2 = c(16n) - c(4n)
+ assert d2 <= limit * d1
+
+ The increment form is what removes the arbitrary additive constant. Any
+ fixed setup cost -- opening the store, writing the skeleton, reading the
+ index once -- appears in both differences and cancels exactly, so there is
+ no slack parameter to pick and no floor for a small-coefficient quadratic
+ to hide under.
+
+ `d1 >= sizes[1]` is a vacuity guard. If a future refactor moves a read onto
+ a path the counters do not see (see the module docstring), every ratio
+ would pass trivially; this fails loudly instead.
+
+ Scale exactly one axis per call and hold the others fixed. Scaling two at
+ once makes legitimate work look quadratic -- an outer-join concat really
+ does produce n_obs x n_var_union cells.
+ """
+ small, mid, large = sizes
+ c_small, c_mid, c_large = (measure(n) for n in (small, mid, large))
+
+ d1 = c_mid - c_small
+ d2 = c_large - c_mid
+ series = (
+ f"{what}, scaling {axis}: "
+ f"{small}->{c_small}, {mid}->{c_mid}, {large}->{c_large} "
+ f"(increments {d1} then {d2}"
+ + (f", ratio {d2 / d1:.1f}" if d1 > 0 else "")
+ + ")"
+ )
+
+ # Half an operation per element added. Work that is exactly one
+ # operation per element -- `export dict` reads each key once -- gives
+ # d1 = mid - small, which is 0.75 * mid at 4x spacing, so a bound of
+ # `mid` would reject correct code. Anything the counters cannot see at
+ # all gives 0 and still fails.
+ floor = (mid - small) / 2
+ assert d1 >= floor, (
+ f"Too little measured work for the growth ratio to mean anything -- "
+ f"the counters are probably not seeing this operation. Expected at "
+ f"least {floor:.0f} more operations between {small} and {mid}. "
+ f"{series}"
+ )
+ assert d2 <= limit * d1, (
+ f"Cost grows faster than linearly in {axis}. Expect an increment "
+ f"ratio near 4.0 for linear work; {limit} is the limit and quadratic "
+ f"would be about 16. {series}"
+ )
+
+
+#: How much a "flat" cost may drift across the whole size span.
+#:
+#: Not 1.0: opening a store, resolving an index and printing a tree all cost
+#: a little more when there is more to describe -- a longer index name, a
+#: wider shape to format. 1.5 across a 64x span leaves room for that and none
+#: at all for reading the data, which would be 64x.
+FLAT_TOLERANCE = 1.5
+
+#: Span for the flat guards. Wider than SIZES, because the claim is stronger
+#: and a wide span is what makes it convincing.
+FLAT_SIZES = (64, 4096)
+
+
+def assert_independent_of(
+ measure: Callable[[int], int],
+ *,
+ what: str,
+ axis: str,
+ sizes: tuple = FLAT_SIZES,
+ tolerance: float = FLAT_TOLERANCE,
+) -> None:
+ """Assert `measure` does not grow with `axis` at all.
+
+ A stronger claim than `assert_grows_linearly`, and the right one wherever
+ the tool promises work proportional to something other than input size:
+
+ * **Inspection.** `view` and `ls` read shapes, dtypes and attributes and
+ never the values behind them, which is the whole reason they return
+ instantly on a store too large to open. Linear growth here would mean
+ the promise had quietly stopped holding.
+ * **Grouping.** Work per row must not depend on how many groups there are.
+ * **Streaming.** At a fixed chunk size, peak memory must not track the
+ size of the input.
+
+ Each of those reads as obvious prose and none of them is checked by a
+ linear-growth guard, which would happily accept a 64x increase.
+ """
+ small, large = sizes[0], sizes[-1]
+ c_small, c_large = measure(small), measure(large)
+
+ span = large / small
+ series = (
+ f"{what}, scaling {axis}: {small}->{c_small}, {large}->{c_large} "
+ f"over a {span:.0f}x span"
+ )
+ ratio = (c_large / c_small) if c_small else float("inf") if c_large else 1.0
+
+ assert ratio <= tolerance, (
+ f"Cost tracks {axis}, and it is supposed to be independent of it. "
+ f"Growing in step with the input would be about {span:.0f}x; "
+ f"{tolerance}x is the limit. {series} (ratio {ratio:.2f})"
+ )
+
+
+def assert_grows_slower_than_input(
+ measure: Callable[[int], int],
+ *,
+ what: str,
+ axis: str,
+ sizes: tuple,
+ at_least: float,
+) -> None:
+ """Assert cost grows at least `at_least` times slower than the input.
+
+ The middle ground between `assert_grows_linearly` and
+ `assert_independent_of`, and the honest shape of most of this tool's
+ streaming: peak memory is not flat -- an index or an indptr is read whole
+ -- but it is far below the input curve, and that margin is the feature.
+
+ Stating the margin as a factor rather than an absolute ceiling keeps the
+ guard meaningful across platforms and interpreters, and makes a drift
+ back towards linear fail long before it becomes a bug report.
+ """
+ small, large = sizes[0], sizes[-1]
+ c_small, c_large = measure(small), measure(large)
+
+ span = large / small
+ growth = (c_large / c_small) if c_small else float("inf")
+ budget = span / at_least
+
+ assert growth <= budget, (
+ f"{what} cost is tracking {axis} too closely. Input grew {span:.0f}x "
+ f"and cost grew {growth:.1f}x; the requirement is at least "
+ f"{at_least:.0f}x better than the input, i.e. no more than "
+ f"{budget:.0f}x. Series: {small}->{c_small}, {large}->{c_large}"
+ )
+
diff --git a/tests/test_benchmark_harness.py b/tests/test_benchmark_harness.py
new file mode 100644
index 0000000..7bc384b
--- /dev/null
+++ b/tests/test_benchmark_harness.py
@@ -0,0 +1,182 @@
+"""The benchmark harness has to be trustworthy before its numbers are.
+
+Only the measurement is tested here, not the cases. If peak RSS were
+attributed to the wrong process every table the project publishes would be
+wrong, and the failure would be invisible -- the numbers would still look
+plausible. Everything else in `benchmarks/` is reporting, and a mistake there
+shows up the moment anyone reads the page.
+"""
+
+from __future__ import annotations
+
+import sys
+
+import pytest
+
+from benchmarks._measure import Measurement, _looks_like_oom, measure
+
+
+#: A child that allocates nothing should cost about what a bare interpreter
+#: costs. Generous, because that depends on the build; far below the ballast
+#: below, which is the point.
+FLOOR_CEILING = 60 * 1024**2
+
+#: Allocated by the parent before measuring, to make an inherited floor
+#: impossible to miss.
+BALLAST = 300 * 1024**2
+
+
+def test_peak_rss_is_the_child_s_own_and_carries_no_floor():
+ """Two ways this number can lie, both of which it has.
+
+ `RUSAGE_CHILDREN` is a running maximum over every child a process has
+ reaped, so a 400 MB case makes every later one look like 400 MB. `wait4`
+ fixes that.
+
+ Worse, and what CI caught: on Linux a forked child inherits its parent's
+ resident pages, and `execve` folds that into the `maxrss` `wait4` reports.
+ A child of a fat parent could not appear small however little it used.
+ Measured under python:3.12-slim before the fix -- parent at 329.6 MB, a
+ no-op child reported 326.4 MB. `posix_spawn` reported 329.5 MB, so
+ dropping `preexec_fn` would not have helped; measuring from a shim
+ reported 8.1 MB.
+
+ macOS resets the high-water mark at exec and shows none of this, so this
+ is mostly a Linux guard -- which is the platform the benchmark runs on.
+
+ The test holds ballast in the parent for its whole duration so that an
+ inherited floor cannot hide.
+ """
+ ballast = bytearray(BALLAST)
+ ballast[::4096] = b"\x01" * len(ballast[::4096]) # make it resident
+
+ baseline = measure([sys.executable, "-c", "pass"])
+ big = measure([sys.executable, "-c", "x = bytearray(400 * 1024 * 1024)"])
+ small = measure([sys.executable, "-c", "x = bytearray(1024)"])
+
+ assert {baseline.status, big.status, small.status} == {"ok"}
+
+ # No floor: a child that allocates nothing looks like nothing, even though
+ # this process is holding 300 MB.
+ assert baseline.maxrss_bytes < FLOOR_CEILING, (
+ f"a no-op child reported {baseline.maxrss_bytes / 1e6:.0f} MB while "
+ f"this process held {BALLAST / 1e6:.0f} MB of ballast. The measured "
+ "peak is inheriting the parent's resident pages, so every benchmark "
+ "row would be floored at roughly the runner's own footprint."
+ )
+
+ # Isolation: the small child resembles the baseline, not the big one that
+ # ran before it.
+ assert big.maxrss_bytes > 300 * 1024**2, big
+ assert small.maxrss_bytes < baseline.maxrss_bytes * 2, (
+ f"peak RSS leaked between children: {small.maxrss_bytes / 1e6:.0f} MB "
+ f"for a 1 KB allocation, against a {baseline.maxrss_bytes / 1e6:.0f} MB "
+ f"baseline, after a {big.maxrss_bytes / 1e6:.0f} MB child"
+ )
+
+ del ballast
+
+
+def test_a_hang_is_reported_as_a_timeout_with_its_output_size(tmp_path):
+ """How the 0.5.1 hang presented: still running, output never growing.
+
+ A harness that recorded this as a slow success would have made the
+ reported bug look like ordinary slowness.
+ """
+ output = tmp_path / "partial.bin"
+ command = [
+ sys.executable,
+ "-c",
+ f"open({str(output)!r}, 'wb').write(b'x' * 1024); import time; time.sleep(60)",
+ ]
+ result = measure(command, output=output, timeout_s=2.0)
+
+ assert result.status == "timeout"
+ assert result.wall_s >= 2.0
+ assert result.output_bytes == 1024, "output size at the kill must be recorded"
+
+
+def test_a_failure_against_the_memory_ceiling_is_recorded_as_oom():
+ """A baseline that cannot cope is a result, not a crashed job.
+
+ Skipped on macOS, which refuses RLIMIT_AS; the ceiling only has to hold
+ on the Linux runner where the benchmark actually runs.
+ """
+ if not sys.platform.startswith("linux"):
+ pytest.skip("RLIMIT_AS is not enforceable on this platform")
+
+ result = measure(
+ [sys.executable, "-c", "x = bytearray(8 * 1024**3)"],
+ memory_limit=512 * 1024**2,
+ )
+ assert result.status == "oom", result
+
+
+@pytest.mark.parametrize(
+ "text",
+ ["MemoryError", "numpy: Unable to allocate 4.0 GiB", "std::bad_alloc"],
+)
+def test_out_of_memory_is_recognised_however_it_is_phrased(text):
+ assert _looks_like_oom(text)
+
+
+def test_an_ordinary_failure_is_not_mistaken_for_an_oom():
+ result = measure([sys.executable, "-c", "raise ValueError('nope')"])
+ assert result.status == "failed"
+ assert "ValueError" in result.stderr_tail
+
+
+def test_the_harness_refuses_to_overwrite_an_existing_output(tmp_path):
+ """A stale output would be sized and reported as this run's work.
+
+ Checked in the parent, not the shim, so it surfaces as an exception the
+ caller must fix rather than as a quiet "failed" row.
+ """
+ existing = tmp_path / "already.h5ad"
+ existing.write_bytes(b"old")
+ with pytest.raises(FileExistsError):
+ measure([sys.executable, "-c", "pass"], output=existing)
+
+
+def test_a_missing_binary_is_a_result_not_a_crash():
+ """`h5ls` ships with the HDF5 tools and is often simply absent."""
+ result = measure(["/nonexistent-binary-for-tests"])
+ assert result.status == "n/a"
+ assert "not installed" in result.stderr_tail
+
+
+def test_maxrss_is_normalised_to_bytes():
+ """`ru_maxrss` is KiB on Linux and bytes on macOS; the field is bytes.
+
+ A missing conversion would be a silent 1024x error in one direction on
+ one platform, which is exactly the kind of thing nobody notices in a
+ table of plausible-looking numbers.
+ """
+ result = measure([sys.executable, "-c", "pass"])
+ assert isinstance(result, Measurement)
+ # Any CPython start-up is well over 1 MB and well under 4 GB.
+ assert 1024**2 < result.maxrss_bytes < 4 * 1024**3, result.maxrss_bytes
+
+
+def test_publishing_keeps_the_page_prose():
+ """A republish must not reduce the docs page to bare tables.
+
+ `publish` rewrites `docs/BENCHMARKS.md` in full on every tag. An earlier
+ version wrote only the rendered results, which would have thrown away
+ everything explaining what the numbers mean the first time it ran.
+ """
+ from benchmarks.report import PAGE_TEMPLATE, build_page
+
+ template = PAGE_TEMPLATE.read_text()
+ assert "" in template, (
+ "the results marker is gone from page_template.md, so results would "
+ "be appended rather than placed"
+ )
+
+ page = build_page("## Results\n\nsome tables\n")
+ assert "some tables" in page
+ assert "Peak RSS is the headline" in page, "the framing was dropped"
+ assert "Running it yourself" in page, "the trailing sections were dropped"
+ assert page.count("\n# ") + page.startswith("# ") == 1, (
+ "the published page must have exactly one H1"
+ )
diff --git a/tests/test_commands_phase2.py b/tests/test_commands_phase2.py
index cf7bd8b..3ea6b99 100644
--- a/tests/test_commands_phase2.py
+++ b/tests/test_commands_phase2.py
@@ -758,3 +758,179 @@ def test_split_rejects_a_bad_zarr_format(tmp_path, sample):
)
assert result.exit_code == 1
assert "must be 2 or 3" in _out(result)
+
+
+# ---------------------------------------------------------------------------
+# concat --merge cost
+#
+# These assert how often a var column is read, not how long the merge takes.
+# The hang reported against 0.5.1 was invisible to every correctness test in
+# this file because all of them use two or three variables, and the defect was
+# quadratic in the number of variables: `tuple(read_str_all(col)[i] for i in
+# where)` re-evaluated `read_str_all` once per target variable. A wall-clock
+# assertion would be flaky on shared CI; a read count is exact.
+
+
+def _count_var_column_reads(monkeypatch):
+ """Count full-column reads performed while merging var."""
+ from adata.core import concat as concat_mod
+
+ calls: list = []
+ original = concat_mod.read_str_all
+
+ def counting(obj, *args, **kwargs):
+ calls.append(getattr(obj, "name", "?"))
+ return original(obj, *args, **kwargs)
+
+ monkeypatch.setattr(concat_mod, "read_str_all", counting)
+ return calls
+
+
+def _two_stores_with_var_columns(tmp_path, n_var: int, n_col: int) -> list:
+ genes = [f"g{i}" for i in range(n_var)]
+ var = pd.DataFrame(
+ {f"col{c}": [f"v{c}-{i}" for i in range(n_var)] for c in range(n_col)},
+ index=genes,
+ )
+ paths = []
+ for name, cells in (("a", ["c1", "c2"]), ("b", ["c3", "c4"])):
+ path = tmp_path / f"{name}.h5ad"
+ ad.AnnData(
+ X=np.ones((2, n_var), dtype="float32"),
+ obs=pd.DataFrame(index=cells),
+ var=var.copy(),
+ ).write_h5ad(path)
+ paths.append(path)
+ return paths
+
+
+@pytest.mark.parametrize("n_var", [4, 64])
+def test_concat_merge_same_reads_each_var_column_once_per_input(
+ tmp_path, monkeypatch, n_var
+):
+ """`--merge same` must not scale with the number of variables.
+
+ Two inputs and three var columns is six reads whatever n_var is. The
+ parametrisation is the whole point: an implementation whose cost grows
+ with n_var fails the second case while passing the first.
+ """
+ a, b = _two_stores_with_var_columns(tmp_path, n_var=n_var, n_col=3)
+ calls = _count_var_column_reads(monkeypatch)
+
+ out = tmp_path / "m.h5ad"
+ result = runner.invoke(
+ app, ["concat", str(a), str(b), "-o", str(out), "--merge", "same"]
+ )
+ assert result.exit_code == 0, _out(result)
+ assert len(calls) == 6, f"{len(calls)} reads for {n_var} vars: {calls[:10]}"
+
+ var = ad.read_h5ad(out).var
+ assert list(var.columns) == ["col0", "col1", "col2"]
+ assert var["col0"].tolist() == [f"v0-{i}" for i in range(n_var)]
+
+
+def test_concat_merge_first_does_not_read_var_column_values(
+ tmp_path, monkeypatch
+):
+ """`first` decides on presence alone, so it reads no column values."""
+ a, b = _two_stores_with_var_columns(tmp_path, n_var=64, n_col=3)
+ calls = _count_var_column_reads(monkeypatch)
+
+ out = tmp_path / "m.h5ad"
+ result = runner.invoke(
+ app, ["concat", str(a), str(b), "-o", str(out), "--merge", "first"]
+ )
+ assert result.exit_code == 0, _out(result)
+ assert calls == [], f"'first' read {len(calls)} columns it did not compare"
+ assert list(ad.read_h5ad(out).var.columns) == ["col0", "col1", "col2"]
+
+
+def test_concat_merge_drop_is_accepted_and_keeps_no_var_columns(tmp_path):
+ """`drop` is the documented default, so it has to be sayable."""
+ a, b = _two_stores_with_var_columns(tmp_path, n_var=4, n_col=2)
+ out = tmp_path / "m.h5ad"
+ result = runner.invoke(
+ app,
+ ["concat", str(a), str(b), "-o", str(out),
+ "--merge", "drop", "--uns-merge", "drop"],
+ )
+ assert result.exit_code == 0, _out(result)
+
+ got = ad.read_h5ad(out)
+ assert list(got.var.columns) == []
+ assert dict(got.uns) == {}
+
+
+# ---------------------------------------------------------------------------
+# concat refuses mismatched matrix encodings
+#
+# The check has been there since concat was written and nothing exercised it.
+# It is the error a user is most likely to meet with real per-sample files,
+# since whether a matrix lands as CSR or CSC depends on how it was made.
+
+
+def _store_with_layout(path, layout, *, cells, layer=None):
+ matrix = sparse.csr_matrix(np.ones((len(cells), 3), dtype="float32"))
+ if layout == "csc":
+ matrix = matrix.tocsc()
+ elif layout == "dense":
+ matrix = matrix.toarray()
+ obj = ad.AnnData(
+ X=matrix,
+ obs=pd.DataFrame(index=cells),
+ var=pd.DataFrame(index=["g1", "g2", "g3"]),
+ )
+ if layer is not None:
+ obj.layers["counts"] = (
+ sparse.csr_matrix(np.ones((len(cells), 3), dtype="float32"))
+ if layer == "csr"
+ else np.ones((len(cells), 3), dtype="float32")
+ )
+ obj.write_h5ad(path)
+ return path
+
+
+@pytest.mark.parametrize(
+ "left,right", [("csr", "csc"), ("csr", "dense"), ("csc", "dense")]
+)
+def test_concat_refuses_mismatched_x_encodings(tmp_path, left, right):
+ a = _store_with_layout(tmp_path / "a.h5ad", left, cells=["c1", "c2"])
+ b = _store_with_layout(tmp_path / "b.h5ad", right, cells=["c3", "c4"])
+ out = tmp_path / "m.h5ad"
+
+ result = runner.invoke(app, ["concat", str(a), str(b), "-o", str(out)])
+ assert result.exit_code == 1
+ text = _out(result)
+ assert "Cannot concatenate 'X'" in text
+ # The error has to name the way out, which until `convert` existed it
+ # could not do.
+ assert "adata convert" in text and "--layout" in text
+ assert not out.exists(), "nothing may be written when the check fails"
+
+
+def test_concat_refuses_mismatched_layer_encodings(tmp_path):
+ a = _store_with_layout(tmp_path / "a.h5ad", "csr", cells=["c1", "c2"],
+ layer="csr")
+ b = _store_with_layout(tmp_path / "b.h5ad", "csr", cells=["c3", "c4"],
+ layer="dense")
+ out = tmp_path / "m.h5ad"
+
+ result = runner.invoke(app, ["concat", str(a), str(b), "-o", str(out)])
+ assert result.exit_code == 1
+ assert "layers/counts" in _out(result)
+
+
+def test_concat_succeeds_once_the_encodings_are_converted(tmp_path):
+ """The suggested fix has to actually work, so the test follows it."""
+ a = _store_with_layout(tmp_path / "a.h5ad", "csr", cells=["c1", "c2"])
+ b = _store_with_layout(tmp_path / "b.h5ad", "csc", cells=["c3", "c4"])
+
+ converted = tmp_path / "b-csr.h5ad"
+ assert runner.invoke(
+ app, ["convert", str(b), "X", "-o", str(converted), "--layout", "csr"]
+ ).exit_code == 0
+
+ out = tmp_path / "m.h5ad"
+ result = runner.invoke(app, ["concat", str(a), str(converted), "-o", str(out)])
+ assert result.exit_code == 0, _out(result)
+ assert ad.read_h5ad(out).shape == (4, 3)
diff --git a/tests/test_convert.py b/tests/test_convert.py
new file mode 100644
index 0000000..b53e0d6
--- /dev/null
+++ b/tests/test_convert.py
@@ -0,0 +1,693 @@
+"""Tests for `adata convert`.
+
+scipy is the oracle throughout. A transpose or a cast is only correct if it
+agrees exactly with what scipy would have produced from the same matrix, and
+"looks plausible" is not a standard worth having for something that rewrites
+people's data.
+"""
+
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+import numpy as np
+import pytest
+from typer.testing import CliRunner
+
+from adata.cli import app
+
+ad = pytest.importorskip("anndata", reason="anndata is required for these tests")
+pd = pytest.importorskip("pandas")
+sparse = pytest.importorskip("scipy.sparse")
+
+runner = CliRunner()
+_ANSI = re.compile(r"\x1b\[[0-9;]*[a-zA-Z]")
+
+
+def _out(result) -> str:
+ text = result.stdout + (result.stderr or "")
+ return " ".join(_ANSI.sub("", text).split())
+
+
+def _matrix(n_obs=40, n_var=25, density=0.2, layout="csr", seed=0, integral=True):
+ rng = np.random.default_rng(seed)
+ matrix = sparse.random(
+ n_obs, n_var, density=density, format="csr", dtype="float64",
+ random_state=rng,
+ )
+ if integral:
+ # Counts, which is what issue #13 is about: float32 holds them exactly.
+ matrix.data = np.round(matrix.data * 100)
+ return matrix.tocsc() if layout == "csc" else matrix
+
+
+def _store(path: Path, matrix, *, layers=None, raw=False) -> Path:
+ obj = ad.AnnData(
+ X=matrix,
+ obs=pd.DataFrame(index=[f"c{i}" for i in range(matrix.shape[0])]),
+ var=pd.DataFrame(index=[f"g{i}" for i in range(matrix.shape[1])]),
+ )
+ for name, value in (layers or {}).items():
+ obj.layers[name] = value
+ if raw:
+ obj.raw = obj
+ obj.write_h5ad(path)
+ return path
+
+
+def _dense(matrix):
+ return matrix.toarray() if sparse.issparse(matrix) else np.asarray(matrix)
+
+
+# ---------------------------------------------------------------------------
+# layout
+
+
+@pytest.mark.parametrize("in_memory", [False, True], ids=["streaming", "in-memory"])
+@pytest.mark.parametrize("source,target", [("csr", "csc"), ("csc", "csr")])
+def test_transpose_matches_scipy_exactly(tmp_path, source, target, in_memory):
+ """Both paths, both directions, against scipy's own conversion.
+
+ Having two implementations is only worth it if they agree; that is the
+ whole reason the streaming one is not checked merely against itself.
+ """
+ matrix = _matrix(layout=source)
+ store = _store(tmp_path / "in.h5ad", matrix)
+
+ out = tmp_path / f"{source}-{target}-{in_memory}.h5ad"
+ argv = ["convert", str(store), "X", "-o", str(out), "--layout", target]
+ if in_memory:
+ argv.append("--in-memory")
+ result = runner.invoke(app, argv)
+ assert result.exit_code == 0, _out(result)
+
+ got = ad.read_h5ad(out).X
+ expected = matrix.tocsc() if target == "csc" else matrix.tocsr()
+ assert got.format == target
+ assert np.array_equal(got.indptr, expected.indptr)
+ assert np.array_equal(got.indices, expected.indices)
+ assert np.array_equal(got.data, expected.data)
+
+
+def test_transposing_twice_is_the_identity(tmp_path):
+ matrix = _matrix()
+ store = _store(tmp_path / "in.h5ad", matrix)
+
+ mid, back = tmp_path / "mid.h5ad", tmp_path / "back.h5ad"
+ assert runner.invoke(
+ app, ["convert", str(store), "X", "-o", str(mid), "--layout", "csc"]
+ ).exit_code == 0
+ assert runner.invoke(
+ app, ["convert", str(mid), "X", "-o", str(back), "--layout", "csr"]
+ ).exit_code == 0
+
+ result = ad.read_h5ad(back).X
+ assert result.format == "csr"
+ assert np.array_equal(result.indptr, matrix.indptr)
+ assert np.array_equal(result.indices, matrix.indices)
+ assert np.array_equal(result.data, matrix.data)
+
+
+def test_the_two_transpose_paths_agree_on_an_awkward_matrix(tmp_path):
+ """Empty rows, a full row, and a single-entry column all in one."""
+ dense = np.zeros((6, 5))
+ dense[0, :] = [1, 2, 3, 4, 5] # full row
+ dense[3, 2] = 7 # lone entry
+ # rows 1, 2, 4, 5 stay empty
+ matrix = sparse.csr_matrix(dense)
+ store = _store(tmp_path / "in.h5ad", matrix)
+
+ outputs = {}
+ for tag, extra in (("stream", []), ("memory", ["--in-memory"])):
+ out = tmp_path / f"{tag}.h5ad"
+ assert runner.invoke(
+ app,
+ ["convert", str(store), "X", "-o", str(out), "--layout", "csc", *extra],
+ ).exit_code == 0
+ outputs[tag] = ad.read_h5ad(out).X
+
+ expected = matrix.tocsc()
+ for tag, got in outputs.items():
+ assert np.array_equal(got.indptr, expected.indptr), tag
+ assert np.array_equal(got.indices, expected.indices), tag
+ assert np.array_equal(got.data, expected.data), tag
+
+
+# ---------------------------------------------------------------------------
+# dtype -- what issue #13 asked for
+
+
+def test_counts_stored_as_float64_convert_to_float32_and_shrink(tmp_path):
+ """The reported case, end to end, including that the file gets smaller.
+
+ An early version produced an output eight times the input, because it
+ created the arrays with a fixed 65,536-element chunk and promoted int32
+ indices to int64. A conversion asked for to halve a file must not
+ enlarge it, so the size is part of the test.
+ """
+ import h5py
+
+ matrix = _matrix(n_obs=200, n_var=120, density=0.1)
+ store = _store(tmp_path / "in.h5ad", matrix)
+ out = tmp_path / "f32.h5ad"
+
+ result = runner.invoke(
+ app, ["convert", str(store), "X", "-o", str(out), "--dtype", "float32"]
+ )
+ assert result.exit_code == 0, _out(result)
+
+ got = ad.read_h5ad(out).X
+ assert got.dtype == np.dtype("float32")
+ assert np.array_equal(got.toarray(), matrix.toarray().astype("float32"))
+
+ def stored(path):
+ with h5py.File(path) as handle:
+ return sum(
+ handle["X"][key].id.get_storage_size()
+ for key in ("data", "indices", "indptr")
+ )
+
+ assert stored(out) < stored(store), (
+ f"X grew from {stored(store):,} to {stored(out):,} bytes on a "
+ "float64 -> float32 conversion"
+ )
+
+
+def test_the_index_dtype_is_preserved_unless_asked(tmp_path):
+ """Defaulting to int64 silently doubled every int32 store's indices."""
+ import h5py
+
+ store = _store(tmp_path / "in.h5ad", _matrix())
+ with h5py.File(store) as handle:
+ source_dtype = handle["X/indices"].dtype
+
+ out = tmp_path / "out.h5ad"
+ assert runner.invoke(
+ app, ["convert", str(store), "X", "-o", str(out), "--dtype", "float32"]
+ ).exit_code == 0
+
+ with h5py.File(out) as handle:
+ assert handle["X/indices"].dtype == source_dtype
+
+
+def test_indices_can_be_narrowed_when_asked(tmp_path):
+ import h5py
+
+ store = _store(tmp_path / "in.h5ad", _matrix())
+ out = tmp_path / "out.h5ad"
+ result = runner.invoke(
+ app,
+ ["convert", str(store), "X", "-o", str(out), "--indices-dtype", "int32"],
+ )
+ assert result.exit_code == 0, _out(result)
+
+ with h5py.File(out) as handle:
+ assert handle["X/indices"].dtype == np.dtype("int32")
+ assert handle["X/indptr"].dtype == np.dtype("int32")
+ assert np.array_equal(ad.read_h5ad(out).X.toarray(), _matrix().toarray())
+
+
+# ---------------------------------------------------------------------------
+# refusing before writing
+
+
+def test_a_lossy_cast_is_refused_and_nothing_is_written(tmp_path):
+ """The store must be left alone, not half converted."""
+ matrix = _matrix(integral=False)
+ matrix.data = matrix.data + 0.123456789012345
+ store = _store(tmp_path / "in.h5ad", matrix)
+ out = tmp_path / "out.h5ad"
+
+ result = runner.invoke(
+ app, ["convert", str(store), "X", "-o", str(out), "--dtype", "float32"]
+ )
+ assert result.exit_code == 1
+ text = _out(result)
+ assert "lossy" in text and "--force" in text
+ assert "do not round-trip" in text
+ assert not out.exists(), "a refused conversion must leave no output behind"
+
+
+def test_force_converts_anyway(tmp_path):
+ matrix = _matrix(integral=False)
+ matrix.data = matrix.data + 0.123456789012345
+ store = _store(tmp_path / "in.h5ad", matrix)
+ out = tmp_path / "out.h5ad"
+
+ result = runner.invoke(
+ app,
+ ["convert", str(store), "X", "-o", str(out), "--dtype", "float32", "--force"],
+ )
+ assert result.exit_code == 0, _out(result)
+ assert ad.read_h5ad(out).X.dtype == np.dtype("float32")
+
+
+def test_a_lossless_cast_says_so(tmp_path):
+ """Worth telling the user: it is the question they were asking."""
+ store = _store(tmp_path / "in.h5ad", _matrix())
+ result = runner.invoke(
+ app,
+ ["convert", str(store), "X", "-o", str(tmp_path / "o.h5ad"),
+ "--dtype", "float32"],
+ )
+ assert result.exit_code == 0
+ assert "round-trips" in _out(result)
+
+
+class _FakeGroup:
+ """Just enough of a sparse group for the index-range check."""
+
+ def __init__(self, nnz: int):
+ self._nnz = nnz
+
+ def __getitem__(self, key):
+ assert key == "indices"
+ return type("D", (), {"shape": (self._nnz,)})()
+
+
+def test_an_index_dtype_too_small_to_address_the_matrix_is_refused():
+ """indices and indptr are bounded by different things, so both are checked.
+
+ `indices` holds coordinates, bounded by the larger dimension; `indptr`
+ holds offsets, bounded by nnz. A matrix can legitimately need int64 for
+ one and not the other.
+ """
+ from adata.core.convert import check_index_dtype
+
+ wide = _FakeGroup(nnz=10)
+ with pytest.raises(ValueError, match="cannot hold this matrix's coordinates"):
+ check_index_dtype(
+ wide, np.dtype("int32"), np.dtype("int64"), (10, 3_000_000_000)
+ )
+
+ many = _FakeGroup(nnz=3_000_000_000)
+ with pytest.raises(ValueError, match="cannot hold this matrix's offsets"):
+ check_index_dtype(many, np.dtype("int64"), np.dtype("int32"), (10, 10))
+
+ # Narrow matrix, huge nnz: int32 coordinates are fine, offsets are not.
+ check_index_dtype(
+ _FakeGroup(nnz=10), np.dtype("int32"), np.dtype("int64"), (10, 10)
+ )
+
+
+def test_indptr_keeps_its_own_width(tmp_path):
+ """A store with int32 indices and int64 indptr must keep both.
+
+ Inferring indptr's dtype from indices narrowed the offsets of any
+ matrix with more than 2^31 nonzeros, which corrupts it silently.
+ """
+ import h5py
+
+ store = _store(tmp_path / "in.h5ad", _matrix())
+ with h5py.File(store, "a") as handle:
+ pointers = handle["X/indptr"][...]
+ coordinates = handle["X/indices"][...]
+ del handle["X/indptr"], handle["X/indices"]
+ handle["X"].create_dataset("indptr", data=pointers.astype("int64"))
+ handle["X"].create_dataset("indices", data=coordinates.astype("int32"))
+
+ out = tmp_path / "out.h5ad"
+ assert runner.invoke(
+ app, ["convert", str(store), "X", "-o", str(out), "--dtype", "float32"]
+ ).exit_code == 0
+
+ with h5py.File(out) as handle:
+ assert handle["X/indptr"].dtype == np.dtype("int64"), "offsets narrowed"
+ assert handle["X/indices"].dtype == np.dtype("int32")
+
+
+# ---------------------------------------------------------------------------
+# density
+
+
+def test_densify_then_sparsify_round_trips(tmp_path):
+ matrix = _matrix(density=0.3)
+ store = _store(tmp_path / "in.h5ad", matrix)
+
+ dense_path, sparse_path = tmp_path / "dense.h5ad", tmp_path / "sparse.h5ad"
+ assert runner.invoke(
+ app,
+ ["convert", str(store), "X", "-o", str(dense_path), "--layout", "dense",
+ "--force"],
+ ).exit_code == 0
+ got_dense = ad.read_h5ad(dense_path).X
+ assert not sparse.issparse(got_dense)
+ assert np.array_equal(_dense(got_dense), matrix.toarray())
+
+ assert runner.invoke(
+ app,
+ ["convert", str(dense_path), "X", "-o", str(sparse_path),
+ "--layout", "csr"],
+ ).exit_code == 0
+ got = ad.read_h5ad(sparse_path).X
+ assert got.format == "csr"
+ assert np.array_equal(got.toarray(), matrix.toarray())
+
+
+def test_densifying_a_sparse_matrix_is_refused_when_it_would_explode(tmp_path):
+ """0.5% dense over a wide matrix is exactly the case that hurts."""
+ matrix = _matrix(n_obs=200, n_var=2000, density=0.005)
+ store = _store(tmp_path / "in.h5ad", matrix)
+ out = tmp_path / "out.h5ad"
+
+ result = runner.invoke(
+ app, ["convert", str(store), "X", "-o", str(out), "--layout", "dense"]
+ )
+ assert result.exit_code == 1
+ text = _out(result)
+ assert "would grow" in text and "--force" in text
+ assert not out.exists()
+
+
+def test_sparsifying_a_mostly_dense_matrix_warns(tmp_path):
+ dense = np.ones((20, 10))
+ store = _store(tmp_path / "in.h5ad", sparse.csr_matrix(dense))
+ dense_path = tmp_path / "dense.h5ad"
+ assert runner.invoke(
+ app,
+ ["convert", str(store), "X", "-o", str(dense_path), "--layout", "dense",
+ "--force"],
+ ).exit_code == 0
+
+ result = runner.invoke(
+ app,
+ ["convert", str(dense_path), "X", "-o", str(tmp_path / "s.h5ad"),
+ "--layout", "csr"],
+ )
+ assert result.exit_code == 0, _out(result)
+ assert "nonzero" in _out(result)
+
+
+# ---------------------------------------------------------------------------
+# selection and plumbing
+
+
+def test_all_converts_x_layers_and_raw(tmp_path):
+ matrix = _matrix()
+ store = _store(
+ tmp_path / "in.h5ad", matrix, layers={"counts": matrix.copy()}, raw=True
+ )
+ out = tmp_path / "out.h5ad"
+
+ result = runner.invoke(
+ app, ["convert", str(store), "--all", "-o", str(out), "--layout", "csc"]
+ )
+ assert result.exit_code == 0, _out(result)
+
+ got = ad.read_h5ad(out)
+ assert got.X.format == "csc"
+ assert got.layers["counts"].format == "csc"
+ assert got.raw is not None and got.raw.X.format == "csc"
+ assert list(got.obs_names) == [f"c{i}" for i in range(matrix.shape[0])]
+
+
+def test_untargeted_elements_are_carried_over_untouched(tmp_path):
+ matrix = _matrix()
+ obj = ad.AnnData(
+ X=matrix,
+ obs=pd.DataFrame(
+ {"group": pd.Categorical(["a", "b"] * (matrix.shape[0] // 2))},
+ index=[f"c{i}" for i in range(matrix.shape[0])],
+ ),
+ var=pd.DataFrame(index=[f"g{i}" for i in range(matrix.shape[1])]),
+ )
+ obj.layers["untouched"] = matrix.copy()
+ obj.obsm["X_pca"] = np.arange(matrix.shape[0] * 3, dtype="float32").reshape(-1, 3)
+ obj.uns["note"] = "keep me"
+ store = tmp_path / "in.h5ad"
+ obj.write_h5ad(store)
+
+ out = tmp_path / "out.h5ad"
+ assert runner.invoke(
+ app, ["convert", str(store), "X", "-o", str(out), "--layout", "csc"]
+ ).exit_code == 0
+
+ got = ad.read_h5ad(out)
+ assert got.X.format == "csc"
+ assert got.layers["untouched"].format == "csr", "an untargeted layer changed"
+ assert got.uns["note"] == "keep me"
+ assert np.array_equal(got.obsm["X_pca"], obj.obsm["X_pca"])
+ assert list(got.obs["group"]) == list(obj.obs["group"])
+
+
+def test_inplace_replaces_the_source(tmp_path):
+ matrix = _matrix()
+ store = _store(tmp_path / "in.h5ad", matrix)
+
+ result = runner.invoke(
+ app, ["convert", str(store), "X", "--inplace", "--dtype", "float32"]
+ )
+ assert result.exit_code == 0, _out(result)
+
+ got = ad.read_h5ad(store)
+ assert got.X.dtype == np.dtype("float32")
+ assert np.array_equal(got.X.toarray(), matrix.toarray().astype("float32"))
+ assert not list(tmp_path.glob("*convert-tmp*")), "temp file left behind"
+
+
+def test_converting_to_zarr_keeps_the_values(tmp_path):
+ matrix = _matrix()
+ store = _store(tmp_path / "in.h5ad", matrix)
+ out = tmp_path / "out.zarr"
+
+ result = runner.invoke(
+ app, ["convert", str(store), "X", "-o", str(out), "--layout", "csc"]
+ )
+ assert result.exit_code == 0, _out(result)
+ got = ad.read_zarr(out).X
+ assert got.format == "csc"
+ assert np.array_equal(got.toarray(), matrix.toarray())
+
+
+# ---------------------------------------------------------------------------
+# argument handling
+
+
+def test_convert_needs_something_to_do(tmp_path):
+ store = _store(tmp_path / "in.h5ad", _matrix())
+ result = runner.invoke(
+ app, ["convert", str(store), "X", "-o", str(tmp_path / "o.h5ad")]
+ )
+ assert result.exit_code == 1
+ assert "at least one of --dtype" in _out(result)
+
+
+def test_convert_needs_an_output(tmp_path):
+ store = _store(tmp_path / "in.h5ad", _matrix())
+ result = runner.invoke(app, ["convert", str(store), "X", "--dtype", "float32"])
+ assert result.exit_code == 1
+ assert "Output file is required" in _out(result)
+
+
+@pytest.mark.parametrize(
+ "flag,value,expected",
+ [
+ ("--dtype", "complex128", "Unknown dtype"),
+ ("--indices-dtype", "float32", "Unknown dtype"),
+ ("--layout", "coo", "--layout must be one of"),
+ ],
+)
+def test_bad_values_are_rejected_by_name(tmp_path, flag, value, expected):
+ store = _store(tmp_path / "in.h5ad", _matrix())
+ result = runner.invoke(
+ app,
+ ["convert", str(store), "X", "-o", str(tmp_path / "o.h5ad"), flag, value],
+ )
+ assert result.exit_code == 1
+ assert expected in _out(result)
+
+
+def test_converting_something_that_is_not_a_matrix_says_so(tmp_path):
+ store = _store(tmp_path / "in.h5ad", _matrix())
+ result = runner.invoke(
+ app,
+ ["convert", str(store), "obs", "-o", str(tmp_path / "o.h5ad"),
+ "--dtype", "float32"],
+ )
+ assert result.exit_code == 1
+ assert "Not a matrix" in _out(result)
+
+
+def test_a_missing_entry_names_the_path(tmp_path):
+ store = _store(tmp_path / "in.h5ad", _matrix())
+ result = runner.invoke(
+ app,
+ ["convert", str(store), "layers/nope", "-o", str(tmp_path / "o.h5ad"),
+ "--dtype", "float32"],
+ )
+ assert result.exit_code == 1
+ assert "layers/nope" in _out(result)
+
+
+# ---------------------------------------------------------------------------
+# what review found
+#
+# Five findings on the first version of this command, four of them able to
+# change or destroy data while reporting success. Each gets a test.
+
+
+def test_an_output_that_names_the_input_is_refused(tmp_path):
+ """Writing over the store being read from destroyed it.
+
+ HDF5 happens to refuse the second open; Zarr does not, and the command
+ completed successfully leaving a store with zero nonzeros where the
+ data had been.
+ """
+ matrix = _matrix()
+ store = _store(tmp_path / "in.h5ad", matrix)
+
+ result = runner.invoke(
+ app, ["convert", str(store), "X", "-o", str(store), "--dtype", "float32"]
+ )
+ assert result.exit_code == 1
+ assert "Output path is the input" in _out(result)
+ assert np.array_equal(ad.read_h5ad(store).X.toarray(), matrix.toarray())
+
+
+def test_an_output_that_aliases_the_input_through_a_relative_path_is_refused(
+ tmp_path,
+):
+ matrix = _matrix()
+ store = _store(tmp_path / "in.h5ad", matrix)
+ alias = tmp_path / "sub" / ".." / "in.h5ad"
+ (tmp_path / "sub").mkdir()
+
+ result = runner.invoke(
+ app, ["convert", str(store), "X", "-o", str(alias), "--dtype", "float32"]
+ )
+ assert result.exit_code == 1
+ assert np.array_equal(ad.read_h5ad(store).X.toarray(), matrix.toarray())
+
+
+def test_an_explicitly_named_obsm_matrix_is_actually_converted(tmp_path):
+ """Only `layers` and `raw` were descended into, so this silently no-opped.
+
+ The command reported success and copied the matrix over unchanged,
+ which is the worst way to get this wrong: the user has no signal.
+ """
+ matrix = _matrix()
+ obj = ad.AnnData(
+ X=matrix,
+ obs=pd.DataFrame(index=[f"c{i}" for i in range(matrix.shape[0])]),
+ var=pd.DataFrame(index=[f"g{i}" for i in range(matrix.shape[1])]),
+ )
+ obj.obsm["X_pca"] = np.ones((matrix.shape[0], 4), dtype="float64")
+ store = tmp_path / "in.h5ad"
+ obj.write_h5ad(store)
+
+ out = tmp_path / "out.h5ad"
+ result = runner.invoke(
+ app,
+ ["convert", str(store), "obsm/X_pca", "-o", str(out), "--dtype", "float32"],
+ )
+ assert result.exit_code == 0, _out(result)
+
+ got = ad.read_h5ad(out)
+ assert got.obsm["X_pca"].dtype == np.dtype("float32")
+ assert np.array_equal(got.obsm["X_pca"], obj.obsm["X_pca"].astype("float32"))
+ assert got.X.dtype == matrix.dtype, "X should not have been touched"
+
+
+def test_densifying_sums_duplicate_coordinates(tmp_path):
+ """A repeated coordinate means the sum, which is what scipy produces.
+
+ Legal in a CSR store and not what anndata writes, so it takes a
+ hand-built file to reach -- but plain assignment kept whichever entry
+ came last and changed the matrix's values on the way to dense.
+ """
+ import h5py
+
+ store = _store(tmp_path / "in.h5ad", sparse.csr_matrix(np.zeros((2, 3))))
+ with h5py.File(store, "a") as handle:
+ for key in ("data", "indices", "indptr"):
+ del handle["X"][key]
+ handle["X"].create_dataset("data", data=np.array([1.0, 2.0]))
+ handle["X"].create_dataset("indices", data=np.array([1, 1]))
+ handle["X"].create_dataset("indptr", data=np.array([0, 2, 2]))
+
+ assert ad.read_h5ad(store).X.toarray()[0, 1] == 3.0, "scipy sums them"
+
+ out = tmp_path / "dense.h5ad"
+ assert runner.invoke(
+ app,
+ ["convert", str(store), "X", "-o", str(out), "--layout", "dense",
+ "--force"],
+ ).exit_code == 0
+ assert np.asarray(ad.read_h5ad(out).X)[0, 1] == 3.0
+
+
+def test_transpose_buckets_are_balanced_by_nonzeros_not_by_coordinate(
+ tmp_path, monkeypatch
+):
+ """A skewed matrix must not land in one bucket.
+
+ Single-cell matrices are skewed -- a few genes carry most of the counts
+ -- so equal-width coordinate bounds defeat the streaming guarantee
+ exactly where it matters. Measured on the matrix below, the largest of
+ three equal-width buckets held 88% of the entries, so the bucket rather
+ than the chunk set the peak.
+ """
+ import adata.core.subset as subset_module
+ from adata.core.convert import describe, transpose_sparse_streaming
+ from adata.storage import open_store
+
+ rng = np.random.default_rng(0)
+ n = 400
+ rows, cols = [], []
+ for row in range(n):
+ for _ in range(20):
+ rows.append(row)
+ cols.append(
+ int(rng.integers(0, 5)) if rng.random() < 0.95
+ else int(rng.integers(0, n))
+ )
+ skewed = sparse.csr_matrix((np.ones(len(rows)), (rows, cols)), shape=(n, n))
+ store = _store(tmp_path / "skew.h5ad", skewed)
+
+ # `transpose_sparse_streaming` imports `_append` from this module when it
+ # runs, so this is the binding it will pick up.
+ per_bucket: dict = {}
+ original = subset_module._append
+
+ def watching_append(dataset, values):
+ name = str(getattr(dataset, "name", "") or getattr(dataset, "path", ""))
+ if "major" in name:
+ per_bucket[name] = per_bucket.get(name, 0) + int(values.size)
+ return original(dataset, values)
+
+ monkeypatch.setattr(subset_module, "_append", watching_append)
+
+ out = tmp_path / "out.h5ad"
+ with open_store(store, "r") as src, open_store(out, "w") as dst:
+ transpose_sparse_streaming(
+ describe(src.root["X"]),
+ dst.root,
+ "X",
+ data_dtype=np.dtype("float64"),
+ index_dtype=np.dtype("int64"),
+ chunk=1000,
+ bucket_entries=1000,
+ )
+
+ sizes = sorted(per_bucket.values(), reverse=True)
+ assert len(sizes) > 1, f"expected several buckets, saw {per_bucket}"
+ assert sizes[0] <= skewed.nnz * 0.5, (
+ f"the largest bucket held {sizes[0]} of {skewed.nnz} nonzeros "
+ f"({sizes[0] / skewed.nnz:.0%}); buckets must be balanced by count, "
+ "not by coordinate range"
+ )
+
+ # Read X back directly: this wrote only the matrix, not a whole store.
+ import h5py
+
+ expected = skewed.tocsc()
+ with h5py.File(out) as handle:
+ got = sparse.csc_matrix(
+ (handle["X/data"][...], handle["X/indices"][...],
+ handle["X/indptr"][...]),
+ shape=tuple(handle["X"].attrs["shape"]),
+ )
+ assert np.array_equal(got.toarray(), expected.toarray()), (
+ "balancing the buckets changed the result"
+ )
diff --git a/tests/test_docs_are_accurate.py b/tests/test_docs_are_accurate.py
index d474a34..f6f4608 100644
--- a/tests/test_docs_are_accurate.py
+++ b/tests/test_docs_are_accurate.py
@@ -100,7 +100,8 @@ def _command_path(tokens: List[str]) -> List[str]:
_GROUPS = {"export", "import"}
_TOP_LEVEL = {
- "view", "ls", "subset", "split", "concat", "create", "export", "import",
+ "view", "ls", "subset", "split", "concat", "convert", "create",
+ "export", "import",
}
_SUBCOMMANDS = {
"export": {"dataframe", "array", "sparse", "dict", "image"},
diff --git a/tests/test_performance.py b/tests/test_performance.py
new file mode 100644
index 0000000..f272055
--- /dev/null
+++ b/tests/test_performance.py
@@ -0,0 +1,1263 @@
+"""Complexity guards: what an operation *costs*, never how long it takes.
+
+Why this file exists
+--------------------
+`concat --merge` hung for hours on a 36,601-var input (REQ-71798). One line in
+`_write_var` re-read a whole var column from disk once per target variable.
+The output would have been correct; only the cost was wrong. 1,029 tests
+missed it, and they were never going to catch it: every fixture in the suite
+is at most a few hundred elements, and nothing measured cost at all.
+
+So this file closes a defect class, not a bug. The instruments are in
+`tests/perf_counters.py`; the rule they enforce is that cost must grow no
+faster than linearly in each axis, measured in operations rather than seconds.
+
+Why not wall clock
+------------------
+A test that can fail because a CI runner was busy does not belong in a merge
+gate. Every number here is deterministic: the same input produces the same
+count on every machine. That is what lets these run in the ordinary test job
+on both interpreters and block a merge when they fail.
+
+Reading a failure
+-----------------
+A failure here says cost grew super-linearly on the named axis. Usually the
+code is wrong. Occasionally the expectation is -- an operation legitimately
+gains work -- and then the new number needs a comment saying why, in the same
+style as the exact counts in `test_commands_phase2.py`.
+
+One axis at a time. Scaling two at once makes legitimate work look quadratic:
+an outer-join concat really does produce n_obs x n_var_union cells.
+"""
+
+from __future__ import annotations
+
+import pathlib
+from pathlib import Path
+from typing import List, Optional
+
+import numpy as np
+import pytest
+from rich.console import Console
+
+from adata.commands.create import create_store
+from adata.commands.export import (
+ export_json,
+ export_mtx,
+ export_npy,
+ export_table,
+)
+from adata.commands.import_data import import_object
+from adata.commands.info import show_info
+from adata.commands.ls import list_store
+from adata.core.concat import concat_on_disk
+from adata.core.subset import subset_h5ad
+
+from tests.perf_counters import (
+ GROWTH_LIMIT,
+ SIZES,
+ assert_grows_linearly,
+ assert_grows_slower_than_input,
+ assert_independent_of,
+ count_allocations,
+ count_scanned_elements,
+ count_io,
+ count_lines,
+)
+
+ad = pytest.importorskip("anndata", reason="anndata builds the fixtures")
+pd = pytest.importorskip("pandas")
+sparse = pytest.importorskip("scipy.sparse")
+
+QUIET = Console(quiet=True)
+
+#: Only Python executed inside the package counts towards the CPU proxy.
+SOURCE_PREFIX = str(pathlib.Path(__file__).resolve().parent.parent / "src" / "adata")
+
+
+# ---------------------------------------------------------------------------
+# fixtures built to scale exactly one axis
+
+
+def _store(
+ path: Path,
+ *,
+ name: str,
+ n_obs: int = 4,
+ n_var: int = 8,
+ n_var_columns: int = 2,
+ n_obs_columns: int = 0,
+ n_categories: int = 0,
+ obs_kind: Optional[str] = None,
+ shared_var: bool = True,
+) -> 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)])
+ if n_categories:
+ obs["ct"] = pd.Categorical(
+ [f"{name}-t{i % n_categories}" for i in range(n_obs)]
+ )
+ if obs_kind == "numeric":
+ obs["col"] = np.arange(n_obs, dtype="int32")
+ elif obs_kind == "masked":
+ obs["col"] = pd.array(
+ [i if i % 2 else None for i in range(n_obs)], dtype="Int32"
+ )
+ elif obs_kind == "string":
+ obs["col"] = [f"{name}-{i}" for i in range(n_obs)]
+ if n_obs_columns:
+ extra = pd.DataFrame(
+ {f"m{c}": np.arange(n_obs, dtype="int32") for c in range(n_obs_columns)},
+ index=obs.index,
+ )
+ obs = pd.concat([obs, extra], axis=1)
+
+ prefix = "" if shared_var else name
+ var = pd.DataFrame(
+ {f"v{c}": [f"{c}-{i}" for i in range(n_var)] for c in range(n_var_columns)},
+ index=[f"{prefix}g{i}" for i in range(n_var)],
+ )
+
+ obj = ad.AnnData(
+ X=sparse.csr_matrix(np.ones((n_obs, n_var), dtype="float32")),
+ obs=obs,
+ var=var,
+ )
+ obj.obsm["X_pca"] = np.zeros((n_obs, 3), dtype="float32")
+ obj.write_h5ad(path)
+ return path
+
+
+def _inputs(tmp_path: Path, tag: str, count: int = 2, **kwargs) -> List[Path]:
+ directory = tmp_path / tag
+ directory.mkdir(parents=True, exist_ok=True)
+ return [
+ _store(directory / f"{chr(ord('a') + i)}.h5ad", name=chr(ord("a") + i), **kwargs)
+ for i in range(count)
+ ]
+
+
+# ---------------------------------------------------------------------------
+# canary
+#
+# Every growth guard below is a comparison between measurements. If a refactor
+# moves a read onto a path the counters cannot see -- `read_direct`,
+# `np.asarray(dataset)`, `asstr()[...]`; see perf_counters' docstring -- the
+# comparisons would all pass on zeros. The vacuity check inside
+# `assert_grows_linearly` catches most of that; this catches the rest, and
+# says plainly what broke.
+
+
+def test_the_counters_see_a_known_operation(tmp_path):
+ """A concat of two 4 x 64 stores must register substantial work.
+
+ If this fails, the hooks have stopped matching the libraries and every
+ other test in this file has become meaningless. Fix the hooks first.
+ """
+ files = _inputs(tmp_path, "canary", n_obs=4, n_var=64)
+ with count_io() as io:
+ concat_on_disk(files, tmp_path / "out.h5ad", QUIET, merge="same")
+
+ assert io.h5_calls > 0, "h5py.Dataset.__getitem__ hook is not firing"
+ assert io.h5_elements >= 64, f"suspiciously few elements read: {io}"
+
+ with count_allocations() as peak:
+ concat_on_disk(files, tmp_path / "out2.h5ad", QUIET, merge="same")
+ assert peak[0] > 0, "tracemalloc recorded no allocation"
+
+ with count_lines(SOURCE_PREFIX) as lines:
+ concat_on_disk(files, tmp_path / "out3.h5ad", QUIET, merge="same")
+ assert lines[0] > 100, f"line tracer saw only {lines[0]} events in src/adata"
+
+
+def test_the_zarr_store_hooks_fire(tmp_path):
+ """Metadata and chunk traffic must be visible, not just array reads.
+
+ An attribute-write storm never touches `Array.__getitem__` -- it is all
+ `LocalStore.set`. That shape of bug has bitten this repo before, so the
+ store hooks get their own canary.
+ """
+ files = _inputs(tmp_path, "zcanary", n_obs=4, n_var=32)
+ with count_io() as io:
+ concat_on_disk(files, tmp_path / "out.zarr", QUIET, merge="same")
+
+ assert io.store_set > 0, "LocalStore.set hook is not firing"
+ assert io.store_get > 0, "LocalStore.get hook is not firing"
+
+
+# ---------------------------------------------------------------------------
+# concat: reads must not grow faster than linearly on any axis
+#
+# The n_var case is the direct regression guard for REQ-71798. The others
+# exist because concat's cost surface is n_obs x n_var x n_columns x n_inputs
+# and a defect can hide in any of them; scaling only the axis that broke last
+# time would be fighting the last war.
+
+
+@pytest.mark.parametrize("merge", [None, "same", "unique", "first", "only"])
+def test_concat_reads_grow_linearly_in_var_count(tmp_path, merge):
+ """The REQ-71798 axis. Every merge strategy, not just the one reported."""
+
+ def measure(n: int) -> int:
+ files = _inputs(tmp_path, f"var{merge}{n}", n_var=n)
+ with count_io() as io:
+ concat_on_disk(files, tmp_path / f"o{merge}{n}.h5ad", QUIET, merge=merge)
+ return io.reads
+
+ assert_grows_linearly(measure, what=f"concat --merge {merge}", axis="n_var")
+
+
+@pytest.mark.parametrize("join", ["inner", "outer"])
+def test_concat_reads_grow_linearly_in_obs_count(tmp_path, join):
+ def measure(n: int) -> int:
+ files = _inputs(tmp_path, f"obs{join}{n}", n_obs=n, n_var=4)
+ with count_io() as io:
+ concat_on_disk(files, tmp_path / f"o{join}{n}.h5ad", QUIET, join=join)
+ return io.reads
+
+ assert_grows_linearly(measure, what=f"concat --join {join}", axis="n_obs")
+
+
+def test_concat_reads_grow_linearly_in_input_count(tmp_path):
+ """Pairwise work across inputs would show up here and nowhere else."""
+
+ def measure(n: int) -> int:
+ files = _inputs(tmp_path, f"in{n}", count=n, n_obs=2, n_var=4)
+ with count_io() as io:
+ concat_on_disk(files, tmp_path / f"oin{n}.h5ad", QUIET, merge="same")
+ return io.reads
+
+ # Sizes are scaled down: n inputs means n files on disk.
+ assert_grows_linearly(
+ measure, what="concat", axis="n_inputs", sizes=(4, 16, 64)
+ )
+
+
+def test_concat_reads_grow_linearly_in_obs_column_count(tmp_path):
+ def measure(n: int) -> int:
+ files = _inputs(tmp_path, f"oc{n}", n_obs=8, n_var=4, n_obs_columns=n)
+ with count_io() as io:
+ concat_on_disk(files, tmp_path / f"ooc{n}.h5ad", QUIET)
+ return io.reads
+
+ assert_grows_linearly(measure, what="concat", axis="n_obs_columns")
+
+
+def test_concat_reads_grow_linearly_in_var_column_count(tmp_path):
+ def measure(n: int) -> int:
+ files = _inputs(tmp_path, f"vc{n}", n_obs=4, n_var=8, n_var_columns=n)
+ with count_io() as io:
+ concat_on_disk(tmp_path / f"vc{n}" and files, tmp_path / f"ovc{n}.h5ad",
+ QUIET, merge="same")
+ return io.reads
+
+ assert_grows_linearly(measure, what="concat --merge same", axis="n_var_columns")
+
+
+# ---------------------------------------------------------------------------
+# concat: the category union
+#
+# This one needs its own instrument. Merging categories reads each category
+# list exactly once however the union is computed, so a read counter sees
+# nothing; and `x not in some_list` is a single bytecode, so the line tracer
+# sees nothing either -- the quadratic lives inside C-level list membership.
+# Counting string comparisons is the only thing that makes it visible.
+
+
+class _CountingStr(str):
+ """A string that tallies its own comparisons.
+
+ `list.__contains__` rich-compares each element, and a subclass's `__eq__`
+ takes priority, so the tally reflects real membership-probe cost.
+ """
+
+ count = 0
+
+ def __eq__(self, other: object) -> bool:
+ _CountingStr.count += 1
+ return str.__eq__(self, other)
+
+ def __ne__(self, other: object) -> bool:
+ return not self.__eq__(other)
+
+ def __hash__(self) -> int:
+ return str.__hash__(self)
+
+
+def _count_category_comparisons(monkeypatch, tmp_path, k: int) -> int:
+ import adata.core.concat as concat_module
+
+ original = concat_module.read_categories
+ monkeypatch.setattr(
+ concat_module,
+ "read_categories",
+ lambda col: [_CountingStr(c) for c in original(col)],
+ )
+
+ files = _inputs(tmp_path, f"cat{k}", n_obs=k, n_var=2, n_categories=k)
+ _CountingStr.count = 0
+ concat_on_disk(files, tmp_path / f"ocat{k}.h5ad", QUIET)
+ return _CountingStr.count
+
+
+def test_concat_unions_categories_without_quadratic_comparisons(
+ tmp_path, monkeypatch
+):
+ """Category merging must use hash lookup, not list membership.
+
+ `if category not in categories` on a list was O(k) per probe and so
+ O(k^2) over the union: 2,096,128 string comparisons at k=1024, and around
+ 5e9 for an obs column with 100k categories -- the same silent hang as
+ REQ-71798, in a different function.
+
+ A dict makes it O(k) with comparisons only on hash collision, so the
+ budget is generous and still three orders of magnitude below the old cost.
+ """
+ k = 1024
+ comparisons = _count_category_comparisons(monkeypatch, tmp_path, k)
+ assert comparisons <= 10 * k, (
+ f"{comparisons} string comparisons to union {k} categories across two "
+ f"inputs. Linear would be near zero (hash collisions only); the "
+ f"list-membership implementation cost 2,096,128."
+ )
+
+
+# ---------------------------------------------------------------------------
+# concat: per-row cost of an obs column
+#
+# Growth ratios cannot see these: building a Python object per row is linear,
+# just with a fat constant. Comparing a column kind against a plain numeric
+# column in the same run calibrates that constant away, so the assertion holds
+# across interpreters and platforms without a hand-tuned byte budget.
+
+
+def _obs_column_cost(tmp_path, kind: str, n_obs: int) -> tuple:
+ files = _inputs(tmp_path, f"k{kind}{n_obs}", n_obs=n_obs, n_var=2, obs_kind=kind)
+ with count_lines(SOURCE_PREFIX) as lines, count_allocations() as peak:
+ concat_on_disk(files, tmp_path / f"ok{kind}{n_obs}.h5ad", QUIET)
+ return lines[0], peak[0]
+
+
+@pytest.mark.perf
+@pytest.mark.parametrize("kind", ["masked", "string"])
+def test_obs_columns_cost_no_more_python_work_per_row_than_a_numeric_one(
+ tmp_path, kind
+):
+ """No Python-level loop over rows when concatenating an obs column.
+
+ `_concat_masked` used to fill a `List[Any]` of length n_obs one element at
+ a time and then walk it twice more. That is invisible to a read counter
+ and to a growth ratio, but it shows up immediately as executed line events
+ per row relative to the numeric path, which has always been vectorised.
+
+ Measured after the fix: masked is 1.00x numeric, string 1.35x (the string
+ path still materialises Python `str` objects via `read_str_all`, which is
+ inherent to how strings are read, not a per-row loop). Before the fix
+ masked was above 1.2x.
+ """
+ n_obs = 4000
+ numeric_lines, _ = _obs_column_cost(tmp_path, "numeric", n_obs)
+ kind_lines, _ = _obs_column_cost(tmp_path, kind, n_obs)
+
+ limit = 1.10 if kind == "masked" else 1.45
+ ratio = kind_lines / numeric_lines
+ assert ratio <= limit, (
+ f"a {kind} obs column executes {ratio:.2f}x the Python lines of a "
+ f"numeric one ({kind_lines} vs {numeric_lines} at n_obs={n_obs}); "
+ f"limit {limit}. Look for a per-element loop over rows."
+ )
+
+
+@pytest.mark.parametrize("kind", ["numeric", "masked", "string"])
+def test_obs_column_allocation_grows_linearly(tmp_path, kind):
+ """Peak allocation must not grow faster than the data itself.
+
+ Note what this does *not* claim: obs columns are read whole, so peak
+ allocation is O(n_obs) and not O(chunk). The streaming guarantee that the
+ README makes holds for X, not for obs annotation. `benchmarks/` reports
+ the real curve; this only stops it getting worse than linear.
+ """
+
+ def measure(n: int) -> int:
+ _, peak = _obs_column_cost(tmp_path, kind, n)
+ return peak
+
+ assert_grows_linearly(
+ measure,
+ what=f"concat allocation, {kind} obs column",
+ axis="n_obs",
+ sizes=(256, 1024, 4096),
+ )
+
+
+# ---------------------------------------------------------------------------
+# subset, split and conversion
+#
+# concat is where the defect was found, but nothing about the defect was
+# specific to concat: every command that aligns one index onto another has the
+# same opportunity to re-read a column per element.
+
+
+@pytest.mark.parametrize("axis", ["obs", "var"])
+def test_subset_by_name_reads_grow_linearly(tmp_path, axis):
+ def measure(n: int) -> int:
+ source = _store(
+ tmp_path / f"s{axis}{n}.h5ad",
+ name="s",
+ n_obs=n if axis == "obs" else 4,
+ n_var=n if axis == "var" else 4,
+ )
+ # A name file, not indices: resolving names onto the stored index is
+ # exactly the kind of alignment that went quadratic in concat.
+ names = [f"sc{i}" for i in range(0, n, 2)] if axis == "obs" else [
+ f"g{i}" for i in range(0, n, 2)
+ ]
+ name_file = tmp_path / f"n{axis}{n}.txt"
+ name_file.write_text("\n".join(names) + "\n")
+ with count_io() as io:
+ subset_h5ad(
+ source,
+ tmp_path / f"os{axis}{n}.h5ad",
+ name_file if axis == "obs" else None,
+ name_file if axis == "var" else None,
+ console=QUIET,
+ )
+ return io.reads
+
+ assert_grows_linearly(
+ measure, what=f"subset --{axis}", axis=f"n_{axis}", sizes=(64, 256, 1024)
+ )
+
+
+def test_subset_by_query_reads_grow_linearly(tmp_path):
+ def measure(n: int) -> int:
+ source = _store(
+ tmp_path / f"q{n}.h5ad", name="q", n_obs=n, n_var=4, obs_kind="numeric"
+ )
+ with count_io() as io:
+ subset_h5ad(
+ source,
+ tmp_path / f"oq{n}.h5ad",
+ None,
+ None,
+ console=QUIET,
+ obs_query="col >= 0",
+ )
+ return io.reads
+
+ assert_grows_linearly(measure, what="subset --obs-query", axis="n_obs")
+
+
+def test_h5ad_to_zarr_conversion_reads_grow_linearly(tmp_path):
+ def measure(n: int) -> int:
+ source = _store(tmp_path / f"c{n}.h5ad", name="c", n_obs=n, n_var=4)
+ # subset is also the conversion path, but it insists on a selection;
+ # selecting every row is the identity and still rewrites the store.
+ keep = tmp_path / f"ck{n}.txt"
+ keep.write_text("\n".join(f"cc{i}" for i in range(n)) + "\n")
+ with count_io() as io:
+ subset_h5ad(
+ source, tmp_path / f"oc{n}.zarr", keep, None, console=QUIET
+ )
+ return io.reads
+
+ assert_grows_linearly(measure, what="h5ad to zarr", axis="n_obs")
+
+
+def test_split_reads_grow_linearly_in_group_count(tmp_path):
+ """Splitting into k groups must stay linear in k at fixed store size.
+
+ n_obs is held constant deliberately. Scaling rows alongside groups would
+ make the total legitimately quadratic -- k passes over 4k rows -- and the
+ guard would be measuring the fixture, not the implementation.
+ """
+ from adata.commands.split import split_store
+
+ n_obs = 512
+
+ def measure(n: int) -> int:
+ source = _store(
+ tmp_path / f"sp{n}.h5ad",
+ name="sp",
+ n_obs=n_obs,
+ n_var=4,
+ n_categories=n,
+ )
+ out = tmp_path / f"osp{n}"
+ with count_io() as io:
+ split_store(source, "ct", out, QUIET, axis="obs", manifest=False)
+ return io.reads
+
+ assert_grows_linearly(
+ measure, what="split --by", axis="n_groups", sizes=(8, 32, 128)
+ )
+
+
+# ---------------------------------------------------------------------------
+# the constants themselves
+
+
+def test_growth_limit_separates_linear_from_super_linear():
+ """Documenting why K is 6, as an executable statement rather than a note.
+
+ At the 4x spacing in SIZES the increment ratio is 4.0 for linear work,
+ about 4.4 for n log n, 8 for n**1.5 and 16 for quadratic.
+ """
+ small, mid, large = SIZES
+ assert mid == small * 4 and large == mid * 4, (
+ "GROWTH_LIMIT is calibrated for 4x spacing; changing SIZES without "
+ "recomputing it invalidates every guard in this file."
+ )
+
+ def ratio(exponent: float) -> float:
+ cost = [float(n) ** exponent for n in SIZES]
+ return (cost[2] - cost[1]) / (cost[1] - cost[0])
+
+ assert ratio(1.0) < GROWTH_LIMIT, "linear work must pass"
+ assert ratio(1.0) == pytest.approx(4.0, rel=0.01)
+ assert ratio(2.0) > GROWTH_LIMIT, "quadratic work must fail"
+ assert ratio(1.5) > GROWTH_LIMIT, "n**1.5 must fail"
+
+
+# ---------------------------------------------------------------------------
+# inspection is free
+#
+# `view` and `ls` reach only `.shape`, `.dtype` and `.attrs`: `axis_len` goes
+# through `element_len`, which reads a shape, and `_array_details` and
+# `_infer_untagged` never touch a value. That is the whole reason they return
+# instantly on a store far too large to open, and it is the one claim in the
+# README that can be stated as an exact number rather than a ratio.
+
+
+def _reads_for_inspection(path: Path, run) -> "object":
+ with count_io() as io:
+ run(path)
+ return io
+
+
+@pytest.mark.parametrize("n_obs", [64, 4096])
+@pytest.mark.parametrize(
+ "label,run",
+ [
+ ("view", lambda p: show_info(p, QUIET, out_console=QUIET)),
+ ("view --types", lambda p: show_info(p, QUIET, show_types=True,
+ out_console=QUIET)),
+ ("ls", lambda p: list_store(p, QUIET)),
+ ("ls --long", lambda p: list_store(p, QUIET, long=True)),
+ ("ls --plain", lambda p: list_store(p, QUIET, plain=True)),
+ ],
+ ids=lambda v: v if isinstance(v, str) else "",
+)
+def test_inspection_reads_no_data_at_all(tmp_path, n_obs, label, run):
+ """Not "grows slowly" -- zero. An exact count, so it needs no tolerance.
+
+ A ratio would be the wrong instrument here: 0 to 0 is not a meaningful
+ ratio, and the moment inspection reads *one* column the answer stops
+ being zero regardless of how the store scales.
+ """
+ source = _store(tmp_path / f"i{label}{n_obs}.h5ad", name="i", n_obs=n_obs,
+ n_var=16, obs_kind="numeric", n_categories=4)
+ io = _reads_for_inspection(source, run)
+
+ assert io.elements == 0, (
+ f"`{label}` read {io.elements} data elements from a {n_obs}-row store; "
+ f"inspection is supposed to touch only shapes and attributes. "
+ f"Largest reader: {max(io.by_name.items(), key=lambda kv: kv[1], default=('-', 0))}"
+ )
+
+
+def test_the_inspection_fixture_really_does_hold_readable_data(tmp_path):
+ """Keeps the test above from passing because there was nothing to read.
+
+ Same store, read by a command that is supposed to read: if this registers
+ nothing either, the fixture or the hooks are broken, not the claim.
+ """
+ source = _store(tmp_path / "control.h5ad", name="i", n_obs=4096, n_var=16,
+ obs_kind="numeric", n_categories=4)
+ with count_io() as io:
+ export_table(source, "obs", None, tmp_path / "control.csv", 10_000, None, QUIET)
+
+ assert io.elements >= 4096, (
+ f"the control read only {io.elements} elements from a 4096-row store, "
+ "so the zero above proves nothing"
+ )
+
+
+# ---------------------------------------------------------------------------
+# grouping work does not depend on the number of groups
+#
+# `split --by` groups rows through `core.select.group_indices`. Its cost is
+# invisible to every other instrument in this file: the chunk is already in
+# memory so no read counter moves, nothing lasting is allocated, and the
+# Python line count per label is constant. Only counting what is handed to
+# numpy shows it.
+
+
+def test_grouping_does_not_rescan_the_column_once_per_group(tmp_path):
+ """One pass over the column, however many distinct values it holds.
+
+ Scanning per label is O(n_rows * n_groups): measured at n_obs=4096 it was
+ 16,384 elements for 4 groups and 1,048,576 for 256 -- exactly n_rows per
+ group. On a million cells split by a thousand samples that is 10^9
+ comparisons, and `split` would appear to hang for the same reason
+ `concat --merge` did.
+
+ n_obs is fixed. The claim is about work per row, not total work.
+ """
+ from adata.core.select import group_indices
+ from adata.storage import open_store
+
+ n_obs = 4096
+
+ def measure(n_groups: int) -> int:
+ source = _store(
+ tmp_path / f"g{n_groups}.h5ad",
+ name="g",
+ n_obs=n_obs,
+ n_var=2,
+ n_categories=n_groups,
+ )
+ with open_store(source, "r") as store, count_scanned_elements() as scanned:
+ groups, order = group_indices(store.root, "obs", "ct")
+ assert len(order) == n_groups, "fixture did not produce the groups asked for"
+ # A floor, not a measurement -- see count_scanned_elements. If the
+ # implementation stops calling numpy by name the count collapses to
+ # zero and the ratio below would pass for the wrong reason.
+ assert scanned[0] >= n_obs, (
+ f"only {scanned[0]} elements scanned for {n_obs} rows; the counter "
+ "is no longer seeing this code path"
+ )
+ return scanned[0]
+
+ assert_independent_of(
+ measure, what="group_indices", axis="n_groups", sizes=(4, 256)
+ )
+
+
+# ---------------------------------------------------------------------------
+# export
+#
+# Each export reads the thing it is asked for and nothing else. The axis
+# differs per subcommand -- rows, elements, nonzeros, keys -- so each says
+# which one it scales and holds the rest fixed.
+
+
+def test_export_dataframe_reads_grow_linearly_in_rows(tmp_path):
+ def measure(n: int) -> int:
+ source = _store(tmp_path / f"ed{n}.h5ad", name="e", n_obs=n, n_var=4,
+ obs_kind="numeric")
+ with count_io() as io:
+ export_table(source, "obs", None, tmp_path / f"ed{n}.csv",
+ 10_000, None, QUIET)
+ return io.reads
+
+ assert_grows_linearly(measure, what="export dataframe", axis="n_obs")
+
+
+def test_export_dataframe_reads_grow_linearly_in_column_count(tmp_path):
+ def measure(n: int) -> int:
+ source = _store(tmp_path / f"ec{n}.h5ad", name="e", n_obs=16, n_var=4,
+ n_obs_columns=n)
+ with count_io() as io:
+ export_table(source, "obs", None, tmp_path / f"ec{n}.csv",
+ 10_000, None, QUIET)
+ return io.reads
+
+ assert_grows_linearly(measure, what="export dataframe", axis="n_obs_columns")
+
+
+def test_export_array_reads_grow_linearly_in_element_count(tmp_path):
+ def measure(n: int) -> int:
+ source = _store(tmp_path / f"ea{n}.h5ad", name="e", n_obs=n, n_var=4)
+ with count_io() as io:
+ export_npy(source, "obsm/X_pca", tmp_path / f"ea{n}.npy",
+ 100_000, QUIET)
+ return io.reads
+
+ assert_grows_linearly(measure, what="export array", axis="n_obs")
+
+
+@pytest.mark.parametrize("in_memory", [False, True])
+def test_export_sparse_reads_grow_linearly_in_nonzeros(tmp_path, in_memory):
+ """Both paths: the streamed one and the one that loads the matrix."""
+
+ def measure(n: int) -> int:
+ source = _store(tmp_path / f"es{in_memory}{n}.h5ad", name="e",
+ n_obs=n, n_var=4)
+ with count_io() as io:
+ export_mtx(source, "X", tmp_path / f"es{in_memory}{n}.mtx",
+ None, 1_000, in_memory, QUIET)
+ return io.reads
+
+ assert_grows_linearly(
+ measure, what=f"export sparse (in_memory={in_memory})", axis="nnz"
+ )
+
+
+def test_export_dict_reads_grow_linearly_in_key_count(tmp_path):
+ def measure(n: int) -> int:
+ source = _store_with_uns_keys(tmp_path / f"ej{n}.h5ad", n)
+ with count_io() as io:
+ export_json(source, "uns", tmp_path / f"ej{n}.json",
+ 100_000, False, QUIET)
+ return io.reads
+
+ assert_grows_linearly(
+ measure, what="export dict", axis="n_keys", sizes=(64, 256, 1024)
+ )
+
+
+def test_export_image_reads_grow_linearly_in_pixels(tmp_path):
+ from adata.commands.export import export_image
+
+ def measure(n: int) -> int:
+ source = _store_with_image(tmp_path / f"ei{n}.h5ad", n)
+ with count_io() as io:
+ export_image(source, "uns/picture", tmp_path / f"ei{n}.png", QUIET)
+ return io.reads
+
+ # n is the side of a square image, so pixels grow as n**2 -- the axis
+ # being scaled is the pixel count, and the sizes below keep it at 4x.
+ assert_grows_linearly(
+ measure, what="export image", axis="pixels", sizes=(16, 32, 64)
+ )
+
+
+# ---------------------------------------------------------------------------
+# import
+
+
+def test_import_dataframe_reads_grow_linearly_in_rows(tmp_path):
+ def measure(n: int) -> int:
+ source = _store(tmp_path / f"id{n}.h5ad", name="m", n_obs=n, n_var=4)
+ csv = tmp_path / f"id{n}.csv"
+ pd.DataFrame(
+ {"score": np.arange(n, dtype="int32")},
+ index=[f"mc{i}" for i in range(n)],
+ ).to_csv(csv)
+ with count_io() as io:
+ import_object(source, "obs", csv, tmp_path / f"od{n}.h5ad",
+ False, None, QUIET)
+ # `work`, not `reads`: import is a write path, and measuring only
+ # reads made this vacuous.
+ return io.work
+
+ assert_grows_linearly(measure, what="import dataframe", axis="n_obs")
+
+
+def test_import_array_reads_grow_linearly_in_element_count(tmp_path):
+ def measure(n: int) -> int:
+ source = _store(tmp_path / f"ia{n}.h5ad", name="m", n_obs=n, n_var=4)
+ npy = tmp_path / f"ia{n}.npy"
+ np.save(npy, np.zeros((n, 3), dtype="float32"))
+ with count_io() as io:
+ import_object(source, "obsm/imported", npy,
+ tmp_path / f"oa{n}.h5ad", False, None, QUIET)
+ return io.work
+
+ assert_grows_linearly(measure, what="import array", axis="n_obs")
+
+
+def test_import_sparse_reads_grow_linearly_in_nonzeros(tmp_path):
+ def measure(n: int) -> int:
+ source = _store(tmp_path / f"is{n}.h5ad", name="m", n_obs=n, n_var=4)
+ mtx = tmp_path / f"is{n}.mtx"
+ sparse_io = pytest.importorskip("scipy.io")
+ sparse_io.mmwrite(
+ str(mtx), sparse.csr_matrix(np.ones((n, 4), dtype="float32"))
+ )
+ with count_io() as io:
+ import_object(source, "layers/imported", mtx,
+ tmp_path / f"os{n}.h5ad", False, None, QUIET)
+ return io.work
+
+ assert_grows_linearly(measure, what="import sparse", axis="nnz")
+
+
+def test_import_dict_reads_grow_linearly_in_key_count(tmp_path):
+ import json
+
+ def measure(n: int) -> int:
+ source = _store(tmp_path / f"ij{n}.h5ad", name="m", n_obs=8, n_var=4)
+ blob = tmp_path / f"ij{n}.json"
+ blob.write_text(json.dumps({f"k{i}": i for i in range(n)}))
+ with count_io() as io:
+ import_object(source, "uns/imported", blob,
+ tmp_path / f"oj{n}.h5ad", False, None, QUIET)
+ return io.work
+
+ assert_grows_linearly(
+ measure, what="import dict", axis="n_keys", sizes=(64, 256, 1024)
+ )
+
+
+def test_import_image_reads_grow_linearly_in_pixels(tmp_path):
+ """Images take their own entry point.
+
+ `import_object` dispatches on extension and `.png` is deliberately not in
+ `EXTENSION_FORMAT`; the CLI's `import image` calls `_import_image`
+ directly, and it edits in place rather than writing a copy.
+ """
+ Image = pytest.importorskip("PIL.Image")
+ from adata.commands.import_data import _import_image
+
+ def measure(n: int) -> int:
+ source = _store(tmp_path / f"ii{n}.h5ad", name="m", n_obs=8, n_var=4)
+ png = tmp_path / f"ii{n}.png"
+ Image.fromarray(np.zeros((n, n, 3), dtype="uint8")).save(png)
+ with count_io() as io:
+ _import_image(source, "uns/picture", png, QUIET)
+ return io.work
+
+ assert_grows_linearly(
+ measure, what="import image", axis="pixels", sizes=(16, 32, 64)
+ )
+
+
+# ---------------------------------------------------------------------------
+# create, and the concat options nothing else covers
+
+
+@pytest.mark.parametrize("from_file", [False, True])
+def test_create_writes_grow_linearly_in_obs_count(tmp_path, from_file):
+ """Both the generated-names path and the name-file path."""
+
+ def measure(n: int) -> int:
+ names = None
+ if from_file:
+ names = tmp_path / f"cn{n}.txt"
+ names.write_text("\n".join(f"c{i}" for i in range(n)) + "\n")
+ output = tmp_path / f"cr{from_file}{n}.h5ad"
+ with count_io() as io:
+ create_store(
+ output,
+ QUIET,
+ n_obs=None if from_file else n,
+ n_var=4,
+ obs_names=names,
+ )
+ return io.work
+
+ assert_grows_linearly(
+ measure, what=f"create (from_file={from_file})", axis="n_obs"
+ )
+
+
+@pytest.mark.parametrize(
+ "option", ["label", "index_unique"], ids=["--label", "--index-unique"]
+)
+def test_concat_option_reads_grow_linearly_in_obs_count(tmp_path, option):
+ """Both build a per-row value, so both are worth a guard of their own."""
+
+ def measure(n: int) -> int:
+ files = _inputs(tmp_path, f"co{option}{n}", n_obs=n, n_var=4)
+ kwargs = {"label": "batch"} if option == "label" else {"index_unique": "-"}
+ with count_io() as io:
+ concat_on_disk(files, tmp_path / f"oco{option}{n}.h5ad", QUIET, **kwargs)
+ return io.reads
+
+ assert_grows_linearly(measure, what=f"concat --{option}", axis="n_obs")
+
+
+def test_split_by_var_reads_grow_linearly_in_group_count(tmp_path):
+ """The var axis takes a different path through `split` than obs does."""
+ from adata.commands.split import split_store
+
+ n_var = 512
+
+ def measure(n: int) -> int:
+ source = _store_with_var_groups(tmp_path / f"sv{n}.h5ad", n_var, n)
+ with count_io() as io:
+ split_store(source, "kind", tmp_path / f"osv{n}", QUIET,
+ axis="var", manifest=False)
+ return io.reads
+
+ assert_grows_linearly(
+ measure, what="split --axis var", axis="n_groups", sizes=(8, 32, 128)
+ )
+
+
+# ---------------------------------------------------------------------------
+# fixture variants the guards above need
+
+
+def _store_with_uns_keys(path: Path, n_keys: int) -> Path:
+ obj = ad.AnnData(
+ X=np.ones((4, 2), dtype="float32"),
+ obs=pd.DataFrame(index=["a", "b", "c", "d"]),
+ var=pd.DataFrame(index=["g1", "g2"]),
+ )
+ obj.uns.update({f"k{i}": i for i in range(n_keys)})
+ obj.write_h5ad(path)
+ return path
+
+
+def _store_with_image(path: Path, side: int) -> Path:
+ obj = ad.AnnData(
+ X=np.ones((4, 2), dtype="float32"),
+ obs=pd.DataFrame(index=["a", "b", "c", "d"]),
+ var=pd.DataFrame(index=["g1", "g2"]),
+ )
+ obj.uns["picture"] = np.zeros((side, side, 3), dtype="uint8")
+ obj.write_h5ad(path)
+ return path
+
+
+def _store_with_var_groups(path: Path, n_var: int, n_groups: int) -> Path:
+ ad.AnnData(
+ X=np.ones((4, n_var), dtype="float32"),
+ obs=pd.DataFrame(index=["a", "b", "c", "d"]),
+ var=pd.DataFrame(
+ {"kind": pd.Categorical([f"k{i % n_groups}" for i in range(n_var)])},
+ index=[f"g{i}" for i in range(n_var)],
+ ),
+ ).write_h5ad(path)
+ return path
+
+
+# ---------------------------------------------------------------------------
+# streaming is bounded by the chunk, not by the input
+#
+# The README's actual claim, and the reason this tool exists. It is only
+# wholly true in one place today, so these guards say what is true rather
+# than what would be nice, and record the measured numbers so that drift is
+# visible. Everything below holds the chunk size fixed and grows the input by
+# 256x; a command that loaded everything would grow 256x with it.
+
+#: 256x span. Wide enough that "flat" and "linear" cannot be confused, and
+#: the reason these carry the `slow` marker: building a 65,536-row store is
+#: most of the ~45 s this section costs. They still gate merges; the marker
+#: is so a local run can say `-m "not slow"`.
+STREAM_SIZES = (256, 65536)
+
+
+def _peak_for(tmp_path: Path, tag: str, n: int, run) -> int:
+ source = _store(tmp_path / f"{tag}{n}.h5ad", name="s", n_obs=n, n_var=8,
+ obs_kind="numeric")
+ with count_allocations() as peak:
+ run(source, n)
+ return peak[0]
+
+
+@pytest.mark.slow
+def test_export_array_peak_memory_is_set_by_the_chunk_not_the_input(tmp_path):
+ """The closest the tool comes to the claim outright.
+
+ Measured at a fixed 1,000-element chunk: 21.5 KiB at 256 rows and
+ 34.3 KiB at 65,536 -- 1.6x for a 256x input. Not flat, so this does not
+ assert flat; the residue is fixed-size bookkeeping that grows with the
+ length of a formatted shape rather than with the data. Stating it as "at
+ least 64x better than the input" is both true and strong: loading
+ everything would be 256x.
+ """
+
+ def measure(n: int) -> int:
+ return _peak_for(
+ tmp_path, "sa", n,
+ lambda p, k: export_npy(p, "obsm/X_pca", tmp_path / f"sa{k}.npy",
+ 1_000, QUIET),
+ )
+
+ assert_grows_slower_than_input(
+ measure,
+ what="export array peak allocation at --chunk 1000",
+ axis="n_obs",
+ sizes=STREAM_SIZES,
+ at_least=64.0,
+ )
+
+
+@pytest.mark.parametrize(
+ "label,factor,run",
+ [
+ # Measured over the 256x span: export sparse 6.5x, export dataframe
+ # 2.2x, subset 46x. The factors below sit roughly midway between the
+ # measurement and linear, so ordinary variation passes and a real
+ # drift towards loading everything fails.
+ (
+ "export sparse",
+ 8.0,
+ lambda p, k, t: export_mtx(p, "X", t / f"ss{k}.mtx", None, 1_000,
+ False, QUIET),
+ ),
+ (
+ "export dataframe",
+ 8.0,
+ lambda p, k, t: export_table(p, "obs", None, t / f"st{k}.csv",
+ 1_000, None, QUIET),
+ ),
+ (
+ "subset",
+ 2.0,
+ lambda p, k, t: subset_h5ad(p, t / f"su{k}.h5ad", None, None,
+ console=QUIET, obs_query="col >= 0",
+ chunk_rows=256),
+ ),
+ ],
+ ids=["export-sparse", "export-dataframe", "subset"],
+)
+@pytest.mark.slow
+def test_streamed_peak_memory_grows_far_slower_than_the_input(
+ tmp_path, label, factor, run
+):
+ """Not flat, and the guards should not pretend otherwise.
+
+ These paths read an index or an indptr whole, so peak allocation does
+ track n_obs -- just far below it. `subset` is the weakest of the three
+ because obs columns are materialised per column; that is a known gap,
+ reported by `benchmarks/` rather than hidden here.
+ """
+
+ def measure(n: int) -> int:
+ return _peak_for(tmp_path, f"st{label[-4:]}", n,
+ lambda p, k: run(p, k, tmp_path))
+
+ assert_grows_slower_than_input(
+ measure,
+ what=f"{label} peak allocation at a fixed chunk",
+ axis="n_obs",
+ sizes=STREAM_SIZES,
+ at_least=factor,
+ )
+
+
+@pytest.mark.slow
+def test_streaming_export_sparse_costs_far_less_than_loading_the_matrix(
+ tmp_path,
+):
+ """`--in-memory` exists as the fast path; the default must earn its place.
+
+ Self-calibrating: both are measured in the same run, so the assertion
+ holds regardless of platform. Measured at 65,536 rows: 777 KiB streamed
+ against 31,763 KiB in memory, a factor of 41.
+ """
+ n = 65_536
+ source = _store(tmp_path / "cmp.h5ad", name="s", n_obs=n, n_var=8)
+
+ with count_allocations() as streamed:
+ export_mtx(source, "X", tmp_path / "streamed.mtx", None, 1_000, False, QUIET)
+ with count_allocations() as loaded:
+ export_mtx(source, "X", tmp_path / "loaded.mtx", None, 1_000, True, QUIET)
+
+ assert streamed[0] <= loaded[0] / 4, (
+ f"streaming export used {streamed[0] // 1024} KiB against "
+ f"{loaded[0] // 1024} KiB for --in-memory, a factor of only "
+ f"{loaded[0] / max(streamed[0], 1):.1f}. The default path is supposed "
+ "to be the one you reach for when the matrix does not fit."
+ )
+
+
+# ---------------------------------------------------------------------------
+# copying a store must stay inside its read budget
+#
+# `copy_dataset` sizes each read to TARGET_READ_BYTES. For variable-length
+# strings it cannot ask the dtype how wide an element is -- h5py reports the
+# itemsize of a pointer -- so it estimated. An estimate is not a bound: at
+# 4 KiB elements the 32 MiB budget became an 827 MB peak, and at 200,000 rows
+# the computed step exceeded the dataset, so the whole array was read at once.
+#
+# This is the one path the rest of this file did not reach, and it is the
+# path every `copy:` task in `subset` and every uns entry goes through.
+
+
+def _vlen_source(path: Path, width: int, n_rows: int):
+ """An HDF5 variable-length string dataset of `n_rows` x `width` bytes."""
+ import h5py
+
+ with h5py.File(path, "w") as handle:
+ handle.create_dataset(
+ "t",
+ data=np.array([("x" * width).encode()] * n_rows, dtype=object),
+ dtype=h5py.string_dtype(),
+ chunks=(1024,),
+ )
+ return path
+
+
+@pytest.mark.parametrize("width", [16, 256, 4096])
+def test_a_read_of_variable_length_strings_stays_inside_the_byte_budget(
+ tmp_path, width
+):
+ """The budget is bytes per read, so a wider element means fewer elements.
+
+ This is the exact invariant, and it costs nothing to check: the step is
+ computed, not measured. Before the fix the element width was assumed to
+ be 64 bytes, so the step was the same 524,288 elements whatever the data
+ actually held -- 2 GiB per read at 4 KiB elements, against a stated
+ 32 MiB budget.
+ """
+ import h5py
+
+ from adata.storage import (
+ TARGET_READ_BYTES,
+ _chunk_step,
+ _sample_element_bytes,
+ )
+
+ n_rows = 50_000
+ source = _vlen_source(tmp_path / f"s{width}.h5", width, n_rows)
+ with h5py.File(source, "r") as handle:
+ dataset = handle["t"]
+ sampled = _sample_element_bytes(dataset, n_rows)
+ step = _chunk_step(dataset.shape, dataset.chunks, dataset.dtype, sampled)
+
+ assert sampled >= min(width, 64), (
+ f"sampled {sampled} bytes for {width}-byte elements; the width is "
+ "being guessed rather than measured"
+ )
+ # A read may round up to a whole source chunk, hence the slack.
+ assert step * width <= TARGET_READ_BYTES * 2, (
+ f"one read would take {step * width / 1e6:.0f} MB of {width}-byte "
+ f"strings ({step} elements), against a {TARGET_READ_BYTES / 1e6:.0f} "
+ "MB budget"
+ )
+
+
+@pytest.mark.slow
+def test_copying_wide_strings_does_not_read_the_whole_array(tmp_path):
+ """And the budget holds in practice, not just in the arithmetic.
+
+ Sized so the array is several times the budget: before the fix the
+ computed step exceeded the row count, so the whole thing was read at once
+ -- 164 MB here, and 827 MB in the 200,000-row case that prompted this.
+ """
+ import h5py
+
+ from adata.storage import TARGET_READ_BYTES, copy_dataset
+
+ width, n_rows = 4096, 40_000 # ~164 MB, about 5x the budget
+ source = _vlen_source(tmp_path / "wide.h5", width, n_rows)
+
+ with h5py.File(source, "r") as src, h5py.File(tmp_path / "out.h5", "w") as dst:
+ with count_allocations() as peak:
+ copy_dataset(src["t"], dst, "t")
+
+ assert peak[0] <= TARGET_READ_BYTES * 3, (
+ f"copying a {width * n_rows / 1e6:.0f} MB array of {width}-byte "
+ f"strings peaked at {peak[0] / 1e6:.0f} MB, against a "
+ f"{TARGET_READ_BYTES / 1e6:.0f} MB read budget. A step larger than "
+ "the array means the whole array is read in one go."
+ )
+
+
+# ---------------------------------------------------------------------------
+# convert
+#
+# The streaming transpose exists because the in-memory one does not scale.
+# That is a claim about peak memory, so it is the peak that is asserted --
+# a correctness test cannot tell the two implementations apart, which is
+# exactly why they are both allowed to exist.
+
+
+def _sparse_store(path: Path, n_obs: int, n_var: int = 32, density: float = 0.2):
+ rng = np.random.default_rng(0)
+ matrix = sparse.random(
+ n_obs, n_var, density=density, format="csr", dtype="float32",
+ random_state=rng,
+ )
+ ad.AnnData(
+ X=matrix,
+ obs=pd.DataFrame(index=[f"c{i}" for i in range(n_obs)]),
+ var=pd.DataFrame(index=[f"g{i}" for i in range(n_var)]),
+ ).write_h5ad(path)
+ return path
+
+
+def _convert(source: Path, out: Path, **kwargs):
+ from adata.commands.convert import convert_store
+
+ convert_store(source, ["X"], out, QUIET, **kwargs)
+
+
+def test_convert_dtype_reads_grow_linearly_in_nonzeros(tmp_path):
+ def measure(n: int) -> int:
+ source = _sparse_store(tmp_path / f"cd{n}.h5ad", n)
+ with count_io() as io:
+ _convert(source, tmp_path / f"od{n}.h5ad", dtype="float64")
+ return io.work
+
+ assert_grows_linearly(measure, what="convert --dtype", axis="nnz")
+
+
+@pytest.mark.parametrize("layout", ["csc", "dense"])
+def test_convert_layout_reads_grow_linearly(tmp_path, layout):
+ def measure(n: int) -> int:
+ source = _sparse_store(tmp_path / f"cl{layout}{n}.h5ad", n)
+ with count_io() as io:
+ _convert(
+ source, tmp_path / f"ol{layout}{n}.h5ad", layout=layout, force=True
+ )
+ return io.work
+
+ assert_grows_linearly(measure, what=f"convert --layout {layout}", axis="n_obs")
+
+
+@pytest.mark.slow
+def test_the_streaming_transpose_does_not_hold_the_matrix(tmp_path):
+ """Peak allocation must not track nnz. This is the whole point of it.
+
+ Measured against the in-memory path in the same run, which does hold the
+ matrix and therefore does grow -- so the comparison shows the difference
+ is real rather than an artefact of how the fixture is built.
+ """
+
+ def peak(n: int, in_memory: bool) -> int:
+ source = _sparse_store(tmp_path / f"tp{in_memory}{n}.h5ad", n, n_var=64)
+ with count_allocations() as measured:
+ _convert(
+ source,
+ tmp_path / f"otp{in_memory}{n}.h5ad",
+ layout="csc",
+ in_memory=in_memory,
+ chunk=4096,
+ )
+ return measured[0]
+
+ small, large = 256, 8192
+ streaming = peak(large, False) / max(1, peak(small, False))
+ loaded = peak(large, True) / max(1, peak(small, True))
+
+ assert streaming < loaded, (
+ f"the streaming transpose grew {streaming:.1f}x over a 32x larger "
+ f"matrix and the in-memory one grew {loaded:.1f}x -- if streaming is "
+ "not the cheaper of the two it has no reason to exist"
+ )
+ assert streaming <= 8.0, (
+ f"streaming transpose peak grew {streaming:.1f}x for 32x the "
+ "nonzeros; it is supposed to be bounded by the chunk"
+ )
+
+
+def test_the_transpose_streams_unless_asked_not_to(tmp_path):
+ """The safe path is the default; --in-memory is opt-in.
+
+ Asserted by watching which function runs, because the two produce
+ identical output and no result can distinguish them.
+ """
+ import adata.core.convert as convert_module
+
+ called: List[str] = []
+ for name in ("transpose_sparse_streaming", "transpose_sparse_in_memory"):
+ original = getattr(convert_module, name)
+
+ def record(*args, _name=name, _original=original, **kwargs):
+ called.append(_name)
+ return _original(*args, **kwargs)
+
+ setattr(convert_module, name, record)
+
+ try:
+ source = _sparse_store(tmp_path / "d.h5ad", 64)
+ _convert(source, tmp_path / "default.h5ad", layout="csc")
+ assert called == ["transpose_sparse_streaming"], called
+
+ called.clear()
+ _convert(source, tmp_path / "asked.h5ad", layout="csc", in_memory=True)
+ assert called == ["transpose_sparse_in_memory"], called
+ finally:
+ for name in ("transpose_sparse_streaming", "transpose_sparse_in_memory"):
+ setattr(
+ convert_module, name, getattr(convert_module, name).__wrapped__
+ if hasattr(getattr(convert_module, name), "__wrapped__")
+ else getattr(convert_module, name)
+ )
diff --git a/uv.lock b/uv.lock
index 6c20152..dd5144a 100644
--- a/uv.lock
+++ b/uv.lock
@@ -486,7 +486,7 @@ wheels = [
[[package]]
name = "pyadata-cli"
-version = "0.5.1"
+version = "0.6.0"
source = { editable = "." }
dependencies = [
{ name = "h5py" },
]