Repository navigation
perf(concat): plan merges analytically instead of materializing an R x S x P provenance map - #414
Merged
Merged
Conversation
…ixture (#406) Design for replacing concat's materialized provenance map and Run list with an analytic, re-iterable RunPlan, plus the bundled test-only fixture work. Corrects three things in the issues as filed, each verified against main: - #403 names only the per-sample-tracks call site. Two further sites remain, and both are larger: _concat.py:529 (pgen_vcf) and :557 (svar) pass real ploidy, so they are (R*S*P, 2) = 64.0 GB at the All of Us chr22 grid against the tracks site's 32.0 GB. concat is unbounded on three backends, not one. - #403's proposed fix (emit runs as an iterator) would break the build: copy_runs iterates runs twice, as does _gather_svar_offsets, so a bare generator yields an empty second pass and silently truncates the output. The replacement must be re-iterable. - #406's first two follow-ons already landed during #357 in c3da6bd: the per-channel oracle masks are cell-level now, and the dense-layout test's docstring already says it is a smoke test rather than a parity pin. The core design claim -- that the run-break pattern depends only on `order` and never on the region index -- was validated against the retained coalesce(provenance(...)) oracle at 2,880 configurations (both axes, ploidy 1-3, 1-3 inputs, four order modes including sorted key-interleave, empty grids) with zero mismatches. That sweep becomes the shipped property test. Validation also showed the region-boundary case needs no special handling: the pending-run carry subsumes it. The enriched fixture VCF was measured, not reasoned about: 7 of 18 vk cells occupied against the current 1 of 18, with both dense channels exercised where dense_snp_range was previously all zeros. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6
Eight tasks turning the approved spec into task-by-task work: RunPlan's analytic planner and its differential sweep against the retained coalesce(provenance(...)) oracle (1-2), the copy_runs/gather_fixed/ _gather_svar_offsets consumer conversions with byte-identical equivalence pins for all three (3-4), the tracemalloc bound that is the only thing failing if a materialized slot-space array comes back (5), the svar2 fixture dedup and enrichment (6-7), and docs (8). Two corrections found while writing the plan against the code: - _concat.py carries two prose claims that this change falsifies and no test covers: _merge_svar2_ranges' docstring says per-sample tracks "still plan in core" so concat is "bounded only for tracks-free datasets" -- the issue this branch closes -- and the region_runs comment warns about diverging `provenance` calls that will no longer exist. Both are now explicit steps in Task 4. (That docstring's :465 line reference was already stale, and its "~204 bytes per run" disagrees with the 184 bytes measured for this codebase's Run; neither figure is carried forward.) - The spec's slot_batches paragraph said the regions axis yields one batch per run. That is wrong: a regions-axis run can span the entire merged grid, so materializing its slot indices would rebuild the 32 GB array the change exists to remove. Batches there are chunked at _SLOT_BATCH_SLOTS = 1 << 20. The spec is amended to match, and the plan flags the divergence so it is not "restored". Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6
Pre-flight scan before execution fingerprinted every svar2_store fixture's _REF and _VCF rather than trusting the plan's list. Three corrections: - `def svar2_store` is defined 8 times, not 7. tests/dataset/ test_svar2_readbound_diffs.py:38 is a genuine exact duplicate the plan omitted from the fold-in set. - tests/dataset/test_svar2_readbound_tracks.py:39 is NOT a duplicate. Its _VCF carries a fourth variant (chr1:10 G>C, 1|1 / 1|0) that the other six lack. Folding it onto the shared 3-variant fixture would change the data its assertions were written against -- failing the task's own zero-assertion-edit gate, or silently weakening that module's coverage. It now has an explicit do-not-modify entry. - svar2_store_dense_snp exists in three modules, not two (readbound_diffs.py:83 was missed), and svar2_store_unsorted (test_write_svar2.py:526) was unlisted. Both are purpose-built, not duplicates. Step 5's verification expected "exactly three lines" while listing four items; it is now an explicit seven-row table of what must survive. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6
The plan cited 1253 passed / 58 skipped / 4 xfailed from memory. Measured on this branch it is 1250 / 61 / 4 (darwin), so the three tasks that gate on "same counts as before" were checking against a number that was never true here. Noted that skip counts are platform-dependent -- CI runs slow/torch tiers this machine does not -- so only a differing failure count signals a regression. Also promoted `pixi run -e dev gen` to a Global Constraint. A fresh worktree without generated test data reports 387 errors and 52 failures from missing fixture files. That looks exactly like real breakage, and every task in this plan runs the suite. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6
Add RunPlan, which streams coalesce(provenance(...))'s run sequence directly from the merge `order` in O(R + S) memory instead of building the (n_slots, 2) provenance map. At the All of Us chr22 grid that map is 32-64 GB and its run list 384-768 GB, so concat cannot run at cohort scale today. Hoists `_default_order` out of `provenance` to module scope so RunPlan can share it. `provenance`/`coalesce` are unchanged and retained as the oracle RunPlan is pinned against by a 2,880-configuration differential sweep (regions/samples x ploidy x n_ds x block/shuffled/sorted-interleave order). Nothing consumes RunPlan yet. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6
…RunPlan Adds RunPlan.n_slots (arithmetic, no materialization) and slot_batches() (destination-contiguous (dst_start, src_ds, src_slots) batches, chunked at 1Mi slots on the regions axis since a single run there can span the whole merged grid). Adds ExplicitRunPlan, which presents a hand-built Run list through the same interface, and as_plan() to normalize either input for the IO layer's single consuming path. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6
copy_runs now accepts RunPlan/ExplicitRunPlan/Sequence[Run] via as_plan, drops its own lengths array (16.0 GB at chr22) by writing run lengths directly into merged[1:] and cumsumming in place, and fills lengths from plan.slot_batches() instead of per-run to avoid per-slot Python iteration on interleaved sample-axis merges. gather_fixed widens the same way. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6
The copy_runs comment cited the per-sample tracks figures (R*S = 2.0e9 slots, no ploidy axis) as though they were the grid-wide worst case, but copy_runs also serves the ploidy-bearing genotype payload (R*S*P = 4.0e9 slots): 32.0 GB for a separate lengths array and 4.0e9 scalar calls for a per-run fill, not half that. gather_fixed's runs: docstring falsely claimed two passes (copied verbatim from copy_runs per the task brief, which was wrong here) -- it is single-pass. Also name ExplicitRunPlan alongside RunPlan in both runs: docstrings to match the accepted type. No logic changes. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6
…ce map Route every provenance()+coalesce() call site in _concat.py through RunPlan/as_plan, and convert _gather_svar_offsets to the streaming plan interface. No stage of a concat merge now materializes an R*S(*P)-sized provenance array or run list, per-sample tracks included. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6
… from review
Address Task 4 review findings: extract the duplicated
RunPlan("regions", [(r, 1) for r, _ in shapes], 1, order=order) idiom into a
_region_plan(shapes, order) helper used by both region-axis sites; correct
_concat_svar2_ranges's docstring to claim only that no stage plans in core
any more (copy_runs/gather_fixed's output offsets array is irreducible, not
a planning artifact); note that the t_runs hoist now only saves object
construction, not run computation, since RunPlan re-derives its runs per
iteration; and parametrize the svar-offsets oracle-equivalence test over
both axes and an interleaved order, plus an absolute output-size assertion.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6
…ot space Adds a tracemalloc-based test on the chr22-shaped interleaved-sample worst case: RunPlan must derive runs from `order` alone without ever building the (n_slots, 2) provenance map or the one-run-per-slot list that coalesce() would otherwise produce. Every existing behavioral test is invisible to this regression since RunPlan and coalesce(provenance(...)) yield identical runs; this is the only test that would catch a future edit reintroducing that materialization. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6
…test.py
Six modules each carried a byte-identical private 2-sample svar2_store
fixture (test_svar2_reconstruct.py, unit/dataset/test_svar2_{store,link}.py,
dataset/test_svar2_readbound_{variants,haps,diffs}.py); consolidate them
into a shared svar2_store_2s fixture in tests/conftest.py. Renamed (not
left as svar2_store) because tests/dataset/conftest.py separately defines
a 3-sample svar2_store that four of these modules would otherwise shadow.
tests/dataset/test_svar2_readbound_tracks.py keeps its own private fixture
untouched: its VCF carries a fourth variant and is not a true duplicate.
svar2_store_dense_snp (3 modules) and svar2_store_unsorted are unrelated,
purpose-built fixtures and are also untouched.
No assertions changed; full suite stays at 1264 passed / 61 skipped /
4 xfailed.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6
…'s comments
The section comment and the fixture docstring both said four of the six
folded-in modules live under tests/dataset/. It is three: the other three are
tests/test_svar2_reconstruct.py and tests/unit/dataset/test_svar2_{store,link}.py.
The section comment contradicted the module list two lines above it.
The neighbouring note about test_svar2_readbound_tracks.py carrying a "fourth
variant" is a different four and is left alone.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6
The 3-sample svar2 fixture's sparse (vk) channel had one non-empty cell in an 18-cell grid (3 regions x 3 samples x ploidy 2), leaving S1 with zero occupied cells anywhere and the dense SNP channel entirely untested. Replace the fixture's VCF with a 7-variant one (measured, not guessed) that gives a 7-of-18 grid with mixed present/absent cells, per-channel-only cells, an ordered sample axis, and both dense channels populated. Rewrite the occupancy guard (renamed to test_fixture_grid_is_non_degenerate) to pin the full grid instead of one cell, fix the empty-cell regression test's now-occupied probe cell, and promote the dense-layout smoke test to a real dense-vs-sparse parity pin: mutation-testing `r_q = np.zeros_like(r_q)` in `_DenseRanges.lookup` showed the haplotype-byte comparison alone doesn't catch a corrupted region index (this bed's nested/disjoint regions let position-based filtering self-heal it), so add a raw vk_snp/vk_indel range comparison against `_SparseRanges.lookup` on the byte-identical sparse store, which the mutation does catch. Fixes #406. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6
…ings Review fix round 1: docstring accuracy only, no logic or assertion changes. - test_concat_svar2.py: two docstrings still described the fixture's old 1-occupied-cell grid and a fixed pass-through count. Restated against the enriched 7-of-18 grid with freshly measured mutation-failure counts (s_maps[d][w]=w: 4/12 fail; r_maps[d][w]=w: 2/12 fail; order=None: 3/12 fail), while keeping each hand-built test's role as the targeted, isolated pin. - test_write_svar2.py / conftest.py: "no variants at all" -> "no sparse variants" (a dense SNP still lives in that region); "3-carrier-call" language corrected to the actual unit (6 carrier calls: 3 samples x both ploids); the dense-vs-sparse raw-range comparison docstring now names its derivation dependency on _SparseRanges.lookup and points to test_lookup_parity_* and test_dense_ranges_matches_fancy_indexing as covering the remaining gap; dropped transient SDD/task references in favor of a plain, timeless reference= API statement. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6
gvl.concat's planner now derives merge runs from the sorted merge order instead of materializing a provenance map over every (region, sample, ploid) slot, so planning memory scales with regions plus samples rather than their product. Document this in write.md and SKILL.md, and correct the merged offsets.npy figure to depend on ploidy: 16 GB for per-sample tracks (no ploidy axis) vs 32 GB for diploid genotypes at the reference 3,734-region by 535,662-sample grid. That array is unchanged by this work — it's the output file's on-disk format, not a planning artifact. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6
- Copy the shared svar2_store_2s fixture into tmp_path before mutating it in test_fingerprint_detects_mutated_store, instead of corrupting the module- scoped store five other test modules also use. - as_plan() now raises TypeError for a non-Sequence, non-plan input (e.g. a one-shot generator) instead of silently list()-materializing it, which would rebuild exactly the run list RunPlan exists to avoid. - Document provenance/coalesce as RunPlan's retained equivalence oracle (not on the production path) in the module docstring and both functions' docstrings, naming the test that checks RunPlan against them. - Correct write.md and SKILL.md to say peak resident memory during a concat copy is roughly double the merged offsets array (merged + each input's own offsets array), not just the merged array alone. - Clarify the tests/conftest.py comment about the readbound_tracks fixture: not a duplicate of _SVAR2_VCF_2S (fourth variant), but is byte-identical to _SVAR2_SLOT_REF/_SVAR2_SLOT_VCF/_svar2_slot_src (fold tracked separately). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #403. Closes #406.
Why
gvl.concatplanned every merge by building aprovenancemap — one row permerged
(region, sample, ploid)slot — and thencoalesce-ing it into runs. Themap is
(R*S*P, 2)int64. At an All of Us chr22 grid (R = 3,734 regions,S = 535,662 samples, ploidy 2) that is 64 GB allocated purely to decide where
bytes go, before a single byte is copied.
What changed
RunPlanderives the same runs analytically from the sorted merge order, nevermaterializing the map. Planning memory is now O(R + S) instead of O(R·S·P).
The key property that makes this work: the run-break pattern depends only on
order, never on the region indexr. On the samples axis ther*S_dtermscancel, leaving
w_{j+1} == w_j + 1, so the break pattern is computed once inO(R + S) and reused for every region.
provenanceandcoalesceare not deleted. They are retained as theequivalence oracle and are now off the production path entirely — zero call sites
remain in
_concat.py.test_run_plan_matches_coalesce_provenance_exhaustivelysweeps 4,800 configurations (including 4-dataset merges and non-permutation
orders) asserting
RunPlanreproducescoalesce(provenance(...))exactly. Zeromismatches. Both functions' docstrings now say why they must not be deleted.
as_plannow raisesTypeErroron a one-shot generator instead of silentlylist()-materializing it — which would have rebuilt the very structure this PRexists to remove.
RunPlanis a re-iterable object, not a generator; consumersiterate it twice.
What this does not claim
The merged
offsets.npyis still one int64 per merged slot — about 16 GB forper-sample tracks, 32 GB for diploid genotypes at that grid. That is a property of
the ragged on-disk format, not of the merge. Peak resident during the copy is
roughly double that (32 / 64 GB), because the copy also holds each input's own
offsets array for the duration. The docs state this explicitly; no speedup is
claimed, only a memory-scaling fix to planning.
Fixture work (#406)
The
svar2_storefixture had one occupied cell in an 18-cell grid, which madeseveral sparse/dense parity tests structurally vacuous — S1 was indistinguishable
from the deliberately-empty S2 column, and
dense_snp_rangewas all zeros, so thedense SNP path had no coverage at all. The fixture now has 7 of 18 occupied,
measured rather than reasoned about: every sparse variant carries exactly one call,
which is what keeps genoray's cost model from routing it to the per-region dense
channel.
This made two things possible that were not before:
reinstated — a transposed sample axis is now detectable.
test_dense_layout_dataset_still_opens_and_readsis promoted from a smoke test toa real dense-vs-sparse parity pin. The original byte-comparison approach was found
to be structurally incapable of catching a corrupted region index (nested bed
regions plus position-based variant filtering self-heal it), so the test asserts on
raw ranges instead.
The duplicated
svar2_storefixture is also hoisted intotests/conftest.py.Testing
pixi run -e dev pytest tests -q— 1265 passed, 61 skipped, 4 xfailed, 0 failedcargo test— 138 passedpyrefly— 0 errors;ruff check/ruff formatclean across 270 filesapi.mdin sync with__all__(MISSING: none); docs build with no new warningsVersioning
No
feat:commits, so commitizen cuts a PATCH bump (0.43.0 → 0.43.1).DATASET_FORMAT_VERSIONis unchanged at2.0.0— the on-disk format does not move.Follow-ups filed
test_svar2_readbound_tracks's byte-identical fixture onto_svar2_slot_srcsvar2_store_2sto session scope (unblocked by this PR; 6 store builds to 1)np.uniqueout ofcopy_runs' per-batch loop🤖 Generated with Claude Code
https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6