refactor(concat,svar2): clear the three follow-ups from the analytic run planning review

python/genvarloader/_dataset/_concat_io.py

@@ -80,10 +80,17 @@ def copy_runs(
     merged = np.empty(n_slots + 1, np.int64)
     merged[0] = 0
     lengths = merged[1:]
-    for dst_start, ds_vec, slots in plan.slot_batches():
+    for dst_start, ds_vec, slots, ds_uniq in plan.slot_batches():
         n = len(slots)
         out = lengths[dst_start : dst_start + n]
-        for d in np.unique(ds_vec):
+        if len(ds_uniq) == 1:
+            # Single-sourced batch -- every region-axis batch and every
+            # ExplicitRunPlan batch. Skip the mask entirely: it would be all-True
+            # over as many as _SLOT_BATCH_SLOTS elements.
+            off = src_offsets[int(ds_uniq[0])]
+            out[:] = off[slots + 1] - off[slots]
+            continue
+        for d in ds_uniq:
             m = ds_vec == d
             off = src_offsets[int(d)]
             sl = slots[m]

python/genvarloader/_dataset/_concat_plan.py

@@ -282,27 +282,39 @@ def n_slots(self) -> int:
 
     def slot_batches(
         self,
-    ) -> "Iterator[tuple[int, NDArray[np.int64], NDArray[np.int64]]]":
-        """Yield ``(dst_start, src_ds, src_slots)`` batches in destination order.
+    ) -> (
+        "Iterator[tuple[int, NDArray[np.int64], NDArray[np.int64], NDArray[np.int64]]]"
+    ):
+        """Yield ``(dst_start, src_ds, src_slots, src_ds_uniq)`` 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.
+        Concatenating every batch's first three fields in order rebuilds
+        :func:`provenance`'s output exactly, which is what pins this method.
+
+        ``src_ds_uniq`` is the sorted distinct values of ``src_ds``. It is
+        yielded rather than left to the consumer because this method already
+        knows it for free: on the regions axis a batch is single-sourced, and on
+        the samples axis ``src_ds`` is loop-invariant, so a consumer calling
+        ``np.unique`` per batch would re-sort the identical array once per
+        region -- 3,734 sorts of a 1.07e6-element array at the chr22 grid.
 
         Yields:
-            ``(dst_start, src_ds, src_slots)``, where the two arrays are int64
-            and equal in length.
+            ``(dst_start, src_ds, src_slots, src_ds_uniq)``. The first three are
+            int64, with ``src_ds`` and ``src_slots`` equal in length.
         """
         if self.axis == "regions":
             for run in self:
+                # Every batch of a region run is single-sourced.
+                uniq = np.array([run.src], np.int64)
                 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),
+                        uniq,
                     )
                     pos += n
                     dst += n
@@ -318,10 +330,13 @@ def slot_batches(
         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)
+        # Loop-invariant: the merged sample axis, hence its source-dataset
+        # grouping, is identical for every region.
+        ds_uniq = np.unique(ds_vec)
         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)
+            yield (r * n_merged * self.ploidy, ds_vec, slots, ds_uniq)
 
     def _segments(self) -> tuple[NDArray[np.int64], NDArray[np.int64]]:
         """Half-open ``[start, stop)`` spans of ``order`` with no break inside."""
@@ -413,14 +428,21 @@ def n_slots(self) -> int:
 
     def slot_batches(
         self,
-    ) -> "Iterator[tuple[int, NDArray[np.int64], NDArray[np.int64]]]":
-        """Yield one ``(dst_start, src_ds, src_slots)`` batch per run."""
+    ) -> (
+        "Iterator[tuple[int, NDArray[np.int64], NDArray[np.int64], NDArray[np.int64]]]"
+    ):
+        """Yield one ``(dst_start, src_ds, src_slots, src_ds_uniq)`` batch per run.
+
+        A run is drawn from exactly one source dataset, so ``src_ds_uniq`` is
+        always a single element here.
+        """
         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),
+                np.array([r.src], np.int64),
             )
 
 

tests/conftest.py

@@ -176,6 +176,33 @@ def _svar2_slot_src(tmp_path_factory) -> tuple[Path, Path]:
     return bcf, ref
 
 
+@pytest.fixture(scope="session")
+def svar2_slot_store(_svar2_slot_src, tmp_path_factory) -> Path:
+    """The ``.svar2`` store built from `_svar2_slot_src`.
+
+    Shared and session-scoped, so consumers must treat it as READ-ONLY: copy it
+    into ``tmp_path`` before mutating anything inside it, or the corruption
+    leaks into every later consumer of this fixture.
+    """
+    from genoray import _core
+
+    bcf, ref = _svar2_slot_src
+    store = tmp_path_factory.mktemp("svar2_slot_store") / "store.svar2"
+    _core.run_conversion_pipeline(
+        str(bcf),
+        str(ref),
+        ["chr1"],
+        str(store),
+        ["S0", "S1"],
+        25_000,
+        2,
+        1,
+        8 * 1024 * 1024,
+    )
+    assert (store / "meta.json").exists(), "svar2 conversion did not finish"
+    return store
+
+
 @pytest.fixture(scope="session")
 def svar2_slot_reference(_svar2_slot_src):
     """Opened ``gvl.Reference`` matching `phased_svar2_gvl`'s tiny chr1.
@@ -192,7 +219,7 @@ def svar2_slot_reference(_svar2_slot_src):
 
 
 @pytest.fixture(scope="session")
-def phased_svar2_gvl(_svar2_slot_src, tmp_path_factory) -> Path:
+def phased_svar2_gvl(svar2_slot_store, tmp_path_factory) -> Path:
     """A gvl dataset written from a ``.svar2`` source -> opens as ``Svar2Haps``.
 
     Reproduces the Phase-0-pinned record class for #315 (see
@@ -215,24 +242,11 @@ def phased_svar2_gvl(_svar2_slot_src, tmp_path_factory) -> Path:
     empirically before picking 80 (green up to ~40 regions, red by 60+).
     """
     import polars as pl
-    from genoray import SparseVar2, _core
+    from genoray import SparseVar2
 
     import genvarloader as gvl
 
-    bcf, ref = _svar2_slot_src
-    store = tmp_path_factory.mktemp("svar2_slot_store") / "store.svar2"
-    _core.run_conversion_pipeline(
-        str(bcf),
-        str(ref),
-        ["chr1"],
-        str(store),
-        ["S0", "S1"],
-        25_000,
-        2,
-        1,
-        8 * 1024 * 1024,
-    )
-    assert (store / "meta.json").exists(), "svar2 conversion did not finish"
+    store = svar2_slot_store
 
     n_regions = 80
     bed = pl.DataFrame(
@@ -325,14 +339,22 @@ def _build_svar2(vcf_text: str, samples: list[str], d: Path, name: str) -> Path:
     return out
 
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="session")
 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.
+
+    Session-scoped, and therefore SHARED: treat it as READ-ONLY. A test that
+    needs to mutate the store must ``shutil.copytree`` it into ``tmp_path`` and
+    mutate the copy. Mutating it in place corrupts every later consumer, and the
+    resulting failures surface in other modules -- when
+    ``test_fingerprint_detects_mutated_store`` appended one byte to a ``.bin``
+    here, it produced 17 failures plus a Rust ``cast_slice`` panic. That is why
+    this fixture was module-scoped until the mutating test was fixed to copy.
     """
     d = tmp_path_factory.mktemp("svar2_2s")
     return _build_svar2(_SVAR2_VCF_2S, ["S0", "S1"], d, "store")

tests/dataset/test_svar2_readbound_tracks.py

@@ -12,58 +12,21 @@
 
 from __future__ import annotations
 
-import subprocess
-from pathlib import Path
 
 import numpy as np
 import pytest
 
-# 40 bp reference (chr1). VCF POS (1-based) -> 0-based: SNP@2 (A>G, low-carrier,
-# routes to var_key), INS@6 (C>CAT), SNP@9 (G>C, carried by 3 haps -> dense/snp
-# per the cost model used in test_svar2_readbound_haps.py), DEL@11 (GTA>G,
-# ilen -2). Exercises both var_key and dense/snp + dense/indel channels.
-_REF = "ACAGTACATGGGTACTAGCTAGGCTAACCGGTTAACCGGT"
-_VCF = """\
-##fileformat=VCFv4.2
-##contig=<ID=chr1,length=40>
-##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">
-#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\t10\t.\tG\tC\t.\t.\t.\tGT\t1|1\t1|0
-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_tracks")
-    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
+# The store comes from the shared session-scoped `svar2_slot_store` fixture in
+# tests/conftest.py. Its reference and VCF used to be duplicated verbatim here:
+# SNP@2 (A>G, low-carrier, routes to var_key), INS@6 (C>CAT), SNP@9 (G>C,
+# carried by 3 haps -> dense/snp per the cost model used in
+# test_svar2_readbound_haps.py), DEL@11 (GTA>G, ilen -2) -- so both var_key and
+# the dense/snp + dense/indel channels are exercised. The fixture is shared and
+# must be treated as read-only.
+#
+# Used under its own name rather than aliased: this module lives under
+# tests/dataset/, where tests/dataset/conftest.py defines a DIFFERENT 3-sample
+# fixture named `svar2_store`, and an alias would silently bind to that instead.
 
 
 def _synthetic_track_inputs(regions, seed=0):
@@ -87,14 +50,14 @@ def _synthetic_track_inputs(regions, seed=0):
         [(0, 40), (2, 2), (20, 25)],  # empty region + a variant-free window
     ],
 )
-def test_readbound_tracks_match_union_oracle(svar2_store, regions):
+def test_readbound_tracks_match_union_oracle(svar2_slot_store, regions):
     import genoray
 
     from tests._oracles.svar2_source import SparseVar2Source
     from tests._oracles.svar2_readbound_inputs import build_readbound_tracks
 
     contig = "chr1"
-    sv = genoray.SparseVar2(str(svar2_store))
+    sv = genoray.SparseVar2(str(svar2_slot_store))
     S, P = sv.n_samples, sv.ploidy
     assert (S, P) == (2, 2)
 
@@ -155,7 +118,7 @@ def test_readbound_tracks_match_union_oracle(svar2_store, regions):
         pytest.fail("data mismatch but no single hap slice differed (offset bug?)")
 
 
-def test_readbound_tracks_match_union_oracle_with_shifts(svar2_store):
+def test_readbound_tracks_match_union_oracle_with_shifts(svar2_slot_store):
     """Non-trivial per-hap jitter shifts must also match byte-for-byte."""
     import genoray
 
@@ -164,7 +127,7 @@ def test_readbound_tracks_match_union_oracle_with_shifts(svar2_store):
 
     contig = "chr1"
     regions = [(0, 40), (5, 20)]
-    sv = genoray.SparseVar2(str(svar2_store))
+    sv = genoray.SparseVar2(str(svar2_slot_store))
     S, P = sv.n_samples, sv.ploidy
     n_q = len(regions) * S
 

tests/unit/dataset/test_concat_plan.py

@@ -257,9 +257,12 @@ def test_slot_batches_reproduce_the_provenance_map():
             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():
+            for dst_start, ds_vec, slots, ds_uniq in plan.slot_batches():
                 n = len(slots)
                 assert len(ds_vec) == n
+                # The yielded uniques must be exactly what a consumer would
+                # have computed itself; copy_runs trusts them to pick files.
+                np.testing.assert_array_equal(ds_uniq, np.unique(ds_vec))
                 got[dst_start : dst_start + n, 0] = ds_vec
                 got[dst_start : dst_start + n, 1] = slots
                 seen[dst_start : dst_start + n] = True
@@ -295,12 +298,12 @@ def test_slot_batches_chunk_large_region_runs(monkeypatch):
     assert len(list(plan)) == 1
 
     batches = list(plan.slot_batches())
-    assert [len(slots) for _, _, slots in batches] == [5, 5, 5, 5, 4]
+    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]
+    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)
+        np.concatenate([slots for _, _, slots, _ in batches]), np.arange(24)
     )
 
 
@@ -314,6 +317,9 @@ def test_explicit_run_plan_round_trips_a_hand_built_list():
     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])
+    # A run is single-sourced, so each batch reports exactly one dataset.
+    np.testing.assert_array_equal(batches[0][3], [0])
+    np.testing.assert_array_equal(batches[1][3], [1])
 
 
 def test_as_plan_rejects_a_generator():

tests/unit/dataset/test_svar2_link.py

@@ -44,8 +44,9 @@ def test_verify_none_link_is_noop(tmp_path: Path):
 
 
 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.
+    # `svar2_store_2s` is a session-scoped fixture shared with five other test
+    # modules -- mutate a private copy, never the shared store itself. Mutating
+    # it in place corrupts every later consumer across the whole session.
     store = shutil.copytree(svar2_store_2s, tmp_path / "store.svar2")
 
     gvl_path = tmp_path / "ds.gvl"