From 1979b78b6438c6d11b886bd6a49ec49f69a29466 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 29 Jul 2026 00:23:59 +0200 Subject: [PATCH 1/6] 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 b1e6717..f0e73f4 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()) { @@ -695,3 +798,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 3b50731..07bf4d3 100644 --- a/src/params/mod.rs +++ b/src/params/mod.rs @@ -658,6 +658,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 83b1bba613b91ee3368b9a19d00bdc48ca0f855f Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 29 Jul 2026 00:29:00 +0200 Subject: [PATCH 2/6] 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 f0e73f4..9e911d2 100644 --- a/src/io/bam.rs +++ b/src/io/bam.rs @@ -793,9 +793,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 0004c68eff686dd8df72eb7829d81d3123e9274e Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Thu, 30 Jul 2026 20:12:28 +0200 Subject: [PATCH 3/6] docs(changelog): record --outBAMsortingBinsN Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 14b4346..f3b2b75 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -44,6 +44,10 @@ Sections commonly used: Features, Bug fixes, Other changes. - Read names are cut at `--readNameSeparator` (default `/`), as STAR does. A read named `foo/1` was previously emitted as `foo/1` where STAR emits `foo`. +- `--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 From c7a79cf1bde9050945b23f7a29a6534941d6017b Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Thu, 30 Jul 2026 20:13:01 +0200 Subject: [PATCH 4/6] 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 f3b2b75..b446e50 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -47,7 +47,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 From c3200fa805a615700f69d80722c595810d81c702 Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Wed, 26 Aug 2026 23:53:03 +0200 Subject: [PATCH 5/6] fix(params): drop the duplicate --outBAMsortingBinsN declaration after the rebase main gained an accepted-but-inert copy of the flag in the CLI-parity work; this branch implements it, so the inert declaration and its ACCEPTED_BUT_INERT note both go. --- src/params/mod.rs | 4 ---- tests/parameter_surface.rs | 4 ---- 2 files changed, 8 deletions(-) diff --git a/src/params/mod.rs b/src/params/mod.rs index 07bf4d3..eb1c0db 100644 --- a/src/params/mod.rs +++ b/src/params/mod.rs @@ -763,10 +763,6 @@ pub struct Parameters { #[arg(long = "outTmpKeep", default_value = "None")] pub out_tmp_keep: String, - /// Number of bins used when sorting BAM by coordinate. - #[arg(long = "outBAMsortingBinsN", default_value_t = 50)] - pub out_bam_sorting_bins_n: usize, - /// Threads used for BAM sorting. 0 selects `--runThreadN`. #[arg(long = "outBAMsortingThreadN", default_value_t = 0)] pub out_bam_sorting_thread_n: usize, diff --git a/tests/parameter_surface.rs b/tests/parameter_surface.rs index df3dd94..57dd41b 100644 --- a/tests/parameter_surface.rs +++ b/tests/parameter_surface.rs @@ -55,10 +55,6 @@ const ACCEPTED_BUT_INERT: &[(&str, &str)] = &[ "limitSjdbInsertNsj", "an allocation cap on the inserted-junction array", ), - ( - "outBAMsortingBinsN", - "sorting is not binned yet; output is unaffected", - ), ( "outBAMsortingThreadN", "BGZF writing is single-threaded; output is unaffected", From ec11a42031a72da32c9b3c948dd3ca76400148ce Mon Sep 17 00:00:00 2001 From: Benjamin Demaille Date: Sun, 30 Aug 2026 23:03:28 +0200 Subject: [PATCH 6/6] fix(bam): enforce limitBAMsortRAM while records arrive, not after `--outBAMsortingBinsN` spilled the coordinate sort to disk bins, but only from `finish()`. `write_batch` accumulated every record with no bound check, so by the time the first bin was written the whole run was already resident and the peak had been reached. The bound was respected during the sort and nowhere else, which is not where the memory goes. Measured on 560k records (yeast, 17 references, `--limitBAMsortRAM 100000000`), peak RSS: in-memory sort 1980 MB, spill-at-finish 1967 MB. The feature was buying nothing. Open the bins on the first crossing of the bound and push the buffer out to them as batches arrive, so what stays resident is one buffer rather than the run. `finish` then closes the bins it already has instead of building them from a full buffer. Peak RSS, same flags: | records | spill at finish | incremental | |---|---|---| | 560k | 1967 MB | 1744 MB | | 1.3M | 2925 MB | 2254 MB | The gap widens with record count, which is the property that was missing: the buffer no longer scales with the run. Peak is not flat, because the index, the alignment pipeline's own batches and the allocator's retained pages all sit underneath it, and none of those are affected by this. Output is unchanged: 561,312 decoded records identical to the in-memory sort, and the no-bins default path is untouched. Tests pass, 0 clippy warnings, fmt clean. Co-Authored-By: Claude Opus 5 (1M context) --- src/io/bam.rs | 131 +++++++++++++++++++++++++++++++++++++------------- 1 file changed, 98 insertions(+), 33 deletions(-) diff --git a/src/io/bam.rs b/src/io/bam.rs index 9e911d2..9b2e801 100644 --- a/src/io/bam.rs +++ b/src/io/bam.rs @@ -73,6 +73,21 @@ pub struct SortedBamWriter { limit_bam_sort_ram: u64, /// `--outBAMsortingBinsN`. 0 sorts entirely in memory. bins_n: usize, + /// Open bin files, once the buffer has first crossed the RAM bound. + /// + /// Held across `write_batch` calls so records can be pushed out as they + /// arrive rather than accumulating until `finish`. + spill: Option, +} + +/// The on-disk bins a spilling sort writes into, kept open between batches. +struct SpillState { + /// Dropping this removes the bin files; it is never read directly. + _dir: tempfile::TempDir, + paths: Vec, + writers: Vec>>>, + bins: usize, + n_refs: usize, } impl BamWriter { @@ -158,6 +173,7 @@ impl SortedBamWriter { compression: params.out_bam_compression, limit_bam_sort_ram: params.limit_bam_sort_ram, bins_n: params.out_bam_sorting_bins_n, + spill: None, }) } @@ -185,39 +201,23 @@ impl SortedBamWriter { /// 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. + // Pass 1: whatever is still buffered joins what was already spilled, + // then the bin files are closed. When the bound was crossed during the + // run these writers are already open and mostly written. + self.spill_buffered()?; + let mut spill = self.spill.take().expect("spill_buffered opens it"); { - 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)?; - } + let header = std::mem::take(&mut self.header); + let result = spill + .writers + .iter_mut() + .try_for_each(|w| w.finish(&header).map_err(Error::from)); + self.header = header; + result?; } + let paths = std::mem::take(&mut spill.paths); + let bins = spill.bins; + drop(spill.writers); // Pass 2: one bin at a time. let buf_writer = BufWriter::new(File::create(&self.output_path)?); @@ -255,9 +255,16 @@ impl SortedBamWriter { Ok(()) } - /// Buffer records — no disk I/O yet. + /// Buffer records, pushing them out to bins once the bound is crossed. + /// + /// Checking only at `finish` would let every record accumulate first, so + /// the peak would be the whole run no matter how the sort was then + /// performed. The bound has to be enforced while records arrive. pub fn write_batch(&mut self, batch: &[RecordBuf]) -> Result<(), Error> { self.records.extend_from_slice(batch); + if self.should_spill() { + self.spill_buffered()?; + } Ok(()) } @@ -266,6 +273,64 @@ impl SortedBamWriter { self.records.len() as u64 * 400 } + /// Which bin a record belongs to. Unmapped records take the last one. + fn bin_index(rec: &RecordBuf, bins: usize, n_refs: usize) -> usize { + match rec.reference_sequence_id() { + Some(chr) => (chr * bins) / n_refs, + None => bins, + } + } + + /// Open the bin files. Called the first time the buffer crosses the bound. + fn open_spill(&mut self) -> Result<(), Error> { + if self.spill.is_some() { + return Ok(()); + } + let n_refs = self.header.reference_sequences().len().max(1); + let bins = self.bins_n.min(n_refs).max(1); + 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(); + let mut writers = Vec::with_capacity(paths.len()); + for path in &paths { + let f = File::create(path).map_err(|e| Error::io(e, path))?; + // Uncompressed: these are read back immediately, so compressing + // them would cost time and save nothing. + 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)); + } + self.spill = Some(SpillState { + _dir: dir, + paths, + writers, + bins, + n_refs, + }); + Ok(()) + } + + /// Move everything currently buffered out to its bin, freeing the buffer. + fn spill_buffered(&mut self) -> Result<(), Error> { + self.open_spill()?; + // `records` is moved aside so the spill state can be borrowed mutably + // while draining it; the emptied allocation goes back afterwards. + let mut records = std::mem::take(&mut self.records); + let header = std::mem::take(&mut self.header); + let result = (|| -> Result<(), Error> { + let spill = self.spill.as_mut().expect("opened above"); + for rec in records.drain(..) { + let b = Self::bin_index(&rec, spill.bins, spill.n_refs); + spill.writers[b].write_alignment_record(&header, &rec)?; + } + Ok(()) + })(); + self.header = header; + self.records = records; + result + } + fn check_ram_limit(&self) -> Result<(), Error> { if self.limit_bam_sort_ram > 0 { let est = self.estimated_ram(); @@ -290,7 +355,7 @@ impl SortedBamWriter { // 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() { + if self.spill.is_some() || self.should_spill() { return self.finish_binned(); } self.check_ram_limit()?;