From 000cd34d5d73d48d7cfd34c7097821fd5135112a Mon Sep 17 00:00:00 2001 From: alejandrogzi Date: Thu, 6 Aug 2026 23:07:26 +0200 Subject: [PATCH 1/2] feat(io): bound coordinate-sort memory with an external merge '--outSAMtype BAM SortedByCoordinate' no longer buffers the whole output in RAM. The sort now fills a '--limitBAMsortRAM' buffer, spills sorted runs beside the output, and k-way merges them on finish, so peak memory is flat in output size. '--limitBAMsortRAM 0' now means 512 MiB instead of unlimited, and runs beyond 64 are merged in balanced passes so a small budget cannot exhaust file descriptors. Output is byte-identical to the previous in-memory sort. --- CHANGELOG.md | 14 + README.md | 8 +- .../content/docs/reference/cli-parameters.md | 2 +- src/io/bam.rs | 735 ++++++++++++++---- 4 files changed, 607 insertions(+), 152 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 087aa49..bbc6613 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -166,6 +166,20 @@ Sections commonly used: Features, Bug fixes, Other changes. - **STARsolo `features.tsv` column 2 now emits the GTF `gene_name`** (symbol), with the STAR gene_id fallback, instead of duplicating the gene_id. +- **`--outSAMtype BAM SortedByCoordinate` no longer buffers the whole + output in RAM.** The coordinate sort is now external: it fills a + `--limitBAMsortRAM` buffer, spills sorted runs beside the output, and + k-way merges them on finish. Peak memory is flat in output size + (measured: 610 MB at a 64 MiB budget for 300 k through 2.4 M records, + versus 648 MB → 2,132 MB before, a growth of 723 B/record that + extrapolated to ~116 GB for a 160 M-record human sample). Output is + byte-identical to the previous in-memory sort — coordinate ties still + resolve to input order, since runs merge with the run index as + tiebreak (verified at 2.4 M records through a 219-run multi-pass + merge). `--limitBAMsortRAM N` now spills above `N` rather than + aborting the run, and `--limitBAMsortRAM 0` means 512 MiB instead of + "unlimited". Runs beyond 64 are merged in balanced passes so a small + budget on a large run cannot exhaust file descriptors. ### Bumps diff --git a/README.md b/README.md index d6536e6..6816821 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ A Rust reimplementation of [STAR](https://github.com/alexdobin/STAR) (Spliced Tr rustar-aligner aims to be a faithful port of STAR, matching the original behavior as closely as possible. It uses the same genome index format, accepts the same `--camelCase` command-line parameters, and produces compatible SAM/BAM output. -**Current status**: End-to-end single-end and paired-end RNA-seq alignment with splice junction detection, two-pass mode, chimeric alignment detection (including multi-junction Tier 3), gene-level quantification, **single-cell quantification (STARsolo: Gene / GeneFull / SJ / Velocyto features, barcode correction, UMI dedup, EmptyDrops_CR cell calling)**, WASP allele-specific filtering, paired-end mate-overlap merging, coverage-track output, adapter/CellRanger clipping, and multi-threaded parallel processing. Default SAM output is byte-identical to STAR's (`NH HI AS nM` attributes); solo count matrices are byte-identical to STARsolo's. Pure-Rust core with an in-tree deterministic RNG (no `rand` dependency). 578 tests passing (553 unit + 25 integration), 0 clippy warnings. See [Performance & Benchmarks](#performance--benchmarks) for a native three-way comparison against STARsolo and CellRanger. +**Current status**: End-to-end single-end and paired-end RNA-seq alignment with splice junction detection, two-pass mode, chimeric alignment detection (including multi-junction Tier 3), gene-level quantification, **single-cell quantification (STARsolo: Gene / GeneFull / SJ / Velocyto features, barcode correction, UMI dedup, EmptyDrops_CR cell calling)**, WASP allele-specific filtering, paired-end mate-overlap merging, coverage-track output, adapter/CellRanger clipping, and multi-threaded parallel processing. Default SAM output is byte-identical to STAR's (`NH HI AS nM` attributes); solo count matrices are byte-identical to STARsolo's. Pure-Rust core with an in-tree deterministic RNG for alignment tie-breaking. 622 tests passing (591 unit + 31 integration), 0 clippy warnings. See [Performance & Benchmarks](#performance--benchmarks) for a native three-way comparison against STARsolo and CellRanger. ## Quick Start @@ -208,7 +208,7 @@ resident; the 16 GB sparse index is stable at ~54 s. - Single-end and paired-end alignment with mate rescue - Read-end alignment mode (`--alignEndsType Local` (default) / `EndToEnd` / `Extend5pOfRead1` / `Extend5pOfReads12` / `Extend3pOfRead1`) -- SAM, unsorted BAM, and coordinate-sorted BAM output (`--outSAMtype SAM`, `BAM Unsorted`, or `BAM SortedByCoordinate`) +- SAM, unsorted BAM, and coordinate-sorted BAM output (`--outSAMtype SAM`, `BAM Unsorted`, or `BAM SortedByCoordinate`). The coordinate sort is external: it buffers up to `--limitBAMsortRAM`, spills sorted runs beside the output, and merges them, so peak memory is independent of output size - Multi-threaded parallel alignment (`--runThreadN`) - GTF-based junction annotation with scoring bonus (`--sjdbGTFfile`) - Two-pass mode for novel junction discovery (`--twopassMode Basic`) @@ -227,8 +227,8 @@ resident; the 16 GB sparse index is stable at ~54 s. - Unmapped read output to FASTQ (`--outReadsUnmapped Fastx` → `Unmapped.out.mate1` / `mate2`) - Gzip-compressed FASTQ input (`--readFilesCommand zcat`) - Read group tags (`--outSAMattrRGline`) -- Deterministic input-order output regardless of thread count (`--outSAMorder Paired` / `PairedKeepInputOrder`, both accepted; rustar always preserves FASTQ order) -- Deterministic in-tree RNG for reproducible tie-breaking (`--runRNGseed`; no `rand` dependency) +- Deterministic input-order output regardless of thread count (`--outSAMorder Paired` / `PairedKeepInputOrder`, both accepted) +- Deterministic in-tree RNG for reproducible alignment tie-breaking (`--runRNGseed`) - SAM optional tags: NH, HI, AS, nM, NM, XS, jM, jI, MD (default `NH HI AS nM` matches STAR; `NM` opt-in) - `--outSAMattributes` control (Standard/All/None/explicit list) - SECONDARY flag (0x100) on multi-mapper alignments diff --git a/docs/src/content/docs/reference/cli-parameters.md b/docs/src/content/docs/reference/cli-parameters.md index f749039..bbdd0f4 100644 --- a/docs/src/content/docs/reference/cli-parameters.md +++ b/docs/src/content/docs/reference/cli-parameters.md @@ -42,7 +42,7 @@ Run `rustar-aligner --help` for the full machine-generated listing. | `--outFileNamePrefix` | `./` | Prefix (path + filename stem) for all output files. | | `--outSAMtype` | `SAM` | `SAM`, `BAM Unsorted`, `BAM SortedByCoordinate`, or `None`. | | `--outBAMcompression` | `1` | BGZF level. `-1`/`0` = uncompressed; `1`–`8` = flate2 levels; `≥9` = max. | -| `--limitBAMsortRAM` | `0` | Max RAM (bytes) for sorted BAM. `0` = unlimited. | +| `--limitBAMsortRAM` | `0` | Max RAM for the coordinate sort; accepts a suffix (`8G`, `512M`). Records beyond it spill to sorted runs beside the output and are merged, so output is unaffected by this value. `0` means 512 MiB. Very small values still work but cost wall time (many spill runs plus an extra merge pass). | | `--outStd` | `None` | Route primary output to stdout: `None`, `SAM`, `BAM_Unsorted`, `BAM_SortedByCoordinate`. | ## Output: SAM/BAM records diff --git a/src/io/bam.rs b/src/io/bam.rs index b1e6717..d181a05 100644 --- a/src/io/bam.rs +++ b/src/io/bam.rs @@ -9,7 +9,7 @@ use noodles::sam::alignment::record_buf::RecordBuf; use noodles::{bam, bgzf, sam}; use std::ffi::CString; use std::fs::File; -use std::io::{BufWriter, Write}; +use std::io::{BufReader, BufWriter, Write}; use std::path::Path; /// Buffer for BAM records built by parallel threads @@ -61,16 +61,350 @@ pub struct BamWriter { header: sam::Header, } -/// BAM file writer that collects all records in memory, sorts by coordinate, -/// then writes a single sorted BAM file on `finish()`. +/// Sort key for coordinate-sorted output. Unmapped records (no reference or no +/// position) sort to the end, matching STAR. +fn sort_key(record: &RecordBuf) -> (usize, usize) { + match (record.reference_sequence_id(), record.alignment_start()) { + (Some(chr), Some(pos)) => (chr, pos.get()), + _ => (usize::MAX, 0), + } +} + +/// Estimate of the heap bytes a buffered `RecordBuf` occupies. /// -/// The header emits `SO:coordinate`. Unmapped records sort to the end. -pub struct SortedBamWriter { +/// Counts the variable-length fields plus a fixed allowance for the struct, its +/// five `Vec`/`BString` headers, and per-allocation allocator slack. Calibrated +/// against measured RSS growth: an unbounded sort of 100 bp single-end records +/// with the default `NH HI AS nM` tags grows at 723 B/record, and the constants +/// below yield 768 B for that shape, so the estimate sits just above actual. +/// +/// Erring high is the safe direction — the sorter spills once the running total +/// crosses the budget, so over-counting spills early. Note that the *realized* +/// sort footprint is about 1.4x the configured budget rather than 1.0x, because +/// the batch being copied in, the spill compressor, and the merge readers all sit +/// outside this accounting. +fn estimated_record_bytes(record: &RecordBuf) -> u64 { + /// `RecordBuf` itself, its `Vec`/`BString` headers, and allocator slack. + const FIXED_OVERHEAD: u64 = 320; + /// `record_buf::Cigar` holds `Vec`; `Op` is a (Kind, usize) pair. + const CIGAR_OP_BYTES: u64 = 16; + /// A tag/value pair in the data map, averaged over the tags STAR emits. + const DATA_FIELD_BYTES: u64 = 56; + + FIXED_OVERHEAD + + record.name().map_or(0, |name| name.len() as u64) + + record.sequence().len() as u64 + + record.quality_scores().len() as u64 + + record.cigar().as_ref().len() as u64 * CIGAR_OP_BYTES + + record.data().len() as u64 * DATA_FIELD_BYTES +} + +/// Coordinate-sort RAM budget when `--limitBAMsortRAM` is 0. +/// +/// `0` previously meant "unlimited", so the default configuration buffered the +/// entire output (measured: 723 B/record, i.e. ~116 GB for a 160 M-record human +/// sample). A fixed modest budget replaces that. +/// +/// Deliberately *not* derived from the genome index size, even though that is +/// STAR's documented rule: scaling the buffer up with the genome enlarges it +/// exactly when RAM is tightest (a 32 GB human index on a 48 GB host would get a +/// multi-GB sort buffer on top). Spilling is measured to be free — 5.83 s with 13 +/// spill runs versus 6.15 s unbounded for 1.2 M records — so there is nothing to +/// buy by sorting more in memory. Output is identical either way; only peak memory +/// and temp-file use change. +const DEFAULT_BAM_SORT_RAM: u64 = 512 << 20; + +/// Maximum spill runs merged at once, bounding open file descriptors. +/// +/// A small `--limitBAMsortRAM` on a large run produces thousands of runs, and +/// opening them all at once hits `EMFILE` (macOS defaults to 256 descriptors; +/// containers are often lower). Above this, runs are merged in balanced passes, +/// costing one extra read/write of the data per `log64(runs)` level. +const MAX_OPEN_RUNS: usize = 64; + +fn resolve_bam_sort_ram(params: &Parameters) -> u64 { + if params.limit_bam_sort_ram > 0 { + params.limit_bam_sort_ram + } else { + DEFAULT_BAM_SORT_RAM + } +} + +/// Directory for coordinate-sort spill files. +/// +/// Deliberately the output directory rather than the system temp dir: `/tmp` is +/// tmpfs on many Linux distributions, so spilling there would keep the records in +/// RAM and defeat the budget entirely. STAR likewise keeps its sort scratch +/// (`_STARtmp`) beside the output. +fn sort_temp_dir(params: &Parameters) -> std::path::PathBuf { + let prefix = params.output_path(""); + match prefix.parent() { + Some(parent) if !parent.as_os_str().is_empty() => parent.to_path_buf(), + _ => std::path::PathBuf::from("."), + } +} + +/// External coordinate sorter with bounded memory. +/// +/// Records accumulate in memory until the RAM budget is reached, at which point +/// the buffer is sorted and written to a spill run (a headerless BGZF stream of +/// BAM records) beside the output. `write_sorted` merges every run plus the +/// in-memory tail with a k-way merge, so peak memory is the budget rather than +/// the whole output. +/// +/// Ordering is identical to a single in-memory sort: runs are created in input +/// order, `sort_by_key` is stable within a run, and the merge breaks ties on run +/// index, so records sharing a coordinate keep their input order. +struct CoordinateSorter { records: Vec, - output_path: std::path::PathBuf, + buffered_bytes: u64, + runs: Vec, header: sam::Header, compression: i32, - limit_bam_sort_ram: u64, + ram_limit: u64, + temp_dir: std::path::PathBuf, + n_records: u64, +} + +impl CoordinateSorter { + fn new(header: sam::Header, params: &Parameters) -> Self { + Self { + records: Vec::new(), + buffered_bytes: 0, + runs: Vec::new(), + header, + compression: params.out_bam_compression, + ram_limit: resolve_bam_sort_ram(params), + temp_dir: sort_temp_dir(params), + n_records: 0, + } + } + + fn push_batch(&mut self, batch: &[RecordBuf]) -> Result<(), Error> { + self.records.reserve(batch.len()); + for record in batch { + self.buffered_bytes += estimated_record_bytes(record); + self.records.push(record.clone()); + } + self.n_records += batch.len() as u64; + if self.buffered_bytes >= self.ram_limit { + self.spill()?; + } + Ok(()) + } + + /// Create an empty spill-run file beside the output. + fn new_run(&self) -> Result { + tempfile::Builder::new() + .prefix("rustar-bamsort-") + .suffix(".tmp") + .tempfile_in(&self.temp_dir) + .map_err(|source| Error::io(source, &self.temp_dir)) + } + + /// Sort the in-memory buffer and append it as a spill run. + fn spill(&mut self) -> Result<(), Error> { + if self.records.is_empty() { + return Ok(()); + } + self.records.sort_by_key(sort_key); + + let temp = self.new_run()?; + // Headerless: `Reader::read_record_buf` ignores the header entirely, so + // runs carry only record blocks and never re-parse reference names. + let mut writer = bam::io::Writer::from(make_bgzf_writer( + BufWriter::new(temp.as_file()), + self.compression, + )); + for record in &self.records { + writer.write_alignment_record(&self.header, record)?; + } + // Finish and flush explicitly rather than on drop, so a failure to write + // the run surfaces here instead of being swallowed and read back short. + writer.try_finish()?; + writer.into_inner().into_inner().flush()?; + log::debug!( + "Coordinate sort: spilled run {} ({} records, ~{} MiB)", + self.runs.len(), + self.records.len(), + self.buffered_bytes >> 20 + ); + self.runs.push(temp.into_temp_path()); + // Keep the capacity: it is the budget, so holding one stable allocation + // for the whole run is both correct and cheaper than releasing it and + // growing back from zero (which leaves a doubling series of abandoned + // blocks behind on every spill). + self.records.clear(); + self.buffered_bytes = 0; + Ok(()) + } + + /// Merge every spill run and the in-memory tail into `out` as a sorted BAM. + fn write_sorted(&mut self, out: W, destination: &str) -> Result<(), Error> { + let spilled_runs = self.runs.len(); + self.reduce_runs()?; + self.records.sort_by_key(sort_key); + + let mut bgzf = make_bgzf_writer(out, self.compression); + write_bam_header_lenient(&mut bgzf, &self.header, Some("coordinate"))?; + let mut writer = bam::io::Writer::from(bgzf); + + let written = if self.runs.is_empty() { + // Nothing spilled: identical to the previous in-memory-only path. + for record in &self.records { + writer.write_alignment_record(&self.header, record)?; + } + self.records.len() as u64 + } else { + let runs = std::mem::take(&mut self.runs); + let written = self.merge(&runs, Some(&self.records), &mut writer)?; + drop(runs); + written + }; + writer.try_finish()?; + writer.into_inner().into_inner().flush()?; + + if written != self.n_records { + return Err(Error::Alignment(format!( + "BAM sort merge wrote {written} records but {} were buffered", + self.n_records + ))); + } + log::info!( + "Sorted BAM written to {destination} ({} records, {spilled_runs} spill run(s), \u{2264}{} MiB sort buffer)", + self.n_records, + self.ram_limit >> 20 + ); + Ok(()) + } + + /// Merge runs in balanced passes until at most `MAX_OPEN_RUNS` remain, so the + /// final merge cannot exhaust the process's file descriptors. + /// + /// Each pass merges consecutive groups of `MAX_OPEN_RUNS` into one run and + /// keeps the groups in order, so a lower run index still means "earlier in the + /// input" and coordinate ties keep resolving to input order. + fn reduce_runs(&mut self) -> Result<(), Error> { + while self.runs.len() > MAX_OPEN_RUNS { + let mut remaining = std::mem::take(&mut self.runs); + let mut reduced = Vec::with_capacity(remaining.len().div_ceil(MAX_OPEN_RUNS)); + while !remaining.is_empty() { + let group: Vec<_> = remaining + .drain(..MAX_OPEN_RUNS.min(remaining.len())) + .collect(); + if group.len() == 1 { + reduced.extend(group); + continue; + } + let temp = self.new_run()?; + let mut writer = bam::io::Writer::from(make_bgzf_writer( + BufWriter::new(temp.as_file()), + self.compression, + )); + self.merge(&group, None, &mut writer)?; + writer.try_finish()?; + writer.into_inner().into_inner().flush()?; + // Dropping `group` here deletes the consumed runs, so scratch use + // does not grow across passes. + drop(group); + reduced.push(temp.into_temp_path()); + } + log::debug!( + "Coordinate sort: reduced spill runs to {} (cap {MAX_OPEN_RUNS})", + reduced.len() + ); + self.runs = reduced; + } + Ok(()) + } + + /// K-way merge `runs` (and optionally an in-memory `tail`, which sorts last on + /// ties) into `writer`. Returns the number of records written. + fn merge( + &self, + runs: &[tempfile::TempPath], + tail: Option<&[RecordBuf]>, + writer: &mut bam::io::Writer>, + ) -> Result { + use std::cmp::Reverse; + use std::collections::BinaryHeap; + + let mut readers = Vec::with_capacity(runs.len()); + for path in runs { + let file = File::open(path).map_err(|source| Error::io(source, path))?; + readers.push(bam::io::Reader::new(BufReader::new(file))); + } + // The unspilled tail is the last run, so it keeps the highest run index + // and therefore loses coordinate ties to everything written before it. + let tail_run = readers.len(); + let mut tail = tail.unwrap_or(&[]).iter(); + // Hoisted: noodles' BAM record decoder ignores this argument, so it must + // not be rebuilt per record. + let decode_header = sam::Header::default(); + + let mut heads: Vec> = Vec::with_capacity(tail_run + 1); + let mut heap: BinaryHeap> = BinaryHeap::new(); + for (run, (reader, path)) in readers.iter_mut().zip(runs).enumerate() { + let record = read_run_record(reader, path, &decode_header)?; + if let Some(record) = &record { + heap.push(Reverse((sort_key(record), run))); + } + heads.push(record); + } + let tail_head = tail.next().cloned(); + if let Some(record) = &tail_head { + heap.push(Reverse((sort_key(record), tail_run))); + } + heads.push(tail_head); + + let mut written = 0u64; + while let Some(Reverse((_, run))) = heap.pop() { + let record = heads[run] + .take() + .ok_or_else(|| Error::Alignment("BAM sort merge lost a record".to_string()))?; + writer.write_alignment_record(&self.header, &record)?; + written += 1; + + let next = if run == tail_run { + tail.next().cloned() + } else { + read_run_record(&mut readers[run], &runs[run], &decode_header)? + }; + if let Some(record) = &next { + heap.push(Reverse((sort_key(record), run))); + } + heads[run] = next; + } + Ok(written) + } +} + +/// Read the next record from a spill run, or `None` at end of run. +/// +/// `header` is passed through to noodles, which ignores it — reference sequence +/// ids are plain indices, so runs need no reference list. +fn read_run_record( + reader: &mut bam::io::Reader>, + path: &Path, + header: &sam::Header, +) -> Result, Error> { + let mut record = RecordBuf::default(); + match reader + .read_record_buf(header, &mut record) + .map_err(|source| Error::io(source, path))? + { + 0 => Ok(None), + _ => Ok(Some(record)), + } +} + +/// BAM file writer that sorts records by coordinate with bounded memory, +/// spilling sorted runs beside the output and merging them on `finish()`. +/// +/// The header emits `SO:coordinate`. Unmapped records sort to the end. +pub struct SortedBamWriter { + sorter: CoordinateSorter, + output_path: std::path::PathBuf, } impl BamWriter { @@ -142,7 +476,8 @@ impl BamWriter { } impl SortedBamWriter { - /// Create a sorted BAM writer. Records are buffered in memory until `finish()`. + /// Create a sorted BAM writer. Records are buffered up to the + /// `--limitBAMsortRAM` budget and spilled to sorted runs beyond it. pub fn create( output_path: &Path, genome: &crate::genome::Genome, @@ -150,88 +485,24 @@ impl SortedBamWriter { ) -> Result { let header = crate::io::sam::build_sam_header(genome, params)?; Ok(Self { - records: Vec::new(), + sorter: CoordinateSorter::new(header, params), output_path: output_path.to_path_buf(), - header, - compression: params.out_bam_compression, - limit_bam_sort_ram: params.limit_bam_sort_ram, }) } - /// Buffer records — no disk I/O yet. + /// Buffer records, spilling a sorted run when the RAM budget is reached. pub fn write_batch(&mut self, batch: &[RecordBuf]) -> Result<(), Error> { - self.records.extend_from_slice(batch); - Ok(()) - } - - /// Estimate memory used by buffered records (rough: 400 bytes/record for 150bp reads). - fn estimated_ram(&self) -> u64 { - self.records.len() as u64 * 400 + self.sorter.push_batch(batch) } - fn check_ram_limit(&self) -> Result<(), Error> { - if self.limit_bam_sort_ram > 0 { - let est = self.estimated_ram(); - if est > self.limit_bam_sort_ram { - return Err(Error::Alignment(format!( - "limitBAMsortRAM={} bytes exceeded: estimated {} bytes for {} records. \ - Increase --limitBAMsortRAM or use --outSAMtype BAM Unsorted.", - self.limit_bam_sort_ram, - est, - self.records.len() - ))); - } - } - Ok(()) - } - - /// Sort all buffered records by coordinate and write a single sorted BAM. - /// - /// Sort key: (reference_sequence_id, alignment_start). - /// Unmapped records (no reference) sort to the end. + /// Merge every run plus the in-memory tail into a coordinate-sorted BAM. pub fn finish(&mut self) -> Result<(), Error> { - self.check_ram_limit()?; - self.records - .sort_by_key(|r| match (r.reference_sequence_id(), r.alignment_start()) { - (Some(chr), Some(pos)) => (chr, pos.get()), - _ => (usize::MAX, 0), - }); - - let buf_writer = BufWriter::new(File::create(&self.output_path)?); - let mut bgzf = make_bgzf_writer(buf_writer, self.compression); - write_bam_header_lenient(&mut bgzf, &self.header, Some("coordinate"))?; - let mut bam_writer = bam::io::Writer::from(bgzf); - for record in &self.records { - bam_writer.write_alignment_record(&self.header, record)?; - } - bam_writer.finish(&self.header)?; - log::info!("Sorted BAM written ({} records)", self.records.len()); - Ok(()) - } - - /// Sort all buffered records and write to stdout (for `--outStd BAM_SortedByCoordinate`). - pub fn finish_to_stdout(&mut self) -> Result<(), Error> { - self.check_ram_limit()?; - self.records - .sort_by_key(|r| match (r.reference_sequence_id(), r.alignment_start()) { - (Some(chr), Some(pos)) => (chr, pos.get()), - _ => (usize::MAX, 0), - }); - - let stdout = std::io::stdout(); - let buf_writer = BufWriter::new(stdout.lock()); - let mut bgzf = make_bgzf_writer(buf_writer, self.compression); - write_bam_header_lenient(&mut bgzf, &self.header, Some("coordinate"))?; - let mut bam_writer = bam::io::Writer::from(bgzf); - for record in &self.records { - bam_writer.write_alignment_record(&self.header, record)?; - } - bam_writer.finish(&self.header)?; - log::info!( - "Sorted BAM written to stdout ({} records)", - self.records.len() - ); - Ok(()) + let file = File::create(&self.output_path) + .map_err(|source| Error::io(source, &self.output_path))?; + self.sorter.write_sorted( + BufWriter::new(file), + &self.output_path.display().to_string(), + ) } } @@ -400,59 +671,27 @@ impl BamStdoutWriter { } } -/// Coordinate-sorted BAM writer that writes to stdout on `finish()`. +/// Coordinate-sorted BAM writer that writes to stdout on `finish()`, with the +/// same bounded-memory spill/merge behavior as [`SortedBamWriter`]. pub struct SortedBamStdoutWriter { - records: Vec, - header: sam::Header, - compression: i32, - limit_bam_sort_ram: u64, + sorter: CoordinateSorter, } impl SortedBamStdoutWriter { pub fn create(genome: &crate::genome::Genome, params: &Parameters) -> Result { let header = crate::io::sam::build_sam_header(genome, params)?; Ok(Self { - records: Vec::new(), - header, - compression: params.out_bam_compression, - limit_bam_sort_ram: params.limit_bam_sort_ram, + sorter: CoordinateSorter::new(header, params), }) } pub fn write_batch(&mut self, batch: &[RecordBuf]) -> Result<(), Error> { - self.records.extend_from_slice(batch); - Ok(()) + self.sorter.push_batch(batch) } pub fn finish(&mut self) -> Result<(), Error> { - if self.limit_bam_sort_ram > 0 { - let est = self.records.len() as u64 * 400; - if est > self.limit_bam_sort_ram { - return Err(Error::Alignment(format!( - "limitBAMsortRAM={} bytes exceeded: estimated {} bytes for {} records.", - self.limit_bam_sort_ram, - est, - self.records.len() - ))); - } - } - self.records - .sort_by_key(|r| match (r.reference_sequence_id(), r.alignment_start()) { - (Some(chr), Some(pos)) => (chr, pos.get()), - _ => (usize::MAX, 0), - }); - let mut bgzf = make_bgzf_writer(BufWriter::new(std::io::stdout()), self.compression); - write_bam_header_lenient(&mut bgzf, &self.header, Some("coordinate"))?; - let mut bam_writer = bam::io::Writer::from(bgzf); - for record in &self.records { - bam_writer.write_alignment_record(&self.header, record)?; - } - bam_writer.finish(&self.header)?; - log::info!( - "Sorted BAM written to stdout ({} records)", - self.records.len() - ); - Ok(()) + self.sorter + .write_sorted(BufWriter::new(std::io::stdout()), "stdout") } } @@ -520,6 +759,32 @@ mod tests { assert!(result.is_ok(), "Finishing BAM file should succeed"); } + #[test] + fn test_bam_missing_quality_is_encoded_as_absent() { + let genome = create_test_genome(); + let params = default_params(); + let temp_file = NamedTempFile::new().unwrap(); + + let record = crate::io::sam::SamWriter::build_unmapped_record( + "read1", + &[0, 1, 2, 3], + &[], + ¶ms, + crate::stats::UnmappedReason::Other, + ) + .unwrap(); + let mut writer = BamWriter::create(temp_file.path(), &genome, ¶ms).unwrap(); + writer.write_batch(&[record]).unwrap(); + writer.finish().unwrap(); + drop(writer); + + let mut reader = bam::io::Reader::new(File::open(temp_file.path()).unwrap()); + reader.read_header().unwrap(); + let record = reader.records().next().unwrap().unwrap(); + assert_eq!(record.sequence().len(), 4); + assert!(record.quality_scores().is_empty()); + } + #[test] fn test_bam_alignment_write() { use cigar::op::{Kind, Op}; @@ -660,38 +925,214 @@ mod tests { ); } - #[test] - fn test_sorted_bam_limit_ram_unlimited() { - let genome = create_test_genome(); + /// A record at `(chr, pos)` named `name`; `None` position means unmapped. + fn placed_record(name: &str, chr: Option, pos: Option) -> RecordBuf { + let mut record = RecordBuf::default(); + *record.name_mut() = Some(name.into()); + *record.reference_sequence_id_mut() = chr; + *record.alignment_start_mut() = pos.map(|pos| pos.try_into().unwrap()); + record + } + + /// `(name, reference_sequence_id, alignment_start)` of a decoded record. + type DecodedRecord = (String, Option, Option); + + /// Decode a sorted BAM into `(name, reference_sequence_id, alignment_start)`. + fn read_sorted(path: &Path) -> Vec { + let mut reader = bam::io::Reader::new(File::open(path).unwrap()); + reader.read_header().unwrap(); + let mut out = Vec::new(); + let mut record = RecordBuf::default(); + let header = sam::Header::default(); + while reader.read_record_buf(&header, &mut record).unwrap() != 0 { + out.push(( + String::from_utf8(record.name().unwrap().to_vec()).unwrap(), + record.reference_sequence_id(), + record.alignment_start().map(|p| p.get()), + )); + } + out + } + + /// A two-chromosome genome, so cross-reference ordering is exercised too. + /// + /// References are long enough to contain every position the fixtures use — + /// records past `LN` would make an out-of-spec BAM that only round-trips + /// because noodles does not validate position against reference length. + fn two_chr_genome() -> Genome { + Genome { + transform_blocks: None, + sequence: vec![0u8; 400].into(), + n_genome: 400, + n_genome_real: 400, + n_chr_real: 2, + chr_name: vec!["chr1".to_string(), "chr2".to_string()], + chr_length: vec![200, 200], + chr_start: vec![0, 200, 400], + } + } + + /// Records in deliberately unsorted input order, spanning two references, + /// including coordinate ties (so tie-order is observable) and unmapped + /// records (which must sort last). + fn shuffled_records() -> Vec { + let mut records = Vec::new(); + for i in 0..200usize { + // Positions cycle so input order and sorted order disagree, and every + // (chr, pos) is hit twice to create ties. + let pos = (i * 37) % 100 + 1; + records.push(placed_record(&format!("r{i}"), Some(i % 2), Some(pos))); + } + for i in 0..10usize { + records.push(placed_record(&format!("u{i}"), None, None)); + } + records + } + + /// Write `records` in batches of `batch` through a `SortedBamWriter` whose + /// sort budget is `ram_limit`, returning the decoded output. + fn sorted_output( + records: &[RecordBuf], + ram_limit: u64, + batch: usize, + ) -> (Vec, usize) { + let genome = two_chr_genome(); + let dir = tempfile::tempdir().unwrap(); let mut params = default_params(); - params.limit_bam_sort_ram = 0; // unlimited - let temp_file = NamedTempFile::new().unwrap(); - let mut writer = SortedBamWriter::create(temp_file.path(), &genome, ¶ms).unwrap(); - let result = writer.finish(); + params.limit_bam_sort_ram = ram_limit; + params.out_file_name_prefix = format!("{}/", dir.path().display()); + + let out = dir.path().join("Aligned.sortedByCoord.out.bam"); + let mut writer = SortedBamWriter::create(&out, &genome, ¶ms).unwrap(); + for chunk in records.chunks(batch) { + writer.write_batch(chunk).unwrap(); + } + // Captured before `finish()`, which consumes and deletes the runs. + let spill_runs = writer.sorter.runs.len(); + writer.finish().unwrap(); + + let decoded = read_sorted(&out); + // Spill files must not outlive the merge. + let leftover = std::fs::read_dir(dir.path()) + .unwrap() + .filter_map(Result::ok) + .filter(|entry| { + entry + .file_name() + .to_string_lossy() + .starts_with("rustar-bamsort-") + }) + .count(); + assert_eq!(leftover, 0, "spill files must be cleaned up"); + (decoded, spill_runs) + } + + #[test] + fn test_sorted_bam_spill_merge_matches_in_memory_sort() { + let records = shuffled_records(); + + // A budget far above the data never spills — the previous behavior. + let (in_memory, in_memory_runs) = sorted_output(&records, 1 << 30, 64); + assert_eq!(in_memory_runs, 0); + assert_eq!(in_memory.len(), records.len()); + + // Tiny budgets force a spill per batch; output must be unchanged. + for batch in [1usize, 7, 64] { + let (spilled, runs) = sorted_output(&records, 1, batch); + assert!( + runs >= records.len() / batch, + "expected a spill run per batch (batch={batch}, runs={runs})" + ); + assert_eq!( + spilled, in_memory, + "spill+merge output must equal the in-memory sort (batch={batch})" + ); + } + } + + #[test] + fn test_sorted_bam_coordinate_ties_keep_input_order_across_runs() { + // Every record shares one coordinate, so the only thing under test is + // whether the merge preserves input order across spill runs. + let records: Vec<_> = (0..64) + .map(|i| placed_record(&format!("r{i:03}"), Some(0), Some(1))) + .collect(); + let expected: Vec<_> = records + .iter() + .map(|r| String::from_utf8(r.name().unwrap().to_vec()).unwrap()) + .collect(); + + for batch in [1usize, 3, 16] { + let (decoded, runs) = sorted_output(&records, 1, batch); + assert!(runs > 1, "batch={batch} should produce several runs"); + let names: Vec<_> = decoded.into_iter().map(|(name, ..)| name).collect(); + assert_eq!( + names, expected, + "tie order must be input order (batch={batch})" + ); + } + } + + /// More runs than `MAX_OPEN_RUNS` must merge in passes rather than opening + /// every run at once — opening them all hits `EMFILE` where the descriptor + /// limit is low (macOS defaults to 256), and a small `--limitBAMsortRAM` on a + /// large run produces thousands of runs. + #[test] + fn test_sorted_bam_merges_more_runs_than_the_descriptor_cap() { + // One record per batch with a 1-byte budget => one spill run each. + let records: Vec<_> = (0..(MAX_OPEN_RUNS * 5 + 3)) + .map(|i| placed_record(&format!("r{i:05}"), Some(i % 2), Some(i % 150 + 1))) + .collect(); + let expected = { + let (decoded, runs) = sorted_output(&records, 1 << 30, records.len()); + assert_eq!(runs, 0); + decoded + }; + + let (decoded, runs) = sorted_output(&records, 1, 1); assert!( - result.is_ok(), - "Sorted BAM with unlimited RAM should succeed" + runs > MAX_OPEN_RUNS, + "fixture must exceed the cap to exercise the reduction pass: {runs}" + ); + assert_eq!( + decoded, expected, + "multi-pass merge must match the in-memory sort exactly" ); } #[test] - fn test_sorted_bam_limit_ram_exceeded() { - let genome = create_test_genome(); + fn test_sorted_bam_empty_output_is_a_valid_bam() { + let (decoded, runs) = sorted_output(&[], 1 << 30, 64); + assert_eq!(runs, 0); + assert!(decoded.is_empty()); + } + + #[test] + fn test_bam_sort_ram_default_is_bounded() { + // `--limitBAMsortRAM 0` must resolve to a fixed budget, never "unlimited", + // and must not scale with the genome (which would enlarge the buffer + // exactly when RAM is tightest). let mut params = default_params(); - params.limit_bam_sort_ram = 1; // 1 byte — will be exceeded by any records - let temp_file = NamedTempFile::new().unwrap(); - let mut writer = SortedBamWriter::create(temp_file.path(), &genome, ¶ms).unwrap(); - // Add a record to trigger the limit - let rec = crate::io::sam::SamWriter::build_unmapped_record( - "r1", - &[0, 1, 2, 3], - &[30; 4], - ¶ms, - crate::stats::UnmappedReason::Other, - ) - .unwrap(); - writer.write_batch(&[rec]).unwrap(); - let result = writer.finish(); - assert!(result.is_err(), "Should fail when RAM limit is exceeded"); + params.limit_bam_sort_ram = 0; + assert_eq!(resolve_bam_sort_ram(¶ms), DEFAULT_BAM_SORT_RAM); + params.genome_dir = std::path::PathBuf::from("/nonexistent/huge/index"); + assert_eq!(resolve_bam_sort_ram(¶ms), DEFAULT_BAM_SORT_RAM); + + params.limit_bam_sort_ram = 12345; + assert_eq!(resolve_bam_sort_ram(¶ms), 12345); + } + + #[test] + fn test_record_size_estimate_tracks_record_contents() { + let small = placed_record("r", Some(0), Some(1)); + let mut large = placed_record("r", Some(0), Some(1)); + *large.sequence_mut() = + noodles::sam::alignment::record_buf::Sequence::from(vec![b'A'; 150]); + *large.quality_scores_mut() = + noodles::sam::alignment::record_buf::QualityScores::from(vec![30u8; 150]); + assert!( + estimated_record_bytes(&large) >= estimated_record_bytes(&small) + 300, + "estimate must account for SEQ and QUAL" + ); } } From fc16d8a31601db253a49dfb3010c81397253cc16 Mon Sep 17 00:00:00 2001 From: alejandrogzi Date: Thu, 6 Aug 2026 23:21:34 +0200 Subject: [PATCH 2/2] test(io): cover bounded-memory coordinate-sort spilling A low --limitBAMsortRAM must spill sorted runs to disk and merge them into a BAM identical to the unbounded in-memory sort, leave no scratch behind, and emit non-decreasing (chrom, pos) order. --- tests/alignment_features.rs | 106 ++++++++++++++++++++++++++++++++++++ 1 file changed, 106 insertions(+) diff --git a/tests/alignment_features.rs b/tests/alignment_features.rs index 62ee32a..815755b 100644 --- a/tests/alignment_features.rs +++ b/tests/alignment_features.rs @@ -2024,6 +2024,112 @@ fn test_wasp_samtag() { "all 10 unique reads overlapping the het SNV should pass WASP (vW:i:1)" ); } +// --------------------------------------------------------------------------- +// Coordinate-sort spilling (--limitBAMsortRAM) +// --------------------------------------------------------------------------- + +/// A low `--limitBAMsortRAM` must spill sorted runs to disk and merge them into +/// a BAM identical to the unbounded in-memory sort — same records, same order. +#[test] +fn test_sorted_bam_spills_to_disk_and_matches_unbounded_sort() { + let tmpdir = TempDir::new().unwrap(); + let genome = build_genome(); + let fasta = write_fasta(&tmpdir, &genome); + let genome_dir = tmpdir.path().join("genome"); + build_index(&fasta, &genome_dir, "7", None); + + // 2000 reads across the genome, emitted in an order that does not match + // coordinate order so the sort is actually doing work. + let fastq_path = tmpdir.path().join("reads.fq"); + { + let mut f = fs::File::create(&fastq_path).unwrap(); + for i in 0..2000usize { + let start = (i * 7919) % (genome.len() - 60); + let seq = &genome[start..start + 50]; + writeln!(f, "@read{i}").unwrap(); + f.write_all(seq).unwrap(); + writeln!(f, "\n+\n{}", "I".repeat(50)).unwrap(); + } + } + + let run = |label: &str, limit: &str| -> (PathBuf, PathBuf) { + let output_dir = tmpdir.path().join(label); + fs::create_dir_all(&output_dir).unwrap(); + let prefix = format!("{}/", output_dir.display()); + cargo_bin_cmd!("rustar-aligner") + .args([ + "--runMode", + "alignReads", + "--genomeDir", + genome_dir.to_str().unwrap(), + "--readFilesIn", + fastq_path.to_str().unwrap(), + "--outSAMtype", + "BAM", + "SortedByCoordinate", + "--limitBAMsortRAM", + limit, + "--outFileNamePrefix", + &prefix, + ]) + .assert() + .success(); + (output_dir.join("Aligned.sortedByCoord.out.bam"), output_dir) + }; + + // 64 KiB forces many spill runs; 1 GiB holds everything in memory. + let (spilled_bam, spilled_dir) = run("out_sort_spill", "65536"); + let (memory_bam, _) = run("out_sort_memory", "1G"); + + let read_bam = |path: &PathBuf| -> Vec<(String, Option, Option, String)> { + let mut reader = bam::io::Reader::new(fs::File::open(path).unwrap()); + reader.read_header().unwrap(); + reader + .records() + .map(|record| { + let record = record.unwrap(); + ( + String::from_utf8(record.name().unwrap().to_vec()).unwrap(), + record.reference_sequence_id().transpose().unwrap(), + record + .alignment_start() + .transpose() + .unwrap() + .map(|p| p.get()), + format!("{:?}", record.cigar().iter().collect::>()), + ) + }) + .collect() + }; + + let spilled = read_bam(&spilled_bam); + let in_memory = read_bam(&memory_bam); + + assert!(!spilled.is_empty(), "expected alignments"); + assert_eq!( + spilled, in_memory, + "spill+merge output must equal the unbounded in-memory sort" + ); + + // Coordinate-sorted, and no spill scratch left behind. + let keys: Vec<_> = spilled + .iter() + .map(|(_, chr, pos, _)| (chr.unwrap_or(usize::MAX), pos.unwrap_or(0))) + .collect(); + assert!( + keys.windows(2).all(|w| w[0] <= w[1]), + "output must be non-decreasing by (chr, pos)" + ); + let leftover: Vec<_> = fs::read_dir(&spilled_dir) + .unwrap() + .filter_map(Result::ok) + .map(|e| e.file_name().to_string_lossy().to_string()) + .filter(|name| name.starts_with("rustar-bamsort-")) + .collect(); + assert!(leftover.is_empty(), "spill files left behind: {leftover:?}"); +} + + // --------------------------------------------------------------------------- // --runMode soloCellFiltering