Stream dense seed context during mate and SA recovery

docs/changelog.md

@@ -5,6 +5,21 @@ your version with `osteosarc --version`. Pin both the package version and your
 snapshot (by name or download date) for reproducible analyses. Full release notes are on
 [GitHub](https://github.com/iskandr/osteosarc/releases).
 
+## 0.15.3
+
+- Stream ordinary seed context during mate/SA recovery instead of retaining every
+  alignment and unrelated query name in memory (#109). Every original seed
+  record, duplicate and diagnostic survives. Unchanged recovered BAMs share the
+  seed cache file; changed outputs are checked against the full record multiset.
+- Partner queries select pointer-bearing seed names in every round, retaining
+  candidates for later recursive links. The changed query scope uses recovery
+  request schema 3; prior cache entries remain intact.
+
+## 0.15.2
+
+- Store repeated query-name filters as verified shared assets and pin upstream
+  acquisition receipts without copying their full contents (#107).
+
 ## 0.15.1
 
 - Optional bounded batches for mate/SA recovery preserve available partners

docs/dense-recovery-spec.md

@@ -0,0 +1,21 @@
+# Dense seed context (#109)
+
+Independent indexed counts found 1,044,612 reads in the failed 13-target 10x
+batch and 1.35–5.67 million reads in individual dense targets. Retry the original
+failures separately; neither a regional record cap nor a subprocess failure is
+evidence against an expressed junction.
+
+Stream seed records, retaining alignment objects only for templates with outgoing
+mate/SA pointers or diagnostics. Include every seed record belonging to a
+pointer-bearing name before resolving leads, and query those names in every round
+so later leads can reuse earlier windows. Preserve all original seed records and
+their multiplicities in a coordinate-ordered output stream; append only matched
+partners. Verify the complete output record multiset. Keep per-record provenance,
+unresolved pointers, caps and timeouts explicit. Version the changed recovery
+request because the partner query-name set changes.
+
+Regression checks cover ordinary context retention, seed partners without their
+own pointers, recursive/cyclic pointers, duplicates, read groups and missing
+QNAME/SEQ diagnostics. Measure memory and successful acquisition on the dense
+pilot inputs before changing Isovar's configurable resource budgets. The pipe
+cleanup failure remains a distinct issue unless reproduced and fixed.

docs/reads.md

@@ -111,6 +111,15 @@ seed extraction. These optional limits leave the default request strategy and
 cache identities unchanged. The receipt records failed batches, visited regions
 and exhausted limits; a failed query never establishes absence.
 
+Ordinary seed context is streamed rather than retained as alignment objects in
+memory. All seed records still survive, with their exact stored values and
+multiplicities. Only names with an outgoing mate/SA pointer enter partner
+queries; partners already present in the seed can satisfy those pointers even
+without pointers of their own. An unchanged recovered BAM shares the cached
+seed file. Counts and per-record provenance still consume memory proportional to
+the number of distinct records. Recovery requests use schema 3, so earlier
+recovery derivatives remain intact and are not reused for the changed queries.
+
 ## Genome builds
 
 ```python

osteosarc/__init__.py

@@ -58,7 +58,7 @@
 from .sv_candidates import load_sv_candidates
 from .timeline import Event, Timeline
 
-__version__ = "0.15.2"
+__version__ = "0.15.3"
 
 # The public API. A few other helpers are importable from here because the OpenVax
 # libraries use them (digest, parse_variants, inspect_alignment, ...); everything

osteosarc/recovery.py

@@ -7,6 +7,8 @@
 import tempfile
 from collections import Counter, defaultdict
 from dataclasses import asdict, dataclass, replace
+from heapq import merge
+from itertools import groupby
 
 from .cache import Cache, digest, file_lock, stable_id
 from .errors import CoordinateError, IntegrityError, OsteosarcError, RecordLimitError
@@ -105,14 +107,55 @@ def _timed_out(error):
         isinstance(error, OsteosarcError) and isinstance(error.__cause__, requests.Timeout))
 
 
+def _seed_context(path, policy):
+    """Keep alignment objects only for templates participating in recovery.
+
+    Counts still cover every context record. A second pass includes partners
+    already in the seed even when those records have no outgoing pointers.
+    """
+    counts, records, names = Counter(), {}, set()
+    for record in read_records(path):
+        counts[record.digest] += 1
+        outgoing, issues = _leads(record, policy)
+        if outgoing:
+            names.add(record.read.query_name)
+        if outgoing or issues:
+            records[record.digest] = record
+    if names:
+        for record in read_records(path, keep=lambda r: r.query_name in names):
+            records[record.digest] = record
+    return counts, records, names
+
+
+def _output_records(seed_path, seed_counts, counts, records, contig_count):
+    """Merge all original context with added partners, without loading the seed."""
+    def position(record):
+        read = record.read
+        return (read.reference_id if read.reference_id >= 0 else contig_count,
+                read.reference_start)
+
+    def order(record):
+        return (*position(record), record.digest)
+
+    def seed_records():
+        # Indexed extraction is coordinate sorted; stabilize ties by identity.
+        for _, group in groupby(read_records(seed_path), key=position):
+            yield from sorted(group, key=lambda r: r.digest)
+
+    extra_counts = counts - seed_counts
+    extra_records = (records[key] for key in sorted(extra_counts, key=lambda k: order(records[k]))
+                     for _ in range(extra_counts[key]))
+    yield from merge(seed_records(), extra_records, key=order)
+
+
 def recover_reads(source, regions, *, policy=None, cache=None, **kwargs):
     """Extend extract_reads with bounded, source/RG/segment-scoped partner recovery.
 
     Returns a ReadSubset with every requested/visited interval, unresolved lead,
     cap, missing sequence and repeated/cyclic pointer in its receipt. A changed
     source/index fails rather than mixing acquisitions. Cache hits resume only
     verified complete intermediates. Seed record overflow fails explicitly.
-    Partner windows select all seed QNAMEs before applying the record cap;
+    Partner windows select pointer-bearing seed QNAMEs before applying the record cap;
     source/RG/segment/placement matching still determines retained partners.
     A partner-query timeout fails unless policy.on_timeout is "incomplete".
     """
@@ -134,7 +177,7 @@ def recover_reads(source, regions, *, policy=None, cache=None, **kwargs):
     seed = extract_reads(source, regions, cache=cache, max_records=policy.max_records, **kwargs)
     # Unset new controls retain existing default request/cache identities.
     recovery_parameters = {k: v for k, v in asdict(policy).items() if k != "on_timeout" and v is not None}
-    request = dict(schema_version=2, operation="recover_reads", seed=seed.receipt["request"],
+    request = dict(schema_version=3, operation="recover_reads", seed=seed.receipt["request"],
                    seed_files=seed.receipt["files"],
                    recovery=recovery_parameters)
     directory = cache.workspace / "derived" / stable_id(request)
@@ -144,18 +187,16 @@ def recover_reads(source, regions, *, policy=None, cache=None, **kwargs):
         with seed.open() as bam:
             header = bam.header.to_dict()
         lengths = {s["SN"]: s["LN"] for s in header["SQ"]}
-        seed_records = list(read_records(seed.path))
-        # Every recursively recovered mate/SA partner has a seed QNAME. Include
-        # all seed names in every round so later leads can reuse a visited window.
+        seed_counts, records, names = _seed_context(seed.path, policy)
+        # Every recursively recovered partner has a pointer-bearing seed QNAME.
+        # Include all such names in every round so later leads can reuse a window.
         # Exact RG/segment/placement matching below remains authoritative.
         partner_kwargs = dict(kwargs, filters=replace(
             kwargs.get("filters") or ReadFilter(),
-            query_names=tuple(sorted({r.read.query_name for r in seed_records
-                                      if r.read.query_name not in (None, "*")}))))
+            query_names=tuple(sorted(names))))
         if policy.partner_timeout is not None:
             partner_kwargs["timeout"] = policy.partner_timeout
-        counts = Counter(r.digest for r in seed_records)
-        records = {r.digest: r for r in seed_records}
+        counts = seed_counts.copy()
         # An earlier query can contain a partner whose pointer is discovered later.
         # Keep bounded candidates even when they are not selected for the output.
         candidates, candidate_counts = defaultdict(dict), counts.copy()
@@ -170,8 +211,10 @@ def matching_candidates(lead):
             key = ((lead["rg"], lead["qname"]), lead["segment"], lead["contig"], lead["start"], lead["reverse"])
             return [r for r in candidates.get(key, {}).values() if _matches(r, lead)]
 
-        add_candidates(seed_records)
-        reasons = {key: {"seed/context region"} for key in records}
+        add_candidates(records.values())
+        # Ordinary context shares one immutable reason instead of millions of sets.
+        seed_reason = frozenset(("seed/context region",))
+        reasons = dict.fromkeys(seed_counts, seed_reason)
         visited = [dict(r, round=0) for r in seed.receipt["resolved_regions"]]
         bases = sum(r["end"] - r["start"] for r in visited)
         if len(visited) > policy.max_intervals or bases > policy.max_bases or sum(counts.values()) > policy.max_records:
@@ -215,7 +258,8 @@ def source_identity(receipt):
                     for r in matches:
                         records[r.digest] = r
                         counts[r.digest] = candidate_counts[r.digest]
-                        reasons.setdefault(r.digest, set()).add("paired mate" if lead["kind"] == "mate" else "SA-linked partner")
+                        reasons[r.digest] = reasons.get(r.digest, frozenset()) | {
+                            "paired mate" if lead["kind"] == "mate" else "SA-linked partner"}
                     handled_leads.add(identity)
                     if len(matches) > 1:
                         problems[identity] = dict(lead=identity, reason="multiple matching placements")
@@ -298,7 +342,8 @@ def source_identity(receipt):
                     found |= multiplicities  # union, not sum across overlapping pointers
                     found_records.update((r.digest, r) for r in matches)
                     for r in matches:
-                        reasons.setdefault(r.digest, set()).add("paired mate" if lead["kind"] == "mate" else "SA-linked partner")
+                        reasons[r.digest] = reasons.get(r.digest, frozenset()) | {
+                            "paired mate" if lead["kind"] == "mate" else "SA-linked partner"}
                     unresolved.pop(identity, None)
                     if len(multiplicities) > 1:
                         problems[identity] = dict(lead=identity, reason="multiple matching placements")
@@ -321,16 +366,18 @@ def source_identity(receipt):
             from pathlib import Path
             work = Path(temporary)
             output = work / "reads.bam"
-            ordered = sorted(records, key=lambda k: (
-                records[k].read.reference_id if records[k].read.reference_id >= 0 else len(lengths),
-                records[k].read.reference_start, k))
-            with pysam.AlignmentFile(output, "wb", header=header) as out:
-                for key in ordered:
-                    for _ in range(counts[key]):
-                        out.write(records[key].read)
-            if record_multiset(output) != counts:
-                raise IntegrityError("Recovery changed original BAM records")
-            pysam.index(str(output))
+            if counts == seed_counts:
+                # Both files live in this cache. Unchanged context need not be
+                # decoded, recompressed, or stored a second time.
+                os.link(seed.path, output)
+                os.link(seed.index_path, work / "reads.bam.bai")
+            else:
+                with pysam.AlignmentFile(output, "wb", header=header) as out:
+                    for record in _output_records(seed.path, seed_counts, counts, records, len(lengths)):
+                        out.write(record.read)
+                if record_multiset(output) != counts:
+                    raise IntegrityError("Recovery changed original BAM records")
+                pysam.index(str(output))
             receipt = dict(request=request, files={name: digest(work / name) for name in ("reads.bam", "reads.bam.bai")},
                            records=sum(counts.values()), scope="bounded_mate_SA_context", complete_template=False,
                            status="incomplete" if failed_queries else "truncated" if limits else "bounded",

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "osteosarc"
-version = "0.15.2"
+version = "0.15.3"
 description = "Reproducible access to the public osteosarc.com dataset"
 readme = "README.md"
 requires-python = ">=3.9"

tests/test_recovery.py

@@ -77,6 +77,71 @@ def test_cap_is_distinct_from_empty_and_cache_reuse_needs_no_commands(split_bam,
     assert extract_reads(split_bam, regions, cache=cache, recovery=policy).receipt == result.receipt
 
 
+def test_dense_context_streams_without_losing_original_records(split_bam, tmp_path):
+    from osteosarc.recovery import _seed_context
+
+    path = tmp_path / "dense-context.bam"
+    with pysam.AlignmentFile(split_bam) as bam:
+        header = bam.header
+    with pysam.AlignmentFile(path, "wb", header=header) as out:
+        for i in range(5000):
+            read = pysam.AlignedSegment(header)
+            read.query_name = "context-%d" % i
+            read.reference_id, read.reference_start = 0, 100 + i
+            read.cigarstring, read.query_sequence = "10M", "ACGTACGTAA"
+            read.query_qualities = None if i % 2 else [30] * 10
+            read.flag = 16 if i % 2 else 0
+            read.set_tag("CB", "cell-%d" % i)
+            out.write(read)
+            if i % 100 == 0:
+                out.write(read)
+    pysam.index(str(path))
+    policy = RecoveryPolicy()
+    counts, active, names = _seed_context(path, policy)
+    assert counts == record_multiset(path)
+    assert not active and not names
+    result = extract_reads(path, [Region("chr1", 0, 10000, "GRCh38")],
+                           cache=tmp_path / "cache", recovery=policy)
+    assert record_multiset(result.path) == counts
+    assert result.receipt["records"] == 5050
+    assert len(result.receipt["reasons"]) == 5000
+    assert all(value == ["seed/context region"] for value in result.receipt["reasons"].values())
+    assert not result.receipt["requested_leads"] and not result.receipt["observations"]
+
+
+def test_seed_partner_without_pointer_and_missing_sequence_survive(split_bam, tmp_path):
+    from osteosarc.recovery import _seed_context
+
+    path = tmp_path / "seed-partner.bam"
+    with pysam.AlignmentFile(split_bam) as bam:
+        seed = next(bam)
+        header = bam.header
+    seed.flag = 0
+    seed.set_tag("SA", "chr1,201,+,10M,60,0;")
+    partner = pysam.AlignedSegment.fromstring(seed.to_string(), header)
+    partner.reference_start, partner.flag = 200, 2048
+    partner.set_tag("SA", None)
+    missing = pysam.AlignedSegment.fromstring(partner.to_string(), header)
+    missing.query_name, missing.reference_start, missing.query_sequence = "missing", 300, None
+    with pysam.AlignmentFile(path, "wb", header=header) as out:
+        out.write(seed)
+        out.write(partner)
+        out.write(partner)
+        out.write(missing)
+    pysam.index(str(path))
+    counts, active, names = _seed_context(path, RecoveryPolicy())
+    assert names == {"split"}
+    assert len(active) == 3 and sum(counts.values()) == 4
+    result = extract_reads(path, [Region("chr1", 0, 400, "GRCh38")],
+                           cache=tmp_path / "cache", recovery=RecoveryPolicy())
+    assert record_multiset(result.path) == counts
+    assert len(result.receipt["acquisition"]) == 1
+    assert not result.receipt["unresolved"]
+    assert any(row["reason"] == "missing SEQ" for row in result.receipt["observations"])
+    assert sum("SA-linked partner" in reasons for reasons in result.receipt["reasons"].values()) == 1
+    assert sum(reasons == ["seed/context region"] for reasons in result.receipt["reasons"].values()) == 2
+
+
 def test_record_limit_fails_instead_of_publishing_partial_success(split_bam, tmp_path):
     from osteosarc import RecordLimitError
     with pytest.raises(RecordLimitError, match="limit"):
@@ -316,7 +381,7 @@ def test_dense_partner_window_filters_names_before_cap(split_bam, tmp_path, wron
     assert not result.receipt["limits"]
     assert not (record_multiset(result.path) - record_multiset(path))
     filters = result.receipt["acquisition"][1]["request"]["filters"]
-    assert filters["query_names"] == ["context", "split"]
+    assert filters["query_names"] == ["split"]
     assert filters["min_mapq"] == 20 and filters["barcodes"] == ["shared-cell"]