diff --git a/src/io/bam.rs b/src/io/bam.rs index 87bcf47..193b449 100644 --- a/src/io/bam.rs +++ b/src/io/bam.rs @@ -2,6 +2,7 @@ use crate::error::Error; use crate::genome::Genome; use crate::io::bgzf_writer::BgzfWriter; +use crate::io::encode::{RecordEncoder, RecordFormat, bam_record_spans}; use crate::params::Parameters; use crate::quant::transcriptome::TranscriptomeIndex; use byteorder::{LittleEndian, WriteBytesExt}; @@ -86,7 +87,7 @@ pub struct BamWriter { /// /// The header emits `SO:coordinate`. Unmapped records sort to the end. pub struct SortedBamWriter { - records: Vec, + records: SortBuffer, output_path: std::path::PathBuf, header: sam::Header, compression: i32, @@ -157,6 +158,29 @@ impl BamWriter { Ok(()) } + /// Encoder producing exactly what [`BamWriter::write_batch`] writes. + pub fn encoder(&self) -> RecordEncoder { + RecordEncoder::new(RecordFormat::Bam, &self.header) + } + + /// Append records serialized by [`BamWriter::encoder`]. + pub fn write_encoded(&mut self, bytes: &[u8]) -> Result<(), Error> { + self.writer.get_mut().write_all(bytes)?; + Ok(()) + } + + /// Write `encoded` if the records were pre-serialized, else `records`. + pub fn write_encoded_or( + &mut self, + encoded: Option<&[u8]>, + records: &[RecordBuf], + ) -> Result<(), Error> { + match encoded { + Some(bytes) => self.write_encoded(bytes), + None => self.write_batch(records), + } + } + /// Flush and close BAM file pub fn finish(&mut self) -> Result<(), Error> { self.writer.get_mut().finish()?; @@ -165,6 +189,67 @@ impl BamWriter { } } +/// Coordinate-sort buffer shared by the two sorted BAM writers. +/// +/// Records are kept BAM-encoded (as the align workers produce them, #223) +/// rather than as `RecordBuf`s, so sorting moves small index entries and +/// writing is a plain copy. The order is the one the `RecordBuf` sort gave: +/// a stable sort on `(reference id, 1-based start)`, with records lacking +/// either sorted last. +struct SortBuffer { + encoder: RecordEncoder, + bytes: Vec, + /// `(sort key, offset, length)` per record, in arrival order. + index: Vec<((usize, usize), usize, usize)>, +} + +impl SortBuffer { + fn new(header: &sam::Header) -> Self { + Self { + encoder: RecordEncoder::new(RecordFormat::Bam, header), + bytes: Vec::new(), + index: Vec::new(), + } + } + + fn len(&self) -> usize { + self.index.len() + } + + fn push_records(&mut self, batch: &[RecordBuf]) -> Result<(), Error> { + let mut buf = Vec::new(); + self.encoder.encode(batch, &mut buf)?; + self.push_encoded(&buf) + } + + fn push_encoded(&mut self, bytes: &[u8]) -> Result<(), Error> { + let base = self.bytes.len(); + for (off, len) in bam_record_spans(bytes)? { + let rec = &bytes[off..off + len]; + // After block_size: refID (i32), pos (i32, 0-based, -1 = none). + let ref_id = i32::from_le_bytes(rec[4..8].try_into().unwrap()); + let pos = i32::from_le_bytes(rec[8..12].try_into().unwrap()); + let key = if ref_id >= 0 && pos >= 0 { + (ref_id as usize, pos as usize + 1) + } else { + (usize::MAX, 0) + }; + self.index.push((key, base + off, len)); + } + self.bytes.extend_from_slice(bytes); + Ok(()) + } + + /// Stable-sort by coordinate and write every record to `out`. + fn write_sorted(&mut self, out: &mut W) -> Result<(), Error> { + self.index.sort_by_key(|&(key, _, _)| key); + for &(_, off, len) in &self.index { + out.write_all(&self.bytes[off..off + len])?; + } + Ok(()) + } +} + impl SortedBamWriter { /// Create a sorted BAM writer. Records are buffered in memory until `finish()`. pub fn create( @@ -174,7 +259,7 @@ impl SortedBamWriter { ) -> Result { let header = crate::io::sam::build_sam_header(genome, params)?; Ok(Self { - records: Vec::new(), + records: SortBuffer::new(&header), output_path: output_path.to_path_buf(), header, compression: params.out_bam_compression, @@ -185,8 +270,17 @@ impl SortedBamWriter { /// Buffer records — no disk I/O yet. pub fn write_batch(&mut self, batch: &[RecordBuf]) -> Result<(), Error> { - self.records.extend_from_slice(batch); - Ok(()) + self.records.push_records(batch) + } + + /// Encoder whose output [`SortedBamWriter::write_encoded`] accepts. + pub fn encoder(&self) -> RecordEncoder { + self.records.encoder.clone() + } + + /// Buffer records serialized by [`SortedBamWriter::encoder`]. + pub fn write_encoded(&mut self, bytes: &[u8]) -> Result<(), Error> { + self.records.push_encoded(bytes) } /// Estimate memory used by buffered records (rough: 400 bytes/record for 150bp reads). @@ -216,20 +310,11 @@ impl SortedBamWriter { /// Unmapped records (no reference) sort to the end. pub fn finish(&mut self) -> Result<(), Error> { self.check_ram_limit()?; - self.records - .sort_by_key(|r| match (r.reference_sequence_id(), r.alignment_start()) { - (Some(chr), Some(pos)) => (chr, pos.get()), - _ => (usize::MAX, 0), - }); - let buf_writer = BufWriter::new(File::create(&self.output_path)?); let mut bgzf = make_bgzf_writer(buf_writer, self.compression, self.threads)?; write_bam_header_lenient(&mut bgzf, &self.header, Some("coordinate"))?; - let mut bam_writer = bam::io::Writer::from(bgzf); - for record in &self.records { - bam_writer.write_alignment_record(&self.header, record)?; - } - bam_writer.get_mut().finish()?; + self.records.write_sorted(&mut bgzf)?; + bgzf.finish()?; log::info!("Sorted BAM written ({} records)", self.records.len()); Ok(()) } @@ -237,20 +322,11 @@ impl SortedBamWriter { /// Sort all buffered records and write to stdout (for `--outStd BAM_SortedByCoordinate`). pub fn finish_to_stdout(&mut self) -> Result<(), Error> { self.check_ram_limit()?; - self.records - .sort_by_key(|r| match (r.reference_sequence_id(), r.alignment_start()) { - (Some(chr), Some(pos)) => (chr, pos.get()), - _ => (usize::MAX, 0), - }); - let buf_writer = BufWriter::new(std::io::stdout()); let mut bgzf = make_bgzf_writer(buf_writer, self.compression, self.threads)?; write_bam_header_lenient(&mut bgzf, &self.header, Some("coordinate"))?; - let mut bam_writer = bam::io::Writer::from(bgzf); - for record in &self.records { - bam_writer.write_alignment_record(&self.header, record)?; - } - bam_writer.get_mut().finish()?; + self.records.write_sorted(&mut bgzf)?; + bgzf.finish()?; log::info!( "Sorted BAM written to stdout ({} records)", self.records.len() @@ -419,6 +495,17 @@ impl BamStdoutWriter { Ok(()) } + /// Encoder producing exactly what [`BamStdoutWriter::write_batch`] writes. + pub fn encoder(&self) -> RecordEncoder { + RecordEncoder::new(RecordFormat::Bam, &self.header) + } + + /// Append records serialized by [`BamStdoutWriter::encoder`]. + pub fn write_encoded(&mut self, bytes: &[u8]) -> Result<(), Error> { + self.writer.get_mut().write_all(bytes)?; + Ok(()) + } + pub fn finish(&mut self) -> Result<(), Error> { self.writer.get_mut().finish()?; Ok(()) @@ -427,7 +514,7 @@ impl BamStdoutWriter { /// Coordinate-sorted BAM writer that writes to stdout on `finish()`. pub struct SortedBamStdoutWriter { - records: Vec, + records: SortBuffer, header: sam::Header, compression: i32, threads: usize, @@ -438,7 +525,7 @@ impl SortedBamStdoutWriter { pub fn create(genome: &crate::genome::Genome, params: &Parameters) -> Result { let header = crate::io::sam::build_sam_header(genome, params)?; Ok(Self { - records: Vec::new(), + records: SortBuffer::new(&header), header, compression: params.out_bam_compression, threads: bgzf_threads(params), @@ -447,8 +534,17 @@ impl SortedBamStdoutWriter { } pub fn write_batch(&mut self, batch: &[RecordBuf]) -> Result<(), Error> { - self.records.extend_from_slice(batch); - Ok(()) + self.records.push_records(batch) + } + + /// Encoder whose output [`SortedBamStdoutWriter::write_encoded`] accepts. + pub fn encoder(&self) -> RecordEncoder { + self.records.encoder.clone() + } + + /// Buffer records serialized by [`SortedBamStdoutWriter::encoder`]. + pub fn write_encoded(&mut self, bytes: &[u8]) -> Result<(), Error> { + self.records.push_encoded(bytes) } pub fn finish(&mut self) -> Result<(), Error> { @@ -463,22 +559,14 @@ impl SortedBamStdoutWriter { ))); } } - self.records - .sort_by_key(|r| match (r.reference_sequence_id(), r.alignment_start()) { - (Some(chr), Some(pos)) => (chr, pos.get()), - _ => (usize::MAX, 0), - }); let mut bgzf = make_bgzf_writer( BufWriter::new(std::io::stdout()), self.compression, self.threads, )?; write_bam_header_lenient(&mut bgzf, &self.header, Some("coordinate"))?; - let mut bam_writer = bam::io::Writer::from(bgzf); - for record in &self.records { - bam_writer.write_alignment_record(&self.header, record)?; - } - bam_writer.get_mut().finish()?; + self.records.write_sorted(&mut bgzf)?; + bgzf.finish()?; log::info!( "Sorted BAM written to stdout ({} records)", self.records.len() @@ -725,4 +813,97 @@ mod tests { let result = writer.finish(); assert!(result.is_err(), "Should fail when RAM limit is exceeded"); } + + /// Records covering every sort-key case: two references, ties, a record + /// with a reference but no position, and fully unmapped records. + fn sort_fixture() -> Vec { + let mk = |name: &str, r: Option, p: Option| { + let mut b = RecordBuf::builder().set_name(name); + if let Some(r) = r { + b = b.set_reference_sequence_id(r); + } + if let Some(p) = p { + b = b.set_alignment_start(p.try_into().unwrap()); + } + b.build() + }; + vec![ + mk("u1", None, None), + mk("a", Some(1), Some(5)), + mk("b", Some(0), Some(7)), + mk("rnopos", Some(0), None), + mk("c", Some(0), Some(7)), + mk("d", Some(0), Some(1)), + mk("u2", None, None), + mk("e", Some(1), Some(1)), + ] + } + + fn two_ref_header() -> sam::Header { + use noodles::sam::header::record::value::{Map, map::ReferenceSequence}; + use std::num::NonZeroUsize; + sam::Header::builder() + .add_reference_sequence( + "chr1", + Map::::new(NonZeroUsize::new(100).unwrap()), + ) + .add_reference_sequence( + "chr2", + Map::::new(NonZeroUsize::new(100).unwrap()), + ) + .build() + } + + #[test] + fn sort_buffer_matches_recordbuf_sort() { + let header = two_ref_header(); + let records = sort_fixture(); + + // Reference: the RecordBuf-level stable sort the writers used before. + let mut expected_records = records.clone(); + expected_records.sort_by_key(|r| match (r.reference_sequence_id(), r.alignment_start()) { + (Some(chr), Some(pos)) => (chr, pos.get()), + _ => (usize::MAX, 0), + }); + let mut expected = bam::io::Writer::from(Vec::new()); + for r in &expected_records { + expected.write_alignment_record(&header, r).unwrap(); + } + + // Fed half as records, half pre-encoded, as the pipelines may. + let mut buf = SortBuffer::new(&header); + buf.push_records(&records[..3]).unwrap(); + let mut enc = Vec::new(); + buf.encoder.encode(&records[3..], &mut enc).unwrap(); + buf.push_encoded(&enc).unwrap(); + assert_eq!(buf.len(), records.len()); + let mut got = Vec::new(); + buf.write_sorted(&mut got).unwrap(); + assert_eq!(got, expected.into_inner()); + } + + #[test] + fn encoder_matches_writer_bytes() { + let header = two_ref_header(); + let records = sort_fixture(); + let mut direct = bam::io::Writer::from(Vec::new()); + for r in &records { + direct.write_alignment_record(&header, r).unwrap(); + } + let mut enc = Vec::new(); + RecordEncoder::new(RecordFormat::Bam, &header) + .encode(&records, &mut enc) + .unwrap(); + assert_eq!(enc, direct.into_inner()); + + let mut direct = sam::io::Writer::new(Vec::new()); + for r in &records { + direct.write_alignment_record(&header, r).unwrap(); + } + let mut enc = Vec::new(); + RecordEncoder::new(RecordFormat::Sam, &header) + .encode(&records, &mut enc) + .unwrap(); + assert_eq!(enc, direct.into_inner()); + } } diff --git a/src/io/encode.rs b/src/io/encode.rs new file mode 100644 index 0000000..9d87680 --- /dev/null +++ b/src/io/encode.rs @@ -0,0 +1,115 @@ +//! Output-record serialization that can run on the align workers (#223). +//! +//! The align pipelines used to hand `RecordBuf`s to the single writer thread, +//! which then encoded every record (SAM text or BAM binary) itself. A +//! [`RecordEncoder`] captures everything that encoding needs (format, header, +//! `--outSAMmode NoQS`), so the rayon workers can serialize each read's records +//! as soon as they are built and the writer thread only appends bytes. +//! +//! The bytes are produced by the same noodles writers the output writers use, +//! record by record, so the concatenation is byte-identical to what the writer +//! thread would have produced. + +use std::sync::Arc; + +use noodles::sam::alignment::record_buf::{QualityScores, RecordBuf}; +use noodles::{bam, sam}; + +use crate::error::Error; + +/// Serialization format of an output stream. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RecordFormat { + /// SAM text lines. + Sam, + /// BAM records (`block_size` prefix included), before BGZF compression. + Bam, +} + +/// Serializes alignment records exactly as an output writer would. +#[derive(Clone)] +pub struct RecordEncoder { + format: RecordFormat, + header: Arc, + strip_quality: bool, + check_cigar: bool, +} + +impl RecordEncoder { + pub fn new(format: RecordFormat, header: &sam::Header) -> Self { + Self { + format, + header: Arc::new(header.clone()), + strip_quality: false, + check_cigar: false, + } + } + + pub fn format(&self) -> RecordFormat { + self.format + } + + /// Drop quality strings before encoding (`--outSAMmode NoQS`). + #[must_use] + pub fn without_quality(mut self) -> Self { + self.strip_quality = true; + self + } + + /// Run the SAM file writer's CIGAR/SEQ length sanity check on every record. + #[must_use] + pub fn with_cigar_check(mut self) -> Self { + self.check_cigar = true; + self + } + + /// Append the serialized form of `records` to `out`. + pub fn encode(&self, records: &[RecordBuf], out: &mut Vec) -> Result<(), Error> { + if records.is_empty() { + return Ok(()); + } + match self.format { + RecordFormat::Sam => self.encode_with(&mut sam::io::Writer::new(out), records), + RecordFormat::Bam => self.encode_with(&mut bam::io::Writer::from(out), records), + } + } + + fn encode_with( + &self, + writer: &mut impl sam::alignment::io::Write, + records: &[RecordBuf], + ) -> Result<(), Error> { + for record in records { + if self.check_cigar { + crate::io::sam::check_cigar_seq_len(record); + } + if self.strip_quality { + let mut stripped = record.clone(); + *stripped.quality_scores_mut() = QualityScores::default(); + writer.write_alignment_record(&self.header, &stripped)?; + } else { + writer.write_alignment_record(&self.header, record)?; + } + } + Ok(()) + } +} + +/// Split a buffer of BAM-encoded records into `(record_start, record_len)` +/// spans, each including its 4-byte `block_size` prefix. +pub fn bam_record_spans(bytes: &[u8]) -> Result, Error> { + let mut spans = Vec::new(); + let mut off = 0; + while off < bytes.len() { + let Some(prefix) = bytes.get(off..off + 4) else { + return Err(Error::Alignment("truncated encoded BAM record".into())); + }; + let len = 4 + u32::from_le_bytes(prefix.try_into().unwrap()) as usize; + if off + len > bytes.len() { + return Err(Error::Alignment("truncated encoded BAM record".into())); + } + spans.push((off, len)); + off += len; + } + Ok(spans) +} diff --git a/src/io/mod.rs b/src/io/mod.rs index fa6e46c..5a53c54 100644 --- a/src/io/mod.rs +++ b/src/io/mod.rs @@ -2,6 +2,7 @@ pub mod bam; pub mod bgzf_writer; +pub mod encode; pub mod fastq; pub mod log; pub mod sam; diff --git a/src/io/sam.rs b/src/io/sam.rs index 189c4bb..3dfd90e 100644 --- a/src/io/sam.rs +++ b/src/io/sam.rs @@ -3,6 +3,7 @@ use crate::align::read_align::PairedAlignment; use crate::align::transcript::{Transcript, cigar_to_string}; use crate::error::Error; use crate::genome::Genome; +use crate::io::encode::{RecordEncoder, RecordFormat}; use crate::io::fastq::{complement_base, decode_base}; use crate::junction::encode_motif; use crate::mapq::calculate_mapq; @@ -32,6 +33,13 @@ use std::path::Path; #[derive(Default)] pub struct BufferedSamRecords { pub records: Vec, + /// `records` (plus any `--chimOutType WithinBAM` supplementary records) + /// already serialized for the main output on an align worker (#223). + /// When set, `records` has been emptied and the writer only appends these + /// bytes. + pub encoded: Option>, + /// Same, for this read's `--quantMode TranscriptomeSAM` records. + pub encoded_transcriptome: Option>, } impl BufferedSamRecords { @@ -39,6 +47,8 @@ impl BufferedSamRecords { pub fn new() -> Self { Self { records: Vec::with_capacity(10000), + encoded: None, + encoded_transcriptome: None, } } @@ -75,6 +85,31 @@ fn insert_unmapped_tags(record: &mut RecordBuf, attrs: SamAttributes, reason: Un data.insert(Tag::new(b'u', b'T'), Value::Character(ut)); } +/// Debug check: panic if a record's CIGAR query length disagrees with its SEQ. +pub(crate) fn check_cigar_seq_len(record: &RecordBuf) { + let cigar_ops = record.cigar().as_ref(); + let cigar_query_len: usize = cigar_ops + .iter() + .filter(|op| op.kind().consumes_read()) + .map(|op| op.len()) + .sum(); + let seq_len = record.sequence().len(); + if cigar_query_len != seq_len && !cigar_ops.is_empty() { + let name = record + .name() + .map(|n| String::from_utf8_lossy(n.as_ref()).to_string()) + .unwrap_or_default(); + panic!( + "[SAM-MISMATCH] read={} cigar_query_len={} seq_len={} flags={:?} cigar={}", + name, + cigar_query_len, + seq_len, + record.flags(), + cigar_to_string(cigar_ops) + ); + } +} + /// SAM file writer pub struct SamWriter { writer: sam::io::Writer>, @@ -170,33 +205,23 @@ impl SamWriter { /// * `batch` - Slice of records to write pub fn write_batch(&mut self, batch: &[RecordBuf]) -> Result<(), Error> { for record in batch { - // Debug: validate CIGAR vs SEQ length before writing - let cigar_ops = record.cigar().as_ref(); - let cigar_query_len: usize = cigar_ops - .iter() - .filter(|op| op.kind().consumes_read()) - .map(|op| op.len()) - .sum(); - let seq_len = record.sequence().len(); - if cigar_query_len != seq_len && !cigar_ops.is_empty() { - let name = record - .name() - .map(|n| String::from_utf8_lossy(n.as_ref()).to_string()) - .unwrap_or_default(); - panic!( - "[SAM-MISMATCH] read={} cigar_query_len={} seq_len={} flags={:?} cigar={}", - name, - cigar_query_len, - seq_len, - record.flags(), - cigar_to_string(cigar_ops) - ); - } + check_cigar_seq_len(record); self.writer.write_alignment_record(&self.header, record)?; } Ok(()) } + /// Encoder producing exactly what [`SamWriter::write_batch`] writes. + pub fn encoder(&self) -> RecordEncoder { + RecordEncoder::new(RecordFormat::Sam, &self.header).with_cigar_check() + } + + /// Append records serialized by [`SamWriter::encoder`]. + pub fn write_encoded(&mut self, bytes: &[u8]) -> Result<(), Error> { + std::io::Write::write_all(self.writer.get_mut(), bytes)?; + Ok(()) + } + /// Build unmapped record (without writing) /// /// # Arguments @@ -882,6 +907,17 @@ impl SamStdoutWriter { } Ok(()) } + + /// Encoder producing exactly what [`SamStdoutWriter::write_batch`] writes. + pub fn encoder(&self) -> RecordEncoder { + RecordEncoder::new(RecordFormat::Sam, &self.header) + } + + /// Append records serialized by [`SamStdoutWriter::encoder`]. + pub fn write_encoded(&mut self, bytes: &[u8]) -> Result<(), Error> { + std::io::Write::write_all(self.writer.get_mut(), bytes)?; + Ok(()) + } } /// Build paired SAM header from genome diff --git a/src/lib.rs b/src/lib.rs index 7b60e9b..910a5d9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -179,6 +179,40 @@ trait AlignmentWriter: Send { fn finish(&mut self) -> Result<(), error::Error> { Ok(()) } + /// Serializer the align workers can run ahead of time so this writer only + /// appends bytes (#223). `None`: records are passed as `RecordBuf`s. + fn encoder(&self) -> Option { + None + } + /// Append records serialized by [`AlignmentWriter::encoder`]. + fn write_encoded(&mut self, _bytes: &[u8]) -> Result<(), error::Error> { + Err(error::Error::Alignment( + "writer does not accept pre-encoded records".into(), + )) + } + /// Write a read's records, pre-encoded or not. + fn write_buffered( + &mut self, + buf: &crate::io::sam::BufferedSamRecords, + ) -> Result<(), error::Error> { + match &buf.encoded { + Some(bytes) => self.write_encoded(bytes), + None => self.write_batch(&buf.records), + } + } +} + +/// Forwards [`AlignmentWriter::encoder`] / `write_encoded` to the writer's +/// inherent methods of the same name. +macro_rules! forward_encoded { + () => { + fn encoder(&self) -> Option { + Some(self.encoder()) + } + fn write_encoded(&mut self, bytes: &[u8]) -> Result<(), error::Error> { + self.write_encoded(bytes) + } + }; } /// Drops quality strings on the way to the real writer (`--outSAMmode NoQS`). @@ -204,6 +238,16 @@ impl AlignmentWriter for NoQsWriter { fn finish(&mut self) -> Result<(), error::Error> { self.0.finish() } + + fn encoder(&self) -> Option { + self.0 + .encoder() + .map(crate::io::encode::RecordEncoder::without_quality) + } + + fn write_encoded(&mut self, bytes: &[u8]) -> Result<(), error::Error> { + self.0.write_encoded(bytes) + } } /// Null writer that discards all output (for two-pass mode pass 1) @@ -225,6 +269,7 @@ impl AlignmentWriter for crate::io::sam::SamWriter { ) -> Result<(), error::Error> { self.write_batch(batch) } + forward_encoded!(); } impl AlignmentWriter for crate::io::bam::BamWriter { @@ -234,6 +279,7 @@ impl AlignmentWriter for crate::io::bam::BamWriter { ) -> Result<(), error::Error> { self.write_batch(batch) } + forward_encoded!(); fn finish(&mut self) -> Result<(), error::Error> { self.finish() } @@ -246,6 +292,7 @@ impl AlignmentWriter for crate::io::bam::SortedBamWriter { ) -> Result<(), error::Error> { self.write_batch(batch) } + forward_encoded!(); fn finish(&mut self) -> Result<(), error::Error> { self.finish() } @@ -258,6 +305,7 @@ impl AlignmentWriter for crate::io::sam::SamStdoutWriter { ) -> Result<(), error::Error> { self.write_batch(batch) } + forward_encoded!(); } impl AlignmentWriter for crate::io::bam::BamStdoutWriter { @@ -267,6 +315,7 @@ impl AlignmentWriter for crate::io::bam::BamStdoutWriter { ) -> Result<(), error::Error> { self.write_batch(batch) } + forward_encoded!(); fn finish(&mut self) -> Result<(), error::Error> { self.finish() } @@ -279,6 +328,7 @@ impl AlignmentWriter for crate::io::bam::SortedBamStdoutWriter { ) -> Result<(), error::Error> { self.write_batch(batch) } + forward_encoded!(); fn finish(&mut self) -> Result<(), error::Error> { self.finish() } @@ -1151,6 +1201,146 @@ where Ok(()) } +/// Output serializers handed to the align workers (#223), so the writer thread +/// only appends bytes instead of encoding every SAM/BAM record itself. +struct PreEncoder { + /// Main output (`Aligned.out.{sam,bam}`, `--outStd`), if it takes bytes. + main: Option, + /// `Aligned.toTranscriptome.out.bam`, if enabled. + transcriptome: Option, + /// Set when `--chimOutType WithinBAM` supplementary records are built and + /// encoded on the workers too (they follow the read's records, as before). + within_bam: Option>, +} + +impl PreEncoder { + /// `None` when records should stay `RecordBuf`s and be encoded on the + /// writer thread. Output bytes are the same either way; only where the + /// encoding CPU is spent changes. + /// + /// * BySJout: records are filtered and re-read from its temp file later. + /// * Spare cores: the decode, writer and BGZF threads run beside the align + /// pool, so while `--runThreadN` leaves a few cores idle, encoding on + /// the writer thread is free and moving it onto the workers only slows + /// alignment. Once the pool covers (nearly) every core, the writer + /// thread is the bottleneck and the workers must share the encoding. + /// SAM text is costlier to format than BAM (and has no BGZF pool to + /// hand work to), so its writer saturates earlier. Measured on 16 + /// cores, 2 M SE reads, median wall, writer-thread vs worker encoding: + /// + /// | threads | BAM | SAM | + /// |---------|---------------|---------------| + /// | 1 | 13.79 / 14.32 | 13.28 / 14.37 | + /// | 8 | 2.99 / 3.20 | 3.04 / 3.19 | + /// | 12 | 2.51 / 2.55 | 2.97 / 2.48 | + /// | 14 | 2.53 / 2.44 | - / 2.49 | + /// | 16 | 2.69 / 2.55 | - / 2.41 | + /// + /// `RUSTAR_PRE_ENCODE=1` / `=0` forces it on / off (parity tests). + fn new( + writer: &W, + tr_writer: Option<&crate::io::bam::BamWriter>, + index: &std::sync::Arc, + params: &Parameters, + ) -> Option> { + if params.out_filter_type == crate::params::OutFilterType::BySJout { + return None; + } + let main = writer.encoder(); + let enabled = match std::env::var("RUSTAR_PRE_ENCODE").as_deref() { + Ok("1") => true, + Ok("0") => false, + _ => { + use crate::io::encode::RecordFormat; + // Cores that must stay idle for writer-thread encoding to win. + let spare = match main.as_ref().map(crate::io::encode::RecordEncoder::format) { + Some(RecordFormat::Sam) => 6, + _ => 4, + }; + let cores = std::thread::available_parallelism().map_or(1, usize::from); + params.run_thread_n.get() + spare > cores + } + }; + if !enabled { + return None; + } + let transcriptome = tr_writer.map(crate::io::bam::BamWriter::encoder); + let within_bam = + (main.is_some() && params.chim_out_within_bam()).then(|| std::sync::Arc::clone(index)); + (main.is_some() || transcriptome.is_some()).then(|| { + std::sync::Arc::new(Self { + main, + transcriptome, + within_bam, + }) + }) + } + + /// Serialize one read's records into `buf.encoded`, emptying `records`. + fn encode_main( + &self, + buf: &mut crate::io::sam::BufferedSamRecords, + chimeric: &[crate::chimeric::ChimericAlignment], + ) -> Result<(), error::Error> { + let Some(enc) = &self.main else { + return Ok(()); + }; + let mut bytes = Vec::new(); + enc.encode(&buf.records, &mut bytes)?; + if let Some(index) = &self.within_bam { + for chim_aln in chimeric { + let supp = crate::chimeric::build_within_bam_records(chim_aln, &index.genome, 255)?; + enc.encode(&supp, &mut bytes)?; + } + } + buf.records.clear(); + buf.encoded = Some(bytes); + Ok(()) + } +} + +/// A per-read pipeline product whose output records can be pre-encoded. +trait PreEncode { + fn pre_encode(&mut self, enc: &PreEncoder) -> Result<(), error::Error>; +} + +impl PreEncode for AlignmentBatchResults { + fn pre_encode(&mut self, enc: &PreEncoder) -> Result<(), error::Error> { + enc.encode_main(&mut self.sam_records, &self.chimeric_alns)?; + if let Some(tr_enc) = &enc.transcriptome { + let mut bytes = Vec::new(); + tr_enc.encode(&self.transcriptome_records, &mut bytes)?; + self.transcriptome_records.clear(); + self.sam_records.encoded_transcriptome = Some(bytes); + } + Ok(()) + } +} + +/// Wrap a batch `align` closure so each finished batch is also serialized for +/// the output writers, in parallel on the rayon pool, before it is handed on. +fn with_pre_encode( + align: impl Fn(u64, Vec) -> BatchOut + Clone + Send + 'static, + enc: Option>, +) -> impl Fn(u64, Vec) -> BatchOut + Clone + Send + 'static { + move |base, batch| { + use rayon::prelude::*; + let mut out = align(base, batch); + if let Some(enc) = &enc { + out.par_iter_mut().for_each(|result| { + let encoded = match result { + Ok(product) => product.pre_encode(enc), + Err(_) => Ok(()), + }; + if let Err(e) = encoded { + *result = Err(e); + } + }); + } + out + } +} + /// Build transcriptome-space records for a single-end read. Projects every /// surviving genome-space alignment onto all compatible transcripts, picks one /// projected alignment at random as the primary (seeded by `per_read_seed`), @@ -1503,6 +1693,7 @@ fn align_reads_single_end( let index_writer = Arc::clone(index); // Shared, 'static parameters for the per-batch aligner tasks spawned below. let params_arc = Arc::new(params.clone()); + let pre_encoder = PreEncoder::new(&*writer, tr_writer.as_deref(), index, params); std::thread::scope(|scope| -> anyhow::Result<()> { let (read_tx, read_rx) = std::sync::mpsc::sync_channel::< Result, error::Error>, @@ -1619,11 +1810,14 @@ fn align_reads_single_end( } // Write SAM/BAM records - writer.write_batch(&batch.sam_records.records)?; + writer.write_buffered(&batch.sam_records)?; // Write transcriptome-space records (if enabled) if let Some(ref mut tw) = tr_writer { - tw.write_batch(&batch.transcriptome_records)?; + tw.write_encoded_or( + batch.sam_records.encoded_transcriptome.as_deref(), + &batch.transcriptome_records, + )?; } // Write chimeric alignments @@ -1636,7 +1830,7 @@ fn align_reads_single_end( )?; } } - if params.chim_out_within_bam() { + if params.chim_out_within_bam() && batch.sam_records.encoded.is_none() { use crate::chimeric::build_within_bam_records; for chim_aln in &batch.chimeric_alns { let supp = build_within_bam_records(chim_aln, &index.genome, 255)?; @@ -2036,7 +2230,7 @@ fn align_reads_single_end( info!("Processed {n} reads..."); } }, - align, + with_pre_encode(align, pre_encoder.clone()), |done| Ok(res_tx.send(done).is_ok()), ) }; @@ -2105,6 +2299,7 @@ fn align_reads_solo( let output_unmapped = emit_sam && params.out_sam_unmapped != params::OutSamUnmapped::None; // Shared, 'static parameters for the per-batch aligner tasks spawned below. let params_arc = Arc::new(params.clone()); + let pre_encoder = PreEncoder::new(&*writer, None, index, params); /// Per-read result for the solo loop (one outcome per quantified feature). struct SoloReadProduct { @@ -2113,6 +2308,11 @@ fn align_reads_solo( sj: Vec, velocyto: Option, } + impl PreEncode for SoloReadProduct { + fn pre_encode(&mut self, enc: &PreEncoder) -> Result<(), error::Error> { + enc.encode_main(&mut self.sam_records, &[]) + } + } info!("STARsolo: aligning cDNA reads and quantifying barcodes..."); // Decode the next batch on a background thread while the current batch is @@ -2162,7 +2362,7 @@ fn align_reads_solo( let mut velo_batch: Vec = Vec::new(); for result in products { let product = result?; - writer.write_batch(&product.sam_records.records)?; + writer.write_buffered(&product.sam_records)?; for (fi, fo) in product.per_feature.into_iter().enumerate() { if let Some(r) = fo.record { feat_records[fi].push(r); @@ -2361,7 +2561,7 @@ fn align_reads_solo( info!("STARsolo: processed {n} reads..."); } }, - align, + with_pre_encode(align, pre_encoder.clone()), consume, ) } @@ -2419,6 +2619,7 @@ fn align_reads_solo_pe( let emit_sam = params.emits_alignments(); let output_unmapped = emit_sam && params.out_sam_unmapped != params::OutSamUnmapped::None; let params_arc = Arc::new(params.clone()); + let pre_encoder = PreEncoder::new(&*writer, None, index, params); struct SoloReadProduct { sam_records: BufferedSamRecords, @@ -2426,6 +2627,11 @@ fn align_reads_solo_pe( sj: Vec, velocyto: Option, } + impl PreEncode for SoloReadProduct { + fn pre_encode(&mut self, enc: &PreEncoder) -> Result<(), error::Error> { + enc.encode_main(&mut self.sam_records, &[]) + } + } info!("STARsolo: aligning 5' paired-end reads and quantifying barcodes..."); std::thread::scope(|scope| -> anyhow::Result<()> { @@ -2463,7 +2669,7 @@ fn align_reads_solo_pe( let mut velo_batch: Vec = Vec::new(); for result in products { let product = result?; - writer.write_batch(&product.sam_records.records)?; + writer.write_buffered(&product.sam_records)?; for (fi, fo) in product.per_feature.into_iter().enumerate() { if let Some(r) = fo.record { feat_records[fi].push(r); @@ -2736,7 +2942,7 @@ fn align_reads_solo_pe( info!("STARsolo: processed {n} read pairs..."); } }, - align, + with_pre_encode(align, pre_encoder.clone()), consume, ) } @@ -2840,6 +3046,7 @@ fn align_reads_paired_end( let index_writer = Arc::clone(index); // Shared, 'static parameters for the per-batch aligner tasks spawned below. let params_arc = Arc::new(params.clone()); + let pre_encoder = PreEncoder::new(&*writer, tr_writer.as_deref(), index, params); std::thread::scope(|scope| -> anyhow::Result<()> { let (read_tx, read_rx) = std::sync::mpsc::sync_channel::< Result, error::Error>, @@ -2955,11 +3162,14 @@ fn align_reads_paired_end( ); } } - writer.write_batch(&batch.sam_records.records)?; + writer.write_buffered(&batch.sam_records)?; if let Some(ref mut tw) = tr_writer { - tw.write_batch(&batch.transcriptome_records)?; + tw.write_encoded_or( + batch.sam_records.encoded_transcriptome.as_deref(), + &batch.transcriptome_records, + )?; } - if params.chim_out_within_bam() { + if params.chim_out_within_bam() && batch.sam_records.encoded.is_none() { use crate::chimeric::build_within_bam_records; for chim_aln in &batch.chimeric_alns { let supp = build_within_bam_records(chim_aln, &index.genome, 255)?; @@ -3515,7 +3725,7 @@ fn align_reads_paired_end( info!("Processed {n} pairs..."); } }, - align, + with_pre_encode(align, pre_encoder.clone()), |done| Ok(res_tx.send(done).is_ok()), ) }; diff --git a/tests/pre_encode_parity.rs b/tests/pre_encode_parity.rs new file mode 100644 index 0000000..defa579 --- /dev/null +++ b/tests/pre_encode_parity.rs @@ -0,0 +1,302 @@ +//! Output parity for #223: records serialized on the align workers +//! (`RUSTAR_PRE_ENCODE=1`) must produce byte-identical files to records +//! serialized on the writer thread (`RUSTAR_PRE_ENCODE=0`). + +use assert_cmd::cargo::cargo_bin_cmd; +use std::fs; +use std::io::Write; +use std::path::Path; +use tempfile::TempDir; + +fn lcg_seq(seed: u32, length: usize) -> Vec { + let mut state = seed; + (0..length) + .map(|_| { + state = state.wrapping_mul(1_103_515_245).wrapping_add(12345); + b"ACGT"[((state >> 16) & 3) as usize] + }) + .collect() +} + +fn rc(seq: &[u8]) -> Vec { + seq.iter() + .rev() + .map(|&b| match b { + b'A' => b'T', + b'T' => b'A', + b'C' => b'G', + _ => b'C', + }) + .collect() +} + +/// Two chromosomes; chr1 carries a GT-AG intron at [10050, 10250). +fn build_genome() -> (Vec, Vec) { + let mut chr1 = lcg_seq(88888, 20000); + chr1[10050..10052].copy_from_slice(b"GT"); + chr1[10248..10250].copy_from_slice(b"AG"); + (chr1, lcg_seq(4242, 12000)) +} + +fn write_fastq(path: &Path, reads: &[Vec]) { + let mut f = fs::File::create(path).unwrap(); + for (i, seq) in reads.iter().enumerate() { + // Varied qualities so --outSAMmode NoQS has something to strip. + let qual: String = (0..seq.len()) + .map(|j| char::from(b"#5?FIJ"[(i + j) % 6])) + .collect(); + writeln!(f, "@r{i}").unwrap(); + f.write_all(seq).unwrap(); + writeln!(f, "\n+\n{qual}").unwrap(); + } +} + +struct Fixture { + _tmp: TempDir, + root: std::path::PathBuf, +} + +fn fixture() -> Fixture { + let tmp = TempDir::new().unwrap(); + let root = tmp.path().to_path_buf(); + let (chr1, chr2) = build_genome(); + + let fasta = root.join("genome.fa"); + let mut f = fs::File::create(&fasta).unwrap(); + for (name, seq) in [("chr1", &chr1), ("chr2", &chr2)] { + writeln!(f, ">{name}").unwrap(); + f.write_all(seq).unwrap(); + writeln!(f).unwrap(); + } + let gtf = root.join("a.gtf"); + fs::write( + >f, + "chr1\tt\texon\t9901\t10050\t.\t+\t.\tgene_id \"G1\"; transcript_id \"T1\";\n\ + chr1\tt\texon\t10251\t10400\t.\t+\t.\tgene_id \"G1\"; transcript_id \"T1\";\n", + ) + .unwrap(); + + let genome_dir = root.join("idx"); + fs::create_dir_all(&genome_dir).unwrap(); + cargo_bin_cmd!("rustar-aligner") + .args(["--runMode", "genomeGenerate", "--genomeSAindexNbases", "7"]) + .arg("--genomeDir") + .arg(&genome_dir) + .arg("--genomeFastaFiles") + .arg(&fasta) + .arg("--sjdbGTFfile") + .arg(>f) + .args(["--sjdbOverhang", "49"]) + .arg("--outFileNamePrefix") + .arg(genome_dir.join("run_")) + .assert() + .success(); + + // SE: genomic (both strands), spliced across the intron, chimeric + // (chr1 + chr2 halves) and random (unmapped) reads. + let mut se = Vec::new(); + for i in 0..120usize { + let s = 200 + i * 150; + let r = chr1[s..s + 60].to_vec(); + se.push(if i % 2 == 0 { r } else { rc(&r) }); + } + for k in 10..40usize { + let mut r = chr1[10050 - k..10050].to_vec(); + r.extend_from_slice(&chr1[10250..10250 + 60 - k]); + se.push(r); + } + for i in 0..20usize { + let mut r = chr1[1000 + i * 300..1030 + i * 300].to_vec(); + r.extend_from_slice(&chr2[500 + i * 400..530 + i * 400]); + se.push(r); + } + for i in 0..20u32 { + se.push(lcg_seq(7 + i, 60)); + } + write_fastq(&root.join("se.fq"), &se); + + // PE: 250 bp fragments, mate 2 reverse-complemented. + let (mut m1, mut m2) = (Vec::new(), Vec::new()); + for i in 0..80usize { + let (src, s) = if i % 3 == 0 { + (&chr2, 100 + i * 120) + } else { + (&chr1, 100 + i * 200) + }; + let frag = &src[s..s + 250]; + m1.push(frag[..60].to_vec()); + m2.push(rc(&frag[190..])); + } + // Chimeric pairs for --chimOutType WithinBAM: mate 1 on chr1, mate 2 on + // chr2 (inter-mate), and mate 1 itself split across chr1 / chr2. + for i in 0..30usize { + let a = &chr1[12000 + i * 200..12000 + i * 200 + 250]; + let b = &chr2[6000 + i * 150..6000 + i * 150 + 250]; + if i % 2 == 0 { + m1.push(a[..100].to_vec()); + m2.push(rc(&b[150..])); + } else { + let mut r = a[..50].to_vec(); + r.extend_from_slice(&b[..50]); + m1.push(r); + m2.push(rc(&b[150..])); + } + } + write_fastq(&root.join("m1.fq"), &m1); + write_fastq(&root.join("m2.fq"), &m2); + + Fixture { _tmp: tmp, root } +} + +/// Run one configuration with pre-encoding forced on and off, from two +/// sibling directories with prefix `./` (the `@PG CL:` line records argv, so +/// both runs need identical arguments), and compare every output file. +fn assert_parity(fx: &Fixture, label: &str, args: &[&str]) { + let mut outputs = Vec::new(); + for mode in ["0", "1"] { + let dir = fx.root.join(format!("{label}_{mode}")); + fs::create_dir_all(&dir).unwrap(); + let out = cargo_bin_cmd!("rustar-aligner") + .current_dir(&dir) + .env("RUSTAR_PRE_ENCODE", mode) + .args(["--genomeDir", "../idx", "--outFileNamePrefix", "./"]) + .args(["--runThreadN", "3"]) + .args(args) + .output() + .unwrap(); + assert!( + out.status.success(), + "{label} (pre-encode {mode}) failed: {}", + String::from_utf8_lossy(&out.stderr) + ); + fs::write(dir.join("stdout.bin"), &out.stdout).unwrap(); + outputs.push(dir); + } + let mut compared = 0; + for entry in fs::read_dir(&outputs[0]).unwrap() { + let path = entry.unwrap().path(); + let name = path.file_name().unwrap().to_str().unwrap().to_string(); + if name.starts_with("Log.") || !path.is_file() { + continue; + } + let a = fs::read(&path).unwrap(); + let b = fs::read(outputs[1].join(&name)) + .unwrap_or_else(|_| panic!("{label}: {name} missing with pre-encoding")); + assert!(a == b, "{label}: {name} differs with pre-encoding"); + compared += 1; + } + assert!(compared >= 2, "{label}: nothing compared"); +} + +#[test] +fn pre_encoded_output_is_byte_identical() { + let fx = fixture(); + let se = ["--readFilesIn", "../se.fq"]; + let pe = ["--readFilesIn", "../m1.fq", "../m2.fq"]; + let uw = ["--outSAMunmapped", "Within"]; + let cat = |parts: &[&[&'static str]]| -> Vec<&'static str> { parts.concat() }; + + let configs: Vec<(&str, Vec<&str>)> = vec![ + ("se_sam", cat(&[&se, &uw])), + ( + "se_bam", + cat(&[&se, &uw, &["--outSAMtype", "BAM", "Unsorted"]]), + ), + ( + "se_sorted", + cat(&[&se, &uw, &["--outSAMtype", "BAM", "SortedByCoordinate"]]), + ), + ("se_noqs", cat(&[&se, &uw, &["--outSAMmode", "NoQS"]])), + ("se_stdout_sam", cat(&[&se, &["--outStd", "SAM"]])), + ( + "se_stdout_bam", + cat(&[ + &se, + &uw, + &[ + "--outSAMtype", + "BAM", + "Unsorted", + "--outStd", + "BAM_Unsorted", + ], + ]), + ), + ( + "se_stdout_sorted", + cat(&[ + &se, + &[ + "--outSAMtype", + "BAM", + "SortedByCoordinate", + "--outStd", + "BAM_SortedByCoordinate", + ], + ]), + ), + ( + "se_transcriptome", + cat(&[ + &se, + &[ + "--quantMode", + "TranscriptomeSAM", + "--outSAMtype", + "BAM", + "Unsorted", + ], + ]), + ), + ( + "se_chimeric", + cat(&[ + &se, + &["--chimSegmentMin", "20", "--outSAMtype", "BAM", "Unsorted"], + ]), + ), + ( + "se_unmapped_fastx", + cat(&[&se, &["--outReadsUnmapped", "Fastx"]]), + ), + ("pe_sam", cat(&[&pe, &uw])), + ( + "pe_sorted", + cat(&[&pe, &uw, &["--outSAMtype", "BAM", "SortedByCoordinate"]]), + ), + ( + "pe_within_bam", + cat(&[ + &pe, + &[ + "--chimSegmentMin", + "20", + "--chimOutType", + "Junctions", + "WithinBAM", + "--outSAMtype", + "BAM", + "Unsorted", + ], + ]), + ), + ( + "pe_transcriptome_noqs", + cat(&[ + &pe, + &[ + "--quantMode", + "TranscriptomeSAM", + "--outSAMmode", + "NoQS", + "--outSAMtype", + "BAM", + "Unsorted", + ], + ]), + ), + ]; + for (label, args) in &configs { + assert_parity(&fx, label, args); + } +}