Skip to content

perf(concat): plan merges analytically instead of materializing an R x S x P provenance map - #414

Merged
d-laub merged 17 commits into
mainfrom
feat/concat-analytic-run-planning
Sep 16, 2026
Merged

d-laub merged 17 commits into
mainfrom
feat/concat-analytic-run-planning

Conversation

@d-laub

@d-laub d-laub commented Sep 16, 2026

Copy link
Copy Markdown
Collaborator

Closes #403. Closes #406.

Why

gvl.concat planned every merge by building a provenance map — one row per
merged (region, sample, ploid) slot — and then coalesce-ing it into runs. The
map 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

RunPlan derives the same runs analytically from the sorted merge order, never
materializing 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 index r. On the samples axis the r*S_d terms
cancel, leaving w_{j+1} == w_j + 1, so the break pattern is computed once in
O(R + S) and reused for every region.

provenance and coalesce are not deleted. They are retained as the
equivalence oracle and are now off the production path entirely — zero call sites
remain in _concat.py. test_run_plan_matches_coalesce_provenance_exhaustively
sweeps 4,800 configurations (including 4-dataset merges and non-permutation
orders) asserting RunPlan reproduces coalesce(provenance(...)) exactly. Zero
mismatches. Both functions' docstrings now say why they must not be deleted.

as_plan now raises TypeError on a one-shot generator instead of silently
list()-materializing it — which would have rebuilt the very structure this PR
exists to remove. RunPlan is a re-iterable object, not a generator; consumers
iterate it twice.

What this does not claim

The merged offsets.npy is still one int64 per merged slot — about 16 GB for
per-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_store fixture had one occupied cell in an 18-cell grid, which made
several sparse/dense parity tests structurally vacuous — S1 was indistinguishable
from the deliberately-empty S2 column, and dense_snp_range was all zeros, so the
dense 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:

The duplicated svar2_store fixture is also hoisted into tests/conftest.py.

Testing

  • pixi run -e dev pytest tests -q — 1265 passed, 61 skipped, 4 xfailed, 0 failed
  • cargo test — 138 passed
  • pyrefly — 0 errors; ruff check / ruff format clean across 270 files
  • api.md in sync with __all__ (MISSING: none); docs build with no new warnings

Versioning

No feat: commits, so commitizen cuts a PATCH bump (0.43.0 → 0.43.1).
DATASET_FORMAT_VERSION is unchanged at 2.0.0 — the on-disk format does not move.

Follow-ups filed

🤖 Generated with Claude Code

https://claude.ai/code/session_015VxRqNngU7Eg1wdb1aEgD6

d-laub and others added 17 commits September 15, 2026 14:18
…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
@d-laub
d-laub merged commit e742aec into main Sep 16, 2026
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

svar2_store fixture has one non-empty cell in 18, weakening the sparse suite concat still plans per-sample tracks in core (R x S)

1 participant