perf(svar2): consume genoray's sparse range stream in the writer

Cargo.toml

@@ -22,18 +22,21 @@ seqpro-core = "0.1"
 # genoray crates are pulled straight from GitHub (no crates.io publish). Cargo finds
 # each package by name inside the repo — no in-repo path is given or needed. Bump `rev`
 # to pull newer genoray code; both crates must share the same rev (one clone, one repo).
-# Currently tag 4.0.1. Keep this rev in step with the `genoray` Python pin: the Rust
+# Currently tag 5.0.0. Keep this rev in step with the `genoray` Python pin: the Rust
 # code here reads svar2 stores that the Python package writes, so a format change that
 # lands in one and not the other is invisible until it corrupts a read.
 # `default-features = false` selects genoray's query-only core; since 4.0.1 that also
 # excludes `tracing`/`tracing-subscriber`, which are gated behind its `conversion`
 # feature (d-laub/genoray#165).
 # The sparse range cache's empty-cell contract also depends on this pin:
-# `gather_haps_readbound_impl` (genoray:src/query/gather.rs:772) must never
-# dereference an empty range's `start`, since that's what makes writing `(0, 0)`
-# for an absent cell safe rather than a corrupt read. Re-check this on any bump.
-svar2-codec = { git = "https://github.com/d-laub/genoray.git", rev = "d66ec0e03d097fa8c338567b0938924be67315db" }
-genoray_core = { git = "https://github.com/d-laub/genoray.git", rev = "d66ec0e03d097fa8c338567b0938924be67315db", package = "genoray", default-features = false }
+# `gather_haps_readbound_impl` (genoray:src/query/gather.rs:869 at 5.0.0) must
+# never dereference an empty range's `start`, since that's what makes writing
+# `(0, 0)` for an absent cell safe rather than a corrupt read. Re-check this on
+# any bump. Re-checked at 5.0.0: every range is consumed as a SLICE bound
+# (`d_snp_pos[ss..se]`, `snp_positions[vs..ve]`), so `len == 0` yields an empty
+# slice and the loop body never runs -- the `start` is never dereferenced.
+svar2-codec = { git = "https://github.com/d-laub/genoray.git", rev = "714f8761574b0d4e934d18c694dcaadacd946a2d" }
+genoray_core = { git = "https://github.com/d-laub/genoray.git", rev = "714f8761574b0d4e934d18c694dcaadacd946a2d", package = "genoray", default-features = false }
 
 [features]
 extension-module = ["pyo3/extension-module"]

pixi.toml

@@ -102,11 +102,13 @@ numba = "==0.59.1"
 pyarrow = ">=21"
 hirola = "==0.3"
 seqpro = "==0.22.0"
-# genoray >=3.4.0 as the prebuilt abi3 wheel from PyPI — one cp310-abi3 wheel covers
-# py310-313 on both platforms. 3.4.0 carries SparseVar2._find_ranges_chunked, the
-# memory-bounded chunked range API _write_from_svar2 consumes (gvl#333). Mirrors
-# the pyproject range, which spans 3.x and 4.x.
-genoray = ">=3.4.0,<5"
+# genoray >=5 as the prebuilt abi3 wheel from PyPI — one cp310-abi3 wheel covers
+# py310-313 on both platforms. 5.0.0 carries
+# SparseVar2._find_ranges_chunked_sparse, the sparse chunked range API
+# _write_from_svar2 consumes (gvl#405), and the keyword-only, struct-grouped
+# _core pipeline entry points the test fixtures call (d-laub/genoray#153, #200).
+# Mirrors the pyproject range.
+genoray = ">=5,<6"
 polars = "==1.37.1"
 loguru = "*"
 natsort = "*"

pyproject.toml

@@ -11,12 +11,13 @@ license = { file = "LICENSE.txt" }
 requires-python = ">=3.10,<3.14" # >= 3.14 blocked by pyarrow/genoray
 dependencies = [
     "seqpro>=0.22",
-    # >=3.4.0 carries SparseVar2._find_ranges_chunked, the memory-bounded chunked
-    # range API _write_from_svar2 consumes to avoid materializing a whole
-    # contig's ranges at once (gvl#333). 4.x is supported: its breaking changes
-    # (SparseVar2.from_vcf(tune=) removal, from_vcf_list(max_mem=) semantics)
-    # are on APIs GVL never calls.
-    "genoray>=3.4.0,<5",
+    # >=5 carries SparseVar2._find_ranges_chunked_sparse, the sparse chunked
+    # range API _write_from_svar2 consumes (gvl#405), and the keyword-only,
+    # struct-grouped _core pipeline entry points the test fixtures call
+    # (d-laub/genoray#153, #200). 5.0.0's breaking changes -- the GENORAY_*
+    # environment variables, SparseVar2.from_vcf(reader_workers=) and
+    # resolve_log_level -- are all on APIs GVL never calls.
+    "genoray>=5,<6",
     "numpy",
     "loguru",
     "natsort",

python/genvarloader/_dataset/_svar2_ranges.py

@@ -38,7 +38,6 @@
     "_DenseRanges",
     "_ranges_reader",
     "_SparseWriter",
-    "nonempty_entries",
     "merge_region_blocks",
 ]
 
@@ -715,7 +714,8 @@ def append_contig(
     ) -> None:
         """Merge one contig's per-chunk blocks into region-major order and append.
 
-        Each block from :func:`nonempty_entries` is already region-major, and
+        Each block from genoray's sparse range stream is already region-major,
+        and
         chunk ``i``'s sample slots lie entirely below chunk ``i + 1``'s, so the
         merged order is fixed by region alone. That makes this a stable counting
         sort with ``O(rc)`` of auxiliary state, not a comparison sort.
@@ -869,72 +869,6 @@ def close(self) -> int:
         return self.n_entries
 
 
-def nonempty_entries(
-    snp: NDArray[np.int64], indel: NDArray[np.int64], slot0: int, ploidy: int
-) -> tuple[NDArray[np.int32], NDArray[np.int32], NDArray[np.void]]:
-    """Extract non-empty cells from a ``(rc, ns, P, 2)`` pair of range blocks.
-
-    Args:
-        snp: SNP ranges, ``(rc, ns, P, 2)`` -- normally a ``transpose(2, 0, 1, 3)``
-            view of a hap-major genoray chunk.
-        indel: Indel ranges, same shape.
-        slot0: Dataset sample slot of this block's first column.
-        ploidy: ``P``.
-
-    Returns:
-        ``(region, cell, entries)``, region-major: ``region`` is **contig-local**
-        and non-decreasing, ``cell`` is ``slot * ploidy + ploid`` and ascends
-        within each region. Split rather than combined into one key because
-        :meth:`_SparseWriter.append_contig` needs the region axis on its own to
-        count, and ``cell`` is what lands on disk -- combining them would only be
-        undone again.
-
-    Raises:
-        ValueError: If ``snp`` and ``indel`` don't share a shape, or their
-            ploidy axis doesn't match ``ploidy``. A caller that transposes the
-            wrong axes (e.g. swapping the region and sample axes) still
-            produces a same-rank ``(a, b, c, 2)`` array, so this is checked
-            explicitly rather than left to fail downstream -- without it, a
-            mis-transposed cache still writes a self-consistent CSR table
-            with no invariant violated, just region/sample-scrambled entries.
-    """
-    if snp.shape != indel.shape:
-        raise ValueError(
-            "svar2 range cache: snp and indel blocks must share a shape, got"
-            f" {snp.shape} and {indel.shape}"
-        )
-    if snp.ndim != 4 or snp.shape[2] != ploidy:
-        raise ValueError(
-            f"svar2 range cache: expected (regions, samples, ploidy={ploidy}, 2)"
-            f" blocks, got shape {snp.shape}"
-        )
-    ne = (snp[..., 1] > snp[..., 0]) | (indel[..., 1] > indel[..., 0])
-    # np.nonzero walks the LOGICAL shape in C order, so (r, slot, ploid) comes
-    # out ascending even though `ne` is NOT C-contiguous: the `>` above inherits
-    # the transposed view's stride permutation, because numpy allocates ufunc
-    # output with NPY_KEEPORDER. Do not "fix" that with ascontiguousarray --
-    # materializing (rc, ns, P) in C order is a strided scatter costing ~11x the
-    # comparison itself (97.7 ms vs 8.8 ms on a 15e6-cell chunk).
-    ri, sj, pj = np.nonzero(ne)
-    ent = np.empty(len(ri), ENTRY_DTYPE)
-    ent["snp_start"] = snp[ri, sj, pj, 0]
-    ent["snp_len"] = snp[ri, sj, pj, 1] - snp[ri, sj, pj, 0]
-    ent["indel_start"] = indel[ri, sj, pj, 0]
-    ent["indel_len"] = indel[ri, sj, pj, 1] - indel[ri, sj, pj, 0]
-    # Overflow guard local to this function: the sole production call site
-    # (`_write_from_svar2`) already guards `n_samples * ploidy < 2**31` by
-    # constructing `_SparseWriter` first, but `nonempty_entries` is a public
-    # module-level function callable independently of that guard.
-    max_cell = (int(slot0) + int(snp.shape[1]) - 1) * int(ploidy) + int(ploidy) - 1
-    if max_cell > np.iinfo(np.int32).max:
-        raise ValueError(
-            "svar2 range cache: slot0 * ploidy overflows int32"
-            f" (max cell id {max_cell})"
-        )
-    cell = (slot0 + sj).astype(np.int64) * ploidy + pj
-    return ri.astype(np.int32), cell.astype(np.int32), ent
-
-
 def merge_region_blocks(
     readers: "list[_RangeLookup]",
     r_maps: "list[NDArray[np.int64]]",

python/genvarloader/_dataset/_write.py

@@ -44,7 +44,7 @@
 from .._variants._utils import path_is_pgen, path_is_vcf
 from ._indexing import s2i
 from ._svar2_link import Svar2Link
-from ._svar2_ranges import ENTRY_DTYPE, _SparseWriter, nonempty_entries
+from ._svar2_ranges import ENTRY_DTYPE, _SparseWriter
 from ._svar_link import SvarLink
 from ._utils import bed_to_regions, regions_to_bed
 
@@ -1267,7 +1267,7 @@ def _write_from_svar2(
             ends = df["chromEnd"].to_numpy()
             # extend_to_length is validated at function entry (False raises); the
             # read-bound kernel sizes haplotype output at read time.
-            stream = svar2._find_ranges_chunked(
+            stream = svar2._find_ranges_chunked_sparse(
                 c, starts, ends, samples=sel, max_mem=max_mem
             )
             dense_snp[lo:hi] = np.asarray(stream.dense_snp_range, np.int64).reshape(
@@ -1290,7 +1290,7 @@ def _write_from_svar2(
                 # append_contig's counting-sort merge assumes chunk i's sample
                 # slots lie entirely below chunk i + 1's (see its docstring): the
                 # merged order is fixed by region alone only because of that.
-                # genoray's `_find_ranges_chunked` happens to yield ascending
+                # genoray's `_find_ranges_chunked_sparse` happens to yield ascending
                 # `sample_start` today, but that is a generator's behaviour in a
                 # separate package, asserted nowhere on either side -- an
                 # out-of-order chunk stream would corrupt the merge silently
@@ -1308,26 +1308,36 @@ def _write_from_svar2(
                         f" {prev_sample_start}."
                     )
                 prev_sample_start = ch.sample_start
-                # Chunks are hap-major (samples, ploidy, regions, 2); transpose to
-                # region-major (regions, samples, ploidy, 2). transpose() is a
-                # view, and nonempty_entries relies on that -- see its comment on
-                # np.nonzero and NPY_KEEPORDER.
-                r, cell, ent = nonempty_entries(
-                    ch.vk_snp_range.transpose(2, 0, 1, 3),
-                    ch.vk_indel_range.transpose(2, 0, 1, 3),
-                    slot0=ch.sample_start,
-                    ploidy=P,
-                )
-                # `nonempty_entries` filters on width, not on genoray's raw start ==
-                # end insertion point (which is what a *dense* cell would carry at
-                # the same coordinate). An all-empty (region, sample, ploid) cell
-                # is therefore never written here, so a real dataset's
+                # The chunk already IS the non-empty set, region-major with
+                # ascending cell_id -- the shape append_contig wants. What used
+                # to stand here was a transpose plus an `np.nonzero` scan over a
+                # dense (samples, ploidy, regions, 2) block, 71% of this kernel
+                # and ~128 GB per All of Us chr22 contig (#405). genoray filters
+                # on the same predicate (`end > start` in either channel) and
+                # emits the same absolute `slot * ploidy + ploid`, so the bytes
+                # written here are unchanged.
+                #
+                # CSR in, region column out: append_contig counts on the region
+                # axis, so expand `region_ptr` rather than make genoray send a
+                # column it would have to build from the same offsets.
+                r = np.repeat(np.arange(rc, dtype=np.int32), np.diff(ch.region_ptr))
+                ent = np.empty(len(ch.cell_id), ENTRY_DTYPE)
+                ent["snp_start"] = ch.snp_start
+                ent["snp_len"] = ch.snp_len
+                ent["indel_start"] = ch.indel_start
+                ent["indel_len"] = ch.indel_len
+                # genoray filters on width, not on its raw start == end insertion
+                # point (which is what a *dense* cell would carry at the same
+                # coordinate). An all-empty (region, sample, ploid) cell is
+                # therefore never written here, so a real dataset's
                 # `_SparseRanges.lookup` always returns (0, 0) for it, never
                 # genoray's insertion point -- unlike `_DenseRanges.lookup`, which
                 # would surface (x, x). That is the one place the two layouts are
-                # not byte-identical (see `_SparseRanges`'s docstring).
+                # not byte-identical (see `_SparseRanges`'s docstring). A cell
+                # non-empty in only ONE channel still carries the other channel's
+                # raw start with length 0, exactly as the dense path produced it.
                 acc_r.append(r)
-                acc_c.append(cell)
+                acc_c.append(ch.cell_id)
                 acc_e.append(ent)
                 np.maximum(keys, ch.max_end_keys, out=keys)
                 pbar.update(rc * ch.n_samples / S)

tests/conftest.py

@@ -185,19 +185,22 @@ def svar2_slot_store(_svar2_slot_src, tmp_path_factory) -> Path:
     leaks into every later consumer of this fixture.
     """
     from genoray import _core
+    from genoray._pipeline_args import FieldSpec, PlanSettings, RegionSpec
 
     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,
+        vcf_path=str(bcf),
+        reference_path=str(ref),
+        output_dir=str(store),
+        regions=RegionSpec(chroms=["chr1"], samples=["S0", "S1"]),
+        fields=FieldSpec(),
+        plan=PlanSettings(
+            chunk_size=25_000,
+            max_threads=1,
+            long_allele_capacity=8 * 1024 * 1024,
+        ),
+        ploidy=2,
     )
     assert (store / "meta.json").exists(), "svar2 conversion did not finish"
     return store
@@ -312,6 +315,7 @@ def _build_svar2(vcf_text: str, samples: list[str], d: Path, name: str) -> Path:
         Path to the finished ``.svar2`` store.
     """
     from genoray import _core
+    from genoray._pipeline_args import FieldSpec, PlanSettings, RegionSpec
 
     ref = d / "ref.fa"
     ref.write_text(f">chr1\n{_SVAR2_REF}\n")
@@ -325,15 +329,17 @@ def _build_svar2(vcf_text: str, samples: list[str], d: Path, name: str) -> Path:
 
     out = d / name
     _core.run_conversion_pipeline(
-        str(bcf),
-        str(ref),
-        ["chr1"],
-        str(out),
-        samples,
-        25_000,
-        2,
-        1,
-        8 * 1024 * 1024,
+        vcf_path=str(bcf),
+        reference_path=str(ref),
+        output_dir=str(out),
+        regions=RegionSpec(chroms=["chr1"], samples=samples),
+        fields=FieldSpec(),
+        plan=PlanSettings(
+            chunk_size=25_000,
+            max_threads=1,
+            long_allele_capacity=8 * 1024 * 1024,
+        ),
+        ploidy=2,
     )
     assert (out / "meta.json").exists(), "conversion did not finish"
     return out

tests/dataset/conftest.py

@@ -63,18 +63,21 @@ def vcf_and_ref(tmp_path_factory) -> tuple[Path, Path]:
 def svar2_store(vcf_and_ref, tmp_path_factory) -> Path:
     bcf, ref = vcf_and_ref
     from genoray import _core
+    from genoray._pipeline_args import FieldSpec, PlanSettings, RegionSpec
 
     out = tmp_path_factory.mktemp("svar2_write") / "store.svar2"
     _core.run_conversion_pipeline(
-        str(bcf),
-        str(ref),
-        ["chr1"],
-        str(out),
-        ["S0", "S1", "S2"],
-        25_000,
-        2,
-        1,
-        8 * 1024 * 1024,
+        vcf_path=str(bcf),
+        reference_path=str(ref),
+        output_dir=str(out),
+        regions=RegionSpec(chroms=["chr1"], samples=["S0", "S1", "S2"]),
+        fields=FieldSpec(),
+        plan=PlanSettings(
+            chunk_size=25_000,
+            max_threads=1,
+            long_allele_capacity=8 * 1024 * 1024,
+        ),
+        ploidy=2,
     )
     assert (out / "meta.json").exists(), "conversion did not finish"
     return out

tests/dataset/test_svar2_contig_naming.py

@@ -40,6 +40,7 @@
 def ensembl_svar2_store(tmp_path_factory) -> Path:
     """A ``.svar2`` store whose only contig is spelled ``1``."""
     from genoray import _core
+    from genoray._pipeline_args import FieldSpec, PlanSettings, RegionSpec
 
     d = tmp_path_factory.mktemp("svar2_contig_naming")
     ref = d / "ref.fa"
@@ -54,15 +55,17 @@ def ensembl_svar2_store(tmp_path_factory) -> Path:
 
     out = d / "store.svar2"
     _core.run_conversion_pipeline(
-        str(bcf),
-        str(ref),
-        ["1"],
-        str(out),
-        ["S0", "S1"],
-        25_000,
-        2,
-        1,
-        8 * 1024 * 1024,
+        vcf_path=str(bcf),
+        reference_path=str(ref),
+        output_dir=str(out),
+        regions=RegionSpec(chroms=["1"], samples=["S0", "S1"]),
+        fields=FieldSpec(),
+        plan=PlanSettings(
+            chunk_size=25_000,
+            max_threads=1,
+            long_allele_capacity=8 * 1024 * 1024,
+        ),
+        ploidy=2,
     )
     assert (out / "meta.json").exists(), "conversion did not finish"
     return out