fix(svar2): normalize contig names on the read path (#336)

python/genvarloader/_dataset/_svar2_haps.py

@@ -31,6 +31,7 @@
 from typing import TYPE_CHECKING, Any, Literal, cast
 
 import numpy as np
+from genoray._contigs import ContigNormalizer
 from genoray._types import POS_TYPE, V_IDX_TYPE
 from numpy.typing import NDArray
 from seqpro.rag import Ragged
@@ -200,6 +201,15 @@ class Svar2Haps(Haps[_H]):
     Populated from ``SparseVar2.available_fields``. These keys are additionally
     advertised in ``available_var_fields`` so users can request them via ``var_fields``.
     """
+    _ds_to_store: list[str | None] = field(init=False)
+    """``ds_contigs`` mapped to the store's spelling, parallel to ``ds_contigs``.
+
+    The dataset and its .svar2 store may name contigs differently (e.g. a UCSC-named
+    ``chr1`` reference/GTF over an Ensembl-named ``1`` store) — the write path
+    reconciles the two, so the read path must too. ``None`` marks a dataset contig
+    with no equivalent in the store; that is only an error if such a contig is
+    actually queried, so it is raised lazily by :meth:`_store_contig`.
+    """
 
     def __post_init__(self):
         # Deliberately does NOT call Haps.__post_init__ (that reads an SVAR1
@@ -209,6 +219,22 @@ def __post_init__(self):
         self.available_var_fields = ["alt", "ilen", "start"] + [
             k for k in self.store_fields if k not in _BUILTIN_VAR_FIELDS
         ]
+        # Resolve once per contig at open, not once per query.
+        self._ds_to_store = ContigNormalizer(self.store_contigs).norm(self.ds_contigs)
+
+    def _store_contig(self, ci: int) -> str:
+        """Return the .svar2 store's spelling of dataset contig index ``ci``.
+
+        The read-bound Rust kernels look contigs up in the store's own keyspace, so
+        they must be handed the store's spelling rather than the dataset's.
+        """
+        contig = self._ds_to_store[ci]
+        if contig is None:
+            raise ValueError(
+                f"Dataset contig {self.ds_contigs[ci]!r} has no equivalent in the"
+                f" .svar2 store, whose contigs are {self.store_contigs}."
+            )
+        return contig
 
     # ---- construction ----
 
@@ -462,7 +488,7 @@ def _reconstruct_spliced(
                 out,
                 g_bounds,
                 self.store,
-                self.ds_contigs[ci],
+                self._store_contig(ci),
                 gi[0],
                 gi[1],
                 gi[2],
@@ -509,7 +535,7 @@ def _haplotype_diffs(
             gi = self._gather_inputs(r_q[qsel], si_q[qsel], regions[qsel], ploidy)
             d = hap_diffs_from_svar2_readbound(
                 self.store,
-                self.ds_contigs[ci],
+                self._store_contig(ci),
                 gi[0],
                 gi[1],
                 gi[2],
@@ -614,7 +640,7 @@ def get_haps_and_shifts(
                 g_total = int(hap_lengths[qsel].sum())
             g_data, g_off = reconstruct_haplotypes_from_svar2_readbound(
                 self.store,
-                self.ds_contigs[ci],
+                self._store_contig(ci),
                 gi[0],
                 gi[1],
                 gi[2],
@@ -718,7 +744,7 @@ def realign_track_block(
             g_total = int(track_ofsts_g[-1])
             out_data, out_off = shift_and_realign_tracks_from_svar2_readbound(
                 self.store,
-                self.ds_contigs[ci],
+                self._store_contig(ci),
                 gi[0],
                 gi[1],
                 gi[2],
@@ -855,7 +881,7 @@ def measure_variant_payload(
             pos, ilen, alt_bytes, str_off, var_off, field_bufs, field_isizes = (
                 decode_variants_from_svar2_readbound(
                     self.store,
-                    self.ds_contigs[ci],
+                    self._store_contig(ci),
                     gi[0],
                     gi[1],
                     gi[2],
@@ -947,7 +973,7 @@ def _reconstruct_variants(
             pos, ilen, alt_bytes, str_off, var_off, field_bufs, field_isizes = (
                 decode_variants_from_svar2_readbound(
                     self.store,
-                    self.ds_contigs[ci],
+                    self._store_contig(ci),
                     gi[0],
                     gi[1],
                     gi[2],
@@ -1092,7 +1118,7 @@ def _reconstruct_variant_windows(
             pos, ilen, alt_bytes, str_off, var_off, field_bufs, field_isizes = (
                 decode_variants_from_svar2_readbound(
                     self.store,
-                    self.ds_contigs[ci],
+                    self._store_contig(ci),
                     gi[0],
                     gi[1],
                     gi[2],

tests/dataset/test_svar2_contig_naming.py

@@ -0,0 +1,112 @@
+"""Regression test for #336: SVAR2 read path must normalize contig naming.
+
+A dataset's contigs come from its BED (``_write._prep_bed``: ``natsorted(bed["chrom"]
+.unique())``), while a ``.svar2`` store carries whatever spelling its source VCF used.
+The two are allowed to differ — a UCSC-named (``chr1``) BED over an Ensembl-named
+(``1``) store is a supported arrangement, and ``gvl.write`` reconciles them.
+
+Before the fix, only the *write* path normalized: ``Svar2Haps`` passed
+``ds_contigs[ci]`` straight into the read-bound Rust kernels, which look contigs up
+in the store's own keyspace. So such a dataset built successfully and then raised
+``ValueError: contig chr1 not in store`` on the first read.
+"""
+
+from __future__ import annotations
+
+import subprocess
+from dataclasses import replace
+from pathlib import Path
+
+import polars as pl
+import pytest
+
+import genvarloader as gvl
+
+# Same 40 bp reference and variants as the other svar2 fixtures, but named "1"
+# (Ensembl) rather than "chr1" (UCSC), so the store's keyspace differs from the BED's.
+_REF = "ACAGTACATGGGTACTAGCTAGGCTAACCGGTTAACCGGT"
+_VCF = """\
+##fileformat=VCFv4.2
+##contig=<ID=1,length=40>
+##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">
+#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tS0\tS1
+1\t3\t.\tA\tG\t.\t.\t.\tGT\t1|0\t0|0
+1\t7\t.\tC\tCAT\t.\t.\t.\tGT\t0|1\t1|1
+1\t12\t.\tGTA\tG\t.\t.\t.\tGT\t1|1\t0|1
+"""
+
+
+@pytest.fixture(scope="module")
+def ensembl_svar2_store(tmp_path_factory) -> Path:
+    """A ``.svar2`` store whose only contig is spelled ``1``."""
+    from genoray import _core
+
+    d = tmp_path_factory.mktemp("svar2_contig_naming")
+    ref = d / "ref.fa"
+    ref.write_text(f">1\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),
+        ["1"],
+        str(out),
+        ["S0", "S1"],
+        25_000,
+        2,
+        1,
+        8 * 1024 * 1024,
+    )
+    assert (out / "meta.json").exists(), "conversion did not finish"
+    return out
+
+
+@pytest.fixture(scope="module")
+def chr_named_dataset(ensembl_svar2_store: Path, tmp_path_factory) -> Path:
+    """A dataset whose regions are ``chr1`` over the ``1``-named store."""
+    from genoray import SparseVar2
+
+    bed = pl.DataFrame(
+        {"chrom": ["chr1", "chr1"], "chromStart": [0, 5], "chromEnd": [20, 15]}
+    )
+    out = tmp_path_factory.mktemp("svar2_contig_naming_ds") / "ds.gvl"
+    # The write path already reconciles chr1 <-> 1; this must succeed.
+    gvl.write(
+        out, bed, variants=SparseVar2(ensembl_svar2_store), samples=None, overwrite=True
+    )
+    return out
+
+
+def test_svar2_read_normalizes_contig_naming(chr_named_dataset: Path):
+    """A chr-named dataset over an Ensembl-named .svar2 store must be readable."""
+    ds = gvl.Dataset.open(chr_named_dataset).with_seqs("variants")
+
+    # The premise: the dataset and its store really do disagree on spelling.
+    assert ds._seqs.ds_contigs == ["chr1"]
+    assert ds._seqs.store_contigs == ["1"]
+
+    # Raised `ValueError: contig chr1 not in store` before the fix.
+    assert ds[0] is not None
+
+
+def test_svar2_unmapped_contig_names_both_spellings(chr_named_dataset: Path):
+    """A dataset contig with no store equivalent raises, naming both spellings.
+
+    Resolution is lazy: an unmappable contig is only an error when that contig is
+    actually queried, so building the resolution table at open must not itself raise.
+    """
+    ds = gvl.Dataset.open(chr_named_dataset).with_seqs("variants")
+
+    # replace() re-runs __post_init__, so the resolution table is rebuilt.
+    haps = replace(ds._seqs, ds_contigs=["chrZ"])
+    assert haps._ds_to_store == [None]  # built without raising
+
+    with pytest.raises(ValueError, match=r"chrZ.*no equivalent"):
+        haps._store_contig(0)