From cd83e08283bb50e157bf2c72b8ed4af31f6d9df4 Mon Sep 17 00:00:00 2001 From: James Kane Date: Thu, 30 Jul 2026 17:26:22 -0500 Subject: [PATCH] CRAM region queries decoded the whole chromosome; skip non-overlapping containers Every CRAM region query cost the same as reading the entire contig, no matter how small the region. Measured on a 30x 1kGP CRAM, a ONE-BASE query: chr21 (45 Mb) 20.92 s chr1 (248 Mb) 116.4 s same query on a BAM: 4-6 ms Cost tracked contig length exactly (248/45 = 5.5x length, 116/21 = 5.5x time) and was completely independent of region size -- a 1 bp and a 1 Mb query on chr21 both took 20.91 s. CAUSE. noodles' `Query::read_next_container` skips a container only when its `reference_sequence_id` does not match; it never consults the requested interval, so it decodes every container of the chromosome and filters records afterwards. Our own `for_each` had replicated the same reference-sequence-only filter. The `.crai` already carries `alignment_start` and `alignment_span` per container -- the information needed to skip was present and unused. On HG00096, chr21 holds 1,140 containers and a point query needs exactly 1. FIX. `cram_container_offsets` selects containers by interval overlap, shared by both read paths. `query` no longer delegates to noodles: it decodes only the containers that can overlap, lazily, one at a time so an early-stopping caller does not pay for the rest. A container the index cannot place (`alignment_start` = None) is KEPT -- skipping is only ever done on positive evidence a container lies outside the interval, because a wrongly skipped container is reads silently missing from a variant call. RESULT, same file, same machine: chr21 1 bp 20.92 s -> 8.3 ms (~2,500x) chr1 1 bp 116.4 s -> 113 ms (~1,030x) chr21 1 Mb 20.91 s -> 438 ms (~48x) `navigator call --contig chr21` >96 min -> 44 s (~130x) The last line is the one that matters: the whole-chromosome cost was paid once per caller chunk AND once per realignment candidate, which is why a single chromosome ran over an hour and a half without finishing. CORRECTNESS. Faster-but-lossy would present as a faster caller rather than a broken one, so equivalence is tested, not assumed: - `cram_query_matches_noodles_query` runs our query and noodles' own Query over the checked-in fixture and compares every field of every record. - `container_offsets_select_only_overlapping_containers` pins the selection boundaries (first/last base of a container overlap, one past either end does not, unbounded intervals keep everything, other references are never selected, unplaceable containers are kept) -- the fixture is a single container and cannot catch an off-by-one here. - On the real 11 GB CRAM, `VERIFY=1 cram_query_probe` compares against noodles record for record: 305 records over 1 kb and 47,898 over 200 kb, byte identical on name, position, flags and sequence. Why this went unnoticed: the ground-truth subject's alignment is a BAM, and its whole genome calls in ~5 minutes. Every CRAM in the workspace -- including all 3,216 1kGP alignments -- has been paying the whole-contig cost. Found while trying to validate archaic segment calls against the hmmix 1000G callset, which needs those CRAMs. Adds `cram_query_probe`, the harness these numbers come from, as the regression tool: it attributes cost per phase (open, first query, warm query, bulk region) so a future slowdown names the part responsible. Co-Authored-By: Claude Opus 5 (1M context) --- .../examples/cram_query_probe.rs | 92 ++++++++ crates/navigator-analysis/src/reader.rs | 206 +++++++++++++++++- 2 files changed, 286 insertions(+), 12 deletions(-) create mode 100644 crates/navigator-analysis/examples/cram_query_probe.rs diff --git a/crates/navigator-analysis/examples/cram_query_probe.rs b/crates/navigator-analysis/examples/cram_query_probe.rs new file mode 100644 index 00000000..82ff1818 --- /dev/null +++ b/crates/navigator-analysis/examples/cram_query_probe.rs @@ -0,0 +1,92 @@ +//! Time the phases of an indexed region read, so a slow BAM/CRAM path can be attributed to the +//! part actually responsible (open, header, first query, warm query, bulk region iteration) rather +//! than to whichever call the wall-clock happened to land in. +//! +//! Written to diagnose a CRAM that took ~500x longer than an equivalent BAM for one region query. +//! +//! ```sh +//! cargo run --release -p navigator-analysis --example cram_query_probe -- \ +//! [span_bp=1000000] +//! ``` + +use std::path::Path; +use std::time::Instant; + +use navigator_analysis::reader::open_indexed; +use noodles::core::Region; + +fn main() { + let a: Vec = std::env::args().collect(); + if a.len() < 5 { + eprintln!("usage: cram_query_probe [span_bp=1000000]"); + std::process::exit(2); + } + let (path, refp, contig) = (Path::new(&a[1]), Path::new(&a[2]), a[3].clone()); + let pos: usize = a[4].parse().expect("pos"); + let span: usize = a.get(5).and_then(|s| s.parse().ok()).unwrap_or(1_000_000); + + let t0 = Instant::now(); + let (header, mut reader) = open_indexed(path, Some(refp)).expect("open"); + println!("open + header : {:>8.2?}", t0.elapsed()); + + // A single-base region: cost here is per-query overhead, not per-record work. + let one = |p: usize| -> Region { format!("{contig}:{p}-{p}").parse().expect("region") }; + + let t = Instant::now(); + let n: usize = reader.query(&header, &one(pos)).expect("q1").count(); + println!("first 1bp query ({n:>4} rec) : {:>8.2?} <- includes any lazy setup", t.elapsed()); + + let t = Instant::now(); + let n: usize = reader.query(&header, &one(pos)).expect("q2").count(); + println!("same query again ({n:>4} rec): {:>8.2?} <- warm: is the cost per-query or one-off?", t.elapsed()); + + let t = Instant::now(); + let n: usize = reader.query(&header, &one(pos + 5_000_000)).expect("q3").count(); + println!("distant 1bp query ({n:>4} rec): {:>8.2?} <- new container: does it re-decode?", t.elapsed()); + + let region: Region = format!("{contig}:{pos}-{}", pos + span).parse().expect("region"); + let t = Instant::now(); + let n: usize = reader.query(&header, ®ion).expect("bulk").count(); + let el = t.elapsed(); + println!( + "{:.1} Mb region ({n} rec) : {:>8.2?} = {:.2?}/Mb", + span as f64 / 1e6, + el, + el / (span as u32 / 1_000_000).max(1) + ); + + // VERIFY=1 checks the container-skipping query against noodles' own (whole-contig) Query on + // REAL data. The checked-in fixture is a single container, so only a large multi-container CRAM + // can catch a container wrongly skipped -- which would present as a faster caller, not a broken + // one. Slow by construction: the oracle is the implementation we replaced. + if std::env::var("VERIFY").is_ok_and(|v| v == "1") { + use noodles::cram; + + let key = |r: &noodles::sam::alignment::RecordBuf| { + ( + r.name().map(|n| n.to_vec()), + r.alignment_start().map(usize::from), + r.flags().bits(), + r.sequence().as_ref().to_vec(), + ) + }; + let mine: Vec<_> = reader.query(&header, ®ion).expect("mine").map(|r| key(&r.expect("rec"))).collect(); + + let repo = navigator_analysis::reader::build_repository(refp).expect("repo"); + let mut oracle = cram::io::indexed_reader::Builder::default() + .set_reference_sequence_repository(repo) + .build_from_path(path) + .expect("noodles open"); + let oh = oracle.read_header().expect("noodles header"); + let t = Instant::now(); + let theirs: Vec<_> = oracle + .query(&oh, ®ion) + .expect("noodles query") + .map(|r| key(&r.expect("rec"))) + .collect(); + println!("\nVERIFY: noodles' own Query took {:?} for the same region", t.elapsed()); + println!(" ours {} records, noodles {} records", mine.len(), theirs.len()); + assert_eq!(mine, theirs, "container skipping changed the records returned"); + println!(" IDENTICAL — container skipping is lossless"); + } +} diff --git a/crates/navigator-analysis/src/reader.rs b/crates/navigator-analysis/src/reader.rs index 29195678..f8a85686 100644 --- a/crates/navigator-analysis/src/reader.rs +++ b/crates/navigator-analysis/src/reader.rs @@ -12,7 +12,8 @@ use std::fs::File; use std::num::NonZeroUsize; use std::path::{Path, PathBuf}; -use noodles::core::Region; +use noodles::core::region::Interval; +use noodles::core::{Position, Region}; use noodles::sam::alignment::RecordBuf; use noodles::{bam, bgzf, cram, fasta, sam}; @@ -216,6 +217,37 @@ pub enum IdxReader { }, } +/// File offsets of the `.crai` containers that can hold records overlapping `interval` on `ref_id`. +/// +/// **This is the whole reason CRAM region queries are usable.** A CRAM container is the unit of +/// decode — you cannot decode part of one — so restricting *which containers get decoded* is the +/// only place a region query can save work. noodles' own `Query` (and, before this, our `for_each`) +/// selected containers by reference sequence alone and then discarded non-overlapping records +/// *after* decoding them, which made every query cost a whole chromosome no matter how small the +/// region: measured at 20.9 s for a 1 bp query on chr21 and 116 s on chr1, against 4–6 ms for the +/// same query on a BAM. chr21 of a 30x WGS holds 1,140 containers and a point query needs exactly +/// one of them. +/// +/// A container whose `alignment_start` is absent is **kept**: that is a container this index cannot +/// place, and dropping it would silently lose records. Skipping is only ever done on positive +/// evidence that the container lies outside the interval. +fn cram_container_offsets(index: &cram::crai::Index, ref_id: usize, interval: Interval) -> Vec { + index + .iter() + .filter(|r| r.reference_sequence_id() == Some(ref_id)) + .filter(|r| match r.alignment_start() { + Some(start) => { + // Span 0 would make an empty range, which intersects nothing; treat it as one base. + let span = r.alignment_span().max(1); + let end = Position::new(usize::from(start).saturating_add(span - 1)).unwrap_or(start); + interval.intersects((start..=end).into()) + } + None => true, + }) + .map(|r| r.offset()) + .collect() +} + /// Open `path` for indexed region queries (autoloads the `.bai`/`.crai`). `reference` is /// required for CRAM. pub fn open_indexed(path: &Path, reference: Option<&Path>) -> Result<(sam::Header, IdxReader), AnalysisError> { @@ -268,10 +300,87 @@ impl IdxReader { RecordBuf::try_from_alignment_record(header, &rec).map_err(|e| AnalysisError::io(&path, e)) }))) } - IdxReader::Cram { inner, path, .. } => { + // Hand-rolled rather than `inner.query(...)`: noodles' `Query` decodes every container + // of the contig and filters records afterwards, so a 1 bp query costs a whole + // chromosome (see [`cram_container_offsets`]). This decodes only the containers that + // can overlap, lazily — one container at a time, so a caller that stops early (a + // `.take(n)` probe, a cancelled walk) does not pay for the rest. + IdxReader::Cram { inner, repo, path } => { + use std::io::{Seek, SeekFrom}; + + use noodles::sam::alignment::Record as _; // alignment_start/_end on cram::Record + let path = path.clone(); - let q = inner.query(header, region).map_err(|e| AnalysisError::io(&path, e))?; - Ok(Box::new(q.map(move |r| r.map_err(|e| AnalysisError::io(&path, e))))) + let repo = repo.clone(); + let ref_id = header.reference_sequences().get_index_of(region.name()).ok_or_else(|| { + AnalysisError::Message(format!( + "contig {} not in {} header", + String::from_utf8_lossy(region.name()), + path.display() + )) + })?; + let interval = region.interval(); + let mut offsets = cram_container_offsets(inner.index(), ref_id, interval).into_iter(); + + let mut pending: std::vec::IntoIter = Vec::new().into_iter(); + let mut container = cram::io::reader::Container::default(); + Ok(Box::new(std::iter::from_fn(move || { + loop { + if let Some(rec) = pending.next() { + return Some(Ok(rec)); + } + // Next container that can overlap; `None` ends the iterator. + let offset = offsets.next()?; + let io_err = |e| AnalysisError::io(&path, e); + if let Err(e) = inner.get_mut().seek(SeekFrom::Start(offset)).map_err(io_err) { + return Some(Err(e)); + } + match inner.read_container(&mut container).map_err(io_err) { + Ok(0) => continue, + Ok(_) => {} + Err(e) => return Some(Err(e)), + } + let compression_header = match container.compression_header().map_err(io_err) { + Ok(h) => h, + Err(e) => return Some(Err(e)), + }; + let mut buf = Vec::new(); + for slice in container.slices() { + let slice = match slice.map_err(io_err) { + Ok(s) => s, + Err(e) => return Some(Err(e)), + }; + let (core, external) = match slice.decode_blocks().map_err(io_err) { + Ok(b) => b, + Err(e) => return Some(Err(e)), + }; + let records = match slice + .records(repo.clone(), header, &compression_header, &core, &external) + .map_err(io_err) + { + Ok(r) => r, + Err(e) => return Some(Err(e)), + }; + for rec in &records { + // Same per-record overlap test noodles applies post-decode — the + // container filter is a coarse prefilter, not a replacement for it. + if let (Some(Ok(start)), Some(Ok(end))) = (rec.alignment_start(), rec.alignment_end()) + { + if !interval.intersects((start..=end).into()) { + continue; + } + } else { + continue; + } + match RecordBuf::try_from_alignment_record(header, rec).map_err(io_err) { + Ok(r) => buf.push(r), + Err(e) => return Some(Err(e)), + } + } + } + pending = buf.into_iter(); + } + }))) } } } @@ -372,14 +481,12 @@ impl IdxReader { })?; let interval = region.interval(); - // Collect the file offsets of this contig's containers before borrowing `inner` - // mutably to seek/read (the `.crai` index borrow can't overlap the read borrow). - let offsets: Vec = inner - .index() - .iter() - .filter(|r| r.reference_sequence_id() == Some(ref_id)) - .map(|r| r.offset()) - .collect(); + // Collect the file offsets of the containers that can overlap the query before + // borrowing `inner` mutably to seek/read (the `.crai` index borrow can't overlap the + // read borrow). Selecting on the interval — not just the contig — is what keeps this + // proportional to the region instead of the chromosome; see + // [`cram_container_offsets`]. + let offsets = cram_container_offsets(inner.index(), ref_id, interval); let mut container = cram::io::reader::Container::default(); for offset in offsets { @@ -578,4 +685,79 @@ mod tests { assert!(!sink.0.is_empty(), "fixture should have chrM records"); assert_eq!(sink.0, via_query, "cram::Record path must match RecordBuf path"); } + + /// [`cram_container_offsets`] decides which containers are decoded at all, so its boundary + /// behaviour *is* the correctness of every CRAM region query: a container wrongly skipped is + /// reads silently missing from a variant call, which no downstream test would attribute to the + /// reader. Checked at the edges, where an off-by-one actually lives. + #[test] + fn container_offsets_select_only_overlapping_containers() { + let p = |n: usize| Position::new(n).unwrap(); + // ref 0 containers spanning [1000,1099], [2000,2099], [3000,3099]; one on ref 1; and one + // the index cannot place. + let idx: cram::crai::Index = vec![ + cram::crai::Record::new(Some(0), Some(p(1000)), 100, 10, 0, 0), + cram::crai::Record::new(Some(0), Some(p(2000)), 100, 20, 0, 0), + cram::crai::Record::new(Some(0), Some(p(3000)), 100, 30, 0, 0), + cram::crai::Record::new(Some(1), Some(p(2000)), 100, 40, 0, 0), + cram::crai::Record::new(Some(0), None, 0, 50, 0, 0), + ]; + let sel = |a: usize, b: usize| cram_container_offsets(&idx, 0, (p(a)..=p(b)).into()); + + // A point inside one container decodes that container — not the contig. This single + // assertion is the difference between 8 ms and 21 s on a real chr21. + assert_eq!(sel(2050, 2050), vec![20, 50]); + // Boundaries: touching the first/last base of a container counts as overlap. + assert_eq!(sel(2099, 2099), vec![20, 50], "last base of a container overlaps"); + assert_eq!(sel(2000, 2000), vec![20, 50], "first base of a container overlaps"); + assert_eq!(sel(2100, 2100), vec![50], "one past the end does not"); + assert_eq!(sel(1999, 1999), vec![50], "one before the start does not"); + // A span crossing several containers takes exactly those it crosses. + assert_eq!(sel(1050, 2050), vec![10, 20, 50]); + // The other reference is never selected, even at identical coordinates. + assert_eq!(cram_container_offsets(&idx, 1, (p(2050)..=p(2050)).into()), vec![40]); + // An unbounded interval keeps every container on the reference. + assert_eq!( + cram_container_offsets(&idx, 0, Region::new(b"x".to_vec(), ..).interval()), + vec![10, 20, 30, 50] + ); + } + + /// Our hand-rolled `query` must return exactly what noodles' own `Query` returns. We replaced + /// it for speed, and a reimplementation that quietly drops records would look like a faster + /// caller rather than a broken one — the failure mode worth a test. + #[test] + fn cram_query_matches_noodles_query() { + let dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures"); + let cram = dir.join("coverage.cram"); + let reference = dir.join("ref.fa"); + + for region in [ + Region::new(b"chrM".to_vec(), ..), + Region::new(b"chrM".to_vec(), Position::new(1).unwrap()..=Position::new(200).unwrap()), + Region::new(b"chrM".to_vec(), Position::new(50).unwrap()..=Position::new(60).unwrap()), + ] { + let (header, mut ours) = open_indexed(&cram, Some(&reference)).expect("open"); + let mine: Vec = ours + .query(&header, ®ion) + .expect("query") + .map(|r| capture(&r.expect("rec"))) + .collect(); + + // noodles' own indexed query, unmodified, as the oracle. + let repo = build_repository(&reference).expect("repo"); + let mut theirs = cram::io::indexed_reader::Builder::default() + .set_reference_sequence_repository(repo) + .build_from_path(&cram) + .expect("noodles open"); + let nheader = theirs.read_header().expect("noodles header"); + let reference_impl: Vec = theirs + .query(&nheader, ®ion) + .expect("noodles query") + .map(|r| capture(&r.expect("rec"))) + .collect(); + + assert_eq!(mine, reference_impl, "region {region:?}: must match noodles' Query exactly"); + } + } }