diff --git a/Cargo.lock b/Cargo.lock index f7e9efd..497fc0e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -602,9 +602,9 @@ checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" [[package]] name = "paraseq" -version = "0.4.14" +version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b7f0c7259499c750423489d3fe955dfc60183f998f203a87a6f2d4c56872a2f4" +checksum = "100aa4a98a2aef79c936248e72b2eb1bb34063464dba3309c17f457790bfb1f9" dependencies = [ "anyhow", "crossbeam-channel", diff --git a/Cargo.toml b/Cargo.toml index 81ef10b..46670fb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,13 +11,13 @@ categories = ["science::bioinformatics", "encoding", "data-structures"] keywords = ["binary", "nucleotide", "sequencing", "genomics", "fastq"] [dependencies] -anyhow = { version = "1.0.103", optional = true } +anyhow = { version = "1.0.104", optional = true } bitnuc-deprec = { package ="bitnuc", version = "0.4.1" } -bitnuc = { version = "0.5.1" } -bytemuck = { version = "1.25.1", features = ["derive", "extern_crate_alloc"] } +bitnuc = "0.5.4" +bytemuck = { version = "1.25.2", features = ["derive", "extern_crate_alloc"] } itoa = "1.0.18" memmap2 = "0.9.11" -paraseq = { version = "0.4.14", optional = true } +paraseq = { version = "0.5.1", optional = true } parking_lot = {version = "0.12.5", optional = true } rand = { version = "0.9.5", features = ["small_rng"] } sucds = "0.8.3" @@ -26,10 +26,9 @@ zstd = { version = "0.13.3", features = ["zstdmt"] } [dev-dependencies] memchr = "2.8.3" -anyhow = "1.0.103" +anyhow = "1.0.104" parking_lot = "0.12.5" -clap = { version = "4.6.2", features = ["derive"] } -paraseq = "0.4.14" +clap = { version = "4.6.6", features = ["derive"] } tempfile = "3.27.0" [features] diff --git a/examples/auto-write.rs b/examples/auto-write.rs index 5a555a3..dbee93f 100644 --- a/examples/auto-write.rs +++ b/examples/auto-write.rs @@ -58,6 +58,10 @@ struct Args { /// Number of threads to use for parallel processing, 0: all available #[clap(short = 'T', long, default_value = "0")] threads: usize, + + /// Disable ordering output records to match input order (small perf gain) + #[clap(long)] + unordered: bool, } impl Args { /// Determines the output format based on the file extension or the provided format @@ -114,6 +118,7 @@ fn main() -> Result<()> { builder.input(&args.input) } .threads(args.threads) + .ordered(!args.unordered) .run()?; Ok(()) diff --git a/src/utils/fastx.rs b/src/utils/fastx.rs index ad7a4a7..66ffb9d 100644 --- a/src/utils/fastx.rs +++ b/src/utils/fastx.rs @@ -36,6 +36,10 @@ enum FastxInput { /// Created by [`BinseqWriterBuilder::encode_fastx`]; configures the input source /// and threading before running the encoding. /// +/// Can be ordered or unordered; +/// though unordered takes better advantage of parallelism, ordering preserves +/// the input order in the fastx. +/// /// # Example /// /// ```rust,no_run @@ -56,6 +60,7 @@ pub struct FastxEncoderBuilder { builder: BinseqWriterBuilder, output: BoxedWrite, input: Option, + ordered: bool, threads: usize, } @@ -66,7 +71,8 @@ impl FastxEncoderBuilder { builder, output, input: None, - threads: 0, // 0 means use all available cores + threads: 0, // 0 means use all available cores + ordered: false, // default to unordered for speed } } @@ -103,6 +109,13 @@ impl FastxEncoderBuilder { self } + /// Set whether the output should be ordered as input (small perf cost, switch off for speed) + #[must_use] + pub fn ordered(mut self, ordered: bool) -> Self { + self.ordered = ordered; + self + } + /// Execute the FASTX encoding, consuming the builder pub fn run(mut self) -> Result<()> { let (r1, r2) = match self.input { @@ -143,7 +156,7 @@ impl FastxEncoderBuilder { let writer = self.builder.build(self.output)?; let paired = writer.is_paired(); - let mut encoder = Encoder::new(writer)?; + let mut encoder = Encoder::new(writer, self.ordered)?; match (paired, r2) { (true, Some(r2)) => r1.process_parallel_paired(r2, &mut encoder, self.threads), (true, None) => r1.process_parallel_interleaved(&mut encoder, self.threads), @@ -193,15 +206,18 @@ struct Encoder { writer: Arc>>>, /// Thread-local writer buffer thread_writer: BinseqWriter>, + /// Whether the output should follow same order as input + ordered: bool, } impl Encoder { /// Create a new encoder with a global writer - pub fn new(writer: BinseqWriter>) -> Result { + pub fn new(writer: BinseqWriter>, ordered: bool) -> Result { let thread_writer = writer.new_headless_buffer()?; Ok(Self { writer: Arc::new(Mutex::new(writer)), thread_writer, + ordered, }) } /// Finish the stream on the global writer @@ -209,6 +225,16 @@ impl Encoder { self.writer.lock().finish()?; Ok(()) } + fn flush_batch(&mut self) -> paraseq::Result<()> { + let mut writer = self.writer.lock(); + if self.ordered { + // can't take full advantage of parallelism if we have to order the output + writer.ingest(&mut self.thread_writer) + } else { + writer.ingest_completed(&mut self.thread_writer) + } + .map_err(IntoProcessError::into_process_error) + } } impl ParallelProcessor for Encoder { @@ -227,12 +253,7 @@ impl ParallelProcessor for Encoder { } fn on_batch_complete(&mut self) -> paraseq::Result<()> { - // Only drain completed blocks mid-stream (keeps CBQ blocks full) - self.writer - .lock() - .ingest_completed(&mut self.thread_writer) - .map_err(IntoProcessError::into_process_error)?; - Ok(()) + self.flush_batch() } fn on_thread_complete(&mut self) -> paraseq::Result<()> { @@ -242,6 +263,10 @@ impl ParallelProcessor for Encoder { .map_err(IntoProcessError::into_process_error)?; Ok(()) } + + fn requires_ordering(&self) -> bool { + self.ordered + } } impl PairedParallelProcessor for Encoder { @@ -265,12 +290,7 @@ impl PairedParallelProcessor for Encoder { } fn on_batch_complete(&mut self) -> paraseq::Result<()> { - // Only drain completed blocks mid-stream (keeps CBQ blocks full) - self.writer - .lock() - .ingest_completed(&mut self.thread_writer) - .map_err(IntoProcessError::into_process_error)?; - Ok(()) + self.flush_batch() } fn on_thread_complete(&mut self) -> paraseq::Result<()> { @@ -280,6 +300,10 @@ impl PairedParallelProcessor for Encoder { .map_err(IntoProcessError::into_process_error)?; Ok(()) } + + fn requires_ordering(&self) -> bool { + self.ordered + } } #[cfg(test)] @@ -293,17 +317,27 @@ mod tests { #[test] fn test_encoder_builder_construction() { - let builder = BinseqWriterBuilder::new(Format::Vbq); + let builder = BinseqWriterBuilder::new(Format::Cbq); let handle = Box::new(Cursor::new(Vec::new())); let encoder_builder = FastxEncoderBuilder::new(builder, handle); assert!(encoder_builder.input.is_none()); assert_eq!(encoder_builder.threads, 0); + assert!(!encoder_builder.ordered); + } + + #[test] + fn test_encoder_builder_ordered_setter() { + let builder = BinseqWriterBuilder::new(Format::Cbq); + let handle = Box::new(Cursor::new(Vec::new())); + let encoder_builder = FastxEncoderBuilder::new(builder, handle).ordered(false); + + assert!(!encoder_builder.ordered); } #[test] fn test_encoder_builder_input_methods() { - let builder = BinseqWriterBuilder::new(Format::Vbq); + let builder = BinseqWriterBuilder::new(Format::Cbq); let handle = Box::new(Cursor::new(Vec::new())); let encoder_builder = FastxEncoderBuilder::new(builder, handle) .input("test.fastq") @@ -315,7 +349,7 @@ mod tests { #[test] fn test_encoder_builder_stdin() { - let builder = BinseqWriterBuilder::new(Format::Vbq); + let builder = BinseqWriterBuilder::new(Format::Cbq); let handle = Box::new(Cursor::new(Vec::new())); let encoder_builder = FastxEncoderBuilder::new(builder, handle).input_stdin(); @@ -324,7 +358,7 @@ mod tests { #[test] fn test_encoder_builder_single() { - let builder = BinseqWriterBuilder::new(Format::Vbq); + let builder = BinseqWriterBuilder::new(Format::Cbq); let handle = Box::new(Cursor::new(Vec::new())); let encoder_builder = FastxEncoderBuilder::new(builder, handle).input(FASTQ_R1_PATH); @@ -336,7 +370,7 @@ mod tests { #[test] fn test_encoder_builder_paired() { - let builder = BinseqWriterBuilder::new(Format::Vbq); + let builder = BinseqWriterBuilder::new(Format::Cbq); let handle = Box::new(Cursor::new(Vec::new())); let encoder_builder = FastxEncoderBuilder::new(builder, handle).input_paired(FASTQ_R1_PATH, FASTQ_R2_PATH); @@ -351,4 +385,65 @@ mod tests { // Run the encoder builder and assert that it is successful assert!(encoder_builder.run().is_ok()); } + + #[derive(Clone, Default)] + struct HeaderCollector { + headers: Arc>>, + } + impl crate::ParallelProcessor for HeaderCollector { + fn process_record(&mut self, record: R) -> crate::Result<()> { + let header = String::from_utf8_lossy(record.sheader()).into_owned(); + self.headers.lock().push((record.index(), header)); + Ok(()) + } + } + + /// Encodes a synthetic multi-threaded FASTQ input with `ordered(true)` and confirms + /// the written BINSEQ records come back out in the same order as the input, even + /// though multiple threads raced to produce them. + #[test] + fn test_encoder_builder_ordered_preserves_record_order() { + use crate::{BinseqReader, ParallelReader as DecodeReader}; + use std::fmt::Write as _; + use std::sync::Arc; + + const N_RECORDS: usize = 4_000; + const SEQ_LEN: usize = 32; + + let mut fastq = String::new(); + for i in 0..N_RECORDS { + let base = b"ACGT"[i % 4] as char; + let seq: String = std::iter::repeat_n(base, SEQ_LEN).collect(); + let qual: String = std::iter::repeat_n('F', SEQ_LEN).collect(); + let _ = writeln!(fastq, "@read_{i:06}\n{seq}\n+\n{qual}"); + } + let temp_dir = tempfile::tempdir().unwrap(); + + let input_path = temp_dir.path().join("input.fastq"); + let output_path = temp_dir.path().join("output.cbq"); + std::fs::write(&input_path, &fastq).unwrap(); + + let builder = BinseqWriterBuilder::new(Format::Cbq).headers(true); + let handle = Box::new(std::fs::File::create(&output_path).unwrap()); + let result = FastxEncoderBuilder::new(builder, handle) + .input(&input_path) + .threads(4) + .ordered(true) + .run(); + assert!(result.is_ok()); + + let reader = BinseqReader::new(&output_path).unwrap(); + let processor = HeaderCollector::default(); + let headers = processor.headers.clone(); + reader.process_parallel(processor, 4).unwrap(); + + let mut results = Arc::try_unwrap(headers).unwrap().into_inner(); + results.sort_by_key(|(idx, _)| *idx); + + assert_eq!(results.len(), N_RECORDS); + for (i, (idx, header)) in results.iter().enumerate() { + assert_eq!(*idx, i as u64); + assert_eq!(header, &format!("read_{i:06}")); + } + } }