perf(align): cut the align batch size and stop cloning every FASTQ record

CHANGELOG.md

@@ -13,6 +13,21 @@ Sections commonly used: Features, Bug fixes, Other changes.
 
 ### Other changes
 
+- **Align batch size 10 000 → 1 000** (`ALIGN_BATCH_SIZE`, was a magic number
+  repeated in the four align pipelines). A batch is live end to end — decoded
+  reads in, SAM records out — and several are in flight at once, so peak RSS
+  scaled with it. Measured on a 50 Mb genome, 2 M SE 100 bp reads, 8 threads,
+  BAM output: wall 3.24 s → 2.43 s, user CPU 25.6 s → 17.5 s, peak RSS 1543 MB
+  → 669 MB. The CPU drop is the point: at 10 000 the batch working set no
+  longer fits in cache next to the per-read alignment scratch. Output is
+  byte-identical at every batch size tried.
+
+- FASTQ decode reads into one reusable `noodles` record instead of going
+  through `Reader::records()`, which allocates a record per call and yields a
+  clone of it — two extra copies of every sequence and quality string per read.
+  The read name is also cut at `--readNameSeparator` before being allocated
+  rather than allocated and then copied shorter.
+
 - `cluster_seeds` reuses its window-bin map across reads on a thread instead
   of rebuilding it per read. Merging two windows re-keys every bin in the
   merged span, so the per-read pre-sizing was only a floor and the map

src/io/fastq.rs

@@ -72,6 +72,11 @@ pub struct PairedRead {
 /// FASTQ reader that handles decompression and base encoding
 pub struct FastqReader {
     inner: fastq::io::Reader<Box<dyn BufRead + Send>>,
+    /// Scratch record reused across reads. `Reader::records()` allocates a fresh
+    /// record buffer per call and clones it out on every `next()`, so going
+    /// through the iterator cost two extra copies of the sequence and quality
+    /// bytes per read. Reading into one long-lived buffer keeps the capacity.
+    buf: fastq::Record,
     /// Signed shift applied to every input quality byte so the rest of the
     /// pipeline always sees Phred+33.
     ///
@@ -128,6 +133,7 @@ impl FastqReader {
 
         Ok(Self {
             inner: fastq_reader,
+            buf: fastq::Record::default(),
             qual_shift: 0,
             name_separators: vec![b'/'],
         })
@@ -171,48 +177,49 @@ impl FastqReader {
 
     /// Get next read with encoded bases
     pub fn next_encoded(&mut self) -> Result<Option<EncodedRead>, Error> {
-        match self.inner.records().next() {
-            Some(Ok(record)) => {
-                let name = std::str::from_utf8(record.name())
-                    .map_err(|e| {
-                        Error::from(std::io::Error::new(
-                            std::io::ErrorKind::InvalidData,
-                            format!("invalid UTF-8 in read name: {e}"),
-                        ))
-                    })?
-                    .to_string();
-
-                let sequence = record.sequence().iter().map(|&b| encode_base(b)).collect();
-
-                let name = match self
-                    .name_separators
-                    .iter()
-                    .filter_map(|&sep| name.as_bytes().iter().position(|&b| b == sep))
-                    .min()
-                {
-                    Some(cut) => name[..cut].to_string(),
-                    None => name,
-                };
-
-                let quality = if self.qual_shift == 0 {
-                    record.quality_scores().to_vec()
-                } else {
-                    record
-                        .quality_scores()
-                        .iter()
-                        .map(|&b| (b as i32 + self.qual_shift).clamp(33, 126) as u8)
-                        .collect()
-                };
-
-                Ok(Some(EncodedRead {
-                    name,
-                    sequence,
-                    quality,
-                }))
-            }
-            Some(Err(e)) => Err(Error::from(e)),
-            None => Ok(None),
+        // Read straight into the reusable buffer (`read_record` clears it first). The `records()` iterator would
+        // allocate a record per call and hand back a clone of it, i.e. two extra
+        // copies of every sequence and quality string.
+        if self.inner.read_record(&mut self.buf).map_err(Error::from)? == 0 {
+            return Ok(None);
         }
+        let record = &self.buf;
+
+        // Cut at the first separator before allocating, so a trimmed name is one
+        // allocation rather than an allocation plus a shortened copy.
+        let raw_name = record.name();
+        let end = self
+            .name_separators
+            .iter()
+            .filter_map(|&sep| raw_name.iter().position(|&b| b == sep))
+            .min()
+            .unwrap_or(raw_name.len());
+        let name = std::str::from_utf8(&raw_name[..end])
+            .map_err(|e| {
+                Error::from(std::io::Error::new(
+                    std::io::ErrorKind::InvalidData,
+                    format!("invalid UTF-8 in read name: {e}"),
+                ))
+            })?
+            .to_string();
+
+        let sequence = record.sequence().iter().map(|&b| encode_base(b)).collect();
+
+        let quality = if self.qual_shift == 0 {
+            record.quality_scores().to_vec()
+        } else {
+            record
+                .quality_scores()
+                .iter()
+                .map(|&b| (b as i32 + self.qual_shift).clamp(33, 126) as u8)
+                .collect()
+        };
+
+        Ok(Some(EncodedRead {
+            name,
+            sequence,
+            quality,
+        }))
     }
 
     /// Read a batch of encoded reads for parallel processing

src/lib.rs

@@ -47,6 +47,31 @@ use noodles::sam::alignment::record::cigar;
 
 use crate::params::{Parameters, RunMode};
 
+/// Reads decoded, aligned and written as one unit by the align pipelines.
+///
+/// The whole batch is live at once — decoded reads on the way in, their SAM
+/// records on the way out — and the pipeline keeps several in flight (a bounded
+/// channel per stage, plus one per rayon worker), so peak RSS scales with this
+/// number times the worker count, not with the input size.
+///
+/// It used to be 10 000. Measured on a 50 Mb genome, 2 M single-end 100 bp
+/// reads, 8 threads, BAM output (median of two runs):
+///
+/// | batch  | wall  | user CPU | peak RSS |
+/// |--------|-------|----------|----------|
+/// | 10 000 | 3.24s | 25.6s    | 1543 MB  |
+/// |  2 500 | 2.69s | 19.6s    |  811 MB  |
+/// |  1 000 | 2.43s | 17.5s    |  669 MB  |
+/// |    500 | 2.19s | 17.0s    |  660 MB  |
+///
+/// Smaller batches are not merely cheaper in memory, they are cheaper in CPU:
+/// a batch's working set has to fit in cache alongside the per-read alignment
+/// scratch, and at 10 000 it does not. The curve flattens below ~1 000, so that
+/// is the value here; going lower trades away the per-batch amortisation for
+/// nothing. Output is byte-identical at every batch size (verified on the 2 M
+/// read run: same BAM payload hash).
+const ALIGN_BATCH_SIZE: usize = 1000;
+
 /// Top-level dispatcher. Called from `main()` after CLI parsing.
 pub fn run(params: &Parameters) -> anyhow::Result<()> {
     info!("rustar-aligner {}", env!("CARGO_PKG_VERSION"));
@@ -1437,7 +1462,7 @@ fn align_reads_single_end<W: AlignmentWriter + ?Sized>(
         params.read_map_number as u64
     };
 
-    let batch_size = 10000;
+    let batch_size = ALIGN_BATCH_SIZE;
     let max_multimaps = params.out_filter_multimap_nmax as usize;
     // `--outSAMtype None` (e.g. quant-only) skips building SAM records.
     let emit_sam = params.emits_alignments();
@@ -2069,7 +2094,7 @@ fn align_reads_solo<W: AlignmentWriter + ?Sized>(
     } else {
         params.read_map_number as u64
     };
-    let batch_size = 10000;
+    let batch_size = ALIGN_BATCH_SIZE;
     let clip5p = params.clip5p(0);
     let clip3p = params.clip3p(0);
     let cr4_clip = params.clip_adapter_type == "CellRanger4";
@@ -2384,7 +2409,7 @@ fn align_reads_solo_pe<W: AlignmentWriter + ?Sized>(
     } else {
         params.read_map_number as u64
     };
-    let batch_size = 10000;
+    let batch_size = ALIGN_BATCH_SIZE;
     // Per-mate clip: mate 1 (--clip5pNbases[0], e.g. 39 to strip the 5' barcode
     // region) and mate 2 ([1], e.g. 0). CellRanger4 adapter clipping is not used
     // by the cellgeni 5' path (it uses clip5pNbases instead), so it is not applied.
@@ -2777,7 +2802,7 @@ fn align_reads_paired_end<W: AlignmentWriter + ?Sized>(
         params.read_map_number as u64
     };
 
-    let batch_size = 10000;
+    let batch_size = ALIGN_BATCH_SIZE;
     let max_multimaps = params.out_filter_multimap_nmax as usize;
     // `--outSAMtype None` (e.g. quant-only) skips building SAM records.
     let emit_sam = params.emits_alignments();

src/ruSTAR.code-workspace

@@ -0,0 +1,10 @@
+{
+    "folders": [
+        {
+            "path": "../../STAR-rs"
+        },
+        {
+            "path": ".."
+        }
+    ]
+}