diff --git a/CHANGELOG.md b/CHANGELOG.md index 14b4346..2237623 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/io/fastq.rs b/src/io/fastq.rs index 0c795de..7e44157 100644 --- a/src/io/fastq.rs +++ b/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>, + /// 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, 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 diff --git a/src/lib.rs b/src/lib.rs index 5086fcb..6f951e2 100644 --- a/src/lib.rs +++ b/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( 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( } 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( } 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( 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(); diff --git a/src/ruSTAR.code-workspace b/src/ruSTAR.code-workspace new file mode 100644 index 0000000..d62934f --- /dev/null +++ b/src/ruSTAR.code-workspace @@ -0,0 +1,10 @@ +{ + "folders": [ + { + "path": "../../STAR-rs" + }, + { + "path": ".." + } + ] +} \ No newline at end of file