diff --git a/CHANGELOG.md b/CHANGELOG.md index 14b4346..b446e50 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 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 diff --git a/src/io/bam.rs b/src/io/bam.rs index b1e6717..9b2e801 100644 --- a/src/io/bam.rs +++ b/src/io/bam.rs @@ -71,6 +71,23 @@ pub struct SortedBamWriter { header: sam::Header, compression: i32, 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 { @@ -155,12 +172,99 @@ 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, + spill: None, }) } - /// Buffer records — no disk I/O yet. + /// 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> { + // 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 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)?); + 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, 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(()) } @@ -169,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(); @@ -190,6 +352,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.spill.is_some() || 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()) { @@ -690,8 +858,77 @@ 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" + ); + } +} + +#[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..eb1c0db 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, @@ -757,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",