feat(bam): --outBAMsortingBinsN, spilling the coordinate sort to disk bins

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

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<SpillState>,
+}
+
+/// 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<std::path::PathBuf>,
+    writers: Vec<bam::io::Writer<bgzf::io::Writer<BufWriter<File>>>>,
+    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<RecordBuf> = 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<std::path::PathBuf> = (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, &params).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());
     }
 }

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,

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",