diff --git a/.dockerignore b/.dockerignore
new file mode 100644
index 0000000..f841322
--- /dev/null
+++ b/.dockerignore
@@ -0,0 +1,18 @@
+# Keep the build context to what `uv sync --locked` actually needs. In
+# particular .venv/ would otherwise ship a host-platform virtualenv into the
+# image, and .claude/worktrees/ holds full copies of this repo.
+#
+# Do NOT add README.md or LICENSE here: pyproject.toml references them via
+# `readme` and `license-files`, so the build fails without them.
+.venv/
+.git/
+.claude/
+.pytest_cache/
+__pycache__/
+*.py[cod]
+.coverage
+coverage.xml
+htmlcov/
+build/
+dist/
+*.egg-info/
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..77e4540 100644
--- a/docs/COMMANDS.md
+++ b/docs/COMMANDS.md
@@ -118,7 +118,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 b37ac86..669fd73 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -22,7 +22,7 @@ Or run it without installing anything:
```bash
docker run --rm -it -v /path/to/data:/data \
- quay.io/cellgeni/adata-cli:latest view /data/your_file.h5ad
+ quay.io/cellgeni/adata-cli:latest adata view /data/your_file.h5ad
```
## Documentation
@@ -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
diff --git a/pyproject.toml b/pyproject.toml
index 1d1ebaf..7767eca 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.0"
+version = "0.5.1"
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..45f1f30 100644
--- a/src/adata/cli.py
+++ b/src/adata/cli.py
@@ -7,7 +7,7 @@
import typer
from adata.commands import (
- MERGE_STRATEGIES,
+ MERGE_CHOICES,
concat_stores,
create_store,
split_store,
@@ -413,12 +413,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 +444,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,
diff --git a/src/adata/commands/__init__.py b/src/adata/commands/__init__.py
index 70e0690..3cc464b 100644
--- a/src/adata/commands/__init__.py
+++ b/src/adata/commands/__init__.py
@@ -5,4 +5,4 @@
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
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/core/concat.py b/src/adata/core/concat.py
index 4982916..f1c4ecc 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(
@@ -945,13 +964,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/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 58e7bad..025bce5 100644
--- a/src/adata/storage/__init__.py
+++ b/src/adata/storage/__init__.py
@@ -440,12 +440,125 @@ def _is_string_src(src: Any) -> bool:
return is_string_dtype(getattr(src, "dtype", None))
-def _chunk_step(shape: Sequence[int], chunks: Optional[Sequence[int]]) -> int:
- if chunks is not None and len(chunks) > 0 and chunks[0]:
- return int(chunks[0])
+#: Byte budget for a single read when streaming a dataset.
+#:
+#: Using a source's chunk height as the read size means inheriting whatever
+#: the writer chose, and a store chunked `(1, n_cols)` is then copied one row
+#: per read. On a local disk that is merely wasteful; on a network filesystem
+#: (Lustre, NFS) every read is a round-trip costing milliseconds, so a
+#: million-row copy spends nearly all of its time waiting. 32 MiB sits
+#: comfortably above a typical 1 MiB Lustre stripe while bounding peak memory.
+TARGET_READ_BYTES = 32 * 1024 * 1024
+
+#: 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 _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)
+ if _is_vlen(dtype):
+ itemsize = element_bytes
+
+ return max(1, width * itemsize)
+
+
+def _chunk_step(
+ 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.
+
+ The step is grown to `TARGET_READ_BYTES`, then rounded down to a whole
+ number of source chunks: reading part of a chunk still costs decompressing
+ all of it, so a step that splits one wastes the remainder.
+ """
if not shape:
return 1
- return max(1, min(1024, int(shape[0])))
+
+ n_rows = int(shape[0])
+ if n_rows <= 0:
+ return 1
+
+ chunk_rows = 0
+ 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, 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
+ # makes TARGET_READ_BYTES a target rather than a cap -- a source whose
+ # own chunk already exceeds the budget (say 1000 x 1e6 float32, chunked
+ # whole) reads that chunk regardless. This matches the previous
+ # behaviour, which used the chunk height verbatim.
+ step = max(chunk_rows, (step // chunk_rows) * chunk_rows)
+
+ return min(step, n_rows)
def copy_dataset(src: Any, dst_group: Any, name: str) -> Any:
@@ -475,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))
+ 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..7f1b367 100644
--- a/tests/test_commands_phase2.py
+++ b/tests/test_commands_phase2.py
@@ -758,3 +758,104 @@ 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) == {}
diff --git a/tests/test_performance.py b/tests/test_performance.py
new file mode 100644
index 0000000..87c0023
--- /dev/null
+++ b/tests/test_performance.py
@@ -0,0 +1,1140 @@
+"""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."
+ )
diff --git a/tests/test_storage.py b/tests/test_storage.py
index 926627f..9c57d17 100644
--- a/tests/test_storage.py
+++ b/tests/test_storage.py
@@ -393,3 +393,77 @@ def test_clamping_handles_one_dimensional_chunks():
assert _clamp_chunks({"chunks": (100,)}, 5)["chunks"] == (5,)
assert _clamp_chunks({}, 5) == {}
+
+
+# ---------------------------------------------------------------------------
+# read sizing
+
+
+def test_read_step_grows_past_a_pathological_source_chunk():
+ """A `(1, n_cols)` source must not be copied one row per read.
+
+ Forwarding the source's chunk height meant inheriting whatever the writer
+ chose. On a network filesystem each read is a round-trip, so a row-chunked
+ million-cell store spent the whole copy waiting on latency.
+ """
+ from adata.storage import TARGET_READ_BYTES, _chunk_step, _row_bytes
+
+ shape, chunks, dtype = (1_000_000, 30_000), (1, 30_000), np.dtype("float32")
+ step = _chunk_step(shape, chunks, dtype)
+
+ assert step > 1, "the source chunk height is no longer used verbatim"
+ assert step * _row_bytes(shape, dtype) >= TARGET_READ_BYTES // 2
+
+
+def test_read_step_stays_a_whole_number_of_source_chunks():
+ """Reading part of a chunk still costs decompressing all of it."""
+ from adata.storage import _chunk_step
+
+ for chunk_rows in (1, 7, 100, 65_536):
+ step = _chunk_step((10_000_000,), (chunk_rows,), np.dtype("int64"))
+ assert step % chunk_rows == 0
+
+
+def test_read_step_never_exceeds_the_dataset():
+ from adata.storage import _chunk_step
+
+ assert _chunk_step((10,), (1,), np.dtype("int64")) == 10
+ assert _chunk_step((), None, np.dtype("int64")) == 1
+ assert _chunk_step((0,), None, np.dtype("int64")) == 1
+
+
+def test_read_step_bounds_peak_memory():
+ """A wide unchunked source used to read 1024 rows however wide they were."""
+ from adata.storage import TARGET_READ_BYTES, _chunk_step, _row_bytes
+
+ cases = [
+ ((1_000_000, 30_000), None, np.dtype("float32")),
+ ((1_000_000, 50), (1, 50), np.dtype("float32")),
+ ((200_000_000,), (65_536,), np.dtype("int64")),
+ ]
+ for shape, chunks, dtype in cases:
+ step = _chunk_step(shape, chunks, dtype)
+ assert step * _row_bytes(shape, dtype) <= 2 * TARGET_READ_BYTES
+
+
+def test_row_bytes_does_not_trust_a_vlen_itemsize():
+ """h5py reports itemsize 8 for vlen str -- that is the pointer, not the text."""
+ import h5py
+
+ from adata.storage import VLEN_ELEMENT_BYTES, _row_bytes
+
+ assert _row_bytes((10,), h5py.string_dtype(encoding="utf-8")) == VLEN_ELEMENT_BYTES
+ assert _row_bytes((10,), np.dtype("O")) == VLEN_ELEMENT_BYTES
+ assert _row_bytes((10, 4), np.dtype("float32")) == 16
+
+
+def test_copy_dataset_of_a_row_chunked_source_round_trips(new_store):
+ """The larger read step must not change what lands on disk."""
+ path, opener = new_store()
+ with opener("a") as root:
+ from adata.storage import create_dataset
+
+ values = np.arange(400, dtype="float32").reshape(100, 4)
+ create_dataset(root["uns"], "src", data=values, chunks=(1, 4))
+ copy_dataset(root["uns"]["src"], root["uns"], "dst")
+ assert np.array_equal(root["uns"]["dst"][...], values)
diff --git a/uv.lock b/uv.lock
index adb6d8d..6c20152 100644
--- a/uv.lock
+++ b/uv.lock
@@ -486,7 +486,7 @@ wheels = [
[[package]]
name = "pyadata-cli"
-version = "0.5.0"
+version = "0.5.1"
source = { editable = "." }
dependencies = [
{ name = "h5py" },
]