CRAM region queries decoded the whole chromosome; skip non-overlapping containers

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 -- \
+//!   <bam/cram> <ref.fa> <contig> <pos> [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<String> = std::env::args().collect();
+    if a.len() < 5 {
+        eprintln!("usage: cram_query_probe <bam/cram> <ref.fa> <contig> <pos> [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, &region).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, &region).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, &region)
+            .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");
+    }
+}

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<u64> {
+    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<RecordBuf> = 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<u64> = 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<Captured> = ours
+                .query(&header, &region)
+                .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<Captured> = theirs
+                .query(&nheader, &region)
+                .expect("noodles query")
+                .map(|r| capture(&r.expect("rec")))
+                .collect();
+
+            assert_eq!(mine, reference_impl, "region {region:?}: must match noodles' Query exactly");
+        }
+    }
 }