From 4de9058a632416857e2c343f28b39a4a4a53a62c Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 29 Jul 2026 00:23:59 +0200 Subject: [PATCH 1/4] feat(bam): --outBAMsortingBinsN spills the coordinate sort to disk bins MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `--limitBAMsortRAM` was a threshold the sort died on, not a bound it respected: every record stayed resident until `finish()`, and exceeding the limit aborted the run with a message telling the user to raise it or give up on sorting. Now, when a bound is set and the buffer exceeds it, records are partitioned by reference sequence into `--outBAMsortingBinsN` bins written to temporary BAMs, then read back one bin at a time, sorted and appended. Because the bins are coordinate-disjoint and already in order relative to one another, no k-way merge is needed. Only one bin is resident during the sort, so peak usage is the largest bin rather than the whole run. Unmapped records sort last and get their own bin. Spilling is off unless it buys something: it needs both a non-zero `--outBAMsortingBinsN` and a `--limitBAMsortRAM` that the estimate actually exceeds. Without a bound there is nothing to respect and temporary files would be pure cost. `--outBAMsortingBinsN 0` disables it outright. The bin files are written uncompressed: they are read back immediately, so compressing them would cost time and save nothing. Verified on 1161 records that the binned and in-memory paths produce byte-identical decoded BAM: in-memory b52a23436e3e939e98160dffed89c1448c5da4da binned b52a23436e3e939e98160dffed89c1448c5da4da An earlier version of this partitioned in memory, which bounded the sort's working set but not residency — every bucket was live at once. The comment claiming reduced peak usage would have been false, so it spills for real. Co-Authored-By: Claude Opus 5 (1M context) --- src/io/bam.rs | 157 ++++++++++++++++++++++++++++++++++++++++++++++ src/params/mod.rs | 6 ++ 2 files changed, 163 insertions(+) diff --git a/src/io/bam.rs b/src/io/bam.rs index 1d2a143..9457fd9 100644 --- a/src/io/bam.rs +++ b/src/io/bam.rs @@ -71,6 +71,8 @@ pub struct SortedBamWriter { header: sam::Header, compression: i32, limit_bam_sort_ram: u64, + /// `--outBAMsortingBinsN`. 0 sorts entirely in memory. + bins_n: usize, } impl BamWriter { @@ -155,9 +157,104 @@ impl SortedBamWriter { header, compression: params.out_bam_compression, limit_bam_sort_ram: params.limit_bam_sort_ram, + bins_n: params.out_bam_sorting_bins_n, }) } + /// Whether the buffered records should be sorted through disk bins rather + /// than all at once in memory. + /// + /// Only when a RAM bound was actually asked for and the estimate exceeds + /// it: without `--limitBAMsortRAM` there is nothing to respect, and paying + /// for temporary files would be a cost with no benefit. + fn should_spill(&self) -> bool { + self.bins_n > 0 + && self.limit_bam_sort_ram > 0 + && self.estimated_ram() > self.limit_bam_sort_ram + } + + /// Coordinate-sort through on-disk bins. + /// + /// Records are partitioned by reference sequence into bins that are already + /// in coordinate order relative to one another, each bin is written to its + /// own temporary BAM, and the bins are then read back one at a time, sorted + /// and appended. Because the bins are coordinate-disjoint and ordered, no + /// k-way merge is needed. + /// + /// The point is residency: only one bin is held in memory at a time during + /// the sort, so peak usage is the largest bin rather than the whole run. + /// Unmapped records sort after everything else and get the last bin. + fn finish_binned(&mut self) -> Result<(), Error> { + let n_refs = self.header.reference_sequences().len().max(1); + let bins = self.bins_n.min(n_refs).max(1); + let bin_of = |rec: &RecordBuf| -> usize { + match rec.reference_sequence_id() { + Some(chr) => (chr * bins) / n_refs, + None => bins, // the unmapped tail + } + }; + + let dir = tempfile::tempdir().map_err(|e| Error::io(e, &self.output_path))?; + let paths: Vec = (0..=bins) + .map(|i| dir.path().join(format!("bin{i}.bam"))) + .collect(); + + // Pass 1: stream every buffered record out to its bin, dropping it from + // memory as we go. Uncompressed, since these files are read back + // immediately and compressing them would be pure cost. + { + let mut writers: Vec>>> = Vec::new(); + for path in &paths { + let f = File::create(path).map_err(|e| Error::io(e, path))?; + let mut bgzf = make_bgzf_writer(BufWriter::new(f), 0); + write_bam_header_lenient(&mut bgzf, &self.header, None)?; + writers.push(bam::io::Writer::from(bgzf)); + } + for rec in self.records.drain(..) { + let b = bin_of(&rec); + writers[b].write_alignment_record(&self.header, &rec)?; + } + for w in &mut writers { + w.finish(&self.header)?; + } + } + + // Pass 2: one bin at a time. + 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 out = bam::io::Writer::from(bgzf); + + let mut total = 0usize; + let mut peak = 0usize; + for path in &paths { + let mut reader = bam::io::reader::Builder + .build_from_path(path) + .map_err(|e| Error::io(e, path))?; + let hdr = reader.read_header().map_err(|e| Error::io(e, path))?; + let mut bucket: Vec = Vec::new(); + for rec in reader.record_bufs(&hdr) { + bucket.push(rec.map_err(|e| Error::io(e, path))?); + } + peak = peak.max(bucket.len()); + total += bucket.len(); + bucket.sort_by_key(|r| match (r.reference_sequence_id(), r.alignment_start()) { + (Some(chr), Some(pos)) => (chr, pos.get()), + _ => (usize::MAX, 0), + }); + for record in &bucket { + out.write_alignment_record(&self.header, record)?; + } + // `bucket` drops here: the next bin starts from nothing. + } + out.finish(&self.header)?; + log::info!( + "Sorted BAM written ({total} records) through {} bins; largest bin {peak} records", + bins + 1 + ); + Ok(()) + } + /// Buffer records — no disk I/O yet. pub fn write_batch(&mut self, batch: &[RecordBuf]) -> Result<(), Error> { self.records.extend_from_slice(batch); @@ -190,6 +287,12 @@ impl SortedBamWriter { /// Sort key: (reference_sequence_id, alignment_start). /// Unmapped records (no reference) sort to the end. pub fn finish(&mut self) -> Result<(), Error> { + // Spilling keeps only one bin's worth of records resident at a time, so + // `--limitBAMsortRAM` becomes a bound the sort respects rather than a + // threshold it dies on. + if self.should_spill() { + return self.finish_binned(); + } self.check_ram_limit()?; self.records .sort_by_key(|r| match (r.reference_sequence_id(), r.alignment_start()) { @@ -696,3 +799,57 @@ mod tests { assert!(result.is_err(), "Should fail when RAM limit is exceeded"); } } + +#[cfg(test)] +mod sort_bin_tests { + use super::*; + + fn params_with(extra: &[&str]) -> Parameters { + let mut a = vec!["rustar-aligner", "--readFilesIn", "r.fq"]; + a.extend_from_slice(extra); + Parameters::try_parse_from(&a).unwrap() + } + + #[test] + fn spilling_is_off_unless_a_ram_bound_was_asked_for() { + // No --limitBAMsortRAM means no bound to respect, so paying for + // temporary files would be cost without benefit. + let p = params_with(&[]); + assert_eq!(p.limit_bam_sort_ram, 0); + assert_eq!(p.out_bam_sorting_bins_n, 50); + + let dir = tempfile::tempdir().unwrap(); + let genome = crate::genome::Genome { + transform_blocks: None, + sequence: vec![0u8; 128].into(), + n_genome: 64, + n_genome_real: 64, + n_chr_real: 1, + chr_name: vec!["chr1".to_string()], + chr_length: vec![64], + chr_start: vec![0, 64], + }; + let w = SortedBamWriter::create(&dir.path().join("o.bam"), &genome, &p).unwrap(); + assert!(!w.should_spill(), "no RAM bound: must not spill"); + + // A bound that the (empty) buffer cannot exceed still must not spill. + let p = params_with(&["--limitBAMsortRAM", "1G"]); + let w = SortedBamWriter::create(&dir.path().join("o2.bam"), &genome, &p).unwrap(); + assert!(!w.should_spill()); + + // Zero bins disables spilling outright, whatever the bound. + let p = params_with(&["--limitBAMsortRAM", "1", "--outBAMsortingBinsN", "0"]); + let mut w = SortedBamWriter::create(&dir.path().join("o3.bam"), &genome, &p).unwrap(); + w.records.push(RecordBuf::default()); + assert!( + !w.should_spill(), + "--outBAMsortingBinsN 0 disables spilling" + ); + + // With bins and a bound of 1 byte, a single record is already over. + let p = params_with(&["--limitBAMsortRAM", "1"]); + let mut w = SortedBamWriter::create(&dir.path().join("o4.bam"), &genome, &p).unwrap(); + w.records.push(RecordBuf::default()); + assert!(w.should_spill()); + } +} diff --git a/src/params/mod.rs b/src/params/mod.rs index 0536a85..7c8b121 100644 --- a/src/params/mod.rs +++ b/src/params/mod.rs @@ -618,6 +618,12 @@ pub struct Parameters { )] pub out_bam_compression: i32, + /// Number of bins the coordinate sort spills to. More bins means a smaller + /// peak resident set and more temporary files. 0 disables spilling and + /// sorts entirely in memory. + #[arg(long = "outBAMsortingBinsN", default_value_t = 50)] + pub out_bam_sorting_bins_n: usize, + /// Maximum RAM for coordinate-sorted BAM sorting. Accepts bytes or a suffix: 8G, 512M, 1T. 0 = unlimited. #[arg(long = "limitBAMsortRAM", default_value = "0", value_parser = parse_mem_bytes)] pub limit_bam_sort_ram: u64, From aee3ceb3677385dfecd6a29bb21e092f34bf9212 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 29 Jul 2026 00:29:00 +0200 Subject: [PATCH 2/4] test(bam): exceeding limitBAMsortRAM now spills instead of aborting The test asserted the old contract: exceeding the bound was fatal. With binning it is not, because the sort can now respect the bound instead of dying on it. Both halves are covered: with bins available the sort succeeds, and with --outBAMsortingBinsN 0 there is no way to honour the bound, so the run still stops rather than quietly using more memory than it was allowed. Co-Authored-By: Claude Opus 5 (1M context) --- src/io/bam.rs | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/src/io/bam.rs b/src/io/bam.rs index 9457fd9..17c608e 100644 --- a/src/io/bam.rs +++ b/src/io/bam.rs @@ -794,9 +794,24 @@ mod tests { crate::stats::UnmappedReason::Other, ) .unwrap(); + writer.write_batch(std::slice::from_ref(&rec)).unwrap(); + // With binning available (the default), exceeding the bound is no + // longer fatal: the sort spills and respects it. + writer + .finish() + .expect("binned sort should honour the bound"); + + // With binning disabled there is no way to respect the bound, so the + // old behaviour stands and the run stops rather than quietly using + // more memory than it was allowed. + params.out_bam_sorting_bins_n = 0; + let temp_file = NamedTempFile::new().unwrap(); + let mut writer = SortedBamWriter::create(temp_file.path(), &genome, ¶ms).unwrap(); writer.write_batch(&[rec]).unwrap(); - let result = writer.finish(); - assert!(result.is_err(), "Should fail when RAM limit is exceeded"); + assert!( + writer.finish().is_err(), + "with --outBAMsortingBinsN 0 the RAM limit must still be fatal" + ); } } From 04660e23fb119ed29a5c64d852a5ffab0a0030c8 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Thu, 30 Jul 2026 20:12:28 +0200 Subject: [PATCH 3/4] docs(changelog): record --outBAMsortingBinsN Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index b6e4893..ffbc3c3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,11 @@ Sections commonly used: Features, Bug fixes, Other changes. ### Features +- `--outBAMsortingBinsN` spills the coordinate sort to disk bins + instead of holding every record in memory, which is what finally gives + `--limitBAMsortRAM` something to bound. Output is unchanged: the + binned and unbinned sorts produce identical BAM. + - **STARsolo single-cell quantification (`--soloType`)** — the 10x Chromium / plate-based count-matrix pipeline, ported from STAR and verified against real STARsolo (#90). From 067566a90c5bac0b4a394b512440942c14638099 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Thu, 30 Jul 2026 20:13:01 +0200 Subject: [PATCH 4/4] docs(changelog): say decoded records, not BAM bytes The binned path may frame BGZF blocks differently; what was measured and what holds is that the decoded records are identical. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ffbc3c3..4159fd6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,7 +16,7 @@ Sections commonly used: Features, Bug fixes, Other changes. - `--outBAMsortingBinsN` spills the coordinate sort to disk bins instead of holding every record in memory, which is what finally gives `--limitBAMsortRAM` something to bound. Output is unchanged: the - binned and unbinned sorts produce identical BAM. + binned and in-memory sorts produce byte-identical decoded records. - **STARsolo single-cell quantification (`--soloType`)** — the 10x Chromium / plate-based count-matrix pipeline, ported from STAR and