diff --git a/docs/source/write.md b/docs/source/write.md index c9516974..a39f6ab7 100644 --- a/docs/source/write.md +++ b/docs/source/write.md @@ -162,6 +162,19 @@ size of the merged dataset at sequential-IO speed, so a merge of a very large da order of a terabyte) could plausibly take hours rather than minutes. It's worth doing against the alternative of re-extracting genotypes from scratch, not as a routine step. +`gvl.concat`'s planning is streaming: it derives the merge plan from the sorted +merge order rather than building a map over every `(region, sample, ploid)` slot, +so its memory stays proportional to regions plus samples rather than to their +product. The merged `offsets.npy` it writes is still one int64 per merged +`(region, sample[, ploid])` slot — about 16 GB for per-sample tracks (no ploidy +axis) or 32 GB for diploid genotypes, at a 3,734-region by 535,662-sample grid — +and that is a property of the ragged on-disk format, not of the merge. Peak +resident memory during the copy is roughly double that figure: the copy also +holds each input's own offsets array for the duration, and since the inputs +partition the merged grid, their combined size is about the same as the merged +array itself — so about 32 GB for per-sample tracks or 64 GB for diploid +genotypes, at the same grid. + For contig-sharded `.svar2` workflows there is a cheaper path at the variant-store layer: `genoray.SparseVar2.concat(output, sources, mode="copy")` merges disjoint-contig `.svar2` stores that share identical samples/ploidy/fields, then a single `gvl.write` call over the merged store diff --git a/docs/superpowers/plans/2026-09-15-concat-analytic-run-planning.md b/docs/superpowers/plans/2026-09-15-concat-analytic-run-planning.md new file mode 100644 index 00000000..9dce1cb6 --- /dev/null +++ b/docs/superpowers/plans/2026-09-15-concat-analytic-run-planning.md @@ -0,0 +1,1425 @@ +# Analytic Run Planning for `concat` Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Remove `concat`'s materialized `(R*S*P, 2)` provenance map and its `list[Run]` — 32–768 GB at the All of Us chr22 grid — by deriving runs analytically from `order`, and strengthen the `svar2_store` test fixture that #357 found nearly degenerate. + +**Architecture:** `provenance()` + `coalesce()` stay in `_concat_plan.py` as the reference implementation, no longer called on the hot path. A new `RunPlan` derives the identical run sequence directly from `order` in `O(R + S)` memory, exploiting the fact that the run-break pattern depends only on `order` and never on the region index. It is re-iterable (not a generator) because `copy_runs` and `_gather_svar_offsets` each iterate runs twice. A `slot_batches()` method lets `copy_runs` build merged offsets with `R` vectorized gathers instead of one tiny numpy call per run. + +**Tech Stack:** Python 3.10–3.13, numpy, pytest, `pixi run -e dev`. No Rust changes. No new dependencies. + +**Spec:** `docs/superpowers/specs/2026-09-15-concat-analytic-run-planning-design.md` + +## Global Constraints + +- **Branch:** `feat/concat-analytic-run-planning`, already created off `main` at `d0e2f498`. One PR targeting `main`. +- **`provenance` and `coalesce` are never deleted.** They are the oracle the new path is verified against. Removing them removes the proof. +- **Behavior-preserving.** Every existing test in `tests/unit/dataset/test_concat_plan.py`, `test_concat_io.py`, and the concat dataset tests must pass **without modification**. If an existing concat test needs editing to go green, that is a signal the change is wrong — stop and report it, do not edit the test. +- **No on-disk format change.** `DATASET_FORMAT_VERSION` stays `2.0.0`. Nothing in `metadata.json` or any payload file changes bytes. +- **No public API change.** `python/genvarloader/_dataset/_concat_plan.py` is private; `gvl.concat`'s signature and semantics are unchanged. +- **No `feat:` commits and no `BREAKING CHANGE` trailers.** This releases as a PATCH (`0.43.0` → `0.43.1`) under commitizen (`version_provider = "pep621"`, `major_version_zero = true`). Use `perf(concat):`, `refactor(concat):`, `test(svar2):`, `docs(concat):`. +- **Commit message trailers**, on every commit: + ``` + Co-Authored-By: Claude Sonnet 5 + Claude-Session: https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6 + ``` +- **prek hooks must be installed** before committing (`prek install`). They run ruff check, ruff format, pyrefly, and commitizen. +- **Lint covers both trees:** `pixi run -e dev ruff check python/ tests/` and `pixi run -e dev ruff format python/ tests/`. `python/` alone misses test-only issues. +- **Docstrings are Google style** (`Args:` / `Returns:` / `Raises:`), enforced by ruff pydocstyle on `python/genvarloader/` only. +- **Test data must be generated once per worktree** before any test run: `pixi run -e dev gen`. Without it a fresh worktree reports ~387 errors and ~52 failures from missing fixture files, which look like real breakage and are not. (Already done in this worktree.) +- **The chr22 reference grid**, used for every size figure in this plan: `R` = 3,734 regions, `S` = 535,662 samples, `P` = 2. `R*S` = 2,000,161,908; `R*S*P` = 4,000,323,816. A `Run` measures 184 bytes on this codebase. + +--- + +## File Structure + +| File | Responsibility | Change | +|---|---|---| +| `python/genvarloader/_dataset/_concat_plan.py` | Pure planning, no IO | Add `RunPlan`, `ExplicitRunPlan`, `as_plan`, `_default_order`, `_SLOT_BATCH_SLOTS`. Keep `provenance`/`coalesce` as the oracle. | +| `python/genvarloader/_dataset/_concat_io.py` | Buffered streaming IO | `copy_runs` and `gather_fixed` take a run source; `copy_runs` builds offsets from `slot_batches()` and cumsums in place. | +| `python/genvarloader/_dataset/_concat.py` | Merge orchestration | Five `provenance`+`coalesce` sites become `RunPlan`; `_gather_svar_offsets` takes a run source. | +| `tests/unit/dataset/test_concat_plan.py` | Planning unit tests | Add the differential sweep, `slot_batches` tests, the memory bound. | +| `tests/unit/dataset/test_concat_io.py` | IO unit tests | Add `RunPlan`-driven equivalence and the cumsum-aliasing pin. | +| `tests/conftest.py` | Tree-wide fixtures | **Create.** Shared `_SVAR2_REF`, `_SVAR2_VCF_2S`, `svar2_store_2s`. | +| `tests/dataset/conftest.py` | Dataset-suite fixtures | Enriched `_VCF`; `_REF` re-exported from `tests/conftest.py`. | +| Six test modules | svar2 read-path tests | Delete local `svar2_store` shadows; consume `svar2_store_2s`. | +| `tests/dataset/test_write_svar2.py` | svar2 write oracle | Rewrite the fixture-occupancy guard against the 7-cell grid. | +| `docs/source/write.md`, `skills/genvarloader/SKILL.md` | User docs | State concat's planning memory, which the change makes true. | + +--- + +### Task 1: `RunPlan.__iter__` — the analytic planner + +**Files:** +- Modify: `python/genvarloader/_dataset/_concat_plan.py` +- Test: `tests/unit/dataset/test_concat_plan.py` + +**Interfaces:** +- Consumes: the existing `Run` NamedTuple, `provenance`, `coalesce` (all already in this file). +- Produces: `class RunPlan` with `__init__(self, axis: str, shape_per_ds: list[tuple[int, int]], ploidy: int, *, order: NDArray[np.int64] | None = None)`, `__iter__(self) -> Iterator[Run]`, and public attributes `axis`, `ploidy`, `shape_per_ds`, `order`. Tasks 2–4 extend and consume it. + +**Background the implementer needs.** A GVL dataset stores ragged arrays over an `(R, S[, P])` C-order grid; the flat slot for `(r, s, p)` is `((r * S) + s) * P + p`. Merging decides, per *merged* flat slot, which input dataset and which *source* flat slot it came from. `order` is an `(n_merged_along_axis, 2)` int64 array giving `(dataset_idx, within_dataset_idx)` per merged position along the axis, in destination order. A **run** is a maximal span of consecutive merged slots over which the source dataset is constant *and* the source slot increases by exactly 1. + +**The property this task exploits.** The break condition between adjacent merged slots is `(source dataset changes) OR (source slot does not increment)`. On the samples axis, between merged samples `j` and `j+1` the source slots are `(r*S_{d_j} + w_j)*P + P-1` and `(r*S_{d_{j+1}} + w_{j+1})*P`. A dataset change is a break unconditionally; if the dataset does not change, the `r*S_d` terms cancel and the condition reduces to `w_{j+1} == w_j + 1`, with no `r` in it. The same cancellation happens across a region boundary. So the entire break pattern is a function of `order` alone. + +**Do not add a special case for the region boundary.** The pending-run carry in `_iter_samples` below already merges runs that span one, and this was verified at 2,880 configurations. Adding an explicit `cross` branch is dead code. + +- [ ] **Step 1: Write the failing differential test** + +Add to `tests/unit/dataset/test_concat_plan.py`. Import `RunPlan` alongside the existing imports. + +```python +def _rand_order(rng, counts): + """A random valid merged order: every (ds, within) slot exactly once.""" + rows = [(d, i) for d, c in enumerate(counts) for i in range(c)] + rng.shuffle(rows) + return np.array(rows, dtype=np.int64).reshape(-1, 2) + + +def _sorted_interleave_order(rng, counts): + """Each input's keys sorted, merged by global key -- the real two-cohort case.""" + keys = [] + for d, c in enumerate(counts): + ks = sorted(rng.choice(10_000, size=c, replace=False)) + keys.extend((k, d, i) for i, k in enumerate(ks)) + keys.sort() + return np.array([(d, i) for _, d, i in keys], dtype=np.int64).reshape(-1, 2) + + +def test_run_plan_matches_coalesce_provenance_exhaustively(): + """RunPlan must reproduce coalesce(provenance(...)) exactly, everywhere. + + This is the correctness proof for the whole change: `provenance` + `coalesce` + are retained purely as this oracle. A plan that is wrong in a way this sweep + misses produces a merged dataset whose slots point at the wrong samples -- + readable, and not obviously corrupt -- so the sweep deliberately includes the + sorted key-interleave order that models the real two-cohort merge, not only + shuffled and block-default orders. + """ + import itertools + + rng = np.random.default_rng(0) + n_checked = 0 + + for axis, ploidy, n_ds in itertools.product( + ("regions", "samples"), (1, 2, 3), (1, 2, 3) + ): + for _ in range(40): + counts = [int(rng.integers(0, 5)) for _ in range(n_ds)] + if axis == "regions": + n_samples = int(rng.integers(0, 5)) + shapes = [(c, n_samples) for c in counts] + else: + n_regions = int(rng.integers(0, 5)) + shapes = [(n_regions, c) for c in counts] + + for mode in ("none", "default", "shuffled", "interleaved"): + if mode == "none": + order = None + elif mode == "default": + order = _default_order(len(counts), counts) + elif mode == "shuffled": + order = _rand_order(rng, counts) + else: + order = _sorted_interleave_order(rng, counts) + + want = coalesce(provenance(axis, shapes, ploidy, order=order)) + got = list(RunPlan(axis, shapes, ploidy, order=order)) + n_checked += 1 + assert got == want, ( + f"axis={axis} ploidy={ploidy} shapes={shapes} mode={mode}\n" + f"want={want[:6]}\ngot={got[:6]}" + ) + + assert n_checked == 2880, "sweep shrank; the oracle coverage is the whole point" +``` + +Import `_default_order` from `genvarloader._dataset._concat_plan` in this test module's import block. + +- [ ] **Step 2: Run it to verify it fails** + +Run: `pixi run -e dev pytest tests/unit/dataset/test_concat_plan.py::test_run_plan_matches_coalesce_provenance_exhaustively -v` +Expected: FAIL with `ImportError: cannot import name 'RunPlan'`. + +- [ ] **Step 3: Hoist `_default_order` to module scope** + +In `_concat_plan.py`, `provenance` currently defines `_default_order` as a closure that captures `n_ds`. Lift it out so `RunPlan` shares it, and take `n_ds` as a parameter: + +```python +def _default_order(n_ds: int, counts: list[int]) -> NDArray[np.int64]: + """Block-concatenation order: dataset 0's whole block, then dataset 1's, etc. + + Args: + n_ds: Number of input datasets. + counts: Positions along the merged axis contributed by each dataset. + + Returns: + An ``(sum(counts), 2)`` int64 array of ``(dataset_idx, within_idx)``. + """ + if n_ds == 0: + return np.zeros((0, 2), np.int64) + ds_col = np.repeat(np.arange(n_ds, dtype=np.int64), counts) + w_col = np.concatenate([np.arange(c, dtype=np.int64) for c in counts]) + return np.stack([ds_col, w_col], axis=1) +``` + +Then in `provenance`, delete the nested `def _default_order(counts)` and change both call sites from `_default_order([...])` to `_default_order(n_ds, [...])`. The two sites are the `axis == "regions"` branch (`order = _default_order([shape_per_ds[d][0] for d in range(n_ds)])`) and the samples branch (`order = _default_order(per_ds_samples)`). + +- [ ] **Step 4: Add `RunPlan` with `__iter__`** + +Append to `_concat_plan.py`: + +```python +class RunPlan: + """Destination-ordered runs, derived from ``order`` without materializing slots. + + Equivalent to ``coalesce(provenance(axis, shape_per_ds, ploidy, order=order))`` + and pinned against it by + ``test_run_plan_matches_coalesce_provenance_exhaustively``, but it never + builds the ``(n_slots, 2)`` map: at the All of Us chr22 grid that map is + 32-64 GB and the run list it compresses to is 384-768 GB, because on an + interleaved sample merge every run is one slot long. + + The whole thing rests on one property: the run-break condition is "source + dataset changes, or source slot does not increment", and on both axes that + reduces to a predicate over ``order`` alone -- the region index ``r`` cancels + out. So the break pattern is computed once, in ``O(R + S)``, and the runs + stream in ``O(1)``. + + **Re-iterable on purpose, not a generator.** ``copy_runs`` iterates runs + twice (once for offsets, once to stream bytes) and so does + ``_gather_svar_offsets``; a one-shot iterator would yield an empty second + pass and silently truncate the output rather than raise. + + Args: + axis: Either ``"regions"`` or ``"samples"``. + shape_per_ds: ``(n_regions, n_samples)`` per input dataset, in input order. + ploidy: Slots per ``(region, sample)`` cell. Pass ``1`` for interval + stores, which have no ploidy axis. + order: ``(n_merged_along_axis, 2)`` int array of ``(dataset_idx, + within_dataset_idx)`` per merged position along ``axis``, in + destination order. ``None`` reproduces block-concatenation. + + Raises: + ValueError: If ``axis`` is not ``"regions"`` or ``"samples"``. + """ + + def __init__( + self, + axis: str, + shape_per_ds: list[tuple[int, int]], + ploidy: int, + *, + order: "NDArray[np.int64] | None" = None, + ) -> None: + if axis not in ("regions", "samples"): + raise ValueError(f'axis must be "regions" or "samples", got {axis!r}') + self.axis = axis + self.ploidy = int(ploidy) + self.shape_per_ds = [(int(r), int(s)) for r, s in shape_per_ds] + n_ds = len(self.shape_per_ds) + if order is None: + counts = [ + r if axis == "regions" else s for r, s in self.shape_per_ds + ] + self.order = _default_order(n_ds, counts) + else: + self.order = np.asarray(order, dtype=np.int64).reshape(-1, 2) + + def _segments(self) -> tuple[NDArray[np.int64], NDArray[np.int64]]: + """Half-open ``[start, stop)`` spans of ``order`` with no break inside.""" + ds, w = self.order[:, 0], self.order[:, 1] + brk = (ds[1:] != ds[:-1]) | (w[1:] != w[:-1] + 1) + starts = np.concatenate([np.zeros(1, np.int64), np.flatnonzero(brk) + 1]) + stops = np.concatenate([starts[1:], np.array([len(self.order)], np.int64)]) + return starts, stops + + def __iter__(self) -> "Iterator[Run]": + if self.axis == "regions": + yield from self._iter_regions() + else: + yield from self._iter_samples() + + def _iter_regions(self) -> "Iterator[Run]": + # Merged region i owns one contiguous S*P block in BOTH source and + # destination, so a segment of `order` maps to exactly one run and no + # carry is needed: adjacent segments are non-contiguous by construction. + cell = self.shape_per_ds[0][1] * self.ploidy if self.shape_per_ds else 0 + if len(self.order) == 0 or cell == 0: + return + ds, w = self.order[:, 0], self.order[:, 1] + for a, b in zip(*self._segments()): + yield Run( + src=int(ds[a]), + src_start=int(w[a]) * cell, + src_stop=(int(w[a]) + int(b - a)) * cell, + dst_start=int(a) * cell, + ) + + def _iter_samples(self) -> "Iterator[Run]": + n_regions = self.shape_per_ds[0][0] if self.shape_per_ds else 0 + n_merged = len(self.order) + if n_regions == 0 or n_merged == 0 or self.ploidy == 0: + return + per_ds_samples = np.asarray([s for _, s in self.shape_per_ds], np.int64) + ds, w = self.order[:, 0], self.order[:, 1] + seg_starts, seg_stops = self._segments() + + # One pending run, extended whenever the next segment continues it in + # both source and destination. This is what makes a run that spans a + # region boundary come out merged, with no special case for that + # boundary -- verified against the oracle at 2,880 configurations. + pending: Run | None = None + for r in range(n_regions): + for a, b in zip(seg_starts, seg_stops): + d = int(ds[a]) + s_d = int(per_ds_samples[d]) + src_start = (r * s_d + int(w[a])) * self.ploidy + src_stop = (r * s_d + int(w[b - 1]) + 1) * self.ploidy + dst_start = (r * n_merged + int(a)) * self.ploidy + if ( + pending is not None + and pending.src == d + and pending.src_stop == src_start + and pending.dst_start + (pending.src_stop - pending.src_start) + == dst_start + ): + pending = Run(d, pending.src_start, src_stop, pending.dst_start) + else: + if pending is not None: + yield pending + pending = Run(d, src_start, src_stop, dst_start) + if pending is not None: + yield pending +``` + +Add `RunPlan` to `__all__`, keeping it sorted: `__all__ = ["CONCAT_CHUNK_BYTES", "Run", "RunPlan", "coalesce", "provenance"]`. Add `from typing import Iterator, NamedTuple` to the imports (the file already has `from typing import NamedTuple`). + +- [ ] **Step 5: Run the differential test** + +Run: `pixi run -e dev pytest tests/unit/dataset/test_concat_plan.py -v` +Expected: PASS, including every pre-existing test in the file unchanged. + +- [ ] **Step 6: Verify the test is not vacuous** + +Temporarily change `_iter_samples`'s `src_stop` from `int(w[b - 1]) + 1` to `int(w[b - 1])`, re-run, and confirm the differential test FAILS. Revert the mutation. This proves the sweep actually compares run contents rather than only lengths. + +Run: `pixi run -e dev pytest tests/unit/dataset/test_concat_plan.py::test_run_plan_matches_coalesce_provenance_exhaustively -q` +Expected before revert: FAIL. Expected after revert: PASS. + +- [ ] **Step 7: Lint and commit** + +```bash +pixi run -e dev ruff check python/ tests/ +pixi run -e dev ruff format python/ tests/ +git add python/genvarloader/_dataset/_concat_plan.py tests/unit/dataset/test_concat_plan.py +git commit -m "perf(concat): derive merge runs from order without materializing slots" +``` + +--- + +### Task 2: `n_slots` and `slot_batches` + +**Files:** +- Modify: `python/genvarloader/_dataset/_concat_plan.py` +- Test: `tests/unit/dataset/test_concat_plan.py` + +**Interfaces:** +- Consumes: `RunPlan` from Task 1. +- Produces: `RunPlan.n_slots -> int` (a property), `RunPlan.slot_batches() -> Iterator[tuple[int, NDArray[np.int64], NDArray[np.int64]]]`, `class ExplicitRunPlan`, `def as_plan(runs) -> RunPlan | ExplicitRunPlan`, and the module constant `_SLOT_BATCH_SLOTS = 1 << 20`. Task 3 and Task 4 consume all of these. + +**Why `slot_batches` exists.** `copy_runs` currently builds its `lengths` array with one numpy slice-assignment *per run*. On an interleaved sample merge every run is one slot long, so that is 2.0e9 tiny numpy calls at chr22 — slower than the byte streaming it feeds. `slot_batches` yields destination-contiguous batches instead: one per region on the samples axis (`n_merged * ploidy` slots, 8.6 MB of int64 at chr22), and fixed-size chunks of a run on the regions axis. A regions-axis run can cover the whole merged grid, so it **must** be chunked — a naive `np.arange(run.src_start, run.src_stop)` there would be 32 GB and defeat the entire change. + +**This deliberately diverges from the spec.** The spec's `slot_batches` paragraph says the regions axis yields "one per run"; that is wrong, for the reason just given, and chunking is the correction. The spec has been amended to match. Do not "restore" the one-batch-per-run form. + +**Why `ExplicitRunPlan` exists.** `tests/unit/dataset/test_concat_io.py` builds small `Run` lists by hand to test IO behavior independently of planning, and that separation is correct. `as_plan` normalizes both inputs so the IO layer has exactly one consuming path. + +- [ ] **Step 1: Write the failing tests** + +Add to `tests/unit/dataset/test_concat_plan.py`: + +```python +def test_slot_batches_reproduce_the_provenance_map(): + """Concatenating the batches must rebuild `provenance` row for row. + + `copy_runs` gathers source lengths through these batches, so a batch whose + slots are right but whose dataset column is wrong would read the correct + offsets out of the wrong file -- a silent wrong merge, not a crash. + """ + for axis, shapes, ploidy in ( + ("regions", [(3, 4), (2, 4)], 2), + ("samples", [(3, 2), (3, 3)], 2), + ("samples", [(2, 1), (2, 2)], 1), + ): + counts = [s for _, s in shapes] if axis == "samples" else [r for r, _ in shapes] + interleaved = _sorted_interleave_order(np.random.default_rng(1), counts) + for order in (None, interleaved): + plan = RunPlan(axis, shapes, ploidy, order=order) + want = provenance(axis, shapes, ploidy, order=order) + got = np.zeros_like(want) + seen = np.zeros(len(want), bool) + for dst_start, ds_vec, slots in plan.slot_batches(): + n = len(slots) + assert len(ds_vec) == n + got[dst_start : dst_start + n, 0] = ds_vec + got[dst_start : dst_start + n, 1] = slots + seen[dst_start : dst_start + n] = True + assert seen.all(), f"{axis} {shapes} left slots uncovered" + np.testing.assert_array_equal(got, want) + + +def test_n_slots_matches_provenance_length(): + for axis, shapes, ploidy in ( + ("regions", [(3, 4), (2, 4)], 2), + ("samples", [(3, 2), (3, 3)], 2), + ("regions", [(0, 4), (2, 4)], 1), + ("samples", [(0, 2), (0, 3)], 2), + ): + plan = RunPlan(axis, shapes, ploidy) + assert plan.n_slots == len(provenance(axis, shapes, ploidy)) + + +def test_slot_batches_chunk_large_region_runs(monkeypatch): + """A regions-axis run can span the whole grid, so batches must be chunked. + + The production cap is 1Mi slots, far above anything a unit test can build, + so this shrinks it and checks the chunking actually happens. Asserting only + `max(sizes) <= cap` would pass vacuously on a single batch -- including + against the unchunked implementation this test exists to rule out. + """ + from genvarloader._dataset import _concat_plan + + monkeypatch.setattr(_concat_plan, "_SLOT_BATCH_SLOTS", 5) + + # One input, identity order -> exactly one run over 4 * 3 * 2 = 24 slots. + plan = RunPlan("regions", [(4, 3)], 2) + assert len(list(plan)) == 1 + + batches = list(plan.slot_batches()) + assert [len(slots) for _, _, slots in batches] == [5, 5, 5, 5, 4] + assert plan.n_slots == 24 + # Destination stays contiguous across every chunk boundary. + assert [dst for dst, _, _ in batches] == [0, 5, 10, 15, 20] + np.testing.assert_array_equal( + np.concatenate([slots for _, _, slots in batches]), np.arange(24) + ) + + +def test_explicit_run_plan_round_trips_a_hand_built_list(): + runs = [Run(0, 0, 2, 0), Run(1, 5, 7, 2)] + plan = as_plan(runs) + assert list(plan) == runs + assert plan.n_slots == 4 + assert as_plan(plan) is plan + batches = list(plan.slot_batches()) + assert [b[0] for b in batches] == [0, 2] + np.testing.assert_array_equal(batches[0][2], [0, 1]) + np.testing.assert_array_equal(batches[1][1], [1, 1]) + + +def test_run_plan_is_re_iterable(): + """copy_runs iterates twice; a generator here would silently truncate output.""" + plan = RunPlan("samples", [(2, 2), (2, 1)], 2) + assert list(plan) == list(plan) + assert len(list(plan)) > 0 +``` + +Add `RunPlan`, `as_plan` and `Run` to the module's import block if not already present. + +- [ ] **Step 2: Run to verify failure** + +Run: `pixi run -e dev pytest tests/unit/dataset/test_concat_plan.py -k "slot_batches or n_slots or explicit or re_iterable" -v` +Expected: FAIL with `AttributeError: 'RunPlan' object has no attribute 'slot_batches'` / `ImportError` for `as_plan`. + +- [ ] **Step 3: Implement** + +Add the constant next to `CONCAT_CHUNK_BYTES` in `_concat_plan.py`: + +```python +_SLOT_BATCH_SLOTS = 1 << 20 +"""Slots per `slot_batches` chunk on the regions axis: 8 MiB of int64 indices. + +A regions-axis run can cover the whole merged grid (4.0e9 slots at chr22), so +emitting one batch per run would rebuild exactly the array this module exists to +avoid. The samples axis is naturally bounded at `n_merged * ploidy` instead. +""" +``` + +Add to `RunPlan`: + +```python + @property + def n_slots(self) -> int: + """Total merged flat slots this plan covers, computed arithmetically.""" + if not self.shape_per_ds: + return 0 + if self.axis == "regions": + return len(self.order) * self.shape_per_ds[0][1] * self.ploidy + return self.shape_per_ds[0][0] * len(self.order) * self.ploidy + + def slot_batches( + self, + ) -> "Iterator[tuple[int, NDArray[np.int64], NDArray[np.int64]]]": + """Yield ``(dst_start, src_ds, src_slots)`` batches in destination order. + + Each batch describes a destination-contiguous span: ``src_ds[i]`` and + ``src_slots[i]`` are the origin of merged slot ``dst_start + i``. + Concatenating every batch in order rebuilds :func:`provenance`'s output + exactly, which is what pins this method. + + Yields: + ``(dst_start, src_ds, src_slots)``, where the two arrays are int64 + and equal in length. + """ + if self.axis == "regions": + for run in self: + pos, dst = run.src_start, run.dst_start + while pos < run.src_stop: + n = min(_SLOT_BATCH_SLOTS, run.src_stop - pos) + yield ( + dst, + np.full(n, run.src, np.int64), + np.arange(pos, pos + n, dtype=np.int64), + ) + pos += n + dst += n + return + + n_regions = self.shape_per_ds[0][0] if self.shape_per_ds else 0 + n_merged = len(self.order) + if n_regions == 0 or n_merged == 0 or self.ploidy == 0: + return + per_ds_samples = np.asarray([s for _, s in self.shape_per_ds], np.int64) + ds, w = self.order[:, 0], self.order[:, 1] + s_d = per_ds_samples[ds] + p = np.arange(self.ploidy, dtype=np.int64) + # `order` is per merged SAMPLE; each contributes `ploidy` adjacent slots. + ds_vec = np.repeat(ds, self.ploidy) + for r in range(n_regions): + base = (r * s_d + w) * self.ploidy + slots = (base[:, None] + p[None, :]).reshape(-1) + yield (r * n_merged * self.ploidy, ds_vec, slots) +``` + +Then append the adapter and the normalizer: + +```python +class ExplicitRunPlan: + """A hand-built run list presented through :class:`RunPlan`'s interface. + + The IO layer takes either this or a :class:`RunPlan`, so its unit tests can + exercise streaming with a two-run list without constructing a merge. + + Args: + runs: Destination-ordered runs. Materialized, so it may be any iterable. + """ + + def __init__(self, runs: "Iterable[Run]") -> None: + self._runs = list(runs) + + def __iter__(self) -> "Iterator[Run]": + return iter(self._runs) + + @property + def n_slots(self) -> int: + """Total merged flat slots covered by the run list.""" + return sum(r.src_stop - r.src_start for r in self._runs) + + def slot_batches( + self, + ) -> "Iterator[tuple[int, NDArray[np.int64], NDArray[np.int64]]]": + """Yield one ``(dst_start, src_ds, src_slots)`` batch per run.""" + for r in self._runs: + n = r.src_stop - r.src_start + yield ( + r.dst_start, + np.full(n, r.src, np.int64), + np.arange(r.src_start, r.src_stop, dtype=np.int64), + ) + + +def as_plan(runs: "RunPlan | ExplicitRunPlan | Sequence[Run]") -> "RunPlan | ExplicitRunPlan": + """Normalize a run source so the IO layer has one consuming path. + + Args: + runs: A plan, or a re-iterable sequence of runs. A one-shot generator is + deliberately not accepted: ``copy_runs`` iterates its runs twice. + + Returns: + ``runs`` itself when it is already a plan, else an + :class:`ExplicitRunPlan` wrapping it. + """ + if isinstance(runs, (RunPlan, ExplicitRunPlan)): + return runs + return ExplicitRunPlan(runs) +``` + +Extend imports to `from typing import Iterable, Iterator, NamedTuple, Sequence` and `__all__` to `["CONCAT_CHUNK_BYTES", "ExplicitRunPlan", "Run", "RunPlan", "as_plan", "coalesce", "provenance"]`. + +- [ ] **Step 4: Run the tests** + +Run: `pixi run -e dev pytest tests/unit/dataset/test_concat_plan.py -v` +Expected: PASS, all tests including Task 1's sweep. + +- [ ] **Step 5: Commit** + +```bash +pixi run -e dev ruff check python/ tests/ +pixi run -e dev ruff format python/ tests/ +git add python/genvarloader/_dataset/_concat_plan.py tests/unit/dataset/test_concat_plan.py +git commit -m "perf(concat): add vectorized slot batches and analytic slot count to RunPlan" +``` + +--- + +### Task 3: `copy_runs` builds offsets from batches + +**Files:** +- Modify: `python/genvarloader/_dataset/_concat_io.py:39-89` (`copy_runs`), `:92-124` (`gather_fixed`) +- Test: `tests/unit/dataset/test_concat_io.py` + +**Interfaces:** +- Consumes: `RunPlan`, `ExplicitRunPlan`, `as_plan` from Task 2; `Run` from `_concat_plan`. +- Produces: `copy_runs(srcs: list[Path], dst: Path, runs: RunPlan | ExplicitRunPlan | Sequence[Run], src_offsets: list[NDArray[np.int64]], itemsize: int) -> NDArray[np.int64]` and `gather_fixed(srcs: list[Path], dst: Path, runs: RunPlan | ExplicitRunPlan | Sequence[Run], record_bytes: int) -> None`. Return values and written bytes are unchanged; only the accepted `runs` type widens. + +**What changes and why.** Today `copy_runs` allocates a separate `lengths` array of `n_slots` int64 (16.0 GB at chr22) and fills it one run at a time. It will instead write lengths directly into `merged[1:]` and cumsum in place, filling from `slot_batches()`. The `n_slots = sum(...)` pre-pass is replaced by `plan.n_slots`. The byte-streaming loop is untouched. + +**The aliasing detail.** `np.cumsum(merged[1:], out=merged[1:])` reads and writes the same memory. numpy supports exact aliasing here, but it is load-bearing and gets its own pin rather than an assumption. + +- [ ] **Step 1: Write the failing tests** + +Add to `tests/unit/dataset/test_concat_io.py`: + +```python +def test_copy_runs_accepts_a_run_plan_and_matches_an_explicit_list(tmp_path): + """Driving copy_runs from RunPlan must be byte-identical to the old path.""" + from genvarloader._dataset._concat_plan import RunPlan, coalesce, provenance + + shapes = [(2, 2), (2, 1)] + axis, ploidy = "samples", 2 + n_src = [r * s * ploidy for r, s in shapes] + + srcs, offsets = [], [] + for d, n in enumerate(n_src): + lens = np.arange(1, n + 1, dtype=np.int64) + off = np.concatenate([[0], np.cumsum(lens)]).astype(np.int64) + payload = np.arange(off[-1], dtype=np.int32) + d * 1000 + p = tmp_path / f"src{d}.bin" + _write_raw(p, payload) + srcs.append(p) + offsets.append(off) + + out_plan = tmp_path / "plan.npy" + out_list = tmp_path / "list.npy" + plan = RunPlan(axis, shapes, ploidy) + merged_plan = copy_runs(srcs, out_plan, plan, offsets, itemsize=4) + merged_list = copy_runs( + srcs, out_list, coalesce(provenance(axis, shapes, ploidy)), offsets, itemsize=4 + ) + + np.testing.assert_array_equal(merged_plan, merged_list) + assert out_plan.read_bytes() == out_list.read_bytes() + + +def test_copy_runs_in_place_cumsum_matches_out_of_place(tmp_path): + """merged[1:] is cumsummed into itself; pin that against the naive form.""" + lens = np.array([3, 0, 5, 2, 7], np.int64) + off = np.concatenate([[0], np.cumsum(lens)]).astype(np.int64) + payload = np.arange(off[-1], dtype=np.int32) + p = tmp_path / "src.bin" + _write_raw(p, payload) + + runs = [Run(0, 0, 5, 0)] + merged = copy_runs([p], tmp_path / "out.npy", runs, [off], itemsize=4) + + want = np.concatenate([[0], np.cumsum(lens)]).astype(np.int64) + np.testing.assert_array_equal(merged, want) + + +def test_gather_fixed_from_a_run_plan_is_byte_identical_to_the_oracle(tmp_path): + """gather_fixed must produce the same bytes from a plan as from a run list.""" + from genvarloader._dataset._concat_plan import RunPlan, coalesce, provenance + + shapes, axis, ploidy = [(2, 2), (3, 2)], "regions", 1 + srcs = [] + for d, (r, sm) in enumerate(shapes): + p = tmp_path / f"g{d}.bin" + _write_raw(p, np.arange(r * sm * ploidy, dtype=np.int32) + d * 100) + srcs.append(p) + + plan = RunPlan(axis, shapes, ploidy) + out_plan, out_list = tmp_path / "plan.bin", tmp_path / "list.bin" + gather_fixed(srcs, out_plan, plan, record_bytes=4) + gather_fixed( + srcs, out_list, coalesce(provenance(axis, shapes, ploidy)), record_bytes=4 + ) + + assert out_plan.read_bytes() == out_list.read_bytes() + assert out_plan.stat().st_size == plan.n_slots * 4 +``` + +- [ ] **Step 2: Run to verify failure** + +Run: `pixi run -e dev pytest tests/unit/dataset/test_concat_io.py -k "run_plan or in_place_cumsum" -v` +Expected: FAIL — `copy_runs` raises `AttributeError` because a `RunPlan` has no `__len__`/is consumed as a list, or `TypeError` on `sum(...)`. + +- [ ] **Step 3: Rewrite `copy_runs`** + +Replace the body of `copy_runs` (`_concat_io.py:39-89`) with: + +```python +def copy_runs( + srcs: list[Path], + dst: Path, + runs: "RunPlan | ExplicitRunPlan | Sequence[Run]", + src_offsets: list[NDArray[np.int64]], + itemsize: int, +) -> NDArray[np.int64]: + """Stream a ragged payload through a run plan and return merged offsets. + + Args: + srcs: Payload file per input dataset (raw, headerless arrays). + dst: Destination payload file, created or truncated. + runs: Destination-ordered runs — a :class:`._concat_plan.RunPlan`, or any + re-iterable sequence of :class:`._concat_plan.Run`. Iterated twice, + so a one-shot generator is not accepted. + src_offsets: Cumulative offsets per input dataset, each of length + ``n_source_slots + 1``, in elements (not bytes). + itemsize: Bytes per element of the payload dtype. + + Returns: + Merged cumulative offsets, length ``total_merged_slots + 1``, in elements. + """ + plan = as_plan(runs) + n_slots = plan.n_slots + + # Lengths are written straight into the output buffer and cumsummed in + # place: a separate `lengths` array would be another n_slots int64s, 16.0 GB + # at the All of Us chr22 grid. Filling happens per destination-contiguous + # BATCH rather than per run, because on an interleaved sample merge every + # run is one slot long and a per-run numpy slice-assignment would be 2.0e9 + # scalar calls -- slower than the byte streaming it feeds. + merged = np.empty(n_slots + 1, np.int64) + merged[0] = 0 + lengths = merged[1:] + for dst_start, ds_vec, slots in plan.slot_batches(): + n = len(slots) + out = lengths[dst_start : dst_start + n] + for d in np.unique(ds_vec): + m = ds_vec == d + off = src_offsets[int(d)] + sl = slots[m] + out[m] = off[sl + 1] - off[sl] + np.cumsum(lengths, out=lengths) + + handles: dict[int, object] = {} + try: + with open(dst, "wb") as fo: + for r in plan: + if r.src not in handles: + handles[r.src] = open(srcs[r.src], "rb") + fi = handles[r.src] + off = src_offsets[r.src] + start = int(off[r.src_start]) * itemsize + stop = int(off[r.src_stop]) * itemsize + if stop > start: + _stream_range(fi, fo, start, stop) + fo.flush() + finally: + for fh in handles.values(): + fh.close() + + return merged +``` + +- [ ] **Step 4: Widen `gather_fixed`** + +In `gather_fixed` (`_concat_io.py:92`), change the `runs` annotation to `"RunPlan | ExplicitRunPlan | Sequence[Run]"`, add `plan = as_plan(runs)` as the first statement, change `for r in runs:` to `for r in plan:`, and update the `runs:` docstring line to match `copy_runs`'. Nothing else changes — it is a single-pass consumer. + +Update the import at `_concat_io.py:19` to: + +```python +from ._concat_plan import ( + CONCAT_CHUNK_BYTES, + ExplicitRunPlan, + Run, + RunPlan, + as_plan, +) +``` + +and add `Sequence` to the `typing` imports (`from typing import Sequence`). + +- [ ] **Step 5: Run the IO tests** + +Run: `pixi run -e dev pytest tests/unit/dataset/test_concat_io.py -v` +Expected: PASS — the nine pre-existing tests must pass **unmodified**, since they pass plain `list[Run]` and `as_plan` accepts that. + +- [ ] **Step 6: Commit** + +```bash +pixi run -e dev ruff check python/ tests/ +pixi run -e dev ruff format python/ tests/ +git add python/genvarloader/_dataset/_concat_io.py tests/unit/dataset/test_concat_io.py +git commit -m "perf(concat): build merged offsets from slot batches, cumsum in place" +``` + +--- + +### Task 4: Route `_concat.py`'s five sites through `RunPlan` + +**Files:** +- Modify: `python/genvarloader/_dataset/_concat.py:203-238` (`_gather_svar_offsets`), `:340-344`, `:529-530`, `:557-558`, `:585-586`, `:631-634` +- Test: existing concat dataset tests (no new test file) + +**Interfaces:** +- Consumes: `RunPlan`, `as_plan` from Task 2; `copy_runs`, `gather_fixed` from Task 3. +- Produces: nothing new. This task removes every `provenance` + `coalesce` call from the merge path. + +**All five sites, exactly.** Three are the unbounded ones (#403); two are already bounded at `(R, 2)` and move for consistency, so there is one obvious way to plan a merge. + +- [ ] **Step 1: Update the import** + +`_concat.py:16` currently reads `from ._concat_plan import Run, coalesce, provenance`. Change to: + +```python +from ._concat_plan import ExplicitRunPlan, Run, RunPlan, as_plan +``` + +`coalesce` and `provenance` stay in `_concat_plan.py` — they are just no longer imported here. + +- [ ] **Step 2: Convert `_gather_svar_offsets`** + +In `_concat.py:203-238`, change the signature's `runs: list[Run]` to `runs: "RunPlan | ExplicitRunPlan | Sequence[Run]"`, and replace the two lines + +```python + out = np.empty(sum(r.src_stop - r.src_start for r in runs), np.int64) + for r in runs: +``` + +with + +```python + out = np.empty(plan.n_slots, np.int64) + for r in plan: +``` + +adding `plan = as_plan(runs)` as the first statement of the function body, above `n_src_slots = ...`. The docstring's reference to the space "``runs`` was coalesced over" should read "the space ``runs`` plans over (see ``_concat_plan``)", since nothing is coalesced any more. + +- [ ] **Step 3: Pin `_gather_svar_offsets` equivalence** + +The spec requires byte-identical output from a plan versus a run list for all three consumers; this is the third. It goes in `tests/unit/dataset/test_concat_io.py` alongside the other two rather than in a new file — it is a gather consumer, even though the function itself lives in `_concat.py`. + +```python +def test_gather_svar_offsets_from_a_run_plan_is_byte_identical(tmp_path): + """The svar offsets gather must agree with the retained oracle. + + This one cannot go through `gather_fixed`: its slot axis is nested inside + two leading planes, so a slot's start and stop live `n_slots` elements + apart. It therefore has its own run-consuming loop, and needs its own pin. + """ + from genvarloader._dataset._concat import _gather_svar_offsets + from genvarloader._dataset._concat_plan import RunPlan, coalesce, provenance + + shapes, axis, ploidy = [(2, 2), (2, 1)], "samples", 2 + + paths = [] + for d, (r, sm) in enumerate(shapes): + n = r * sm * ploidy + p = tmp_path / f"ds{d}" + (p / "genotypes").mkdir(parents=True) + starts = np.arange(n, dtype=np.int64) + d * 100 + planes = np.stack([starts, starts + 1]) + _write_raw(p / "genotypes" / "offsets.npy", planes) + paths.append(p) + + out_plan, out_list = tmp_path / "plan", tmp_path / "list" + out_plan.mkdir() + out_list.mkdir() + + _gather_svar_offsets( + paths, out_plan, RunPlan(axis, shapes, ploidy), shapes, ploidy + ) + _gather_svar_offsets( + paths, + out_list, + coalesce(provenance(axis, shapes, ploidy)), + shapes, + ploidy, + ) + + assert (out_plan / "offsets.npy").read_bytes() == ( + out_list / "offsets.npy" + ).read_bytes() +``` + +Run: `pixi run -e dev pytest tests/unit/dataset/test_concat_io.py -k gather_svar_offsets -v` +Expected: PASS once Step 2's conversion is in place. + +- [ ] **Step 4: Convert the region_runs site (`:340-344`)** + +Replace: + +```python + region_runs = ( + coalesce(provenance("regions", [(r, 1) for r, _ in shapes], 1, order=order)) + if axis == "regions" + else None + ) +``` + +with: + +```python + region_runs = ( + RunPlan("regions", [(r, 1) for r, _ in shapes], 1, order=order) + if axis == "regions" + else None + ) +``` + +Leave the explanatory comment above it intact, and leave both `assert region_runs is not None` guards in place. + +- [ ] **Step 5: Convert the three unbounded sites** + +At `:529-530` (the `pgen_vcf` backend) and `:557-558` (the `svar` backend), replace each two-line pair + +```python + prov = provenance(axis, shapes, ploidy, order=order) + runs = coalesce(prov) +``` + +with the single line + +```python + runs = RunPlan(axis, shapes, ploidy, order=order) +``` + +At `:585-586` (per-sample tracks), replace + +```python + t_prov = provenance(axis, shapes, 1, order=order) + t_runs = coalesce(t_prov) +``` + +with + +```python + t_runs = RunPlan(axis, shapes, 1, order=order) +``` + +Keep `t_runs` hoisted above `for name in ref.tracks:` — it is built once regardless of track count, and that must not regress. + +- [ ] **Step 6: Convert the annot-track site (`:631-634`)** + +Replace + +```python + a_prov = provenance( + "regions", [(r, 1) for r, _ in shapes], 1, order=order + ) + a_runs = coalesce(a_prov) +``` + +with + +```python + a_runs = RunPlan( + "regions", [(r, 1) for r, _ in shapes], 1, order=order + ) +``` + +- [ ] **Step 7: Correct the two comments this change makes false** + +`_concat.py` documents the old behavior in prose in two places, and both become wrong here. Neither is caught by any test. + +First, the `region_runs` comment at `:331-339` warns about "a future edit that changes one `provenance` call and not the other". There is no `provenance` call left after Step 3. Change that clause to "a future edit that changes one `RunPlan` construction and not the other". + +Second, and more important, `_merge_svar2_ranges`'s docstring at `:259-265` currently reads: + +``` + After this change an svar2 ``concat`` with no per-sample tracks holds nothing + ``R x S``-sized: the ``(R*S*P, 2)`` ``provenance`` array (64 GB at All of Us + chr22) and the ``list[Run]`` ``coalesce`` builds from it (~204 bytes per run, + which on an interleaved sample merge degenerates to one run per slot) are both + gone from this path. Per-sample tracks still plan in core at + ``_concat.py:465`` -- 32 GB plus a ~424 GB run list at the same projection -- + so ``concat`` is bounded only for tracks-free datasets. Tracked separately. +``` + +Every claim in its second half is now false: per-sample tracks no longer plan in core, `concat` is no longer bounded only for tracks-free datasets, and the thing "tracked separately" is the issue this branch closes. (The `:465` line reference was already stale before this change — the site is at `:585` — and `~204 bytes per run` disagrees with the 184 bytes measured for this codebase's `Run`. Do not preserve either figure.) Replace the whole paragraph with: + +``` + No stage of this merge holds anything ``R x S``-sized. The + ``(R*S*P, 2)`` ``provenance`` array (64 GB at the All of Us chr22 grid) and + the ``list[Run]`` ``coalesce`` built from it (184 bytes per run, degenerating + to one run per slot on an interleaved sample merge) are gone from every path + in this module, per-sample tracks included: ``RunPlan`` derives the same runs + from ``order`` alone. See ``_concat_plan.RunPlan``. +``` + +- [ ] **Step 8: Confirm no planning calls remain on the merge path** + +Run: `pixi run -e dev python -c "import re; s=open('python/genvarloader/_dataset/_concat.py').read(); print('provenance:', s.count('provenance(')); print('coalesce:', s.count('coalesce('))"` +Expected: both counts `0`. + +- [ ] **Step 9: Run the concat suite** + +Run: `pixi run -e dev pytest tests/unit/dataset tests/dataset -k concat -v` +Expected: PASS, with **zero test files modified** in this task. If a concat test fails, the change is not behavior-preserving — report it rather than editing the test. + +- [ ] **Step 10: Run the full tree** + +Run: `pixi run -e dev pytest tests -q` +Expected: PASS. The measured baseline on this branch is **1250 passed, 61 skipped, 4 xfailed** (macOS/darwin, `pixi run -e dev pytest tests -q`, after `pixi run -e dev gen`). Counts are platform-dependent — CI runs slow/torch tiers this does not — so treat a differing skip count on another machine as environment, and a differing *failure* count as a real regression. + +- [ ] **Step 11: Commit** + +```bash +pixi run -e dev ruff check python/ tests/ +pixi run -e dev ruff format python/ tests/ +pixi run -e dev typecheck +git add python/genvarloader/_dataset/_concat.py +git commit -m "perf(concat): plan every merge through RunPlan, dropping the provenance map" +``` + +--- + +### Task 5: Pin the memory bound + +**Files:** +- Test: `tests/unit/dataset/test_concat_plan.py` + +**Interfaces:** +- Consumes: `RunPlan` from Tasks 1–2. +- Produces: nothing consumed by later tasks. + +**Why this is its own test rather than a comment.** The whole change is invisible to every behavioral test: `RunPlan` and `coalesce(provenance(...))` produce identical output, so a future edit that quietly reintroduces a materialized slot-space array would keep the suite green while restoring the 32 GB allocation. This test is the only thing that fails in that case. + +- [ ] **Step 1: Write the failing test** + +```python +def test_run_plan_never_materializes_the_slot_space(): + """Peak allocation must stay O(R + S), not O(R * S * P). + + The old path on this grid allocates a (800_000, 2) int64 provenance map + (12.8 MB) and coalesces it into 800_000 one-slot Runs (~147 MB at 184 bytes + each). Both are invisible to every behavioural test, because RunPlan yields + exactly the same runs -- so this bound is what stops a future edit from + silently restoring them. + """ + import tracemalloc + + shapes = [(200, 1000), (200, 1000)] + # Interleaved samples: the worst case, one run per slot. + order = np.array( + [(d, i) for i in range(1000) for d in (0, 1)], dtype=np.int64 + ).reshape(-1, 2) + + tracemalloc.start() + try: + plan = RunPlan("samples", shapes, 2, order=order) + n_runs = sum(1 for _ in plan) + peak = tracemalloc.get_traced_memory()[1] + finally: + tracemalloc.stop() + + assert plan.n_slots == 200 * 2000 * 2 + assert n_runs == plan.n_slots // 2, "interleaved merge should give one run per cell" + assert peak < (1 << 20), f"peak {peak} bytes: the slot space is being materialized" +``` + +The `n_runs` expectation: with `ploidy=2` and fully interleaved samples, each merged sample's two ploidy slots are contiguous and every sample boundary breaks, so there is one run per `(region, sample)` cell — `200 * 2000 = 400_000` runs covering `800_000` slots. + +- [ ] **Step 2: Run it** + +Run: `pixi run -e dev pytest tests/unit/dataset/test_concat_plan.py::test_run_plan_never_materializes_the_slot_space -v` +Expected: PASS (Tasks 1–2 already satisfy it). + +- [ ] **Step 3: Verify it is not vacuous** + +Temporarily add `_ = provenance("samples", shapes, 2, order=order)` immediately after the `RunPlan(...)` construction inside the `tracemalloc` block and re-run. +Expected: FAIL on the `peak < (1 << 20)` assertion. Remove the line and confirm PASS. + +- [ ] **Step 4: Commit** + +```bash +pixi run -e dev ruff check python/ tests/ +pixi run -e dev ruff format python/ tests/ +git add tests/unit/dataset/test_concat_plan.py +git commit -m "test(concat): pin RunPlan's peak allocation below the materialized slot space" +``` + +--- + +### Task 6: Consolidate the seven `svar2_store` fixtures + +**Files:** +- Create: `tests/conftest.py` +- Modify: `tests/test_svar2_reconstruct.py`, `tests/unit/dataset/test_svar2_store.py`, `tests/unit/dataset/test_svar2_link.py`, `tests/dataset/test_svar2_readbound_variants.py`, `tests/dataset/test_svar2_readbound_haps.py`, `tests/dataset/test_svar2_readbound_diffs.py`, `tests/dataset/conftest.py` +- **Do not modify:** `tests/dataset/test_svar2_readbound_tracks.py` (see "The situation" below) + +**Interfaces:** +- Consumes: nothing from earlier tasks. Independent of Tasks 1–5. +- Produces: `tests/conftest.py` exporting `_SVAR2_REF` and the `svar2_store_2s` fixture. Task 7 adds the enriched 3-sample VCF on top. + +**This task changes no test content.** It is pure deduplication, and its gate is that the full suite stays green with **zero assertion edits**. Do not enrich anything here — that is Task 7. + +**The situation, as measured.** `svar2_store` is defined **eight** times. `tests/dataset/conftest.py:52` builds the **3-sample** store (`S0, S1, S2`) from `vcf_and_ref` — that is Task 7's fixture, untouched here. The other seven are **2-sample** (`S0, S1`), each with a private copy of the same 40 bp `_REF` and a `_VCF`, each re-running `samtools faidx` + `bcftools view` + `bcftools index`. + +**Only six of those seven are duplicates.** Their `_VCF` bodies are byte-identical (the same three variants at chr1:3, :7, :12): + +| module | line | fold in? | +|---|---|---| +| `tests/test_svar2_reconstruct.py` | 34 | yes | +| `tests/unit/dataset/test_svar2_store.py` | 28 | yes | +| `tests/unit/dataset/test_svar2_link.py` | 59 | yes | +| `tests/dataset/test_svar2_readbound_variants.py` | 36 | yes | +| `tests/dataset/test_svar2_readbound_haps.py` | 34 | yes | +| `tests/dataset/test_svar2_readbound_diffs.py` | 38 | yes | +| `tests/dataset/test_svar2_readbound_tracks.py` | 39 | **NO** | + +`test_svar2_readbound_tracks.py` carries a **four**-variant `_VCF` — it adds `chr1 10 . G C ... 1|1 1|0` to the other three. Folding it onto the shared fixture would silently change the data its assertions were written against, which is exactly what this task's zero-assertion-edit gate forbids. Leave its `_VCF` and its `svar2_store` fixture exactly where they are. It may take `_SVAR2_REF` from the shared module (its `_REF` *is* identical), but nothing else. + +**Also leave alone**, all purpose-built rather than duplicated: `svar2_store_dense_snp` in **three** modules (`test_svar2_readbound_variants.py:140`, `test_svar2_readbound_haps.py:280`, `test_svar2_readbound_diffs.py:83`) and `svar2_store_unsorted` (`test_write_svar2.py:526`). + +Four of the six folded shadows live in `tests/dataset/`, where they silently shadow `tests/dataset/conftest.py`'s 3-sample fixture of the same name. Renaming the shared 2-sample fixture to `svar2_store_2s` makes that distinction explicit instead of accidental — that is the point of the rename, not cosmetics. + +- [ ] **Step 1: Create `tests/conftest.py`** + +```python +"""Fixtures shared across the whole test tree.""" + +import subprocess +from pathlib import Path + +import pytest + +_SVAR2_REF = "ACAGTACATGGGTACTAGCTAGGCTAACCGGTTAACCGGT" +"""40 bp synthetic chr1 backing every .svar2 fixture. + +Positions used by the VCFs below: 3=A, 7=C, 9=T, 12=GTA, 17=A, 19=C, 30=G. +""" + +_SVAR2_VCF_2S = """\ +##fileformat=VCFv4.2 +##contig= +##FORMAT= +#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tS0\tS1 +chr1\t3\t.\tA\tG\t.\t.\t.\tGT\t1|0\t0|0 +chr1\t7\t.\tC\tCAT\t.\t.\t.\tGT\t0|1\t1|1 +chr1\t12\t.\tGTA\tG\t.\t.\t.\tGT\t1|1\t0|1 +""" + + +def _build_svar2(vcf_text: str, samples: list[str], d: Path, name: str) -> Path: + """Write a VCF + FASTA under ``d`` and convert them to a .svar2 store. + + Args: + vcf_text: Full VCF text, including header. + samples: Sample names to convert, in the order genoray should store them. + d: Directory to build in. + name: Basename of the resulting store directory. + + Returns: + Path to the finished ``.svar2`` store. + """ + from genoray import _core + + ref = d / "ref.fa" + ref.write_text(f">chr1\n{_SVAR2_REF}\n") + subprocess.run(["samtools", "faidx", str(ref)], check=True) + + vcf = d / "in.vcf" + vcf.write_text(vcf_text) + bcf = d / "in.bcf" + subprocess.run(["bcftools", "view", "-Ob", "-o", str(bcf), str(vcf)], check=True) + subprocess.run(["bcftools", "index", str(bcf)], check=True) + + out = d / name + _core.run_conversion_pipeline( + str(bcf), str(ref), ["chr1"], str(out), samples, + 25_000, 2, 1, 8 * 1024 * 1024, + ) + assert (out / "meta.json").exists(), "conversion did not finish" + return out + + +@pytest.fixture(scope="module") +def svar2_store_2s(tmp_path_factory) -> Path: + """A two-sample (S0, S1) .svar2 store. + + Named apart from the dataset suite's three-sample ``svar2_store`` on purpose: + three of this fixture's former copies lived in ``tests/dataset/`` and silently + shadowed ``tests/dataset/conftest.py``'s same-named three-sample fixture. + """ + d = tmp_path_factory.mktemp("svar2_2s") + return _build_svar2(_SVAR2_VCF_2S, ["S0", "S1"], d, "store") +``` + +- [ ] **Step 2: Delete the six shadows and repoint their consumers** + +In each of the **six** modules marked "fold in" above — `tests/test_svar2_reconstruct.py`, `tests/unit/dataset/test_svar2_store.py`, `tests/unit/dataset/test_svar2_link.py`, `tests/dataset/test_svar2_readbound_variants.py`, `tests/dataset/test_svar2_readbound_haps.py`, `tests/dataset/test_svar2_readbound_diffs.py` (**not** `test_svar2_readbound_tracks.py`): + +1. Delete the module-level `_REF` and `_VCF` constants. +2. Delete the whole `svar2_store` fixture function and its `@pytest.fixture` decorator. +3. Rename every `svar2_store` **parameter** in that module's test functions to `svar2_store_2s`, and every in-body use of the name. +4. Leave `svar2_store_dense_snp` (and its own `_VCF_*` constant) untouched in the two modules that define it. If it references `_REF`, change that reference to `_SVAR2_REF` imported from `tests.conftest` — do **not** delete the constant it needs. +5. Drop any now-unused `import subprocess` / `from pathlib import Path` / `from genoray import _core` left behind. `ruff check` will name them. + +- [ ] **Step 3: Repoint `tests/dataset/conftest.py`** + +Delete its local `_REF` and import the shared one instead, leaving its `_VCF` (the 3-sample one) and its `vcf_and_ref` / `svar2_store` / `svar1_store` fixtures in place: + +```python +from tests.conftest import _SVAR2_REF as _REF +``` + +If that import path does not resolve under this repo's pytest configuration, define `tests/dataset/conftest.py`'s `_REF` as `_REF = _SVAR2_REF` by importing from the `conftest` module pytest already loaded, or simply keep `_REF` inline in `tests/dataset/conftest.py` with a comment pointing at `tests/conftest.py::_SVAR2_REF` as the canonical copy. Pick whichever the test run actually accepts and say which in the task report — do not leave a broken import. + +- [ ] **Step 4: Run the full tree** + +Run: `pixi run -e dev pytest tests -q` +Expected: PASS with the same counts as before this task (**1250 passed, 61 skipped, 4 xfailed** on darwin). **No assertion may have been edited to achieve this.** If a test fails, the consolidation changed behavior — most likely a module silently picked up the 3-sample fixture — and the fix is to repoint it, never to weaken the assertion. + +- [ ] **Step 5: Confirm the duplication is actually gone** + +Run: `pixi run -e dev python -c "import subprocess; print(subprocess.run(['grep','-rn','def svar2_store','tests/'],capture_output=True,text=True).stdout)"` + +Expected: exactly **seven** definitions remain, and no more — + +| definition | where | +|---|---| +| `svar2_store_2s` | `tests/conftest.py` (new, shared) | +| `svar2_store` | `tests/dataset/conftest.py` (3-sample) | +| `svar2_store` | `tests/dataset/test_svar2_readbound_tracks.py` (4-variant, deliberately kept) | +| `svar2_store_dense_snp` | `test_svar2_readbound_variants.py`, `test_svar2_readbound_haps.py`, `test_svar2_readbound_diffs.py` | +| `svar2_store_unsorted` | `test_write_svar2.py` | + +Six definitions were removed. If any other `def svar2_store(` survives, the dedup is incomplete; if one of the seven above is gone, something was folded in that should not have been. + +- [ ] **Step 6: Commit** + +```bash +pixi run -e dev ruff check python/ tests/ +pixi run -e dev ruff format python/ tests/ +git add tests/ +git commit -m "test(svar2): hoist the duplicated svar2_store fixture into tests/conftest.py" +``` + +--- + +### Task 7: Enrich the three-sample fixture and rewrite its guards + +**Files:** +- Modify: `tests/dataset/conftest.py` (the `_VCF` constant), `tests/dataset/test_write_svar2.py:624-690` (the occupancy guard), `tests/dataset/test_write_svar2.py:737-758` (the dense-layout test's docstring) + +**Interfaces:** +- Consumes: the consolidated fixture layout from Task 6. +- Produces: nothing consumed by later tasks. + +**What is already done, so do not redo it.** #406 lists three follow-ons; two landed during #357 in commit `c3da6bda` and must be left alone: +- The per-channel oracle masks at `test_write_svar2.py:~117-133` already read `present = (widths_snp > 0) | (widths_indel > 0)` — cell-level, with a comment explaining why per-channel masking is wrong. +- `test_dense_layout_dataset_still_opens_and_reads` already says in its docstring that it is a smoke test, not a parity pin. + +**The problem that remains.** The 3-sample fixture's vk (sparse) channel has **one** non-empty cell in an 18-cell grid (3 regions × 3 samples × ploidy 2). genoray's `cost_model.rs::choose_representation` routes a variant whose carrier-call count crosses a bit-cost threshold to the per-region **dense** channel, which has no sample axis; the fixture's SNP has one carrier and stays sparse, both indels have three carrier calls and route dense. So S1 has zero cells in the vk view anywhere, structurally indistinguishable from S2's deliberately empty column. Its `dense_snp_range` is also `[[0,0],[0,0],[0,0]]` — the dense **SNP** channel is never exercised at all. + +**The replacement was built and measured, not reasoned about.** Every sparse variant below carries exactly **one** call, which is what keeps it out of the dense channel; the two 3-carrier variants exist to populate both dense channels. + +- [ ] **Step 1: Replace `_VCF` in `tests/dataset/conftest.py`** + +```python +_VCF = """\ +##fileformat=VCFv4.2 +##contig= +##FORMAT= +#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tS0\tS1\tS2 +chr1\t3\t.\tA\tG\t.\t.\t.\tGT\t1|0\t0|0\t0|0 +chr1\t7\t.\tC\tCAT\t.\t.\t.\tGT\t0|1\t0|0\t0|0 +chr1\t9\t.\tT\tC\t.\t.\t.\tGT\t0|0\t0|1\t0|0 +chr1\t12\t.\tGTA\tG\t.\t.\t.\tGT\t0|0\t1|0\t0|0 +chr1\t17\t.\tA\tC\t.\t.\t.\tGT\t0|1\t0|0\t0|0 +chr1\t19\t.\tC\tCGG\t.\t.\t.\tGT\t1|1\t1|1\t1|1 +chr1\t30\t.\tG\tA\t.\t.\t.\tGT\t1|1\t1|1\t1|1 +""" +``` + +Regions are unchanged: `[0, 20)`, `[5, 15)`, `[25, 40)`. `_REF` is unchanged. + +- [ ] **Step 2: Run the suite to see what the enrichment breaks** + +Run: `pixi run -e dev pytest tests/dataset/test_write_svar2.py -q` +Expected: FAIL — at minimum `test_fixture_has_empty_cells`, which pins the old one-cell grid. Record every failure; each one is a test that was passing because the fixture was weak, and each is fixed in the following steps rather than weakened. + +- [ ] **Step 3: Rewrite the occupancy guard** + +Replace the body of the fixture guard at `tests/dataset/test_write_svar2.py:624-690` (the test asserting `sorted_samples == ["S0", "S1", "S2"]`, reshaping `nonempty` to `(3, S, P)`, and pinning `grid[0, S0, 0]`) with: + +```python + svar2 = SparseVar2(svar2_store) + sorted_samples = sorted(svar2.available_samples) + assert sorted_samples == ["S0", "S1", "S2"], "fixture lost its third sample" + + d = svar2._find_ranges( + "chr1", + np.array([0, 5, 25]), + np.array([20, 15, 40]), + samples=sorted_samples, + ) + snp = np.asarray(d["vk_snp_range"], np.int64) # (R*S*P, 2) + indel = np.asarray(d["vk_indel_range"], np.int64) + w_snp = (snp[:, 1] - snp[:, 0]).reshape(3, len(sorted_samples), svar2.ploidy) + w_indel = (indel[:, 1] - indel[:, 0]).reshape(w_snp.shape) + occ = (w_snp > 0) | (w_indel > 0) + + s0, s1, s2 = (sorted_samples.index(s) for s in ("S0", "S1", "S2")) + + # Row-major (R, S, P) -- pinned by the layout oracle in + # test_write_svar2_emits_cache, which asserts this same reshape against the + # cache memmaps. Assert each structure SEPARATELY: a single `not occ.all()` + # is a disjunction that stays green when any one of them regresses alone. + assert not occ[:, s2].any(), "S2 is no longer all-reference; the empty COLUMN is gone" + assert not occ[2].any(), "region [25, 40) now holds sparse variants; the empty ROW is gone" + + # The measured grid. Every sparse variant carries exactly one call, which is + # what keeps genoray's cost model from routing it to the per-region dense + # channel; if a cost-model change pushes any of them dense this fails loudly + # rather than letting the sparse/dense parity tests pass against a thinner + # table while appearing green. + expected = np.zeros_like(occ) + expected[0, s0] = [True, True] + expected[0, s1] = [True, True] + expected[1, s0] = [False, True] + expected[1, s1] = [True, True] + np.testing.assert_array_equal(occ, expected) + assert int(occ.sum()) == 7, "fixture occupancy changed; update this pin deliberately" + + # Non-vacuity, kept from the original guard: an all-empty vk grid would make + # every sparse/dense parity test built on this fixture pass trivially. + assert occ.any(), ( + "vk channel is entirely empty: genoray routed every variant to the " + "dense channel, so all sparse-cache parity tests on this fixture are " + "now vacuous" + ) + + # Mixed per-channel emptiness inside a PRESENT cell -- the property the old + # one-cell grid could not express, and the reason two shipped oracles were + # able to mask per channel without failing. + assert (w_snp[0, s0, 0] > 0) and (w_indel[0, s0, 0] == 0) + assert (w_snp[0, s0, 1] > 0) and (w_indel[0, s0, 1] > 0) + assert (w_snp[0, s1, 0] == 0) and (w_indel[0, s1, 0] > 0) + + # An ABSENT cell inside a non-empty region: region 1 holds variants, but + # (S0, ploid 0) has none, so lookup must return (0, 0) there rather than a + # neighbour's range. + assert not occ[1, s0, 0] + + # The sample axis is now ORDERED, not just occupied: S0 and S1 differ at + # (region 1, ploid 0), so a transposed sample axis is detectable. During + # #357 this guard had to be withdrawn as unsatisfiable. + assert not np.array_equal(occ[:, s0], occ[:, s1]) + + # Both dense channels are exercised. dense_snp_range was all zeros before + # this fixture was enriched, so the dense SNP path had no coverage here. + dense_snp = np.asarray(d["dense_snp_range"], np.int64) + dense_indel = np.asarray(d["dense_indel_range"], np.int64) + assert (dense_snp[:, 1] > dense_snp[:, 0]).any(), "dense SNP channel is empty again" + assert (dense_indel[:, 1] > dense_indel[:, 0]).any(), "dense indel channel is empty" +``` + +Rename the test to `test_fixture_grid_is_non_degenerate` and update its docstring to describe the 7-of-18 grid. Keep the module's existing import style (`from genoray import SparseVar2` inside the test, as the current version does). + +- [ ] **Step 4: Fix the remaining fallout from Step 2** + +For each other failure recorded in Step 2, fix it by making the assertion express the **new, stronger** reality — never by loosening it. Expect `test_write_svar2_emits_cache`'s `n_empty_seen > 0` and its `n * 28 < bed.height * S * P * 32` size assertion to still hold (fill rises from 1/18 to 7/18, well under the 28/32 break-even). If any test now needs a *weaker* assertion to pass, stop and report it — that means the enrichment broke a real invariant rather than a weak test. + +- [ ] **Step 5: Promote the dense-layout test** + +At `tests/dataset/test_write_svar2.py:737-758`, the docstring justifies the test being a smoke test by "this fixture's grid has only one non-empty cell (see `svar2_store`), so it can't distinguish a correct dense reader from one that always returns the same wrong answer." That reason is now false. Update the docstring to state the grid is 7-of-18 and that the test *does* pin dense-vs-sparse parity, and add an assertion comparing the dense-layout dataset's reads against the sparse-layout dataset's reads over the full `(region, sample)` grid rather than only asserting it opens. + +- [ ] **Step 6: Mutation-check the promotion** + +Temporarily inject `r_q = np.zeros_like(r_q)` at the top of `_DenseRanges.lookup` in `python/genvarloader/_dataset/_svar2_ranges.py` and run: + +Run: `pixi run -e dev pytest tests/dataset/test_write_svar2.py::test_dense_layout_dataset_still_opens_and_reads -q` +Expected: FAIL. Revert the injection and confirm PASS. If it still passes with the mutation, the promotion did not happen — the assertion is not reaching the dense reader. + +- [ ] **Step 7: Run the full tree** + +Run: `pixi run -e dev pytest tests -q` +Expected: PASS. Note the new counts in the task report; they will differ from the 1250-passed baseline if assertions were added. + +- [ ] **Step 8: Commit** + +```bash +pixi run -e dev ruff check python/ tests/ +pixi run -e dev ruff format python/ tests/ +git add tests/ +git commit -m "test(svar2): enrich the fixture grid from 1 to 7 occupied cells" +``` + +--- + +### Task 8: Documentation + +**Files:** +- Modify: `docs/source/write.md:158-163`, `skills/genvarloader/SKILL.md:196-199` + +**Interfaces:** +- Consumes: the finished behavior from Tasks 1–5. +- Produces: nothing. + +**Why docs change at all.** `gvl.concat` is public API (`python/genvarloader/__init__.py:73`), and both docs currently claim its cost "is dominated by I/O rather than computation" without mentioning memory — which was false before this change, since planning allocated 32–768 GB at cohort scale. The claim becomes true here, and the repo's docs-audit rule requires user-facing docs to stay true for any user-facing change. `docs/source/changelog.md` is auto-generated from commit messages and does **not** count as documentation; do not hand-edit it. + +- [ ] **Step 1: Update `docs/source/write.md`** + +After the existing paragraph ending "not as a routine step." (`write.md:158-163`), add: + +```markdown +`gvl.concat`'s planning is streaming: it derives the merge plan from the sorted +merge order rather than building a map over every `(region, sample, ploid)` slot, +so its memory stays proportional to regions plus samples rather than to their +product. The merged `offsets.npy` it writes is still one int64 per merged slot — +about 16 GB for a 3,734-region by 535,662-sample grid — and that is a property of +the ragged on-disk format, not of the merge. +``` + +- [ ] **Step 2: Update `skills/genvarloader/SKILL.md`** + +In the `**Cost:**` paragraph at `SKILL.md:196-199`, after "it moves roughly the full size of the merged dataset.", insert: + +```markdown +Planning is streaming — memory scales with regions plus samples, not their +product — so the resident cost is the merged offsets array (one int64 per merged +`(region, sample[, ploid])` slot), not the merge plan. +``` + +- [ ] **Step 3: Verify the docs build** + +Run: `pixi run -e docs doc` +Expected: builds without new warnings. If the `docs` environment is stale, run `pixi install -e docs` first. + +- [ ] **Step 4: Check `api.md` is still in sync with `__all__`** + +Run: `pixi run -e dev python -c "import genvarloader as g; api=open('docs/source/api.md').read(); print('MISSING:', [n for n in g.__all__ if n not in api] or 'none')"` +Expected: `MISSING: none`. This change adds no public symbol, so this is a regression check, not new work. + +- [ ] **Step 5: Commit** + +```bash +git add docs/source/write.md skills/genvarloader/SKILL.md +git commit -m "docs(concat): state that merge planning is streaming, not slot-sized" +``` + +--- + +## Final Verification + +After Task 8, before opening the PR: + +- [ ] `pixi run -e dev test` — the full pytest + cargo gate, whole tree. +- [ ] `pixi run -e dev ruff check python/ tests/` — clean. +- [ ] `pixi run -e dev ruff format python/ tests/` — reports no reformatting. +- [ ] `pixi run -e dev typecheck` — pyrefly reports 0 errors. **Confirm it is not vacuous:** this worktree lives under the git-ignored `.claude/worktrees/`, which has previously made pyrefly check zero files and report a meaningless "0 errors". A real run reports a non-zero count of suppressed errors and warnings alongside the zero. +- [ ] `git log --oneline origin/main..HEAD` — every commit is `perf:`, `refactor:`, `test:`, or `docs:`. No `feat:`, no `BREAKING CHANGE`. +- [ ] `pixi run -e dev python -c "import re; s=open('python/genvarloader/_dataset/_concat_plan.py').read(); assert 'def provenance' in s and 'def coalesce' in s, 'the oracle was deleted'"` — the oracle survives. diff --git a/docs/superpowers/specs/2026-09-15-concat-analytic-run-planning-design.md b/docs/superpowers/specs/2026-09-15-concat-analytic-run-planning-design.md new file mode 100644 index 00000000..69ba589c --- /dev/null +++ b/docs/superpowers/specs/2026-09-15-concat-analytic-run-planning-design.md @@ -0,0 +1,368 @@ +# Analytic run planning for `concat` — design + +**Issues:** #403 (primary), #406 (bundled, test-only) +**Status:** approved for planning, 2026-09-15 +**Target:** `main`, one PR + +## Problem + +`gvl.concat` decides, for each *merged* flat slot, which input dataset and which +source slot it comes from. Today it does that by materializing the whole mapping +and then compressing it: + +```python +prov = provenance(axis, shapes, ploidy, order=order) # (n_slots, 2) int64 +runs = coalesce(prov) # list[Run] +``` + +Both intermediates are sized by the `(R, S[, P])` grid, which is exactly the +quantity #357 just removed from the svar2 range cache. At the All of Us chr22 +grid (`R` = 3,734, `S` = 535,662, `P` = 2): + +| intermediate | shape | bytes | +|---|---|---| +| `provenance`, genotypes (`ploidy=2`) | `(4,000,323,816, 2)` int64 | **64.0 GB** | +| `provenance`, tracks (`ploidy=1`) | `(2,000,161,908, 2)` int64 | **32.0 GB** | +| `coalesce` list, genotypes | 4.00e9 × 184 B + pointers | **768 GB** | +| `coalesce` list, tracks | 2.00e9 × 184 B + pointers | **384 GB** | +| `copy_runs` `lengths` temp, tracks | 2.00e9 int64 | **16.0 GB** | + +`Run` is a 4-field `NamedTuple` measured at 184 B on this codebase (72 B tuple + +112 B for its four `int` objects). A run is a maximal span over which the source +dataset is constant *and* the source slot increases by exactly 1, so on an +interleaved sample merge — two cohorts whose sample IDs interleave in sorted +order, the common case — every run is one slot long and the list degenerates to +one `Run` per slot. + +### Scope correction to #403 + +#403 names only the per-sample-tracks call site and states that an svar2 concat +without per-sample tracks holds nothing `R x S`-sized. That is true for the svar2 +range cache specifically, but **three** call sites remain, and the two #403 does +not name are the larger ones: + +| site | call | at chr22 | +|---|---|---| +| `_concat.py:529` (`pgen_vcf`) | `provenance(axis, shapes, ploidy, ...)` | 64.0 GB | +| `_concat.py:557` (`svar`) | `provenance(axis, shapes, ploidy, ...)` | 64.0 GB | +| `_concat.py:585` (per-sample tracks) | `provenance(axis, shapes, 1, ...)` | 32.0 GB | + +Two further sites are already bounded and are **out of scope**: `_concat.py:341` +(`region_runs`, `(R, 2)`) and `_concat.py:631` (annot tracks, region axis, +`(R, 2)`). Both pass `[(r, 1) for r, _ in shapes]` with `ploidy=1`, so they are R +rows, not `R*S`. + +So `concat` is unbounded on all three backends that store per-slot ragged +payloads, not on the tracks path alone. + +### Correction to #403's proposed fix + +#403 proposes emitting runs "as an iterator instead of a list," noting that +`copy_runs` consumes runs in destination order and never indexes backwards. That +is true but insufficient: `copy_runs` iterates `runs` **twice** — once at +`_concat_io.py:59-66` to build `lengths`, once at `:75-83` to stream bytes — and +`_gather_svar_offsets` (`_concat.py:232-233`) does the same. A bare generator +would yield an empty second pass and silently produce a truncated or zero-length +output rather than an error. The replacement must be **re-iterable**, not a +one-shot iterator. + +## Design + +### The key property + +For both axes, the run-break pattern is a function of `order` alone and is +**independent of the region index `r`**. The break condition between adjacent +merged slots is `(source dataset changes) OR (source slot does not increment)`. + +*Samples axis, within a region.* Between merged samples `j` and `j+1`, the source +slots are `(r*S_{d_j} + w_j)*P + P-1` and `(r*S_{d_{j+1}} + w_{j+1})*P`. If the +dataset changes it is a break unconditionally. If it does not, the `r*S_d` terms +cancel and the condition reduces to `w_{j+1} == w_j + 1` — no `r`. + +*Samples axis, across a region boundary.* Continuity from `(r, S-1, P-1)` to +`(r+1, 0, 0)` requires `d_0 == d_{S-1}` and `S_d + w_0 == w_{S-1} + 1`. The `r` +terms cancel here too. + +*Regions axis.* Merged region `i` occupies one contiguous `S*P` block in both +source and destination, so breaks are found by scanning `order` (length R) alone. + +Therefore the entire run structure is derivable from `order` in `O(R + S)` work +and memory, and the runs themselves stream in `O(1)`. + +### `RunPlan` + +Replace the `provenance` → `coalesce` pipeline on the hot path with a small +re-iterable object in `_concat_plan.py`: + +```python +class RunPlan: + def __init__(self, axis, shape_per_ds, ploidy, *, order=None): ... + def __iter__(self) -> Iterator[Run]: ... + def slot_batches(self) -> Iterator[tuple[int, NDArray[np.int64], NDArray[np.int64]]]: ... + @property + def n_slots(self) -> int: ... +``` + +**`__iter__`** yields `Run`s in destination order, recomputing from `order` on +each call, so the two-pass consumers work unchanged. Implementation carries a +single pending `Run` and extends it when the next segment is contiguous in both +source and destination. That carry subsumes the region-boundary case entirely — +**no special-casing of `cross` is needed**, which the validation below confirms. + +**`n_slots`** is computed arithmetically (`R * S_merged * ploidy`, or +`len(order) * S * ploidy` on the regions axis), removing `copy_runs`' +`sum(r.src_stop - r.src_start for r in runs)` pre-pass. + +**`slot_batches`** exists because of a second bottleneck found during validation. +`copy_runs` builds its `lengths` array with one numpy slice-assignment *per run*. +On an interleaved merge that is one tiny numpy call per slot — 2.0e9 of them at +chr22, far slower than the streaming it feeds. `slot_batches` instead yields +destination-contiguous batches as `(dst_start, ds_vec, slot_vec)`, so `copy_runs` +fills lengths with `R` vectorized gathers (3,734 at chr22) instead of `R*S` +scalar ones. On the samples axis a batch is one region — `S*P` int64 = 8.6 MB at +chr22. On the regions axis a batch is a **fixed-size chunk of a run**, capped at +`_SLOT_BATCH_SLOTS = 1 << 20` (8 MiB of int64): one batch per run would not do, +because a regions-axis run can span the entire merged grid, and materializing its +slot indices would rebuild the 32 GB array this whole change exists to remove. + +### Consumer changes + +1. **`copy_runs`** (`_concat_io.py:39`) — take `RunPlan`; build merged offsets + from `slot_batches()`; write lengths directly into `merged[1:]` and cumsum in + place, dropping the separate `lengths` array (16.0 GB at chr22); stream bytes + from `__iter__`. +2. **`gather_fixed`** (`_concat_io.py:92`) — take `RunPlan`; single pass, only the + type annotation and docstring change. +3. **`_gather_svar_offsets`** (`_concat.py:203`) — take `RunPlan`; replace its + `sum(...)` pre-pass with `plan.n_slots`. Its two materialized planes are + `n_slots` int64 each and stay as they are: they are the same order of + magnitude as the offsets array `copy_runs` must build anyway, and the + docstring already justifies them. +4. **`_concat.py`** — the three unbounded sites construct `RunPlan` instead of + calling `provenance` + `coalesce`. The two bounded sites (`:341`, `:631`) also + move over, for one obvious way to do it, but gain nothing. + +### What remains after the change + +| quantity | before | after | +|---|---|---| +| provenance map | 32–64 GB | 0 | +| run list | 384–768 GB | one live `Run` | +| `copy_runs` `lengths` | 16.0 GB | 0 | +| merged `offsets.npy` | 16.0 GB | 16.0 GB — **irreducible** | + +The merged offsets array is the output file itself, required by the ragged +format. Removing it is a format change and explicitly not in scope. + +**Planning is not free.** On a fully interleaved sample merge the plan still +yields one `Run` per slot — 2.0e9 of them at chr22, at roughly a microsecond +each. That is real time, but it is dominated by the per-run file IO those runs +describe (one seek per run), and it replaces an allocation that cannot complete +at all. If planning ever becomes the measured bottleneck, a batched array-valued +run form is the follow-up; it is not proposed here. + +### `provenance` and `coalesce` are retained + +They stay in `_concat_plan.py`, no longer called on the hot path, documented as +the reference implementation that `RunPlan` is verified against. They are already +tested, they are small, and they are what makes the new path provably correct. +Deleting them would delete the oracle. `_concat_plan` is private +(`python/genvarloader/_dataset/_concat_plan.py`), so nothing here is a public API +change. + +## Testing + +### Differential property test (the primary gate) + +```python +list(RunPlan(axis, shapes, ploidy, order=order)) == coalesce( + provenance(axis, shapes, ploidy, order=order) +) +``` + +Exact equality, swept over: both axes; `ploidy` 1–3; 1–3 input datasets; region +and sample counts 0–4 inclusive (so empty grids are covered); and four `order` +modes — `None`, explicit block-default, uniformly shuffled, and a sorted +key-interleave that models the real two-cohort merge. + +This prototype has already been run at 2,880 configurations with **zero +mismatches**, including every degenerate case in the sweep. The shipped test is +that sweep, seeded, as a normal pytest — not a hypothesis strategy, so failures +are reproducible by config rather than by shrink. + +### Memory bound + +`tracemalloc` peak over `RunPlan` construction plus one full iteration on a +synthetic grid whose `provenance` map provably exceeds the bound (e.g. `R = 200`, +`S = 2000`, `ploidy = 2` → 800k slots, 12.8 MB under the old path). Assert peak +stays below a small multiple of `O(R + S)`. The point is to fail loudly if +someone reintroduces a materialized slot-space array. + +### Consumer equivalence + +For each of `copy_runs`, `gather_fixed`, `_gather_svar_offsets`: byte-identical +output when driven by `RunPlan` versus by `coalesce(provenance(...))` on small +fixtures, across both axes. + +### Existing suite + +The full tree is the regression gate: `pixi run -e dev test`. Concat has existing +coverage that must stay green without modification — if a concat test needs +editing to pass, that is a signal the change is not behavior-preserving. + +## #406 — `svar2_store` fixture (bundled, test-only) + +Independent of the above; bundled into the same PR because it is test-only and +the project prefers one PR per initiative. + +### Scope correction to #406 + +#406 refers to "the `svar2_store` fixture." There are **seven** definitions of +it, each carrying its own copy-pasted `_REF` (byte-identical 40 bp string) and +`_VCF` (the same three variants: SNP at 3, insertion at 7, deletion at 12), +differing only in whether sample `S2` is present: + +| file | samples | +|---|---| +| `tests/dataset/conftest.py:52` | `S0, S1, S2` — the 18-cell grid #406 measured | +| `tests/test_svar2_reconstruct.py:34` | `S0, S1` | +| `tests/unit/dataset/test_svar2_store.py:28` | `S0, S1` | +| `tests/unit/dataset/test_svar2_link.py:59` | `S0, S1` | +| `tests/dataset/test_svar2_readbound_variants.py:36` | `S0, S1` | +| `tests/dataset/test_svar2_readbound_haps.py:34` | `S0, S1` | +| `tests/dataset/test_svar2_readbound_tracks.py:39` | `S0, S1` | + +Enriching one of seven leaves the other six degenerate, so consolidation comes +first. + +**Two of #406's three follow-ons are already done.** Verified against `main` at +`d0e2f498`: + +- The per-channel oracle masks it reports at `test_write_svar2.py:117-119` and + `:497-499` were fixed during #357 by `c3da6bda` ("test(svar2): fix oracle + masking bug, strengthen weak assertions, add divergence pin"). The site now + reads `present = (widths_snp > 0) | (widths_indel > 0)` — cell-level — with a + comment explaining why per-channel masking is wrong. Nothing to do. +- `test_dense_layout_dataset_still_opens_and_reads` (`:737`) no longer overclaims: + its docstring already says "This is a round-trip smoke test, not a parity pin" + and points at `test_dense_ranges_matches_fancy_indexing` as what actually pins + `_DenseRanges.lookup`. It does, however, name the one-non-empty-cell grid as its + reason, so that docstring becomes stale when the fixture is enriched and must be + updated — and the test can then be promoted to a real parity pin. + +So #406's live scope is the degenerate grid, the seven duplicates, and the +withdrawn sample-axis ordering guard. + +**A further weakness the issue does not mention:** the current fixture's +`dense_snp_range` is `[[0,0],[0,0],[0,0]]` — the per-region dense *SNP* channel +is never exercised at all. Only `dense_indel_range` is populated. + +### Change + +1. Hoist one shared `_REF`, `_VCF`, `vcf_and_ref`, and `svar2_store` into the + top-level `tests/conftest.py`; delete the six local shadows. This also removes + six redundant `samtools faidx` + `bcftools view` + `bcftools index` subprocess + triples from the suite. +2. Replace `_VCF` with the enriched genotypes below. Every sparse variant carries + exactly **one** call, which is what keeps genoray's + `cost_model.rs::choose_representation` from routing it to the per-region dense + channel; the two 3-carrier variants exist to populate the dense channel, + including the SNP half that is currently unexercised. `_REF` is unchanged + (`ACAGTACATGGGTACTAGCTAGGCTAACCGGTTAACCGGT`, 40 bp), as are the three regions + `[0,20)`, `[5,15)`, `[25,40)`. + + ``` + #CHROM POS ID REF ALT QUAL FILTER INFO FORMAT S0 S1 S2 + chr1 3 . A G . . . GT 1|0 0|0 0|0 + chr1 7 . C CAT . . . GT 0|1 0|0 0|0 + chr1 9 . T C . . . GT 0|0 0|1 0|0 + chr1 12 . GTA G . . . GT 0|0 1|0 0|0 + chr1 17 . A C . . . GT 0|1 0|0 0|0 + chr1 19 . C CGG . . . GT 1|1 1|1 1|1 + chr1 30 . G A . . . GT 1|1 1|1 1|1 + ``` + + This was not reasoned about — it was built and measured. Realized vk widths + (`snp`/`indel` per `(region, sample, ploid)`), against the current fixture's + 1-of-18: + + | | ploid 0 | ploid 1 | + |---|---|---| + | r0 S0 | snp1/ind0 | snp1/ind1 | + | r0 S1 | snp0/ind1 | snp1/ind0 | + | r0 S2 | empty | empty | + | r1 S0 | empty | snp0/ind1 | + | r1 S1 | snp0/ind1 | snp1/ind0 | + | r1 S2 | empty | empty | + | r2 (all) | empty | empty | + + **7 of 18 cells occupied.** `dense_snp_range` becomes `[[0,0],[0,0],[0,1]]` + and `dense_indel_range` `[[0,1],[0,0],[1,1]]`, so both dense channels are + exercised where previously `dense_snp_range` was all zeros. + + Every property #406 asks for is satisfied and was checked mechanically: + empty column (S2) kept; empty row (region 2) kept; ≥2 non-empty sample + columns; a present cell with one empty channel (r0 S0 p0, r0 S1 p0); a present + cell with both channels (r0 S0 p1); an absent cell *inside* a non-empty region + (r1 S0 p0); S0's and S1's occupancy patterns differ, so a transposed sample + axis is detectable; dense channel non-empty. + +3. Rewrite `test_fixture_has_empty_cells` (`test_write_svar2.py:~645-690`) against + the new grid. Its structure is right and must be preserved: assert the empty + column and empty row **separately** (a single `not grid.all()` is a disjunction + that stays green when either regresses alone), and keep the non-vacuity pin so + a future genoray cost-model change that pushed every variant dense fails loudly + instead of making every sparse/dense parity test pass against an empty table. + Replace the "exactly one cell is occupied" pin with the 7-cell pattern above. +4. Restore the sample-axis ordering guard that had to be withdrawn during #357 as + unsatisfiable — it is satisfiable now that S0 and S1 differ at (r1, p0). +5. Update `test_dense_layout_dataset_still_opens_and_reads`'s docstring, which + currently justifies itself by the one-non-empty-cell grid, and promote it to a + real parity pin now that the grid can distinguish a correct dense reader from + one returning a constant. Mutation check: injecting + `r_q = np.zeros_like(r_q)` into `_DenseRanges.lookup` must turn it red. + +### Gate + +`pixi run -e dev test` over the full tree. Every test that changes meaning must +change *because* the fixture got stronger, and each such change is called out in +its commit message. + +## #404 — explicitly not in scope + +Analyzed and left deferred; the reasoning is recorded as a comment on the issue. +Summary: the concern was that the probe's share of batch wall rises with thread +count (0.86% single-threaded, 1.68% at 8 threads) and would cross the 2% gate on +a wide machine. Fitting Amdahl's law to the two recorded batch figures +(`T(1) = 171 ms`, `T(8) = 88 ms`) gives serial 76.1 ms and parallel 94.9 ms, so +the batch wall asymptotes at 76.1 ms — above the 73.9 ms break-even the gate +implies. Worst-row share at infinite threads is 1.94%, under gate. Thread count +alone never trips it. Caveat: a two-point fit to a two-parameter model has no +residual and cannot be validated against its own data; it is a strong prior, not +a measurement. What would make #404 live is the *serial* 76.1 ms shrinking from +other work, and the existing benchmark gate already catches that. + +## Out of scope + +- The merged `offsets.npy` itself (16.0 GB at chr22) — a format change. +- A batched array-valued run form — only if planning is measured as the + bottleneck. +- `_concat.py:341` and `:631` beyond the mechanical move to `RunPlan` — already + bounded at `(R, 2)`. +- Any change to `Run`'s field set or to `_concat_io`'s IO strategy + (`CONCAT_CHUNK_BYTES`, the 16 MiB NFSv3 knee) — unrelated. +- #404, #405. + +## Risks + +**The one real correctness risk is a silent wrong merge.** A plan that is wrong +in a way the exactness test misses produces a dataset whose slots point at the +wrong samples — readable, not obviously corrupt. The differential test against +the retained oracle is the mitigation, which is why the oracle is kept rather +than deleted, and why the sweep includes the interleaved order mode that models +the real failure case rather than only shuffled and default orders. + +**Secondary:** `copy_runs`' in-place cumsum into `merged[1:]` aliases the array it +reads. `np.cumsum(x, out=x)` is documented as safe for exact aliasing, but the +change must be pinned by a test comparing against the out-of-place form rather +than assumed. diff --git a/python/genvarloader/_dataset/_concat.py b/python/genvarloader/_dataset/_concat.py index b2e6bc04..ee6ea21e 100644 --- a/python/genvarloader/_dataset/_concat.py +++ b/python/genvarloader/_dataset/_concat.py @@ -13,7 +13,7 @@ from .._atomic import atomic_dir from .._fasta_cache import fingerprint as _bounded_fingerprint from ._concat_io import copy_runs, gather_fixed, link_or_copy_buffered -from ._concat_plan import Run, coalesce, provenance +from ._concat_plan import ExplicitRunPlan, Run, RunPlan, as_plan from ._concat_validate import ( ConcatInput, load_inputs, @@ -200,10 +200,15 @@ def _assert_annot_track_matches(name: str, src_dirs: list[Path]) -> None: ) +def _region_plan(shapes: list[tuple[int, int]], order: "NDArray[np.int64]") -> RunPlan: + """Build the region-only (``n_samples=1, ploidy=1``) merge plan shared by both region-axis sites.""" + return RunPlan("regions", [(r, 1) for r, _ in shapes], 1, order=order) + + def _gather_svar_offsets( paths: list[Path], out_dir: Path, - runs: list[Run], + runs: "RunPlan | ExplicitRunPlan | Sequence[Run]", shapes: list[tuple[int, int]], ploidy: int, ) -> None: @@ -211,7 +216,7 @@ def _gather_svar_offsets( Stored as two leading planes (starts, then stops), each a flat ``(R*S*P,)`` int64 array over the same ``(region, sample, ploid)`` C-order flat-slot - space ``runs`` was coalesced over (see ``_concat_plan``). ``gather_fixed`` + space ``runs`` plans over (see ``_concat_plan``). ``gather_fixed`` streams fixed-stride *records* through a single byte-addressed file and can't express this layout: the slot axis is nested *inside* the two leading planes, so a given slot's start and stop live ``R*S*P`` elements @@ -222,6 +227,7 @@ def _gather_svar_offsets( backend) is the correct call here, not a shortcut around the "don't materialize a full array" constraint. """ + plan = as_plan(runs) n_src_slots = [r * s * ploidy for r, s in shapes] planes = [] for plane in (0, 1): @@ -229,8 +235,8 @@ def _gather_svar_offsets( for p, n in zip(paths, n_src_slots): arr = np.fromfile(p / "genotypes" / "offsets.npy", dtype=np.int64) srcs.append(arr.reshape(2, -1)[plane]) - out = np.empty(sum(r.src_stop - r.src_start for r in runs), np.int64) - for r in runs: + out = np.empty(plan.n_slots, np.int64) + for r in plan: n = r.src_stop - r.src_start out[r.dst_start : r.dst_start + n] = srcs[r.src][r.src_start : r.src_stop] planes.append(out) @@ -256,13 +262,15 @@ def _concat_svar2_ranges( input, remapped into the merged keyspace, ordered, and appended. Legacy dense inputs feed the same merge through ``_DenseRanges.entries_for_regions``. - After this change an svar2 ``concat`` with no per-sample tracks holds nothing - ``R x S``-sized: the ``(R*S*P, 2)`` ``provenance`` array (64 GB at All of Us - chr22) and the ``list[Run]`` ``coalesce`` builds from it (~204 bytes per run, - which on an interleaved sample merge degenerates to one run per slot) are both - gone from this path. Per-sample tracks still plan in core at - ``_concat.py:465`` -- 32 GB plus a ~424 GB run list at the same projection -- - so ``concat`` is bounded only for tracks-free datasets. Tracked separately. + No stage of this merge plans in core any more. The ``(R*S*P, 2)`` + ``provenance`` array (64 GB at the All of Us chr22 grid) and the + ``list[Run]`` ``coalesce`` built from it (184 bytes per run, degenerating + to one run per slot on an interleaved sample merge) are gone from every + path in this module, per-sample tracks included: ``RunPlan`` derives the + same runs from ``order`` alone. What remains is the output offsets array + itself -- the merged file this function (or ``copy_runs``/``gather_fixed`` + on its behalf) writes -- which is irreducible, not a planning artifact. + See ``_concat_plan.RunPlan``. ``dense_snp_range``/``dense_indel_range`` are per-region only (sample- and ploidy-independent), and cannot be sparsified -- genoray's ``dense_abs_row`` @@ -332,16 +340,16 @@ def _concat_svar2_ranges( # dense_snp_range/dense_indel_range gather need "which input contributed # each merged region, in merged order", and it is the same value either # way -- computing it twice would let a future edit that changes one - # `provenance` call and not the other silently order the dense_* gather - # differently from the cell_* copy. Only meaningful on `axis == "regions"`: - # on `axis == "samples"`, `order` maps merged *sample* slots, not regions, - # so this must stay unevaluated there (dense_snp_range/dense_indel_range - # are instead linked from input #0 below, unchanged across inputs). - region_runs = ( - coalesce(provenance("regions", [(r, 1) for r, _ in shapes], 1, order=order)) - if axis == "regions" - else None - ) + # `_region_plan` call and not the other silently order the dense_* gather + # differently from the cell_* copy. `_region_plan` is the same helper the + # annot-track region-axis site below calls, one scope up -- a single + # obvious way to build a region-only plan, instead of the two independent + # `RunPlan("regions", ...)` constructions this used to be. Only meaningful + # on `axis == "regions"`: on `axis == "samples"`, `order` maps merged + # *sample* slots, not regions, so this must stay unevaluated there + # (dense_snp_range/dense_indel_range are instead linked from input #0 + # below, unchanged across inputs). + region_runs = _region_plan(shapes, order) if axis == "regions" else None out_span = n_samples * ploidy if axis == "regions" and all(isinstance(rd, _SparseRanges) for rd in readers): @@ -526,8 +534,7 @@ def concat( ) meta["variants_fingerprint"] = variants_fingerprint(paths[0]) - prov = provenance(axis, shapes, ploidy, order=order) - runs = coalesce(prov) + runs = RunPlan(axis, shapes, ploidy, order=order) src_offsets = [ np.fromfile(p / "genotypes" / "offsets.npy", dtype=np.int64) for p in paths @@ -554,8 +561,7 @@ def concat( # global array. Gather the R/S axes in the same merged `order` as # every other store; values copy verbatim (they're absolute indices # into the shared external store, not slot-relative). - prov = provenance(axis, shapes, ploidy, order=order) - runs = coalesce(prov) + runs = RunPlan(axis, shapes, ploidy, order=order) _gather_svar_offsets(paths, geno, runs, shapes, ploidy) shape = [2, n_regions, len(samples), ploidy] (geno / "svar_meta.json").write_text( @@ -582,8 +588,12 @@ def concat( # `order` as everything else -- a track store is indexed by the same # (region, sample) grid as the genotypes/offsets stores above. if ref.tracks: - t_prov = provenance(axis, shapes, 1, order=order) - t_runs = coalesce(t_prov) + # Hoisted once per `concat` call, not per track: a `RunPlan` re-derives + # its runs on every `for r in plan` rather than caching a list, so this + # only saves the object construction, not run computation -- the + # re-derivation across tracks is the design's deliberate memory-for-CPU + # trade, not something this hoist avoids. + t_runs = RunPlan(axis, shapes, 1, order=order) for name in ref.tracks: src_dirs = [p / "intervals" / name for p in paths] out_t = tmp / "intervals" / name @@ -628,10 +638,7 @@ def concat( src_dirs[0] / f"{fname}.npy", out_a / f"{fname}.npy" ) else: - a_prov = provenance( - "regions", [(r, 1) for r, _ in shapes], 1, order=order - ) - a_runs = coalesce(a_prov) + a_runs = _region_plan(shapes, order) a_offsets = [ np.fromfile(d / "offsets.npy", dtype=np.int64) for d in src_dirs ] diff --git a/python/genvarloader/_dataset/_concat_io.py b/python/genvarloader/_dataset/_concat_io.py index 204415db..44340757 100644 --- a/python/genvarloader/_dataset/_concat_io.py +++ b/python/genvarloader/_dataset/_concat_io.py @@ -12,11 +12,18 @@ import errno import shutil from pathlib import Path +from typing import Sequence import numpy as np from numpy.typing import NDArray -from ._concat_plan import CONCAT_CHUNK_BYTES, Run +from ._concat_plan import ( + CONCAT_CHUNK_BYTES, + ExplicitRunPlan, + Run, + RunPlan, + as_plan, +) __all__ = ["copy_runs", "gather_fixed", "link_or_copy_buffered"] @@ -39,7 +46,7 @@ def _stream_range(fi, fo, start: int, stop: int) -> None: def copy_runs( srcs: list[Path], dst: Path, - runs: list[Run], + runs: "RunPlan | ExplicitRunPlan | Sequence[Run]", src_offsets: list[NDArray[np.int64]], itemsize: int, ) -> NDArray[np.int64]: @@ -48,7 +55,10 @@ def copy_runs( Args: srcs: Payload file per input dataset (raw, headerless arrays). dst: Destination payload file, created or truncated. - runs: Destination-ordered runs from :func:`._concat_plan.coalesce`. + runs: Destination-ordered runs — a :class:`._concat_plan.RunPlan`, a + :class:`._concat_plan.ExplicitRunPlan`, or any re-iterable sequence of + :class:`._concat_plan.Run`. Iterated twice, so a one-shot generator is + not accepted. src_offsets: Cumulative offsets per input dataset, each of length ``n_source_slots + 1``, in elements (not bytes). itemsize: Bytes per element of the payload dtype. @@ -56,23 +66,34 @@ def copy_runs( Returns: Merged cumulative offsets, length ``total_merged_slots + 1``, in elements. """ - n_slots = sum(r.src_stop - r.src_start for r in runs) - lengths = np.empty(n_slots, np.int64) - for r in runs: - off = src_offsets[r.src] - n = r.src_stop - r.src_start - lengths[r.dst_start : r.dst_start + n] = ( - off[r.src_start + 1 : r.src_stop + 1] - off[r.src_start : r.src_stop] - ) - + plan = as_plan(runs) + n_slots = plan.n_slots + + # Lengths are written straight into the output buffer and cumsummed in + # place: a separate `lengths` array would be another n_slots int64s, 32.0 GB + # at the All of Us chr22 grid for the ploidy-bearing genotype payload + # (R*S*P = 4.0e9 slots; half that for per-sample tracks, which have no + # ploidy axis). Filling happens per destination-contiguous BATCH rather + # than per run, because on an interleaved sample merge every run is one + # slot long and a per-run numpy slice-assignment would be 4.0e9 scalar + # calls -- slower than the byte streaming it feeds. merged = np.empty(n_slots + 1, np.int64) merged[0] = 0 - np.cumsum(lengths, out=merged[1:]) + lengths = merged[1:] + for dst_start, ds_vec, slots in plan.slot_batches(): + n = len(slots) + out = lengths[dst_start : dst_start + n] + for d in np.unique(ds_vec): + m = ds_vec == d + off = src_offsets[int(d)] + sl = slots[m] + out[m] = off[sl + 1] - off[sl] + np.cumsum(lengths, out=lengths) handles: dict[int, object] = {} try: with open(dst, "wb") as fo: - for r in runs: + for r in plan: if r.src not in handles: handles[r.src] = open(srcs[r.src], "rb") fi = handles[r.src] @@ -92,7 +113,7 @@ def copy_runs( def gather_fixed( srcs: list[Path], dst: Path, - runs: list[Run], + runs: "RunPlan | ExplicitRunPlan | Sequence[Run]", record_bytes: int, ) -> None: """Gather fixed-size records through a run plan. @@ -103,13 +124,16 @@ def gather_fixed( Args: srcs: Source file per input dataset. dst: Destination file, created or truncated. - runs: Destination-ordered runs. + runs: Destination-ordered runs — a :class:`._concat_plan.RunPlan`, a + :class:`._concat_plan.ExplicitRunPlan`, or any re-iterable sequence of + :class:`._concat_plan.Run`. Iterated once. record_bytes: Bytes per slot. """ + plan = as_plan(runs) handles: dict[int, object] = {} try: with open(dst, "wb") as fo: - for r in runs: + for r in plan: if r.src not in handles: handles[r.src] = open(srcs[r.src], "rb") _stream_range( diff --git a/python/genvarloader/_dataset/_concat_plan.py b/python/genvarloader/_dataset/_concat_plan.py index 9a50bce5..4cdc8139 100644 --- a/python/genvarloader/_dataset/_concat_plan.py +++ b/python/genvarloader/_dataset/_concat_plan.py @@ -3,23 +3,65 @@ A GVL dataset stores parallel ragged arrays over an ``(R, S[, P])`` C-order grid; the flat slot for ``(r, s, p)`` is ``((r * S) + s) * P + p``. Merging two or more datasets means deciding, for each *merged* flat slot, which input dataset and which -*source* flat slot it comes from. That mapping is the provenance map, and -:func:`coalesce` compresses it into maximal contiguous runs so the IO layer can move -large byte ranges instead of individual slots. +*source* flat slot it comes from. + +:class:`RunPlan` is the production path: it derives the destination-ordered run +list analytically from the sorted merge ``order`` alone, without ever +materializing a map over every merged slot. :func:`provenance` and +:func:`coalesce` describe that same mapping the naive way — a full +``(n_slots, 2)`` provenance map, compressed into runs by brute-force scanning — +and are kept **only** as the equivalence oracle +``test_run_plan_matches_coalesce_provenance_exhaustively`` checks +:class:`RunPlan` against in ``tests/unit/dataset/test_concat_plan.py``. Neither +function is called from the IO layer (``_concat.py``/``_concat_io.py``); do not +delete them. """ from __future__ import annotations -from typing import NamedTuple +from typing import Iterable, Iterator, NamedTuple, Sequence import numpy as np from numpy.typing import NDArray -__all__ = ["CONCAT_CHUNK_BYTES", "Run", "coalesce", "provenance"] +__all__ = [ + "CONCAT_CHUNK_BYTES", + "ExplicitRunPlan", + "Run", + "RunPlan", + "as_plan", + "coalesce", + "provenance", +] CONCAT_CHUNK_BYTES = 16 << 20 """Buffered-IO chunk size. 16 MiB is the measured knee on NFSv3; 64 MiB is no better.""" +_SLOT_BATCH_SLOTS = 1 << 20 +"""Slots per `slot_batches` chunk on the regions axis: 8 MiB of int64 indices. + +A regions-axis run can cover the whole merged grid (4.0e9 slots at chr22), so +emitting one batch per run would rebuild exactly the array this module exists to +avoid. The samples axis is naturally bounded at `n_merged * ploidy` instead. +""" + + +def _default_order(n_ds: int, counts: list[int]) -> NDArray[np.int64]: + """Block-concatenation order: dataset 0's whole block, then dataset 1's, etc. + + Args: + n_ds: Number of input datasets. + counts: Positions along the merged axis contributed by each dataset. + + Returns: + An ``(sum(counts), 2)`` int64 array of ``(dataset_idx, within_idx)``. + """ + if n_ds == 0: + return np.zeros((0, 2), np.int64) + ds_col = np.repeat(np.arange(n_ds, dtype=np.int64), counts) + w_col = np.concatenate([np.arange(c, dtype=np.int64) for c in counts]) + return np.stack([ds_col, w_col], axis=1) + class Run(NamedTuple): """A maximal contiguous span of merged slots drawn from one source dataset. @@ -46,6 +88,12 @@ def provenance( ) -> NDArray[np.int64]: """Map each merged flat slot to its ``(dataset, source flat slot)`` origin. + Not on the production path: :class:`RunPlan` derives the same runs + analytically without ever materializing this map. Retained solely as the + equivalence oracle ``test_run_plan_matches_coalesce_provenance_exhaustively`` + (in ``tests/unit/dataset/test_concat_plan.py``) checks :class:`RunPlan` + against — do not delete. + Without ``order``, merged positions along ``axis`` are laid out as dataset 0's whole block, then dataset 1's, etc. (block-concatenation) — the default is correct only when the inputs' key ranges do not interleave along ``axis``. @@ -79,11 +127,6 @@ def provenance( n_ds = len(shape_per_ds) - def _default_order(counts: list[int]) -> NDArray[np.int64]: - ds_col = np.repeat(np.arange(n_ds, dtype=np.int64), counts) - w_col = np.concatenate([np.arange(c, dtype=np.int64) for c in counts]) - return np.stack([ds_col, w_col], axis=1) - if axis == "regions": # Merged region i's S*P slots are contiguous in both source and # destination: they're `order[i, 1] * S*P .. +S*P` in dataset @@ -91,7 +134,7 @@ def _default_order(counts: list[int]) -> NDArray[np.int64]: n_samples = shape_per_ds[0][1] cell = n_samples * ploidy if order is None: - order = _default_order([shape_per_ds[d][0] for d in range(n_ds)]) + order = _default_order(n_ds, [shape_per_ds[d][0] for d in range(n_ds)]) else: order = np.asarray(order, dtype=np.int64) @@ -110,7 +153,7 @@ def _default_order(counts: list[int]) -> NDArray[np.int64]: n_regions = shape_per_ds[0][0] per_ds_samples = [s for _, s in shape_per_ds] if order is None: - order = _default_order(per_ds_samples) + order = _default_order(n_ds, per_ds_samples) else: order = np.asarray(order, dtype=np.int64) @@ -135,6 +178,13 @@ def _default_order(counts: list[int]) -> NDArray[np.int64]: def coalesce(prov: NDArray[np.int64]) -> list[Run]: """Compress a provenance map into maximal contiguous runs. + Not on the production path: :class:`RunPlan` derives the same runs + analytically without ever materializing a provenance map to compress. + Retained solely as the equivalence oracle + ``test_run_plan_matches_coalesce_provenance_exhaustively`` (in + ``tests/unit/dataset/test_concat_plan.py``) checks :class:`RunPlan` + against — do not delete. + A run is a maximal span of consecutive merged slots over which the source dataset is constant and the source slot increases by exactly 1. Iterating the returned runs in order walks the destination sequentially *and* each source @@ -165,3 +215,247 @@ def coalesce(prov: NDArray[np.int64]) -> list[Run]: ) for a, b in zip(starts, stops) ] + + +class RunPlan: + """Destination-ordered runs, derived from ``order`` without materializing slots. + + Equivalent to ``coalesce(provenance(axis, shape_per_ds, ploidy, order=order))`` + and pinned against it by + ``test_run_plan_matches_coalesce_provenance_exhaustively``, but it never + builds the ``(n_slots, 2)`` map: at the All of Us chr22 grid that map is + 32-64 GB and the run list it compresses to is 384-768 GB, because on an + interleaved sample merge every run is one slot long. + + The whole thing rests on one property: the run-break condition is "source + dataset changes, or source slot does not increment", and on both axes that + reduces to a predicate over ``order`` alone -- the region index ``r`` cancels + out. So the break pattern is computed once, in ``O(R + S)``, and the runs + stream in ``O(1)``. + + **Re-iterable on purpose, not a generator.** ``copy_runs`` iterates runs + twice (once for offsets, once to stream bytes) and so does + ``_gather_svar_offsets``; a one-shot iterator would yield an empty second + pass and silently truncate the output rather than raise. + + Args: + axis: Either ``"regions"`` or ``"samples"``. + shape_per_ds: ``(n_regions, n_samples)`` per input dataset, in input order. + ploidy: Slots per ``(region, sample)`` cell. Pass ``1`` for interval + stores, which have no ploidy axis. + order: ``(n_merged_along_axis, 2)`` int array of ``(dataset_idx, + within_dataset_idx)`` per merged position along ``axis``, in + destination order. ``None`` reproduces block-concatenation. + + Raises: + ValueError: If ``axis`` is not ``"regions"`` or ``"samples"``. + """ + + def __init__( + self, + axis: str, + shape_per_ds: list[tuple[int, int]], + ploidy: int, + *, + order: "NDArray[np.int64] | None" = None, + ) -> None: + if axis not in ("regions", "samples"): + raise ValueError(f'axis must be "regions" or "samples", got {axis!r}') + self.axis = axis + self.ploidy = int(ploidy) + self.shape_per_ds = [(int(r), int(s)) for r, s in shape_per_ds] + n_ds = len(self.shape_per_ds) + if order is None: + counts = [r if axis == "regions" else s for r, s in self.shape_per_ds] + self.order = _default_order(n_ds, counts) + else: + self.order = np.asarray(order, dtype=np.int64).reshape(-1, 2) + + @property + def n_slots(self) -> int: + """Total merged flat slots this plan covers, computed arithmetically.""" + if not self.shape_per_ds: + return 0 + if self.axis == "regions": + return len(self.order) * self.shape_per_ds[0][1] * self.ploidy + return self.shape_per_ds[0][0] * len(self.order) * self.ploidy + + def slot_batches( + self, + ) -> "Iterator[tuple[int, NDArray[np.int64], NDArray[np.int64]]]": + """Yield ``(dst_start, src_ds, src_slots)`` batches in destination order. + + Each batch describes a destination-contiguous span: ``src_ds[i]`` and + ``src_slots[i]`` are the origin of merged slot ``dst_start + i``. + Concatenating every batch in order rebuilds :func:`provenance`'s output + exactly, which is what pins this method. + + Yields: + ``(dst_start, src_ds, src_slots)``, where the two arrays are int64 + and equal in length. + """ + if self.axis == "regions": + for run in self: + pos, dst = run.src_start, run.dst_start + while pos < run.src_stop: + n = min(_SLOT_BATCH_SLOTS, run.src_stop - pos) + yield ( + dst, + np.full(n, run.src, np.int64), + np.arange(pos, pos + n, dtype=np.int64), + ) + pos += n + dst += n + return + + n_regions = self.shape_per_ds[0][0] if self.shape_per_ds else 0 + n_merged = len(self.order) + if n_regions == 0 or n_merged == 0 or self.ploidy == 0: + return + per_ds_samples = np.asarray([s for _, s in self.shape_per_ds], np.int64) + ds, w = self.order[:, 0], self.order[:, 1] + s_d = per_ds_samples[ds] + p = np.arange(self.ploidy, dtype=np.int64) + # `order` is per merged SAMPLE; each contributes `ploidy` adjacent slots. + ds_vec = np.repeat(ds, self.ploidy) + for r in range(n_regions): + base = (r * s_d + w) * self.ploidy + slots = (base[:, None] + p[None, :]).reshape(-1) + yield (r * n_merged * self.ploidy, ds_vec, slots) + + def _segments(self) -> tuple[NDArray[np.int64], NDArray[np.int64]]: + """Half-open ``[start, stop)`` spans of ``order`` with no break inside.""" + ds, w = self.order[:, 0], self.order[:, 1] + brk = (ds[1:] != ds[:-1]) | (w[1:] != w[:-1] + 1) + starts = np.concatenate([np.zeros(1, np.int64), np.flatnonzero(brk) + 1]) + stops = np.concatenate([starts[1:], np.array([len(self.order)], np.int64)]) + return starts, stops + + def __iter__(self) -> "Iterator[Run]": + if self.axis == "regions": + yield from self._iter_regions() + else: + yield from self._iter_samples() + + def _iter_regions(self) -> "Iterator[Run]": + # Merged region i owns one contiguous S*P block in BOTH source and + # destination, so a segment of `order` maps to exactly one run and no + # carry is needed: adjacent segments are non-contiguous by construction. + cell = self.shape_per_ds[0][1] * self.ploidy if self.shape_per_ds else 0 + if len(self.order) == 0 or cell == 0: + return + ds, w = self.order[:, 0], self.order[:, 1] + for a, b in zip(*self._segments()): + yield Run( + src=int(ds[a]), + src_start=int(w[a]) * cell, + src_stop=(int(w[a]) + int(b - a)) * cell, + dst_start=int(a) * cell, + ) + + def _iter_samples(self) -> "Iterator[Run]": + n_regions = self.shape_per_ds[0][0] if self.shape_per_ds else 0 + n_merged = len(self.order) + if n_regions == 0 or n_merged == 0 or self.ploidy == 0: + return + per_ds_samples = np.asarray([s for _, s in self.shape_per_ds], np.int64) + ds, w = self.order[:, 0], self.order[:, 1] + seg_starts, seg_stops = self._segments() + + # One pending run, extended whenever the next segment continues it in + # both source and destination. This is what makes a run that spans a + # region boundary come out merged, with no special case for that + # boundary -- verified against the oracle at 2,880 configurations. + pending: Run | None = None + for r in range(n_regions): + for a, b in zip(seg_starts, seg_stops): + d = int(ds[a]) + s_d = int(per_ds_samples[d]) + src_start = (r * s_d + int(w[a])) * self.ploidy + src_stop = (r * s_d + int(w[b - 1]) + 1) * self.ploidy + dst_start = (r * n_merged + int(a)) * self.ploidy + if ( + pending is not None + and pending.src == d + and pending.src_stop == src_start + and pending.dst_start + (pending.src_stop - pending.src_start) + == dst_start + ): + pending = Run(d, pending.src_start, src_stop, pending.dst_start) + else: + if pending is not None: + yield pending + pending = Run(d, src_start, src_stop, dst_start) + if pending is not None: + yield pending + + +class ExplicitRunPlan: + """A hand-built run list presented through :class:`RunPlan`'s interface. + + The IO layer takes either this or a :class:`RunPlan`, so its unit tests can + exercise streaming with a two-run list without constructing a merge. + + Args: + runs: Destination-ordered runs. Materialized, so it may be any iterable. + """ + + def __init__(self, runs: "Iterable[Run]") -> None: + self._runs = list(runs) + + def __iter__(self) -> "Iterator[Run]": + return iter(self._runs) + + @property + def n_slots(self) -> int: + """Total merged flat slots covered by the run list.""" + return sum(r.src_stop - r.src_start for r in self._runs) + + def slot_batches( + self, + ) -> "Iterator[tuple[int, NDArray[np.int64], NDArray[np.int64]]]": + """Yield one ``(dst_start, src_ds, src_slots)`` batch per run.""" + for r in self._runs: + n = r.src_stop - r.src_start + yield ( + r.dst_start, + np.full(n, r.src, np.int64), + np.arange(r.src_start, r.src_stop, dtype=np.int64), + ) + + +def as_plan( + runs: "RunPlan | ExplicitRunPlan | Sequence[Run]", +) -> "RunPlan | ExplicitRunPlan": + """Normalize a run source so the IO layer has one consuming path. + + Args: + runs: A plan, or a re-iterable sequence of runs. A one-shot generator is + deliberately not accepted: ``copy_runs`` iterates its runs twice, and + a one-shot iterator would silently exhaust on the first pass and + truncate the second rather than raise. + + Returns: + ``runs`` itself when it is already a plan, else an + :class:`ExplicitRunPlan` wrapping it. + + Raises: + TypeError: if ``runs`` is neither a plan nor a :class:`~typing.Sequence` + (e.g. a generator). This matches the declared parameter type -- + ``RunPlan | ExplicitRunPlan | Sequence[Run]`` -- rather than + silently ``list()``-materializing whatever is iterable, which for a + one-shot generator over the full slot space would rebuild exactly + the run list this module exists to avoid (768 GB in the degenerate + interleaved case at the All of Us chr22 grid). + """ + if isinstance(runs, (RunPlan, ExplicitRunPlan)): + return runs + if not isinstance(runs, Sequence): + raise TypeError( + "as_plan() requires a RunPlan, ExplicitRunPlan, or Sequence[Run], " + f"got {type(runs).__name__}. A one-shot generator is not accepted: " + "the plan must be re-iterable because consumers (e.g. copy_runs) " + "iterate it more than once, and a one-shot iterator would silently " + "truncate rather than raise." + ) + return ExplicitRunPlan(runs) diff --git a/skills/genvarloader/SKILL.md b/skills/genvarloader/SKILL.md index 0ba4f696..c0232e3f 100644 --- a/skills/genvarloader/SKILL.md +++ b/skills/genvarloader/SKILL.md @@ -194,10 +194,13 @@ re-verifies it, so an out-of-band rewrite of the variant index raises instead of misresolving. **Cost:** `gvl.concat` streams bytes rather than re-deriving anything, so its cost tracks I/O, not -compute — it moves roughly the full size of the merged dataset. Only worth it against the -alternative of re-extracting genotypes from scratch, not as a routine step. See "Merging datasets" -in `docs/source/write.md` for the full cost discussion and the `genoray.SparseVar2.concat` -store-level alternative for contig-sharded `.svar2` workflows. +compute — it moves roughly the full size of the merged dataset. Planning is streaming — memory +scales with regions plus samples, not their product — so the resident cost is the merged offsets +array (one int64 per merged `(region, sample[, ploid])` slot) plus the inputs' own offsets arrays +(together roughly twice that figure, since the inputs partition the merged grid), not the merge +plan. Only worth it against the alternative of re-extracting genotypes from scratch, not as a +routine step. See "Merging datasets" in `docs/source/write.md` for the full cost discussion and the +`genoray.SparseVar2.concat` store-level alternative for contig-sharded `.svar2` workflows. Source: `python/genvarloader/_dataset/_concat.py`, `_concat_validate.py`. diff --git a/tests/conftest.py b/tests/conftest.py index 3d15c2c4..2435b1f5 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -254,6 +254,90 @@ def phased_svar2_gvl(_svar2_slot_src, tmp_path_factory) -> Path: return out +# --- shared 2-sample .svar2 fixture ----------------------------------------- +# +# Consolidates six byte-identical private `svar2_store` fixtures that used to +# live in tests/test_svar2_reconstruct.py, tests/unit/dataset/test_svar2_store.py, +# tests/unit/dataset/test_svar2_link.py, and tests/dataset/test_svar2_readbound_ +# {variants,haps,diffs}.py. Named `svar2_store_2s` (not `svar2_store`) because +# tests/dataset/conftest.py separately defines a 3-sample `svar2_store` -- three +# of the folded-in modules live under tests/dataset/ and would otherwise +# silently shadow that fixture. tests/dataset/test_svar2_readbound_tracks.py +# keeps its own private fixture: its VCF carries a fourth variant and is not a +# duplicate of `_SVAR2_VCF_2S` above. It IS byte-identical to the +# `_SVAR2_SLOT_REF` / `_SVAR2_SLOT_VCF` / `_svar2_slot_src` fixture defined +# elsewhere in this file; folding it into that fixture is tracked separately +# and deliberately out of scope here. + +# 40 bp reference (chr1). VCF POS (1-based) -> 0-based: SNP@2 (A>G), INS@6 +# (C>CAT), DEL@11 (GTA>G, ilen -2). Genotypes exercise both samples and both +# ploids. +_SVAR2_REF = "ACAGTACATGGGTACTAGCTAGGCTAACCGGTTAACCGGT" + +_SVAR2_VCF_2S = """\ +##fileformat=VCFv4.2 +##contig= +##FORMAT= +#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tS0\tS1 +chr1\t3\t.\tA\tG\t.\t.\t.\tGT\t1|0\t0|0 +chr1\t7\t.\tC\tCAT\t.\t.\t.\tGT\t0|1\t1|1 +chr1\t12\t.\tGTA\tG\t.\t.\t.\tGT\t1|1\t0|1 +""" + + +def _build_svar2(vcf_text: str, samples: list[str], d: Path, name: str) -> Path: + """Write a VCF + FASTA under ``d`` and convert them to a .svar2 store. + + Args: + vcf_text: Full VCF text, including header. + samples: Sample names to convert, in the order genoray should store them. + d: Directory to build in. + name: Basename of the resulting store directory. + + Returns: + Path to the finished ``.svar2`` store. + """ + from genoray import _core + + ref = d / "ref.fa" + ref.write_text(f">chr1\n{_SVAR2_REF}\n") + subprocess.run(["samtools", "faidx", str(ref)], check=True) + + vcf = d / "in.vcf" + vcf.write_text(vcf_text) + bcf = d / "in.bcf" + subprocess.run(["bcftools", "view", "-Ob", "-o", str(bcf), str(vcf)], check=True) + subprocess.run(["bcftools", "index", str(bcf)], check=True) + + out = d / name + _core.run_conversion_pipeline( + str(bcf), + str(ref), + ["chr1"], + str(out), + samples, + 25_000, + 2, + 1, + 8 * 1024 * 1024, + ) + assert (out / "meta.json").exists(), "conversion did not finish" + return out + + +@pytest.fixture(scope="module") +def svar2_store_2s(tmp_path_factory) -> Path: + """A two-sample (S0, S1) .svar2 store. + + Named apart from the dataset suite's three-sample ``svar2_store`` on purpose: + six of this fixture's former copies were private duplicates, three of which + lived in ``tests/dataset/`` and silently shadowed ``tests/dataset/conftest.py``'s + same-named three-sample fixture. + """ + d = tmp_path_factory.mktemp("svar2_2s") + return _build_svar2(_SVAR2_VCF_2S, ["S0", "S1"], d, "store") + + # --- 1kg datasets (slow tier) ------------------------------------------------ diff --git a/tests/dataset/conftest.py b/tests/dataset/conftest.py index afa34766..c26b1acd 100644 --- a/tests/dataset/conftest.py +++ b/tests/dataset/conftest.py @@ -9,26 +9,37 @@ import pyBigWig import genvarloader as gvl +from tests.conftest import _SVAR2_REF as _REF SEQLEN = 20 -# 40 bp reference (chr1). VCF POS (1-based) -> 0-based: SNP@2 (A>G), INS@6 (C>CAT), -# DEL@11 (GTA>G, ilen -2). Genotypes exercise both samples and both ploids. -# S2 is 0|0 everywhere: an ENTIRELY EMPTY sample column, so the sparse cache's -# "cell not present" branch is exercised. Variants stop at 0-based 13, so a -# region past that (see test_write_svar2_emits_cache) is an entirely empty ROW. -# Mirrors tests/test_svar2_reconstruct.py's svar2_store fixture (which keeps only -# S0/S1), so the matched .svar (SVAR1) store built from the same VCF is still a -# valid parity oracle. -_REF = "ACAGTACATGGGTACTAGCTAGGCTAACCGGTTAACCGGT" +# 40 bp reference (chr1). VCF POS (1-based) -> 0-based: SNP@2 (A>G), INS@6 +# (C>CAT), SNP@8 (T>C), DEL@11 (GTA>G, ilen -2), SNP@16 (A>C) each carry a +# single ALT haplotype call (one ploid, in one sample), which keeps genoray's +# cost model from routing them to the per-region dense channel. INS@18 +# (C>CGG) and SNP@29 (G>A) are 1|1 in every sample -- 6 carrier calls (3 +# samples x both ploids) -- crossing the dense threshold, so each populates +# one of the two dense channels (indel and SNP respectively). +# Over regions [0,20), [5,15), [25,40): S2 is 0|0 at every sparse variant, an +# ENTIRELY EMPTY sample column in the vk view, so the sparse cache's "cell not +# present" branch is still exercised. The resulting vk grid (region x sample x +# ploid) has 7 of 18 cells occupied -- see test_fixture_grid_is_non_degenerate +# in test_write_svar2.py for the exact layout. Shares its reference sequence +# (``_SVAR2_REF``) with tests/conftest.py's svar2_store_2s fixture, but the two +# VCFs are no longer the same variants, so this store is not a parity oracle +# for that fixture. _VCF = """\ ##fileformat=VCFv4.2 ##contig= ##FORMAT= #CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tS0\tS1\tS2 chr1\t3\t.\tA\tG\t.\t.\t.\tGT\t1|0\t0|0\t0|0 -chr1\t7\t.\tC\tCAT\t.\t.\t.\tGT\t0|1\t1|1\t0|0 -chr1\t12\t.\tGTA\tG\t.\t.\t.\tGT\t1|1\t0|1\t0|0 +chr1\t7\t.\tC\tCAT\t.\t.\t.\tGT\t0|1\t0|0\t0|0 +chr1\t9\t.\tT\tC\t.\t.\t.\tGT\t0|0\t0|1\t0|0 +chr1\t12\t.\tGTA\tG\t.\t.\t.\tGT\t0|0\t1|0\t0|0 +chr1\t17\t.\tA\tC\t.\t.\t.\tGT\t0|1\t0|0\t0|0 +chr1\t19\t.\tC\tCGG\t.\t.\t.\tGT\t1|1\t1|1\t1|1 +chr1\t30\t.\tG\tA\t.\t.\t.\tGT\t1|1\t1|1\t1|1 """ diff --git a/tests/dataset/test_concat_svar2.py b/tests/dataset/test_concat_svar2.py index 4a9dddba..03784438 100644 --- a/tests/dataset/test_concat_svar2.py +++ b/tests/dataset/test_concat_svar2.py @@ -431,20 +431,23 @@ def test_concat_svar2_samples_axis_remaps_nonzero_local_slot(tmp_path: Path): """The samples-axis s_map (scatter-inverse of `order`) must be exercised at a non-trivial local index, not just index 0. - The shared `svar2_store` fixture (used by `svar2_shards_by_samples`) has only - ONE occupied sparse cell across its whole (3 region, 3 sample, ploidy 2) grid, - and that cell sits at local sample slot 0 in both its owning shard and the - merged order -- so it stays in place under a completely broken s_map that - forgets to invert `order` at all (`s_maps[d][w] = w`). Verified directly: that - exact one-line mutation in `_concat_svar2_ranges` left the whole - `test_concat_svar2.py` suite (all 9 tests, including both characterization - tests) passing. - - This test builds two minimal, from-scratch `svar2_ranges/` directories (no - `gvl.write`, no VCF) whose one occupied cell each sits at a local index that - is NOT its merged index, and drives `_concat_svar2_ranges` directly so the - scatter-inverse itself is what's under test, not diluted by an otherwise-empty - grid. + The shared `svar2_store` fixture (used by `svar2_shards_by_samples`) has a + 7-of-18 occupied grid over (3 region, 3 sample, ploidy 2) -- see + `test_fixture_grid_is_non_degenerate` in `test_write_svar2.py` -- with + occupied cells at non-trivial local sample slots, so a completely broken + s_map that forgets to invert `order` at all (`s_maps[d][w] = w`) is now + also caught incidentally by the characterization tests: verified directly, + that exact one-line mutation in `_concat_svar2_ranges` now fails 4 of the + 12 tests in `test_concat_svar2.py` (this one plus 3 others, including both + characterization tests). + + This test remains the TARGETED pin for the s_map regardless: it builds two + minimal, from-scratch `svar2_ranges/` directories (no `gvl.write`, no VCF) + whose one occupied cell each sits at a local index that is NOT its merged + index, and drives `_concat_svar2_ranges` directly, so the scatter-inverse + itself is what's under test -- isolated from the shared fixture's other + axes and the rest of the write/read pipeline, and unambiguous about which + mutation broke it, rather than relying on incidental fallout elsewhere. """ import json @@ -520,34 +523,37 @@ def test_concat_svar2_regions_axis_remaps_and_orders_correctly(tmp_path: Path): """Regions-axis twin of `..._samples_axis_remaps_nonzero_local_slot`. `axis="regions"` has TWO independent merge implementations, each with its - own way to get the region order wrong, and the shared fixture's regions - shards (`svar2_shards_by_regions`) only ever put their one occupied cell in - merged region 0 -- where every ordering agrees -- so neither bug is visible - to `..._regions_reads_like_single_write` or - `..._regions_fast_path_matches_general_merge`. Verified directly, by - injecting each bug into `_concat_svar2_ranges` and re-running the whole - `test_concat_svar2.py` suite (10 tests, both characterization tests - included): + own way to get the region order wrong. The shared fixture's regions shards + (`svar2_shards_by_regions`) are built from the enriched 7-of-18 grid (see + `test_fixture_grid_is_non_degenerate` in `test_write_svar2.py`), so both + bugs are now caught incidentally by other tests in this file too -- + verified directly, by injecting each bug into `_concat_svar2_ranges` and + re-running the whole `test_concat_svar2.py` suite (12 tests total, both + characterization tests included): - the general merge path's `r_maps` (used whenever any input is `_DenseRanges`, so the fast path is unavailable) must be the scatter-inverse of `order` (source -> merged). Setting - `r_maps[d][w] = w` (un-inverted / identity) left all 10 tests passing. + `r_maps[d][w] = w` (un-inverted / identity) now fails 2 of the 12 tests: + this one and `..._regions_fast_path_matches_general_merge`. - the fast path (all-sparse inputs) reorders region-CSR blocks via - `provenance("regions", ..., order=order)`; passing `order=None` (block - concat: all of input #0's regions, then all of input #1's) also left - all 10 tests passing. + `_region_plan`'s `RunPlan(..., order=order)`; passing `order=None` + (block concat: all of input #0's regions, then all of input #1's) now + fails 3 of the 12 tests: this one, `..._regions_fast_path_matches_general_merge`, + and the `..._regions_reads_like_single_write` characterization test. Both bugs preserve every per-region entry COUNT and never collide two entries onto the same `(region, cell_id)` pair, so `region_ptr`'s shape is unaffected, nothing raises, and the dataset opens cleanly -- they only move entries to the wrong region or reorder them within one. This test + remains the TARGETED pin regardless of that incidental coverage: it hand-builds two from-scratch shards (no `gvl.write`) with entries placed so the correct interleaved order visibly disagrees with both bugs' output, and drives `_concat_svar2_ranges` directly against both shards sparse (fast path) and with one rewritten dense (general path via `rewrite_as_dense`, the existing dense-input test oracle), so both merge - implementations are actually exercised. + implementations are actually exercised in isolation, unambiguous about + which one broke. """ import json diff --git a/tests/dataset/test_svar2_readbound_diffs.py b/tests/dataset/test_svar2_readbound_diffs.py index 8b097945..5504e0c3 100644 --- a/tests/dataset/test_svar2_readbound_diffs.py +++ b/tests/dataset/test_svar2_readbound_diffs.py @@ -19,51 +19,7 @@ import numpy as np import pytest -# Same fixtures as tests/dataset/test_svar2_readbound_haps.py: 40 bp reference -# (chr1). VCF POS (1-based) -> 0-based: SNP@2 (A>G), INS@6 (C>CAT), DEL@11 -# (GTA>G, ilen -2). Genotypes exercise both samples and both ploids. -_REF = "ACAGTACATGGGTACTAGCTAGGCTAACCGGTTAACCGGT" -_VCF = """\ -##fileformat=VCFv4.2 -##contig= -##FORMAT= -#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tS0\tS1 -chr1\t3\t.\tA\tG\t.\t.\t.\tGT\t1|0\t0|0 -chr1\t7\t.\tC\tCAT\t.\t.\t.\tGT\t0|1\t1|1 -chr1\t12\t.\tGTA\tG\t.\t.\t.\tGT\t1|1\t0|1 -""" - - -@pytest.fixture(scope="module") -def svar2_store(tmp_path_factory) -> Path: - from genoray import _core - - d = tmp_path_factory.mktemp("svar2_readbound_diffs") - ref = d / "ref.fa" - ref.write_text(f">chr1\n{_REF}\n") - subprocess.run(["samtools", "faidx", str(ref)], check=True) - - vcf = d / "in.vcf" - vcf.write_text(_VCF) - bcf = d / "in.bcf" - subprocess.run(["bcftools", "view", "-Ob", "-o", str(bcf), str(vcf)], check=True) - subprocess.run(["bcftools", "index", str(bcf)], check=True) - - out = d / "store" - _core.run_conversion_pipeline( - str(bcf), - str(ref), - ["chr1"], - str(out), - ["S0", "S1"], - 25_000, - 2, - 1, - 8 * 1024 * 1024, - ) - assert (out / "meta.json").exists(), "conversion did not finish" - return out - +from tests.conftest import _SVAR2_REF # Fixture whose cost model routes a SNP into the DENSE/snp table (not var_key) — # see test_svar2_readbound_haps.py for the routing rationale. @@ -85,7 +41,7 @@ def svar2_store_dense_snp(tmp_path_factory) -> Path: d = tmp_path_factory.mktemp("svar2_readbound_diffs_dense_snp") ref = d / "ref.fa" - ref.write_text(f">chr1\n{_REF}\n") + ref.write_text(f">chr1\n{_SVAR2_REF}\n") subprocess.run(["samtools", "faidx", str(ref)], check=True) vcf = d / "in.vcf" @@ -151,17 +107,17 @@ def _implied_diffs(regions, ref_arr, ref_offsets, sv, contig) -> np.ndarray: [(0, 40), (2, 2), (20, 25)], # empty region + a variant-free window ], ) -def test_readbound_diffs_matches_implied_haps(svar2_store, regions): +def test_readbound_diffs_matches_implied_haps(svar2_store_2s, regions): import genoray from tests._oracles.svar2_readbound_inputs import build_readbound_diffs contig = "chr1" - ref_bytes = _REF.encode() + ref_bytes = _SVAR2_REF.encode() ref_arr = np.frombuffer(ref_bytes, np.uint8) ref_offsets = np.array([0, len(ref_bytes)], np.int64) - sv = genoray.SparseVar2(str(svar2_store)) + sv = genoray.SparseVar2(str(svar2_store_2s)) assert (sv.n_samples, sv.ploidy) == (2, 2) implied = _implied_diffs(regions, ref_arr, ref_offsets, sv, contig) @@ -181,7 +137,7 @@ def test_readbound_diffs_dense_snp_matches_implied_haps(svar2_store_dense_snp): from tests._oracles.svar2_readbound_inputs import build_readbound_diffs contig = "chr1" - ref_bytes = _REF.encode() + ref_bytes = _SVAR2_REF.encode() ref_arr = np.frombuffer(ref_bytes, np.uint8) ref_offsets = np.array([0, len(ref_bytes)], np.int64) diff --git a/tests/dataset/test_svar2_readbound_haps.py b/tests/dataset/test_svar2_readbound_haps.py index 9f3450e8..a95e9e15 100644 --- a/tests/dataset/test_svar2_readbound_haps.py +++ b/tests/dataset/test_svar2_readbound_haps.py @@ -16,49 +16,7 @@ import numpy as np import pytest -# 40 bp reference (chr1). VCF POS (1-based) -> 0-based: SNP@2 (A>G), INS@6 (C>CAT), -# DEL@11 (GTA>G, ilen -2). Genotypes exercise both samples and both ploids. -_REF = "ACAGTACATGGGTACTAGCTAGGCTAACCGGTTAACCGGT" -_VCF = """\ -##fileformat=VCFv4.2 -##contig= -##FORMAT= -#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tS0\tS1 -chr1\t3\t.\tA\tG\t.\t.\t.\tGT\t1|0\t0|0 -chr1\t7\t.\tC\tCAT\t.\t.\t.\tGT\t0|1\t1|1 -chr1\t12\t.\tGTA\tG\t.\t.\t.\tGT\t1|1\t0|1 -""" - - -@pytest.fixture(scope="module") -def svar2_store(tmp_path_factory) -> Path: - from genoray import _core - - d = tmp_path_factory.mktemp("svar2_readbound") - ref = d / "ref.fa" - ref.write_text(f">chr1\n{_REF}\n") - subprocess.run(["samtools", "faidx", str(ref)], check=True) - - vcf = d / "in.vcf" - vcf.write_text(_VCF) - bcf = d / "in.bcf" - subprocess.run(["bcftools", "view", "-Ob", "-o", str(bcf), str(vcf)], check=True) - subprocess.run(["bcftools", "index", str(bcf)], check=True) - - out = d / "store" - _core.run_conversion_pipeline( - str(bcf), - str(ref), - ["chr1"], - str(out), - ["S0", "S1"], - 25_000, - 2, - 1, - 8 * 1024 * 1024, - ) - assert (out / "meta.json").exists(), "conversion did not finish" - return out +from tests.conftest import _SVAR2_REF @pytest.mark.parametrize( @@ -69,18 +27,18 @@ def svar2_store(tmp_path_factory) -> Path: [(0, 40), (2, 2), (20, 25)], # empty region + a variant-free window ], ) -def test_readbound_matches_union_oracle(svar2_store, regions): +def test_readbound_matches_union_oracle(svar2_store_2s, regions): import genoray from tests._oracles.svar2_source import SparseVar2Source from tests._oracles.svar2_readbound_inputs import build_readbound_haps contig = "chr1" - ref_bytes = _REF.encode() + ref_bytes = _SVAR2_REF.encode() ref_arr = np.frombuffer(ref_bytes, np.uint8) ref_offsets = np.array([0, len(ref_bytes)], np.int64) - sv = genoray.SparseVar2(str(svar2_store)) + sv = genoray.SparseVar2(str(svar2_store_2s)) S, P = sv.n_samples, sv.ploidy assert (S, P) == (2, 2) @@ -131,7 +89,7 @@ def test_readbound_matches_union_oracle(svar2_store, regions): pytest.fail("data mismatch but no single hap slice differed (offset bug?)") -def test_readbound_haps_noncontiguous_ref_raises(svar2_store): +def test_readbound_haps_noncontiguous_ref_raises(svar2_store_2s): """A non-C-contiguous ``ref_`` view must surface as ``ValueError``, not a Rust panic. @@ -152,17 +110,17 @@ def test_readbound_haps_noncontiguous_ref_raises(svar2_store): contig = "chr1" regions = [(0, 40)] - ref_bytes = _REF.encode() + ref_bytes = _SVAR2_REF.encode() ref_offsets = np.array([0, len(ref_bytes)], np.int64) - # A strided (non-contiguous) view carrying the same bytes as `_REF`: double up + # A strided (non-contiguous) view carrying the same bytes as `_SVAR2_REF`: double up # each byte, then stride over every other one to recover the original values. doubled = np.repeat(np.frombuffer(ref_bytes, np.uint8), 2) ref_strided = doubled[::2] assert ref_strided.flags["C_CONTIGUOUS"] is False assert bytes(ref_strided) == ref_bytes - sv = genoray.SparseVar2(str(svar2_store)) + sv = genoray.SparseVar2(str(svar2_store_2s)) S, P = sv.n_samples, sv.ploidy d = sv._find_ranges( @@ -211,7 +169,7 @@ def test_readbound_haps_noncontiguous_ref_raises(svar2_store): ) -def test_readbound_matches_union_oracle_with_shifts(svar2_store): +def test_readbound_matches_union_oracle_with_shifts(svar2_store_2s): """Non-trivial per-hap jitter shifts must also match byte-for-byte.""" import genoray @@ -220,11 +178,11 @@ def test_readbound_matches_union_oracle_with_shifts(svar2_store): contig = "chr1" regions = [(0, 40), (5, 20)] - ref_bytes = _REF.encode() + ref_bytes = _SVAR2_REF.encode() ref_arr = np.frombuffer(ref_bytes, np.uint8) ref_offsets = np.array([0, len(ref_bytes)], np.int64) - sv = genoray.SparseVar2(str(svar2_store)) + sv = genoray.SparseVar2(str(svar2_store_2s)) S, P = sv.n_samples, sv.ploidy n_q = len(regions) * S rng = np.random.default_rng(0) @@ -282,7 +240,7 @@ def svar2_store_dense_snp(tmp_path_factory) -> Path: d = tmp_path_factory.mktemp("svar2_readbound_dense_snp") ref = d / "ref.fa" - ref.write_text(f">chr1\n{_REF}\n") + ref.write_text(f">chr1\n{_SVAR2_REF}\n") subprocess.run(["samtools", "faidx", str(ref)], check=True) vcf = d / "in.vcf" @@ -321,7 +279,7 @@ def test_readbound_dense_snp_matches_union_oracle(svar2_store_dense_snp): from tests._oracles.svar2_readbound_inputs import build_readbound_haps contig = "chr1" - ref_bytes = _REF.encode() + ref_bytes = _SVAR2_REF.encode() ref_arr = np.frombuffer(ref_bytes, np.uint8) ref_offsets = np.array([0, len(ref_bytes)], np.int64) @@ -372,8 +330,8 @@ def test_readbound_dense_snp_matches_union_oracle(svar2_store_dense_snp): ) -def _svar2_haps_dataset(tmp_path: Path, svar2_store: Path): - """Build a full gvl Dataset over the ``svar2_store`` fixture and return its +def _svar2_haps_dataset(tmp_path: Path, svar2_store_2s: Path): + """Build a full gvl Dataset over the ``svar2_store_2s`` fixture and return its haplotypes view (Svar2Haps-backed). Lifted/adapted from ``test_svar2_dataset.py::_open_pair`` -- this file has no @@ -387,14 +345,14 @@ def _svar2_haps_dataset(tmp_path: Path, svar2_store: Path): import genvarloader as gvl bed = pl.DataFrame({"chrom": ["chr1"], "chromStart": [0], "chromEnd": [40]}) - ref = svar2_store.parent / "ref.fa" + ref = svar2_store_2s.parent / "ref.fa" d = tmp_path / "ds.gvl" - gvl.write(d, bed, variants=SparseVar2(svar2_store), samples=None, overwrite=True) + gvl.write(d, bed, variants=SparseVar2(svar2_store_2s), samples=None, overwrite=True) return gvl.Dataset.open(d, reference=ref).with_seqs("haplotypes") def test_deterministic_haps_read_skips_pre_reconstruct_diffs( - tmp_path: Path, svar2_store: Path, monkeypatch: pytest.MonkeyPatch + tmp_path: Path, svar2_store_2s: Path, monkeypatch: pytest.MonkeyPatch ): """A deterministic (shifts=0) haplotypes read must NOT call the separate hap_diffs readbound kernel -- reconstruct sizes itself internally. Guards the @@ -410,13 +368,13 @@ def counting(*a, **k): monkeypatch.setattr(m, "hap_diffs_from_svar2_readbound", counting) - ds2 = _svar2_haps_dataset(tmp_path, svar2_store) + ds2 = _svar2_haps_dataset(tmp_path, svar2_store_2s) ds2[:, :] assert calls["diffs"] == 0 -def _svar2_spliced_dataset(tmp_path: Path, svar2_store: Path): - """A spliced (2-exon transcript) view over the ``svar2_store`` fixture. +def _svar2_spliced_dataset(tmp_path: Path, svar2_store_2s: Path): + """A spliced (2-exon transcript) view over the ``svar2_store_2s`` fixture. Splicing is the path where sizing and reconstruction are separate calls -- the plan builder needs per-query haplotype lengths before the kernel can be @@ -437,9 +395,9 @@ def _svar2_spliced_dataset(tmp_path: Path, svar2_store: Path): "exon_number": [1, 2], } ) - ref = svar2_store.parent / "ref.fa" + ref = svar2_store_2s.parent / "ref.fa" d = tmp_path / "spliced.gvl" - gvl.write(d, bed, variants=SparseVar2(svar2_store), samples=None, overwrite=True) + gvl.write(d, bed, variants=SparseVar2(svar2_store_2s), samples=None, overwrite=True) return ( gvl.Dataset.open(d, reference=ref) .with_settings(splice_info=("transcript_id", "exon_number")) @@ -448,7 +406,7 @@ def _svar2_spliced_dataset(tmp_path: Path, svar2_store: Path): def test_spliced_haps_read_gathers_inputs_once_per_contig( - tmp_path: Path, svar2_store: Path, monkeypatch: pytest.MonkeyPatch + tmp_path: Path, svar2_store_2s: Path, monkeypatch: pytest.MonkeyPatch ): """Regression for #349: the sizing and reconstruct passes share one gather. @@ -469,7 +427,7 @@ def counting(self, *a, **k): monkeypatch.setattr(m.Svar2Haps, "_gather_inputs", counting) - ds = _svar2_spliced_dataset(tmp_path, svar2_store) + ds = _svar2_spliced_dataset(tmp_path, svar2_store_2s) out = ds[:, :] assert calls["gather"] == 1 @@ -477,7 +435,7 @@ def counting(self, *a, **k): def test_spliced_haps_read_runs_the_rust_gather_once( - tmp_path: Path, svar2_store: Path, monkeypatch: pytest.MonkeyPatch + tmp_path: Path, svar2_store_2s: Path, monkeypatch: pytest.MonkeyPatch ): """Regression for #349: the two FFI passes share one ``gather_haps_readbound``. @@ -499,7 +457,7 @@ def counting(*a, **k): monkeypatch.setattr(m, "gather_svar2_readbound", counting) - ds = _svar2_spliced_dataset(tmp_path, svar2_store) + ds = _svar2_spliced_dataset(tmp_path, svar2_store_2s) out = ds[:, :] assert calls["gather"] == 1 # one contig group, one gather for both passes @@ -507,7 +465,7 @@ def counting(*a, **k): def test_spliced_haps_read_is_byte_identical_without_the_shared_gather( - tmp_path: Path, svar2_store: Path, monkeypatch: pytest.MonkeyPatch + tmp_path: Path, svar2_store_2s: Path, monkeypatch: pytest.MonkeyPatch ): """Sharing the gather must not change a single output byte. @@ -517,17 +475,17 @@ def test_spliced_haps_read_is_byte_identical_without_the_shared_gather( """ import genvarloader._dataset._svar2_haps as m - shared = _svar2_spliced_dataset(tmp_path, svar2_store)[:, :] + shared = _svar2_spliced_dataset(tmp_path, svar2_store_2s)[:, :] monkeypatch.setattr(m.Svar2Haps, "_readbound_gather", lambda self, ci, gi, P: None) - per_pass = _svar2_spliced_dataset(tmp_path / "again", svar2_store)[:, :] + per_pass = _svar2_spliced_dataset(tmp_path / "again", svar2_store_2s)[:, :] np.testing.assert_array_equal(shared.data, per_pass.data) np.testing.assert_array_equal(shared.offsets, per_pass.offsets) def test_readbound_gather_matches_the_unfused_call_and_guards_its_shape( - svar2_store: Path, + svar2_store_2s: Path, ): """The handle is a pure caching detail, and a mismatched one must raise. @@ -545,7 +503,7 @@ def test_readbound_gather_matches_the_unfused_call_and_guards_its_shape( from tests._oracles.svar2_readbound_inputs import readbound_diff_inputs - svar2 = SparseVar2(svar2_store) + svar2 = SparseVar2(svar2_store_2s) contig = svar2.contigs[0] args = readbound_diff_inputs(svar2, contig, [(0, 13), (20, 40)]) @@ -564,7 +522,7 @@ def test_readbound_gather_matches_the_unfused_call_and_guards_its_shape( def test_svar2_haps_n_variants_is_a_readonly_zero_stride_view( - tmp_path: Path, svar2_store: Path + tmp_path: Path, svar2_store_2s: Path ): """#355 mutation-kill: assert on a REAL ``Svar2Haps.n_variants``, not a standalone ``np.broadcast_to`` call. @@ -576,7 +534,7 @@ def test_svar2_haps_n_variants_is_a_readonly_zero_stride_view( goes back to a real ``np.zeros`` allocation: shape/dtype are preserved by a plain ``np.zeros`` too, but strides/writeability are not. """ - ds = _svar2_haps_dataset(tmp_path, svar2_store) + ds = _svar2_haps_dataset(tmp_path, svar2_store_2s) seqs = ds._seqs n_variants = seqs.n_variants diff --git a/tests/dataset/test_svar2_readbound_variants.py b/tests/dataset/test_svar2_readbound_variants.py index d5527ffd..2389255e 100644 --- a/tests/dataset/test_svar2_readbound_variants.py +++ b/tests/dataset/test_svar2_readbound_variants.py @@ -18,49 +18,7 @@ import numpy as np import pytest -# 40 bp reference (chr1). VCF POS (1-based) -> 0-based: SNP@2 (A>G), INS@6 (C>CAT), -# DEL@11 (GTA>G, ilen -2). Genotypes exercise both samples and both ploids. -_REF = "ACAGTACATGGGTACTAGCTAGGCTAACCGGTTAACCGGT" -_VCF = """\ -##fileformat=VCFv4.2 -##contig= -##FORMAT= -#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tS0\tS1 -chr1\t3\t.\tA\tG\t.\t.\t.\tGT\t1|0\t0|0 -chr1\t7\t.\tC\tCAT\t.\t.\t.\tGT\t0|1\t1|1 -chr1\t12\t.\tGTA\tG\t.\t.\t.\tGT\t1|1\t0|1 -""" - - -@pytest.fixture(scope="module") -def svar2_store(tmp_path_factory) -> Path: - from genoray import _core - - d = tmp_path_factory.mktemp("svar2_readbound_variants") - ref = d / "ref.fa" - ref.write_text(f">chr1\n{_REF}\n") - subprocess.run(["samtools", "faidx", str(ref)], check=True) - - vcf = d / "in.vcf" - vcf.write_text(_VCF) - bcf = d / "in.bcf" - subprocess.run(["bcftools", "view", "-Ob", "-o", str(bcf), str(vcf)], check=True) - subprocess.run(["bcftools", "index", str(bcf)], check=True) - - out = d / "store" - _core.run_conversion_pipeline( - str(bcf), - str(ref), - ["chr1"], - str(out), - ["S0", "S1"], - 25_000, - 2, - 1, - 8 * 1024 * 1024, - ) - assert (out / "meta.json").exists(), "conversion did not finish" - return out +from tests.conftest import _SVAR2_REF def _assert_variants_match(oracle, rb) -> None: @@ -103,14 +61,14 @@ def _assert_variants_match(oracle, rb) -> None: [(0, 40), (2, 2), (20, 25)], # empty region + a variant-free window ], ) -def test_readbound_variants_match_decode_oracle(svar2_store, regions): +def test_readbound_variants_match_decode_oracle(svar2_store_2s, regions): import genoray from tests._oracles.svar2_readbound_inputs import build_readbound_variants contig = "chr1" - sv = genoray.SparseVar2(str(svar2_store)) + sv = genoray.SparseVar2(str(svar2_store_2s)) S, P = sv.n_samples, sv.ploidy assert (S, P) == (2, 2) @@ -142,7 +100,7 @@ def svar2_store_dense_snp(tmp_path_factory) -> Path: d = tmp_path_factory.mktemp("svar2_readbound_variants_dense_snp") ref = d / "ref.fa" - ref.write_text(f">chr1\n{_REF}\n") + ref.write_text(f">chr1\n{_SVAR2_REF}\n") subprocess.run(["samtools", "faidx", str(ref)], check=True) vcf = d / "in.vcf" diff --git a/tests/dataset/test_write_svar2.py b/tests/dataset/test_write_svar2.py index d74e6f35..ebc5ebfa 100644 --- a/tests/dataset/test_write_svar2.py +++ b/tests/dataset/test_write_svar2.py @@ -29,8 +29,10 @@ def test_write_svar2_emits_cache(svar2_store: Path, tmp_path: Path): svar2 = SparseVar2(svar2_store) bed = pl.DataFrame( { - # [25, 40) holds no variants at all: an entirely empty region row, - # which the sparse layout must round-trip as (0, 0) everywhere. + # [25, 40) holds no sparse variants: an entirely empty region row in + # the vk view, which the sparse layout must round-trip as (0, 0) + # everywhere. (The dense SNP at 0-based 29 lives here, but it's read + # through dense_snp_range, not the vk cache this row exercises.) "chrom": ["chr1", "chr1", "chr1"], "chromStart": [0, 5, 25], "chromEnd": [20, 15, 40], @@ -621,24 +623,24 @@ def test_write_svar2_duplicate_store_samples_raises( gvl.write(tmp_path / "ds.gvl", bed, variants=svar2, overwrite=True) -def test_fixture_has_empty_cells(svar2_store: Path, tmp_path: Path): - """The fixture must NOT be 100% fill, or the sparse cache is untested. +def test_fixture_grid_is_non_degenerate(svar2_store: Path, tmp_path: Path): + """The fixture's vk (sparse) grid must be a real 7-of-18 grid, not degenerate. Empty cells are the entire point of the sparse layout (#357): if every (region, sample, ploid) window holds a variant, the "cell is absent" branch never executes and a sparse/dense divergence there is invisible. S2 is 0|0 - everywhere (empty column) and [25, 40) holds no variants (empty row). + everywhere (empty column) and [25, 40) holds no sparse variants (empty row). genoray also routes each variant to the per-sample SPARSE (vk) channel or the per-region DENSE channel by carrier-call count (see - `choose_representation` in genoray's cost model), and only the sparse - channel reaches this grid. This fixture's single-carrier SNP stays sparse - but both three-carrier indels route dense, so the grid this test inspects - is already 1/18 fill by construction, not 8/8 -- exactly one cell - (region 0, S0, ploid 0) is occupied. A cost-model shift that pushed that - last SNP dense too would empty the grid entirely and make every - sparse/dense parity test built on this fixture pass vacuously; guard - against that directly rather than assuming any occupancy at all. + `choose_representation` in genoray's cost model): every variant in this + fixture that carries a single ALT haplotype call stays sparse, and the two + all-1|1 variants (6 carrier calls each: 3 samples x both ploids) route + dense, one per dense channel. The resulting vk grid is 7 of + 18 cells occupied (region x sample x ploid over 3 regions, 3 samples, + ploidy 2) -- a mix of present/absent cells, present cells with only one + channel filled, and an ordered (non-symmetric) sample axis, none of which + the old 1-of-18 grid could express (#406). """ from genoray import SparseVar2 @@ -654,39 +656,69 @@ def test_fixture_has_empty_cells(svar2_store: Path, tmp_path: Path): ) snp = np.asarray(d["vk_snp_range"], np.int64) # (R*S*P, 2) indel = np.asarray(d["vk_indel_range"], np.int64) - nonempty = (snp[:, 1] > snp[:, 0]) | (indel[:, 1] > indel[:, 0]) + w_snp = (snp[:, 1] - snp[:, 0]).reshape(3, len(sorted_samples), svar2.ploidy) + w_indel = (indel[:, 1] - indel[:, 0]).reshape(w_snp.shape) + occ = (w_snp > 0) | (w_indel > 0) + + s0, s1, s2 = (sorted_samples.index(s) for s in ("S0", "S1", "S2")) - S, P = len(sorted_samples), svar2.ploidy - assert len(nonempty) == 3 * S * P # Row-major (R, S, P) -- pinned by the layout oracle in # test_write_svar2_emits_cache, which asserts this same reshape against the - # cache memmaps. Assert each empty structure SEPARATELY: a single - # `not nonempty.all()` is a disjunction that stays green when either one - # regresses alone, which is exactly the regression this guard exists to catch. - grid = nonempty.reshape(3, S, P) - s2 = sorted_samples.index("S2") - assert not grid[:, s2].any(), ( + # cache memmaps. Assert each structure SEPARATELY: a single `not occ.all()` + # is a disjunction that stays green when any one of them regresses alone. + assert not occ[:, s2].any(), ( "S2 is no longer all-reference; the empty COLUMN is gone" ) - assert not grid[2].any(), ( - "region [25, 40) now holds variants; the empty ROW is gone" + assert not occ[2].any(), ( + "region [25, 40) now holds sparse variants; the empty ROW is gone" + ) + + # The measured grid. Every sparse variant carries exactly one call, which is + # what keeps genoray's cost model from routing it to the per-region dense + # channel; if a cost-model change pushes any of them dense this fails loudly + # rather than letting the sparse/dense parity tests pass against a thinner + # table while appearing green. + expected = np.zeros_like(occ) + expected[0, s0] = [True, True] + expected[0, s1] = [True, True] + expected[1, s0] = [False, True] + expected[1, s1] = [True, True] + np.testing.assert_array_equal(occ, expected) + assert int(occ.sum()) == 7, ( + "fixture occupancy changed; update this pin deliberately" ) - # Non-vacuity. genoray routes each variant to the per-sample SPARSE (vk) - # channel or the per-region DENSE channel by carrier-call count, and only - # the sparse channel lands in this grid: the fixture's single-carrier SNP - # stays sparse, both three-carrier indels route dense. Exactly one cell is - # therefore occupied -- (region 0, S0, ploid 0). If a cost-model change - # pushed that last variant dense too, this grid would be ALL empty and - # every sparse/dense parity test built on this fixture would pass - # trivially against an empty table, green and meaningless. Pin it. - assert grid.any(), ( + + # Non-vacuity, kept from the original guard: an all-empty vk grid would make + # every sparse/dense parity test built on this fixture pass trivially. + assert occ.any(), ( "vk channel is entirely empty: genoray routed every variant to the " "dense channel, so all sparse-cache parity tests on this fixture are " "now vacuous" ) - assert grid[0, sorted_samples.index("S0"), 0], ( - "the fixture's one sparse-channel cell (region 0, S0, ploid 0) is gone" - ) + + # Mixed per-channel emptiness inside a PRESENT cell -- the property the old + # one-cell grid could not express, and the reason two shipped oracles were + # able to mask per channel without failing. + assert (w_snp[0, s0, 0] > 0) and (w_indel[0, s0, 0] == 0) + assert (w_snp[0, s0, 1] > 0) and (w_indel[0, s0, 1] > 0) + assert (w_snp[0, s1, 0] == 0) and (w_indel[0, s1, 0] > 0) + + # An ABSENT cell inside a non-empty region: region 1 holds variants, but + # (S0, ploid 0) has none, so lookup must return (0, 0) there rather than a + # neighbour's range. + assert not occ[1, s0, 0] + + # The sample axis is now ORDERED, not just occupied: S0 and S1 differ at + # (region 1, ploid 0), so a transposed sample axis is detectable. During + # #357 this guard had to be withdrawn as unsatisfiable. + assert not np.array_equal(occ[:, s0], occ[:, s1]) + + # Both dense channels are exercised. dense_snp_range was all zeros before + # this fixture was enriched, so the dense SNP path had no coverage here. + dense_snp = np.asarray(d["dense_snp_range"], np.int64) + dense_indel = np.asarray(d["dense_indel_range"], np.int64) + assert (dense_snp[:, 1] > dense_snp[:, 0]).any(), "dense SNP channel is empty again" + assert (dense_indel[:, 1] > dense_indel[:, 0]).any(), "dense indel channel is empty" def test_sparse_writer_rejects_oversized_grid(tmp_path): @@ -743,21 +775,48 @@ def test_dense_layout_dataset_still_opens_and_reads( on MAJOR only, so a bump would make new GVL refuse every old dataset). The dense reader is what keeps old datasets openable. - This is a round-trip smoke test, not a parity pin: this fixture's grid has - only one non-empty cell (see `svar2_store`), so it can't distinguish a - correct dense reader from one that always returns the same wrong answer. - `test_dense_ranges_matches_fancy_indexing` in - `tests/unit/dataset/test_svar2_ranges.py` is what actually pins - `_DenseRanges.lookup` against dense fancy-indexing over a randomized grid; - this test's job is only to confirm the dense layout still opens and reads - through the full `Dataset.open` -> `with_seqs` stack. - - Deviation from the brief: `Dataset.open` (not `gvl.write`) is what takes - `reference=` -- `with_seqs("haplotypes")` raises `ValueError` without one, - which the brief's snippet omitted. Added `vcf_and_ref` for the FASTA path. + This fixture's vk grid is 7 of 18 cells occupied (see + `test_fixture_grid_is_non_degenerate`), with a mix of present and absent + cells and a non-symmetric sample axis, so the read-for-read comparison of + the dense-layout dataset against the sparse-layout dataset below is no + longer vacuous the way it was against the old 1-of-18 grid. + + It is, however, still not sufficient on its own to pin `_DenseRanges.lookup` + end-to-end: this bed's regions happen to nest ([0,20) is a superset of + [5,15), and disjoint from [25,40)), and haplotype construction filters + variant calls by absolute position. Confirmed by mutation-testing + `r_q = np.zeros_like(r_q)` at the top of `_DenseRanges.lookup` (#406): + every region query silently reads region 0's row, but the rendered + haplotype bytes still matched dataset `a` byte-for-byte, because the wrong + (region-0) variants either coincide with the right ones (shared position + 6) or land outside the query window and get filtered as if absent. So the + byte-level loop below is kept as an end-to-end smoke check, but the + load-bearing parity pin is the raw-range comparison that follows: it calls + `_DenseRanges.lookup` (the exact code the mutation targets) directly and + compares its output against `_SparseRanges.lookup` on the same store, + which is immune to the positional self-healing above because it compares + indices, not rendered bytes. + + That comparison pins the DENSE reader, not the sparse one, and the two + sides are less independent than they look: `rewrite_as_dense` builds + `dense_ds`'s vk arrays by calling `_SparseRanges.lookup` with this exact + `unravel_index(arange(R * S))` query, so a `_SparseRanges.lookup` bug would + reproduce on both sides of this comparison and cancel out. That blind spot + is covered elsewhere: `test_lookup_parity_*` in + `tests/unit/dataset/test_svar2_ranges.py` pins `_SparseRanges.lookup` + against a randomized independent reference, and + `test_dense_ranges_matches_fancy_indexing` in the same module pins + `_DenseRanges.lookup` against dense fancy-indexing directly (and was + independently confirmed to fail under this same mutation) -- so between + the three tests, both readers are pinned against something other than + each other. Defense in depth, not a single self-consistent check. + + `Dataset.open`, not `gvl.write`, takes `reference=`; `with_seqs("haplotypes")` + raises `ValueError` without one, hence `vcf_and_ref`. """ from genoray import SparseVar2 + from genvarloader._dataset._svar2_ranges import _ranges_reader from tests._oracles.svar2_dense_layout import rewrite_as_dense _bcf, ref = vcf_and_ref @@ -779,6 +838,19 @@ def test_dense_layout_dataset_still_opens_and_reads( np.asarray(b[r, s].to_padded(b"N")), ) + # The load-bearing check: compare _DenseRanges.lookup's raw output against + # _SparseRanges.lookup on the byte-identical sparse store, over the full + # (region, sample) grid. Unlike the byte comparison above, this cannot be + # satisfied by coincidence -- see the docstring. + sparse_reader = _ranges_reader(sparse_ds / "genotypes" / "svar2_ranges") + dense_reader = _ranges_reader(dense_ds / "genotypes" / "svar2_ranges") + R, S, P = sparse_reader.n_regions, sparse_reader.n_samples, sparse_reader.ploidy + r_q, si_q = np.unravel_index(np.arange(R * S), (R, S)) + sparse_snp, sparse_indel = sparse_reader.lookup(r_q, si_q, P) + dense_snp, dense_indel = dense_reader.lookup(r_q, si_q, P) + np.testing.assert_array_equal(dense_snp, sparse_snp) + np.testing.assert_array_equal(dense_indel, sparse_indel) + def test_write_svar2_empty_cell_is_zero_not_insertion_point( svar2_store: Path, tmp_path: Path @@ -792,12 +864,13 @@ def test_write_svar2_empty_cell_is_zero_not_insertion_point( dense cell at the same coordinate would instead carry genoray's real insertion point (x, x) for x > 0. - This fixture's grid (bed [0,20)/[5,15)/[25,40) over samples S0/S1/S2) has - exactly one non-empty cell: (region 0, S0, ploid 0). Region 0 / S1 / ploid 0 - is empty in both channels, and genoray's own dense insertion point there is - (1, 1) -- confirmed directly via `_find_ranges` below -- so this test would - fail loudly (assert (1, 1) == (0, 0)) if the sparse writer ever stored - empty cells verbatim instead of collapsing them to (0, 0). + This fixture's grid (bed [0,20)/[5,15)/[25,40) over samples S0/S1/S2) has 7 + of 18 non-empty cells (see `test_fixture_grid_is_non_degenerate`). (region + 1, S0, ploid 0) is one of the empty ones, and genoray's own dense + insertion point there is (1, 1) in the SNP channel -- confirmed directly + via `_find_ranges` below -- so this test would fail loudly + (assert (1, 1) == (0, 0)) if the sparse writer ever stored empty cells + verbatim instead of collapsing them to (0, 0). """ from genoray import SparseVar2 @@ -824,20 +897,20 @@ def test_write_svar2_empty_cell_is_zero_not_insertion_point( samples=sorted_samples, ) P = svar2.ploidy - s1 = sorted_samples.index("S1") + s0 = sorted_samples.index("S0") exp_snp = np.asarray(d["vk_snp_range"], np.int64).reshape(bed.height, -1, P, 2) exp_indel = np.asarray(d["vk_indel_range"], np.int64).reshape(bed.height, -1, P, 2) - assert tuple(exp_snp[0, s1, 0]) != (0, 0), ( + assert tuple(exp_snp[1, s0, 0]) != (0, 0), ( "fixture assumption broken: genoray's dense insertion point for the" - " empty (region 0, S1, ploid 0) cell is expected to be nonzero" + " empty (region 1, S0, ploid 0) cell is expected to be nonzero" ) - assert exp_indel[0, s1, 0, 1] == exp_indel[0, s1, 0, 0], ( + assert exp_indel[1, s0, 0, 1] == exp_indel[1, s0, 0, 0], ( "fixture assumption broken: expected the indel channel empty too" ) reader = _ranges_reader(out / "genotypes" / "svar2_ranges") got_snp, got_indel = reader.lookup( - np.array([0], np.int64), np.array([s1], np.int64), P + np.array([1], np.int64), np.array([s0], np.int64), P ) np.testing.assert_array_equal(got_snp[0], (0, 0)) np.testing.assert_array_equal(got_indel[0], (0, 0)) diff --git a/tests/test_svar2_reconstruct.py b/tests/test_svar2_reconstruct.py index 58b21165..54aa35f3 100644 --- a/tests/test_svar2_reconstruct.py +++ b/tests/test_svar2_reconstruct.py @@ -10,55 +10,12 @@ from __future__ import annotations -import subprocess -from pathlib import Path - import numpy as np -import pytest + +from tests.conftest import _SVAR2_REF # 40 bp reference (chr1). VCF POS (1-based) -> 0-based: SNP@2 (A>G), INS@6 (C>CAT), # DEL@11 (GTA>G, ilen -2). Genotypes exercise both samples and both ploids. -_REF = "ACAGTACATGGGTACTAGCTAGGCTAACCGGTTAACCGGT" -_VCF = """\ -##fileformat=VCFv4.2 -##contig= -##FORMAT= -#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tS0\tS1 -chr1\t3\t.\tA\tG\t.\t.\t.\tGT\t1|0\t0|0 -chr1\t7\t.\tC\tCAT\t.\t.\t.\tGT\t0|1\t1|1 -chr1\t12\t.\tGTA\tG\t.\t.\t.\tGT\t1|1\t0|1 -""" - - -@pytest.fixture(scope="module") -def svar2_store(tmp_path_factory) -> Path: - from genoray import _core - - d = tmp_path_factory.mktemp("svar2") - ref = d / "ref.fa" - ref.write_text(f">chr1\n{_REF}\n") - subprocess.run(["samtools", "faidx", str(ref)], check=True) - - vcf = d / "in.vcf" - vcf.write_text(_VCF) - bcf = d / "in.bcf" - subprocess.run(["bcftools", "view", "-Ob", "-o", str(bcf), str(vcf)], check=True) - subprocess.run(["bcftools", "index", str(bcf)], check=True) - - out = d / "store" - _core.run_conversion_pipeline( - str(bcf), - str(ref), - ["chr1"], - str(out), - ["S0", "S1"], - 25_000, - 2, - 1, - 8 * 1024 * 1024, - ) - assert (out / "meta.json").exists(), "conversion did not finish" - return out def _consensus(ref: bytes, pos, ilen, alleles, q_start: int, q_end: int) -> bytes: @@ -90,16 +47,16 @@ def _consensus(ref: bytes, pos, ilen, alleles, q_start: int, q_end: int) -> byte return bytes(out) -def test_svar2_two_source_matches_decode_oracle(svar2_store): +def test_svar2_two_source_matches_decode_oracle(svar2_store_2s): import genoray from tests._oracles.svar2_source import SparseVar2Source contig = "chr1" q_start, q_end = 0, 40 regions = [(q_start, q_end)] - ref_bytes = _REF.encode() + ref_bytes = _SVAR2_REF.encode() - sv = genoray.SparseVar2(str(svar2_store)) + sv = genoray.SparseVar2(str(svar2_store_2s)) S, P = sv.n_samples, sv.ploidy assert (S, P) == (2, 2) diff --git a/tests/unit/dataset/test_concat_io.py b/tests/unit/dataset/test_concat_io.py index 60dcbb92..1bba06e7 100644 --- a/tests/unit/dataset/test_concat_io.py +++ b/tests/unit/dataset/test_concat_io.py @@ -140,3 +140,146 @@ def test_copy_runs_writes_nothing_for_no_runs(tmp_path): merged = copy_runs([], dst, [], [], itemsize=4) assert merged.tolist() == [0] assert dst.read_bytes() == b"" + + +def test_copy_runs_accepts_a_run_plan_and_matches_an_explicit_list(tmp_path): + """Driving copy_runs from RunPlan must be byte-identical to the old path.""" + from genvarloader._dataset._concat_plan import RunPlan, coalesce, provenance + + shapes = [(2, 2), (2, 1)] + axis, ploidy = "samples", 2 + n_src = [r * s * ploidy for r, s in shapes] + + srcs, offsets = [], [] + for d, n in enumerate(n_src): + lens = np.arange(1, n + 1, dtype=np.int64) + off = np.concatenate([[0], np.cumsum(lens)]).astype(np.int64) + payload = np.arange(off[-1], dtype=np.int32) + d * 1000 + p = tmp_path / f"src{d}.bin" + _write_raw(p, payload) + srcs.append(p) + offsets.append(off) + + out_plan = tmp_path / "plan.npy" + out_list = tmp_path / "list.npy" + plan = RunPlan(axis, shapes, ploidy) + merged_plan = copy_runs(srcs, out_plan, plan, offsets, itemsize=4) + merged_list = copy_runs( + srcs, out_list, coalesce(provenance(axis, shapes, ploidy)), offsets, itemsize=4 + ) + + np.testing.assert_array_equal(merged_plan, merged_list) + assert out_plan.read_bytes() == out_list.read_bytes() + + +def test_copy_runs_in_place_cumsum_matches_out_of_place(tmp_path): + """merged[1:] is cumsummed into itself; pin that against the naive form.""" + lens = np.array([3, 0, 5, 2, 7], np.int64) + off = np.concatenate([[0], np.cumsum(lens)]).astype(np.int64) + payload = np.arange(off[-1], dtype=np.int32) + p = tmp_path / "src.bin" + _write_raw(p, payload) + + runs = [Run(0, 0, 5, 0)] + merged = copy_runs([p], tmp_path / "out.npy", runs, [off], itemsize=4) + + want = np.concatenate([[0], np.cumsum(lens)]).astype(np.int64) + np.testing.assert_array_equal(merged, want) + + +def test_gather_fixed_from_a_run_plan_is_byte_identical_to_the_oracle(tmp_path): + """gather_fixed must produce the same bytes from a plan as from a run list.""" + from genvarloader._dataset._concat_plan import RunPlan, coalesce, provenance + + shapes, axis, ploidy = [(2, 2), (3, 2)], "regions", 1 + srcs = [] + for d, (r, sm) in enumerate(shapes): + p = tmp_path / f"g{d}.bin" + _write_raw(p, np.arange(r * sm * ploidy, dtype=np.int32) + d * 100) + srcs.append(p) + + plan = RunPlan(axis, shapes, ploidy) + out_plan, out_list = tmp_path / "plan.bin", tmp_path / "list.bin" + gather_fixed(srcs, out_plan, plan, record_bytes=4) + gather_fixed( + srcs, out_list, coalesce(provenance(axis, shapes, ploidy)), record_bytes=4 + ) + + assert out_plan.read_bytes() == out_list.read_bytes() + assert out_plan.stat().st_size == plan.n_slots * 4 + + +@pytest.mark.parametrize( + ("axis", "shapes", "order"), + [ + # axis="samples" requires equal n_regions (index 0) across datasets. + ("samples", [(2, 2), (2, 1)], None), + # Interleaved merged-sample order: ds0 sample0, ds1 sample0, ds0 + # sample1 -- built the same way test_concat_plan.py's interleaved-order + # tests do, with an explicit (dataset_idx, within_idx) literal. + ("samples", [(2, 2), (2, 1)], np.array([[0, 0], [1, 0], [0, 1]])), + # axis="regions" requires equal n_samples (index 1) across datasets. + ("regions", [(2, 2), (3, 2)], None), + # Interleaved merged-region order over ds0's 2 regions and ds1's 3. + ( + "regions", + [(2, 2), (3, 2)], + np.array([[1, 0], [0, 0], [1, 1], [0, 1], [1, 2]]), + ), + ], + ids=[ + "samples-block", + "samples-interleaved", + "regions-block", + "regions-interleaved", + ], +) +def test_gather_svar_offsets_from_a_run_plan_is_byte_identical( + tmp_path, axis, shapes, order +): + """The svar offsets gather must agree with the retained oracle. + + This one cannot go through `gather_fixed`: its slot axis is nested inside + two leading planes, so a slot's start and stop live `n_slots` elements + apart. It therefore has its own run-consuming loop, and needs its own pin. + Parametrized over both axes and both an `order=None` block merge and an + interleaved `order`, since every production call site always passes an + explicit `order` and interleaving is the case the design calls load-bearing. + """ + from genvarloader._dataset._concat import _gather_svar_offsets + from genvarloader._dataset._concat_plan import RunPlan, coalesce, provenance + + ploidy = 2 + + paths = [] + for d, (r, sm) in enumerate(shapes): + n = r * sm * ploidy + p = tmp_path / f"ds{d}" + (p / "genotypes").mkdir(parents=True) + starts = np.arange(n, dtype=np.int64) + d * 100 + planes = np.stack([starts, starts + 1]) + _write_raw(p / "genotypes" / "offsets.npy", planes) + paths.append(p) + + out_plan, out_list = tmp_path / "plan", tmp_path / "list" + out_plan.mkdir() + out_list.mkdir() + + _gather_svar_offsets( + paths, out_plan, RunPlan(axis, shapes, ploidy, order=order), shapes, ploidy + ) + _gather_svar_offsets( + paths, + out_list, + coalesce(provenance(axis, shapes, ploidy, order=order)), + shapes, + ploidy, + ) + + plan_bytes = (out_plan / "offsets.npy").read_bytes() + assert plan_bytes == (out_list / "offsets.npy").read_bytes() + # A shared bug that writes one plane instead of two would still pass the + # byte-identity check above (both paths would be wrong the same way), so + # also pin the absolute size against a freshly built plan's `n_slots`. + expected_n_slots = RunPlan(axis, shapes, ploidy, order=order).n_slots + assert len(plan_bytes) == 2 * expected_n_slots * 8 diff --git a/tests/unit/dataset/test_concat_plan.py b/tests/unit/dataset/test_concat_plan.py index 088ba0ae..870b5140 100644 --- a/tests/unit/dataset/test_concat_plan.py +++ b/tests/unit/dataset/test_concat_plan.py @@ -6,6 +6,9 @@ from genvarloader._dataset._concat_plan import ( CONCAT_CHUNK_BYTES, Run, + RunPlan, + _default_order, + as_plan, coalesce, provenance, ) @@ -168,3 +171,197 @@ def test_provenance_samples_interleaved_order_covers_all_slots(): n = r.src_stop - r.src_start covered[r.dst_start : r.dst_start + n] += 1 assert (covered == 1).all() + + +def _rand_order(rng, counts): + """A random valid merged order: every (ds, within) slot exactly once.""" + rows = [(d, i) for d, c in enumerate(counts) for i in range(c)] + rng.shuffle(rows) + return np.array(rows, dtype=np.int64).reshape(-1, 2) + + +def _sorted_interleave_order(rng, counts): + """Each input's keys sorted, merged by global key -- the real two-cohort case.""" + keys = [] + for d, c in enumerate(counts): + ks = sorted(rng.choice(10_000, size=c, replace=False)) + keys.extend((k, d, i) for i, k in enumerate(ks)) + keys.sort() + return np.array([(d, i) for _, d, i in keys], dtype=np.int64).reshape(-1, 2) + + +def test_run_plan_matches_coalesce_provenance_exhaustively(): + """RunPlan must reproduce coalesce(provenance(...)) exactly, everywhere. + + This is the correctness proof for the whole change: `provenance` + `coalesce` + are retained purely as this oracle. A plan that is wrong in a way this sweep + misses produces a merged dataset whose slots point at the wrong samples -- + readable, and not obviously corrupt -- so the sweep deliberately includes the + sorted key-interleave order that models the real two-cohort merge, not only + shuffled and block-default orders. + """ + import itertools + + rng = np.random.default_rng(0) + n_checked = 0 + + for axis, ploidy, n_ds in itertools.product( + ("regions", "samples"), (1, 2, 3), (1, 2, 3) + ): + for _ in range(40): + counts = [int(rng.integers(0, 5)) for _ in range(n_ds)] + if axis == "regions": + n_samples = int(rng.integers(0, 5)) + shapes = [(c, n_samples) for c in counts] + else: + n_regions = int(rng.integers(0, 5)) + shapes = [(n_regions, c) for c in counts] + + for mode in ("none", "default", "shuffled", "interleaved"): + if mode == "none": + order = None + elif mode == "default": + order = _default_order(len(counts), counts) + elif mode == "shuffled": + order = _rand_order(rng, counts) + else: + order = _sorted_interleave_order(rng, counts) + + want = coalesce(provenance(axis, shapes, ploidy, order=order)) + got = list(RunPlan(axis, shapes, ploidy, order=order)) + n_checked += 1 + assert got == want, ( + f"axis={axis} ploidy={ploidy} shapes={shapes} mode={mode}\n" + f"want={want[:6]}\ngot={got[:6]}" + ) + + assert n_checked == 2880, "sweep shrank; the oracle coverage is the whole point" + + +def test_slot_batches_reproduce_the_provenance_map(): + """Concatenating the batches must rebuild `provenance` row for row. + + `copy_runs` gathers source lengths through these batches, so a batch whose + slots are right but whose dataset column is wrong would read the correct + offsets out of the wrong file -- a silent wrong merge, not a crash. + """ + for axis, shapes, ploidy in ( + ("regions", [(3, 4), (2, 4)], 2), + ("samples", [(3, 2), (3, 3)], 2), + ("samples", [(2, 1), (2, 2)], 1), + ): + counts = [s for _, s in shapes] if axis == "samples" else [r for r, _ in shapes] + interleaved = _sorted_interleave_order(np.random.default_rng(1), counts) + for order in (None, interleaved): + plan = RunPlan(axis, shapes, ploidy, order=order) + want = provenance(axis, shapes, ploidy, order=order) + got = np.zeros_like(want) + seen = np.zeros(len(want), bool) + for dst_start, ds_vec, slots in plan.slot_batches(): + n = len(slots) + assert len(ds_vec) == n + got[dst_start : dst_start + n, 0] = ds_vec + got[dst_start : dst_start + n, 1] = slots + seen[dst_start : dst_start + n] = True + assert seen.all(), f"{axis} {shapes} left slots uncovered" + np.testing.assert_array_equal(got, want) + + +def test_n_slots_matches_provenance_length(): + for axis, shapes, ploidy in ( + ("regions", [(3, 4), (2, 4)], 2), + ("samples", [(3, 2), (3, 3)], 2), + ("regions", [(0, 4), (2, 4)], 1), + ("samples", [(0, 2), (0, 3)], 2), + ): + plan = RunPlan(axis, shapes, ploidy) + assert plan.n_slots == len(provenance(axis, shapes, ploidy)) + + +def test_slot_batches_chunk_large_region_runs(monkeypatch): + """A regions-axis run can span the whole grid, so batches must be chunked. + + The production cap is 1Mi slots, far above anything a unit test can build, + so this shrinks it and checks the chunking actually happens. Asserting only + `max(sizes) <= cap` would pass vacuously on a single batch -- including + against the unchunked implementation this test exists to rule out. + """ + from genvarloader._dataset import _concat_plan + + monkeypatch.setattr(_concat_plan, "_SLOT_BATCH_SLOTS", 5) + + # One input, identity order -> exactly one run over 4 * 3 * 2 = 24 slots. + plan = RunPlan("regions", [(4, 3)], 2) + assert len(list(plan)) == 1 + + batches = list(plan.slot_batches()) + assert [len(slots) for _, _, slots in batches] == [5, 5, 5, 5, 4] + assert plan.n_slots == 24 + # Destination stays contiguous across every chunk boundary. + assert [dst for dst, _, _ in batches] == [0, 5, 10, 15, 20] + np.testing.assert_array_equal( + np.concatenate([slots for _, _, slots in batches]), np.arange(24) + ) + + +def test_explicit_run_plan_round_trips_a_hand_built_list(): + runs = [Run(0, 0, 2, 0), Run(1, 5, 7, 2)] + plan = as_plan(runs) + assert list(plan) == runs + assert plan.n_slots == 4 + assert as_plan(plan) is plan + batches = list(plan.slot_batches()) + assert [b[0] for b in batches] == [0, 2] + np.testing.assert_array_equal(batches[0][2], [0, 1]) + np.testing.assert_array_equal(batches[1][1], [1, 1]) + + +def test_as_plan_rejects_a_generator(): + """A one-shot generator must raise, not silently list()-materialize. + + That materialization is exactly the 768 GB the analytic RunPlan exists to + avoid in the degenerate interleaved case. + """ + + def gen(): + yield Run(0, 0, 2, 0) + + with pytest.raises(TypeError, match="re-iterable"): + as_plan(gen()) + + +def test_run_plan_is_re_iterable(): + """copy_runs iterates twice; a generator here would silently truncate output.""" + plan = RunPlan("samples", [(2, 2), (2, 1)], 2) + assert list(plan) == list(plan) + assert len(list(plan)) > 0 + + +def test_run_plan_never_materializes_the_slot_space(): + """Peak allocation must stay O(R + S), not O(R * S * P). + + The old path on this grid allocates a (800_000, 2) int64 provenance map + (12.8 MB) and coalesces it into 800_000 one-slot Runs (~147 MB at 184 bytes + each). Both are invisible to every behavioural test, because RunPlan yields + exactly the same runs -- so this bound is what stops a future edit from + silently restoring them. + """ + import tracemalloc + + shapes = [(200, 1000), (200, 1000)] + # Interleaved samples: the worst case, one run per slot. + order = np.array( + [(d, i) for i in range(1000) for d in (0, 1)], dtype=np.int64 + ).reshape(-1, 2) + + tracemalloc.start() + try: + plan = RunPlan("samples", shapes, 2, order=order) + n_runs = sum(1 for _ in plan) + peak = tracemalloc.get_traced_memory()[1] + finally: + tracemalloc.stop() + + assert plan.n_slots == 200 * 2000 * 2 + assert n_runs == plan.n_slots // 2, "interleaved merge should give one run per cell" + assert peak < (1 << 20), f"peak {peak} bytes: the slot space is being materialized" diff --git a/tests/unit/dataset/test_svar2_link.py b/tests/unit/dataset/test_svar2_link.py index fb0c2950..5649f923 100644 --- a/tests/unit/dataset/test_svar2_link.py +++ b/tests/unit/dataset/test_svar2_link.py @@ -10,7 +10,7 @@ from __future__ import annotations -import subprocess +import shutil from pathlib import Path import pytest @@ -22,18 +22,6 @@ make_svar2_link, ) -# Same tiny fixture recipe as tests/test_svar2_reconstruct.py::svar2_store. -_REF = "ACAGTACATGGGTACTAGCTAGGCTAACCGGTTAACCGGT" -_VCF = """\ -##fileformat=VCFv4.2 -##contig= -##FORMAT= -#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tS0\tS1 -chr1\t3\t.\tA\tG\t.\t.\t.\tGT\t1|0\t0|0 -chr1\t7\t.\tC\tCAT\t.\t.\t.\tGT\t0|1\t1|1 -chr1\t12\t.\tGTA\tG\t.\t.\t.\tGT\t1|1\t0|1 -""" - def test_resolve_prefers_override(tmp_path: Path): real = tmp_path / "cohort.svar2" @@ -55,48 +43,21 @@ def test_verify_none_link_is_noop(tmp_path: Path): _verify_svar2_fingerprint(tmp_path, None) # must not raise -@pytest.fixture(scope="module") -def svar2_store(tmp_path_factory) -> Path: - from genoray import _core - - d = tmp_path_factory.mktemp("svar2_link") - ref = d / "ref.fa" - ref.write_text(f">chr1\n{_REF}\n") - subprocess.run(["samtools", "faidx", str(ref)], check=True) - - vcf = d / "in.vcf" - vcf.write_text(_VCF) - bcf = d / "in.bcf" - subprocess.run(["bcftools", "view", "-Ob", "-o", str(bcf), str(vcf)], check=True) - subprocess.run(["bcftools", "index", str(bcf)], check=True) - - out = d / "store.svar2" - _core.run_conversion_pipeline( - str(bcf), - str(ref), - ["chr1"], - str(out), - ["S0", "S1"], - 25_000, - 2, - 1, - 8 * 1024 * 1024, - ) - assert (out / "meta.json").exists(), "conversion did not finish" - return out - +def test_fingerprint_detects_mutated_store(svar2_store_2s: Path, tmp_path: Path): + # `svar2_store_2s` is a module-scoped fixture shared with five other test + # modules -- mutate a private copy, never the shared store itself. + store = shutil.copytree(svar2_store_2s, tmp_path / "store.svar2") -def test_fingerprint_detects_mutated_store(svar2_store: Path, tmp_path: Path): gvl_path = tmp_path / "ds.gvl" gvl_path.mkdir() - link = make_svar2_link(gvl_path, svar2_store) - _verify_svar2_fingerprint(svar2_store, link) # must not raise + link = make_svar2_link(gvl_path, store) + _verify_svar2_fingerprint(store, link) # must not raise - bin_files = sorted(svar2_store.rglob("*.bin")) + bin_files = sorted(store.rglob("*.bin")) assert bin_files, "expected at least one .bin file in a real .svar2 store" with open(bin_files[0], "ab") as f: f.write(b"\x00") with pytest.raises(ValueError): - _verify_svar2_fingerprint(svar2_store, link) + _verify_svar2_fingerprint(store, link) diff --git a/tests/unit/dataset/test_svar2_store.py b/tests/unit/dataset/test_svar2_store.py index 8f18d5f0..b147cf94 100644 --- a/tests/unit/dataset/test_svar2_store.py +++ b/tests/unit/dataset/test_svar2_store.py @@ -5,56 +5,11 @@ from __future__ import annotations -import subprocess from pathlib import Path -import pytest - from genvarloader.genvarloader import Svar2Store # compiled extension -_REF = "ACAGTACATGGGTACTAGCTAGGCTAACCGGTTAACCGGT" -_VCF = """\ -##fileformat=VCFv4.2 -##contig= -##FORMAT= -#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tS0\tS1 -chr1\t3\t.\tA\tG\t.\t.\t.\tGT\t1|0\t0|0 -chr1\t7\t.\tC\tCAT\t.\t.\t.\tGT\t0|1\t1|1 -chr1\t12\t.\tGTA\tG\t.\t.\t.\tGT\t1|1\t0|1 -""" - - -@pytest.fixture(scope="module") -def svar2_store(tmp_path_factory) -> Path: - from genoray import _core - - d = tmp_path_factory.mktemp("svar2_store") - ref = d / "ref.fa" - ref.write_text(f">chr1\n{_REF}\n") - subprocess.run(["samtools", "faidx", str(ref)], check=True) - - vcf = d / "in.vcf" - vcf.write_text(_VCF) - bcf = d / "in.bcf" - subprocess.run(["bcftools", "view", "-Ob", "-o", str(bcf), str(vcf)], check=True) - subprocess.run(["bcftools", "index", str(bcf)], check=True) - - out = d / "cohort.svar2" - _core.run_conversion_pipeline( - str(bcf), - str(ref), - ["chr1"], - str(out), - ["S0", "S1"], - 25_000, - 2, - 1, - 8 * 1024 * 1024, - ) - assert (out / "meta.json").exists(), "conversion did not finish" - return out - -def test_store_opens_contigs(svar2_store: Path): - store = Svar2Store(str(svar2_store), ["chr1"], n_samples=2, ploidy=2) +def test_store_opens_contigs(svar2_store_2s: Path): + store = Svar2Store(str(svar2_store_2s), ["chr1"], n_samples=2, ploidy=2) assert store.contigs() == ["chr1"]