From c6567d0844f43f8dadcc3b5e1a220dc5005b4d12 Mon Sep 17 00:00:00 2001 From: Nils Homer Date: Tue, 24 Mar 2026 00:28:16 -0700 Subject: [PATCH] feat(sort): add fgumi merge command with loser tree Add `fgumi merge` subcommand for k-way merging of pre-sorted BAM files, supporting coordinate, queryname, and template-coordinate sort orders. Features: - CLI follows fgumi conventions (-i, -o, --threads, --order) - Merges headers from all inputs (read groups, program records) - Validates reference dictionaries match across inputs - Detects output aliasing input paths - Accepts single input (acts as copy, matching samtools merge) - Supports -b/--input-list for file-of-filenames input Includes a loser tree (tournament tree) data structure for the merge, providing log2(k) comparisons per element vs 2*log2(k) for a binary heap. Based on the incremental insertion approach from Apache DataFusion. Also fixes chunk file naming collision during sort consolidation by replacing chunk_files.len() with a monotonic counter. --- src/commands/merge.rs | 459 +++++++++++++++++++++++++++++++++++++ src/commands/mod.rs | 1 + src/commands/sort.rs | 20 +- src/lib/sort/loser_tree.rs | 386 +++++++++++++++++++++++++++++++ src/lib/sort/mod.rs | 1 + src/lib/sort/raw.rs | 343 ++++++++++++++++++++++++++- src/main.rs | 27 ++- 7 files changed, 1217 insertions(+), 20 deletions(-) create mode 100644 src/commands/merge.rs create mode 100644 src/lib/sort/loser_tree.rs diff --git a/src/commands/merge.rs b/src/commands/merge.rs new file mode 100644 index 000000000..3bccb280a --- /dev/null +++ b/src/commands/merge.rs @@ -0,0 +1,459 @@ +//! Merge pre-sorted BAM files into a single sorted BAM. +//! +//! Performs a k-way merge of BAM files that are already sorted in the same +//! order, producing a single merged output that preserves the sort order. +//! +//! Similar to `samtools merge`, but supports template-coordinate order and +//! uses the same high-performance merge infrastructure as `fgumi sort`. + +use std::collections::HashSet; +use std::path::PathBuf; + +use anyhow::{Result, bail}; +use clap::Parser; +use fgumi_lib::bam_io::create_bam_reader; +use fgumi_lib::logging::OperationTimer; +use fgumi_lib::sort::RawExternalSorter; +use fgumi_lib::validation::validate_file_exists; +use log::info; +use noodles::sam::Header; + +use crate::commands::command::Command; +use crate::commands::sort::SortOrderArg; + +/// Merge pre-sorted BAM files. +/// +/// Performs a k-way merge of multiple BAM files that are already sorted in +/// the same order, similar to `samtools merge`. Input files must all be +/// sorted in the specified order. +#[derive(Debug, Parser)] +#[command( + name = "merge", + about = "\x1b[38;5;72m[ALIGNMENT]\x1b[0m \x1b[36mMerge pre-sorted BAM files into a single sorted BAM\x1b[0m", + long_about = r#" +Merge pre-sorted BAM files into a single sorted BAM. + +Performs a k-way merge of multiple BAM files that are already sorted in the +same order, producing a single merged output that preserves the sort order. +Similar to `samtools merge`, but supports template-coordinate order. + +Input files must all be sorted in the specified sort order. + +EXAMPLES: + + # Merge coordinate-sorted BAMs + fgumi merge -o merged.bam sorted1.bam sorted2.bam sorted3.bam + + # Merge template-coordinate sorted BAMs + fgumi merge -o merged.bam --order template-coordinate tc1.bam tc2.bam + + # Merge from a file listing input BAMs (one per line) + fgumi merge -o merged.bam -b input_list.txt --order queryname + + # Merge with multiple threads + fgumi merge -o merged.bam -@ 4 sorted1.bam sorted2.bam + +"# +)] +pub struct Merge { + /// Output BAM file. + #[arg(short = 'o', long = "output")] + pub output: PathBuf, + + /// Input BAM files to merge (positional). + #[arg(required_unless_present = "input_list")] + pub inputs: Vec, + + /// File containing a list of input BAM paths, one per line. + /// + /// Can be combined with positional inputs. + #[arg(short = 'b', long = "input-list")] + pub input_list: Option, + + /// Sort order of the input files. + #[arg(long = "order", value_enum, default_value = "template-coordinate")] + pub order: SortOrderArg, + + /// Number of threads for parallel operations. + /// + /// Used for multi-threaded BGZF compression. + #[arg(short = '@', short_alias = 't', long = "threads", default_value = "1")] + pub threads: usize, + + /// Compression level for output BAM (1-12). + /// + /// Level 1 is fastest with larger files. + /// Level 6 (default) balances speed and file size. + /// Level 12 produces smallest files but is slowest. + #[arg(long = "compression-level", default_value_t = 6)] + pub compression_level: u32, + + /// Cell barcode tag for template-coordinate merge. + /// + /// When merging in template-coordinate order, this tag is included in the + /// sort key so that templates from different cells at the same locus are + /// not interleaved. Only used for template-coordinate merge. + #[arg(short = 'c', long = "cell-tag", default_value = "CB")] + pub cell_tag: String, +} + +impl Command for Merge { + fn execute(&self, _command_line: &str) -> Result<()> { + let mut input_paths: Vec = self.inputs.clone(); + + if let Some(ref list_path) = self.input_list { + validate_file_exists(list_path, "Input list")?; + let contents = std::fs::read_to_string(list_path)?; + for line in contents.lines() { + let line = line.trim(); + if !line.is_empty() && !line.starts_with('#') { + input_paths.push(PathBuf::from(line)); + } + } + } + + if input_paths.is_empty() { + bail!("No input files specified"); + } + + for path in &input_paths { + validate_file_exists(path, "Input BAM")?; + } + + // Check output doesn't alias any input + if let Ok(output_canon) = std::fs::canonicalize(&self.output) { + for path in &input_paths { + if let Ok(input_canon) = std::fs::canonicalize(path) { + if output_canon == input_canon { + bail!( + "Output file '{}' is the same as input file '{}'", + self.output.display(), + path.display() + ); + } + } + } + } + + let cell_tag = crate::commands::sort::parse_cell_tag(self.order, &self.cell_tag)?; + + let timer = OperationTimer::new("Merging BAMs"); + + info!("Starting Merge"); + info!("Inputs: {} files", input_paths.len()); + for path in &input_paths { + info!(" {}", path.display()); + } + info!("Output: {}", self.output.display()); + info!("Sort order: {:?}", self.order); + if let Some(ct) = cell_tag { + info!("Cell tag: {}{}", ct[0] as char, ct[1] as char); + } + info!("Threads: {}", self.threads); + + // Read and merge headers from all inputs + let header = merge_headers(&input_paths)?; + + let mut sorter = RawExternalSorter::new(self.order.into()) + .threads(self.threads) + .output_compression(self.compression_level); + + if let Some(ct) = cell_tag { + sorter = sorter.cell_tag(ct); + } + + let records_merged = sorter.merge_bams(&input_paths, &header, &self.output)?; + + info!("=== Summary ==="); + info!("Records merged: {records_merged}"); + info!("Output: {}", self.output.display()); + + timer.log_completion(records_merged); + Ok(()) + } +} + +/// Merge headers from multiple BAM files. +/// +/// Uses the first input's reference sequences and header line as the base. +/// Combines read groups and program records from all inputs, with earlier +/// inputs taking precedence for duplicate IDs (matching `samtools merge`). +/// Validates that all inputs share the same reference sequences (names and order). +fn merge_headers(input_paths: &[PathBuf]) -> Result
{ + if input_paths.is_empty() { + bail!("No input files to merge headers from"); + } + + // Read all headers once + let headers: Vec
= input_paths + .iter() + .map(|path| { + let (_, header) = create_bam_reader(path, 1)?; + Ok(header) + }) + .collect::>>()?; + + let first_header = &headers[0]; + + if headers.len() == 1 { + return Ok(first_header.clone()); + } + + // Verify reference sequences match across all inputs + let first_refs = first_header.reference_sequences(); + for (i, (path, header)) in input_paths[1..].iter().zip(headers[1..].iter()).enumerate() { + let other_refs = header.reference_sequences(); + if first_refs.len() != other_refs.len() { + bail!( + "Reference sequence count mismatch: {} has {} references, {} has {}", + input_paths[0].display(), + first_refs.len(), + path.display(), + other_refs.len() + ); + } + for ((name1, _), (name2, _)) in first_refs.iter().zip(other_refs.iter()) { + if name1 != name2 { + bail!( + "Reference sequence mismatch at input {}: '{}' has '{}', '{}' has '{}'", + i + 2, + input_paths[0].display(), + String::from_utf8_lossy(name1.as_ref()), + path.display(), + String::from_utf8_lossy(name2.as_ref()), + ); + } + } + } + + let mut builder = Header::builder(); + + // Reference sequences from first input + for (name, seq) in first_header.reference_sequences() { + builder = builder.add_reference_sequence(name.clone(), seq.clone()); + } + + // Header line from first input + if let Some(hdr) = first_header.header() { + builder = builder.set_header(hdr.clone()); + } + + // Collect read groups and programs from all inputs (first input wins ties) + let mut rg_ids = HashSet::new(); + let mut pg_ids = HashSet::new(); + + for header in &headers { + for (id, rg) in header.read_groups() { + if rg_ids.insert(id.clone()) { + builder = builder.add_read_group(id.clone(), rg.clone()); + } + } + + for (id, pg) in header.programs().as_ref() { + if pg_ids.insert(id.clone()) { + builder = builder.add_program(id.clone(), pg.clone()); + } + } + } + + // Comments from first input only (matching samtools behavior) + for comment in first_header.comments() { + builder = builder.add_comment(comment.clone()); + } + + Ok(builder.build()) +} + +#[cfg(test)] +mod tests { + use super::*; + use bstr::BString; + use noodles::sam::header::record::value::Map; + use noodles::sam::header::record::value::map::ReadGroup; + use noodles::sam::header::record::value::map::read_group::tag as rg_tag; + use std::num::NonZeroUsize; + + /// Create a BAM with the given read group IDs and return its path. + fn write_bam_with_read_groups(dir: &std::path::Path, name: &str, rg_ids: &[&str]) -> PathBuf { + use noodles::sam::header::record::value::map::{Program, ReferenceSequence}; + + let mut header_builder = Header::builder(); + + // Add one reference sequence so we can write mapped records + let map = Map::::new(NonZeroUsize::new(200_000_000).unwrap()); + header_builder = header_builder.add_reference_sequence(BString::from("chr1"), map); + + for rg_id in rg_ids { + let rg = Map::::builder() + .insert(rg_tag::LIBRARY, format!("Lib_{rg_id}")) + .build() + .expect("valid RG"); + header_builder = header_builder.add_read_group(BString::from(*rg_id), rg); + } + + // Add a program record + let pg = Map::::default(); + header_builder = header_builder.add_program(BString::from("test"), pg); + + let header = header_builder.build(); + + // Write a minimal BAM with this header (no records needed for header tests) + let path = dir.join(format!("{name}.bam")); + let file = std::fs::File::create(&path).unwrap(); + let mut writer = noodles::bam::io::Writer::new(file); + writer.write_header(&header).unwrap(); + // write EOF + drop(writer); + path + } + + #[test] + fn test_merge_headers_single_input() { + let dir = tempfile::tempdir().unwrap(); + let bam = write_bam_with_read_groups(dir.path(), "single", &["RG1", "RG2"]); + + let header = merge_headers(std::slice::from_ref(&bam)).unwrap(); + + // Should be the same header (same RGs) + let rg_ids: Vec = + header.read_groups().iter().map(|(id, _)| id.to_string()).collect(); + assert_eq!(rg_ids.len(), 2); + assert!(rg_ids.contains(&"RG1".to_string())); + assert!(rg_ids.contains(&"RG2".to_string())); + } + + #[test] + fn test_merge_headers_combines_read_groups() { + let dir = tempfile::tempdir().unwrap(); + let bam_a = write_bam_with_read_groups(dir.path(), "a", &["RG1"]); + let bam_b = write_bam_with_read_groups(dir.path(), "b", &["RG2"]); + + let header = merge_headers(&[bam_a, bam_b]).unwrap(); + + let rg_ids: Vec = + header.read_groups().iter().map(|(id, _)| id.to_string()).collect(); + assert_eq!(rg_ids.len(), 2, "expected 2 read groups, got {rg_ids:?}"); + assert!(rg_ids.contains(&"RG1".to_string())); + assert!(rg_ids.contains(&"RG2".to_string())); + } + + #[test] + fn test_merge_headers_deduplicates_read_groups() { + let dir = tempfile::tempdir().unwrap(); + // Both BAMs have RG1, but with different library names + let bam_a = write_bam_with_read_groups(dir.path(), "a", &["RG1"]); + let bam_b = write_bam_with_read_groups(dir.path(), "b", &["RG1", "RG2"]); + + let header = merge_headers(&[bam_a, bam_b]).unwrap(); + + let rg_ids: Vec = + header.read_groups().iter().map(|(id, _)| id.to_string()).collect(); + // RG1 from first input wins, RG2 from second input is added => 2 total + assert_eq!(rg_ids.len(), 2, "expected 2 unique read groups, got {rg_ids:?}"); + assert!(rg_ids.contains(&"RG1".to_string())); + assert!(rg_ids.contains(&"RG2".to_string())); + + // Verify RG1's library comes from the first input + let (_, rg1) = header + .read_groups() + .iter() + .find(|(id, _)| >::as_ref(id) == b"RG1") + .unwrap(); + let lib = rg1.other_fields().get(&rg_tag::LIBRARY).map(|v| v.to_string()); + assert_eq!(lib, Some("Lib_RG1".to_string())); + } + + /// Create a BAM with the given reference sequence names and return its path. + fn write_bam_with_refs(dir: &std::path::Path, name: &str, ref_names: &[&str]) -> PathBuf { + use noodles::sam::header::record::value::map::ReferenceSequence; + + let mut header_builder = Header::builder(); + + for ref_name in ref_names { + let map = Map::::new(NonZeroUsize::new(200_000_000).unwrap()); + header_builder = header_builder.add_reference_sequence(BString::from(*ref_name), map); + } + + let header = header_builder.build(); + + let path = dir.join(format!("{name}.bam")); + let file = std::fs::File::create(&path).unwrap(); + let mut writer = noodles::bam::io::Writer::new(file); + writer.write_header(&header).unwrap(); + drop(writer); + path + } + + #[test] + fn test_merge_headers_rejects_different_ref_count() { + let dir = tempfile::tempdir().unwrap(); + let bam_a = write_bam_with_refs(dir.path(), "a", &["chr1", "chr2"]); + let bam_b = write_bam_with_refs(dir.path(), "b", &["chr1"]); + + let result = merge_headers(&[bam_a, bam_b]); + assert!(result.is_err()); + let msg = result.unwrap_err().to_string(); + assert!(msg.contains("Reference sequence count mismatch"), "unexpected error: {msg}"); + } + + #[test] + fn test_merge_headers_rejects_different_ref_names() { + let dir = tempfile::tempdir().unwrap(); + let bam_a = write_bam_with_refs(dir.path(), "a", &["chr1", "chr2"]); + let bam_b = write_bam_with_refs(dir.path(), "b", &["chr1", "chrX"]); + + let result = merge_headers(&[bam_a, bam_b]); + assert!(result.is_err()); + let msg = result.unwrap_err().to_string(); + assert!(msg.contains("Reference sequence mismatch"), "unexpected error: {msg}"); + } + + #[test] + fn test_merge_output_aliasing_input() { + let dir = tempfile::tempdir().unwrap(); + let bam = write_bam_with_read_groups(dir.path(), "input", &["RG1"]); + + let merge = Merge { + output: bam.clone(), + inputs: vec![bam.clone()], + input_list: None, + order: SortOrderArg::Coordinate, + threads: 1, + compression_level: 6, + cell_tag: "CB".to_string(), + }; + + let result = merge.execute("test"); + assert!(result.is_err()); + let msg = result.unwrap_err().to_string(); + assert!(msg.contains("is the same as input file"), "unexpected error: {msg}"); + } + + #[test] + fn test_merge_empty_bams_produces_output() { + let dir = tempfile::tempdir().unwrap(); + let bam_a = write_bam_with_read_groups(dir.path(), "a", &["RG1"]); + let bam_b = write_bam_with_read_groups(dir.path(), "b", &["RG2"]); + let output = dir.path().join("merged.bam"); + + let merge = Merge { + output: output.clone(), + inputs: vec![bam_a, bam_b], + input_list: None, + order: SortOrderArg::Coordinate, + threads: 1, + compression_level: 6, + cell_tag: "CB".to_string(), + }; + + let result = merge.execute("test"); + assert!(result.is_ok(), "merge failed: {:?}", result.unwrap_err()); + assert!(output.exists(), "output BAM was not created"); + + // Verify the output is a valid BAM by reading its header + let (_, header) = create_bam_reader(&output, 1).unwrap(); + let rg_ids: Vec = + header.read_groups().iter().map(|(id, _)| id.to_string()).collect(); + assert_eq!(rg_ids.len(), 2); + } +} diff --git a/src/commands/mod.rs b/src/commands/mod.rs index ca9049a4d..4aff097ed 100644 --- a/src/commands/mod.rs +++ b/src/commands/mod.rs @@ -66,6 +66,7 @@ pub mod extract; pub mod fastq; pub mod filter; pub mod group; +pub mod merge; pub mod review; pub mod simplex; #[cfg(feature = "simulate")] diff --git a/src/commands/sort.rs b/src/commands/sort.rs index ed9c5ff8a..8fc986f0c 100644 --- a/src/commands/sort.rs +++ b/src/commands/sort.rs @@ -188,7 +188,7 @@ pub struct Sort { } /// Parse memory size string (e.g., "512M", "1G", "2G"). -fn parse_memory(s: &str) -> Result { +pub(crate) fn parse_memory(s: &str) -> Result { let s = s.trim().to_uppercase(); if s.is_empty() { @@ -215,6 +215,17 @@ fn parse_memory(s: &str) -> Result { Ok((num * multiplier as f64) as usize) } +/// Parse the cell tag for template-coordinate sort/verify, returning `None` +/// for other sort orders. +pub(crate) fn parse_cell_tag(order: SortOrderArg, cell_tag: &str) -> Result> { + if matches!(order, SortOrderArg::TemplateCoordinate) { + let tag = string_to_tag(cell_tag, "cell-tag")?; + Ok(Some([tag.as_ref()[0], tag.as_ref()[1]])) + } else { + Ok(None) + } +} + /// Summary of sort-order verification: `(total_records, violations, first_violation)`. type VerifySummary = (u64, u64, Option<(u64, String)>); @@ -276,12 +287,7 @@ impl Sort { /// Parse the cell tag for template-coordinate sort/verify, returning `None` /// for other sort orders. fn parse_cell_tag(&self) -> Result> { - if matches!(self.order, SortOrderArg::TemplateCoordinate) { - let tag = string_to_tag(&self.cell_tag, "cell-tag")?; - Ok(Some([tag.as_ref()[0], tag.as_ref()[1]])) - } else { - Ok(None) - } + parse_cell_tag(self.order, &self.cell_tag) } /// Execute sort mode: read, sort, and write output. diff --git a/src/lib/sort/loser_tree.rs b/src/lib/sort/loser_tree.rs new file mode 100644 index 000000000..9054f2e44 --- /dev/null +++ b/src/lib/sort/loser_tree.rs @@ -0,0 +1,386 @@ +//! Loser tree (tournament tree) for efficient k-way merge. +//! +//! A loser tree supports k-way merge with `log2(k)` comparisons per element, +//! compared to `2·log2(k)` for a binary heap. Each internal node stores the +//! *loser* of the comparison at that level, and updating after replacing the +//! winner only requires walking leaf-to-root with one comparison per level. +//! +//! # Algorithm +//! +//! Uses the incremental insertion approach (no power-of-2 padding required): +//! - Array of `k` entries: `losers[0]` = overall winner, `losers[1..k]` = internal losers +//! - Leaf mapping: source `i` maps to internal node `(k + i) / 2` +//! - Parent navigation: `parent(node) = node / 2` +//! - Build: insert sources `0..k` one at a time, each sifting up from leaf to root +//! +//! Based on the approach used in Apache `DataFusion` (`arrow-rs`). +//! +//! # References +//! +//! - Knuth, *The Art of Computer Programming*, Vol. 3, Section 5.4.1 +//! - Apache `DataFusion` `SortPreservingMergeStream` (PR #4301) +//! - Grafana blog: "The Loser Tree Data Structure" + +use std::cmp::Ordering; + +/// Sentinel value indicating an unoccupied internal node during build. +const EMPTY: usize = usize::MAX; + +/// A loser tree for k-way merge of sorted sequences. +/// +/// The tree has `k` leaves (one per source). Internal nodes store the source +/// index of the loser at each tournament level. `losers[0]` holds the overall +/// winner (the source with the current minimum key). +pub struct LoserTree { + /// `losers[0]` = winner, `losers[1..k]` = internal loser nodes. + losers: Vec, + /// Current key for each source. + keys: Vec, + /// Number of sources. + k: usize, + /// Per-source active flag. Inactive sources always lose. + active: Vec, + /// Number of sources still active. + num_active: usize, +} + +impl LoserTree { + /// Create a new loser tree from initial keys (one per source). + /// + /// # Panics + /// + /// Panics if `initial_keys` is empty. + #[must_use] + pub fn new(initial_keys: Vec) -> Self { + let k = initial_keys.len(); + assert!(k > 0, "loser tree requires at least one source"); + + let mut tree = Self { + losers: vec![EMPTY; k], + keys: initial_keys, + k, + active: vec![true; k], + num_active: k, + }; + tree.build(); + tree + } + + /// Build the tournament by inserting each source from `0` to `k-1`. + /// + /// Each source sifts up from its leaf node toward the root. At occupied + /// internal nodes, the loser stays and the winner continues up. At empty + /// nodes (value = `EMPTY`), the winner is written and sifting stops. + fn build(&mut self) { + for i in 0..self.k { + let mut winner = i; + let mut node = self.leaf_to_node(i); + + while node > 0 && self.losers[node] != EMPTY { + let stored = self.losers[node]; + if self.is_greater(winner, stored) { + self.losers[node] = winner; + winner = stored; + } + node >>= 1; + } + + self.losers[node] = winner; + } + } + + /// Sift source `source` up from its leaf, replaying the tournament. + /// Used after replacing or removing the winner. Cost: `log2(k)` comparisons. + fn replay(&mut self, source: usize) { + let mut winner = source; + let mut node = self.leaf_to_node(source); + + while node > 0 { + let challenger = self.losers[node]; + if self.is_greater(winner, challenger) { + self.losers[node] = winner; + winner = challenger; + } + node >>= 1; + } + + self.losers[0] = winner; + } + + /// Map source index to its internal node index. + #[inline] + fn leaf_to_node(&self, source: usize) -> usize { + (self.k + source) >> 1 + } + + /// Returns true if source `a` is greater than source `b` (a loses to b). + /// Inactive sources are always greater. Ties broken by source index. + #[inline] + fn is_greater(&self, a: usize, b: usize) -> bool { + match (self.active[a], self.active[b]) { + (true, true) => { + matches!(self.keys[a].cmp(&self.keys[b]).then_with(|| a.cmp(&b)), Ordering::Greater) + } + (true, false) => false, + (false, true) => true, + (false, false) => a > b, + } + } + + /// Get the index of the current winner (source with minimum key). + #[inline] + #[must_use] + pub fn winner(&self) -> usize { + self.losers[0] + } + + /// Check if the winner is still an active source. + #[inline] + #[must_use] + pub fn winner_is_active(&self) -> bool { + self.num_active > 0 && self.active[self.losers[0]] + } + + /// Replace the winner's key with a new key and replay the tournament. + /// Costs `log2(k)` comparisons. + #[inline] + pub fn replace_winner(&mut self, new_key: K) { + let w = self.losers[0]; + self.keys[w] = new_key; + self.replay(w); + } + + /// Mark the winner's source as exhausted and replay the tournament. + /// Costs `log2(k)` comparisons. + pub fn remove_winner(&mut self) { + let w = self.losers[0]; + self.active[w] = false; + self.num_active -= 1; + self.replay(w); + } + + /// Get a reference to the winner's current key. + #[inline] + #[must_use] + pub fn winner_key(&self) -> &K { + &self.keys[self.losers[0]] + } + + /// Number of sources in the tree (including exhausted ones). + #[inline] + #[must_use] + pub fn len(&self) -> usize { + self.k + } + + /// Whether the tree has no sources. + /// + /// Always returns `false` since the constructor requires at least one source. + /// Present for clippy's `len_without_is_empty` lint. + #[inline] + #[must_use] + pub fn is_empty(&self) -> bool { + self.k == 0 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_single_source() { + let mut tree = LoserTree::new(vec![5]); + assert_eq!(tree.winner(), 0); + assert_eq!(*tree.winner_key(), 5); + assert!(tree.winner_is_active()); + + tree.replace_winner(10); + assert_eq!(tree.winner(), 0); + assert_eq!(*tree.winner_key(), 10); + + tree.remove_winner(); + assert!(!tree.winner_is_active()); + } + + #[test] + fn test_two_sources() { + let mut tree = LoserTree::new(vec![3, 1]); + assert_eq!(tree.winner(), 1); + assert_eq!(*tree.winner_key(), 1); + + tree.replace_winner(5); + assert_eq!(tree.winner(), 0); + assert_eq!(*tree.winner_key(), 3); + } + + #[test] + fn test_three_sources() { + let tree = LoserTree::new(vec![30, 10, 20]); + assert_eq!(tree.winner(), 1); + } + + #[test] + fn test_merge_three_sorted_sequences() { + let sources = [vec![1, 4, 7], vec![2, 5, 8], vec![3, 6, 9]]; + let mut indices = [0usize; 3]; + let initial_keys: Vec = sources.iter().map(|s| s[0]).collect(); + let mut tree = LoserTree::new(initial_keys); + + let mut result = Vec::new(); + while tree.winner_is_active() { + let w = tree.winner(); + result.push(*tree.winner_key()); + indices[w] += 1; + if indices[w] < sources[w].len() { + tree.replace_winner(sources[w][indices[w]]); + } else { + tree.remove_winner(); + } + } + + assert_eq!(result, vec![1, 2, 3, 4, 5, 6, 7, 8, 9]); + } + + #[test] + fn test_merge_with_duplicates_stable() { + let sources = [vec![1, 3, 3], vec![1, 2, 3]]; + let mut indices = [0usize; 2]; + let initial_keys: Vec = sources.iter().map(|s| s[0]).collect(); + let mut tree = LoserTree::new(initial_keys); + + let mut result = Vec::new(); + let mut winner_sources = Vec::new(); + while tree.winner_is_active() { + let w = tree.winner(); + result.push(*tree.winner_key()); + winner_sources.push(w); + indices[w] += 1; + if indices[w] < sources[w].len() { + tree.replace_winner(sources[w][indices[w]]); + } else { + tree.remove_winner(); + } + } + + assert_eq!(result, vec![1, 1, 2, 3, 3, 3]); + assert_eq!(winner_sources[0], 0); + assert_eq!(winner_sources[1], 1); + } + + #[test] + fn test_many_sources_descending_keys() { + let keys: Vec = (0..16).rev().collect(); + let mut tree = LoserTree::new(keys); + + let mut result = Vec::new(); + while tree.winner_is_active() { + result.push(*tree.winner_key()); + tree.remove_winner(); + } + + assert_eq!(result, (0..16).collect::>()); + } + + #[test] + fn test_non_power_of_two() { + let keys = vec![10, 30, 20, 50, 40]; + let mut tree = LoserTree::new(keys); + + let mut result = Vec::new(); + while tree.winner_is_active() { + result.push(*tree.winner_key()); + tree.remove_winner(); + } + + assert_eq!(result, vec![10, 20, 30, 40, 50]); + } + + #[test] + fn test_merge_longer_sequences() { + let sources = [vec![1, 5, 9, 13], vec![2, 6, 10], vec![3, 7, 11, 14, 15], vec![4, 8, 12]]; + let mut indices = [0usize; 4]; + let initial_keys: Vec = sources.iter().map(|s| s[0]).collect(); + let mut tree = LoserTree::new(initial_keys); + + let mut result = Vec::new(); + while tree.winner_is_active() { + let w = tree.winner(); + result.push(*tree.winner_key()); + indices[w] += 1; + if indices[w] < sources[w].len() { + tree.replace_winner(sources[w][indices[w]]); + } else { + tree.remove_winner(); + } + } + + assert_eq!(result, (1..=15).collect::>()); + } + + #[test] + fn test_large_fan_in() { + let n: i32 = 64; + let sources: Vec> = (0..n).map(|i| vec![i]).collect(); + let initial_keys: Vec = sources.iter().map(|s| s[0]).collect(); + let mut tree = LoserTree::new(initial_keys); + + let mut result = Vec::new(); + while tree.winner_is_active() { + result.push(*tree.winner_key()); + tree.remove_winner(); + } + + assert_eq!(result, (0..n).collect::>()); + } + + #[test] + fn test_all_same_keys_stable() { + let keys = vec![42, 42, 42, 42]; + let mut tree = LoserTree::new(keys); + + let mut winners = Vec::new(); + while tree.winner_is_active() { + winners.push(tree.winner()); + tree.remove_winner(); + } + + assert_eq!(winners, vec![0, 1, 2, 3]); + } + + #[test] + fn test_k_equals_7() { + let keys = vec![70, 10, 50, 30, 60, 20, 40]; + let mut tree = LoserTree::new(keys); + + let mut result = Vec::new(); + while tree.winner_is_active() { + result.push(*tree.winner_key()); + tree.remove_winner(); + } + + assert_eq!(result, vec![10, 20, 30, 40, 50, 60, 70]); + } + + #[test] + fn test_interleaved_merge() { + let sources = [vec![1, 3, 5, 7, 9], vec![2, 4, 6, 8, 10]]; + let mut indices = [0usize; 2]; + let initial_keys: Vec = sources.iter().map(|s| s[0]).collect(); + let mut tree = LoserTree::new(initial_keys); + + let mut result = Vec::new(); + while tree.winner_is_active() { + let w = tree.winner(); + result.push(*tree.winner_key()); + indices[w] += 1; + if indices[w] < sources[w].len() { + tree.replace_winner(sources[w][indices[w]]); + } else { + tree.remove_winner(); + } + } + + assert_eq!(result, (1..=10).collect::>()); + } +} diff --git a/src/lib/sort/mod.rs b/src/lib/sort/mod.rs index 5de22f18c..d05195082 100644 --- a/src/lib/sort/mod.rs +++ b/src/lib/sort/mod.rs @@ -32,6 +32,7 @@ pub use fgumi_raw_bam as bam_fields; pub mod external; pub mod inline_buffer; pub mod keys; +pub mod loser_tree; pub mod pipeline; pub mod radix; pub mod raw; diff --git a/src/lib/sort/raw.rs b/src/lib/sort/raw.rs index 6cee25356..011e6e606 100644 --- a/src/lib/sort/raw.rs +++ b/src/lib/sort/raw.rs @@ -681,6 +681,144 @@ impl RawExternalSorter { } } + /// Merge multiple pre-sorted BAM files into a single sorted BAM. + /// + /// Each input BAM must already be sorted in the order specified by + /// `self.sort_order`. The output preserves the sort order. + /// + /// # Errors + /// + /// Returns an error if any input cannot be opened, or writing fails. + pub fn merge_bams(&self, inputs: &[PathBuf], header: &Header, output: &Path) -> Result { + use crate::sort::inline_buffer::extract_coordinate_key_inline; + use crate::sort::keys::{RawCoordinateKey, RawQuerynameKey, RawSortKey, SortContext}; + + info!("Starting k-way merge of {} BAM files", inputs.len()); + + let mut readers = Self::open_bam_prefetch_readers(inputs)?; + let output_header = self.create_output_header(header); + + match self.sort_order { + SortOrder::TemplateCoordinate => { + let lib_lookup = LibraryLookup::from_header(header); + let cell_tag = self.cell_tag; + self.run_merge_loop(&mut readers, &output_header, output, |bam| { + extract_template_key_inline(bam, &lib_lookup, cell_tag.as_ref()) + }) + } + SortOrder::Coordinate => { + #[allow(clippy::cast_possible_truncation)] + let nref = header.reference_sequences().len() as u32; + self.run_merge_loop(&mut readers, &output_header, output, |bam| RawCoordinateKey { + sort_key: extract_coordinate_key_inline(bam, nref), + }) + } + SortOrder::Queryname => { + let ctx = SortContext::from_header(header); + self.run_merge_loop(&mut readers, &output_header, output, |bam| { + RawQuerynameKey::extract(bam, &ctx) + }) + } + } + } + + /// Open background prefetch readers for multiple BAM files. + fn open_bam_prefetch_readers(inputs: &[PathBuf]) -> Result> { + inputs + .iter() + .map(|path| { + let (reader, _header) = create_raw_bam_reader(path, 1)?; + Ok(RawReadAheadReader::new(reader)) + }) + .collect() + } + + /// K-way merge loop: extract keys on the merge thread, write to output. + fn run_merge_loop( + &self, + readers: &mut [RawReadAheadReader], + output_header: &Header, + output: &Path, + extract_key: impl Fn(&[u8]) -> K, + ) -> Result { + use crate::sort::loser_tree::LoserTree; + + const OUTPUT_BUFFER_SIZE: usize = 2048; + + // Initialize: collect first record + key from each reader + let mut initial_keys: Vec = Vec::with_capacity(readers.len()); + let mut records: Vec> = Vec::with_capacity(readers.len()); + let mut source_map: Vec = Vec::with_capacity(readers.len()); + + for (idx, reader) in readers.iter_mut().enumerate() { + if let Some(raw_record) = reader.next() { + let record_bytes = raw_record.as_ref().to_vec(); + initial_keys.push(extract_key(&record_bytes)); + records.push(record_bytes); + source_map.push(idx); + } + } + + if initial_keys.is_empty() { + info!("Merge complete: 0 records merged"); + let writer = crate::bam_io::create_raw_bam_writer( + output, + output_header, + self.threads, + self.output_compression, + )?; + writer.finish()?; + return Ok(0); + } + + let mut tree = LoserTree::new(initial_keys); + let k = tree.len(); + + let mut writer = crate::bam_io::create_raw_bam_writer( + output, + output_header, + self.threads, + self.output_compression, + )?; + + let mut records_merged = 0u64; + let mut output_buffer: Vec> = Vec::with_capacity(OUTPUT_BUFFER_SIZE); + + while tree.winner_is_active() { + let winner = tree.winner(); + let record_bytes = std::mem::take(&mut records[winner]); + + output_buffer.push(record_bytes); + records_merged += 1; + + if output_buffer.len() >= OUTPUT_BUFFER_SIZE { + for rec in output_buffer.drain(..) { + writer.write_raw_record(&rec)?; + } + } + + // Fetch next record from the same source + let reader_idx = source_map[winner]; + if let Some(raw_record) = readers[reader_idx].next() { + let next_rec = raw_record.as_ref().to_vec(); + let new_key = extract_key(&next_rec); + records[winner] = next_rec; + tree.replace_winner(new_key); + } else { + tree.remove_winner(); + } + } + + for rec in output_buffer { + writer.write_raw_record(&rec)?; + } + + writer.finish()?; + + info!("Merge complete: {records_merged} records merged (loser tree, k={k})",); + Ok(records_merged) + } + /// Sort by coordinate order using optimized radix sort for large arrays. fn sort_coordinate( &self, @@ -841,7 +979,7 @@ impl RawExternalSorter { /// Similar to `sort_coordinate_optimized` but uses `IndexingBamWriter` to /// build the BAI index incrementally during write. Uses single-threaded /// compression for accurate virtual position tracking. - #[allow(clippy::cast_possible_truncation)] + #[allow(clippy::cast_possible_truncation, clippy::too_many_lines)] fn sort_coordinate_with_index( &self, reader: crate::bam_io::RawBamReaderAuto, @@ -1843,6 +1981,7 @@ pub use super::SortStats as RawSortStats; #[cfg(test)] mod tests { use super::*; + use crate::sam::builder::SamBuilder; use noodles::sam::header::record::value::map::ReadGroup; // ======================================================================== @@ -2269,4 +2408,206 @@ mod tests { let observed = count_bam_records(&output); assert_eq!(observed, expected, "chunk filename collision likely lost data"); } + + // ======================================================================== + // merge_bams tests + // ======================================================================== + + /// Helper: create a `SamBuilder` with `num_pairs` pairs at non-overlapping positions, + /// write an unsorted BAM, sort it with the given order, and return the sorted path. + /// `start_offset` shifts all positions so different inputs have distinct read names/positions. + fn create_sorted_bam( + dir: &Path, + prefix: &str, + num_pairs: usize, + start_offset: usize, + sort_order: SortOrder, + ) -> (PathBuf, Vec) { + let mut builder = SamBuilder::new(); + let mut names = Vec::with_capacity(num_pairs); + for i in 0..num_pairs { + let name = format!("{prefix}_read{i:04}"); + names.push(name.clone()); + let _ = builder + .add_pair() + .name(&name) + .start1((start_offset + i * 200) + 1) + .start2((start_offset + i * 200) + 101) + .build(); + } + let unsorted = dir.join(format!("{prefix}_unsorted.bam")); + let sorted = dir.join(format!("{prefix}_sorted.bam")); + builder.write_bam(&unsorted).unwrap(); + RawExternalSorter::new(sort_order).output_compression(0).sort(&unsorted, &sorted).unwrap(); + (sorted, names) + } + + /// Collect all read names from a BAM file as strings. + fn collect_read_names(path: &Path) -> Vec { + use crate::sort::read_ahead::RawReadAheadReader; + let (reader, _) = create_raw_bam_reader(path, 1).unwrap(); + RawReadAheadReader::new(reader) + .map(|rec| { + let name_bytes = fgumi_raw_bam::fields::read_name(rec.as_ref()); + String::from_utf8(name_bytes.to_vec()).unwrap() + }) + .collect() + } + + /// Collect (`ref_id`, pos) tuples for every record in a BAM. + fn collect_positions(path: &Path) -> Vec<(i32, i32)> { + use crate::sort::read_ahead::RawReadAheadReader; + let (reader, _) = create_raw_bam_reader(path, 1).unwrap(); + RawReadAheadReader::new(reader) + .map(|rec| { + let bytes = rec.as_ref(); + (fgumi_raw_bam::fields::ref_id(bytes), fgumi_raw_bam::fields::pos(bytes)) + }) + .collect() + } + + /// Helper to build a merge header from the `SamBuilder` default header. + fn default_merge_header() -> Header { + SamBuilder::new().header.clone() + } + + #[test] + fn test_merge_bams_coordinate_sort() { + let dir = tempfile::tempdir().unwrap(); + let (bam_a, _) = create_sorted_bam(dir.path(), "a", 10, 0, SortOrder::Coordinate); + let (bam_b, _) = create_sorted_bam(dir.path(), "b", 10, 10_000, SortOrder::Coordinate); + + let merged = dir.path().join("merged.bam"); + let header = default_merge_header(); + let count = RawExternalSorter::new(SortOrder::Coordinate) + .output_compression(0) + .merge_bams(&[bam_a, bam_b], &header, &merged) + .unwrap(); + + // 10 pairs * 2 records * 2 inputs = 40 + assert_eq!(count, 40); + assert_eq!(count_bam_records(&merged), 40); + + // Verify coordinate sort order: (ref_id, pos) is non-decreasing + let positions = collect_positions(&merged); + for w in positions.windows(2) { + assert!(w[0] <= w[1], "coordinate sort violated: {:?} > {:?}", w[0], w[1]); + } + } + + #[test] + fn test_merge_bams_template_coordinate_sort() { + let dir = tempfile::tempdir().unwrap(); + let (bam_a, _) = create_sorted_bam(dir.path(), "a", 10, 0, SortOrder::TemplateCoordinate); + let (bam_b, _) = + create_sorted_bam(dir.path(), "b", 10, 10_000, SortOrder::TemplateCoordinate); + + let merged = dir.path().join("merged.bam"); + let header = default_merge_header(); + let count = RawExternalSorter::new(SortOrder::TemplateCoordinate) + .output_compression(0) + .merge_bams(&[bam_a, bam_b], &header, &merged) + .unwrap(); + + assert_eq!(count, 40); + assert_eq!(count_bam_records(&merged), 40); + } + + #[test] + fn test_merge_bams_queryname_sort() { + let dir = tempfile::tempdir().unwrap(); + let (bam_a, _) = create_sorted_bam(dir.path(), "a", 10, 0, SortOrder::Queryname); + let (bam_b, _) = create_sorted_bam(dir.path(), "b", 10, 10_000, SortOrder::Queryname); + + let merged = dir.path().join("merged.bam"); + let header = default_merge_header(); + let count = RawExternalSorter::new(SortOrder::Queryname) + .output_compression(0) + .merge_bams(&[bam_a, bam_b], &header, &merged) + .unwrap(); + + assert_eq!(count, 40); + assert_eq!(count_bam_records(&merged), 40); + + // Verify queryname sort order: read names are non-decreasing + let names = collect_read_names(&merged); + for w in names.windows(2) { + assert!(w[0] <= w[1], "queryname sort violated: {:?} > {:?}", w[0], w[1]); + } + } + + #[test] + fn test_merge_bams_single_input() { + let dir = tempfile::tempdir().unwrap(); + let (bam_a, _) = create_sorted_bam(dir.path(), "a", 15, 0, SortOrder::Coordinate); + + let merged = dir.path().join("merged.bam"); + let header = default_merge_header(); + let count = RawExternalSorter::new(SortOrder::Coordinate) + .output_compression(0) + .merge_bams(&[bam_a], &header, &merged) + .unwrap(); + + // 15 pairs * 2 = 30 + assert_eq!(count, 30); + assert_eq!(count_bam_records(&merged), 30); + } + + #[test] + fn test_merge_bams_preserves_all_records() { + let dir = tempfile::tempdir().unwrap(); + let (bam_a, names_a) = create_sorted_bam(dir.path(), "a", 5, 0, SortOrder::Queryname); + let (bam_b, names_b) = create_sorted_bam(dir.path(), "b", 5, 10_000, SortOrder::Queryname); + + let merged = dir.path().join("merged.bam"); + let header = default_merge_header(); + RawExternalSorter::new(SortOrder::Queryname) + .output_compression(0) + .merge_bams(&[bam_a, bam_b], &header, &merged) + .unwrap(); + + let merged_names: std::collections::HashSet = + collect_read_names(&merged).into_iter().collect(); + + // Every expected name (from both inputs) must appear in the merged output + for name in names_a.iter().chain(names_b.iter()) { + assert!(merged_names.contains(name), "read name {name:?} missing from merged output"); + } + } + + #[test] + fn test_merge_bams_many_inputs() { + let dir = tempfile::tempdir().unwrap(); + let k = 8; + let pairs_per_input = 5; + let mut inputs = Vec::with_capacity(k); + + for i in 0..k { + let (bam, _) = create_sorted_bam( + dir.path(), + &format!("in{i}"), + pairs_per_input, + i * 50_000, + SortOrder::Coordinate, + ); + inputs.push(bam); + } + + let merged = dir.path().join("merged.bam"); + let header = default_merge_header(); + let count = RawExternalSorter::new(SortOrder::Coordinate) + .output_compression(0) + .merge_bams(&inputs, &header, &merged) + .unwrap(); + + let expected = (k * pairs_per_input * 2) as u64; // 8 * 5 * 2 = 80 + assert_eq!(count, expected); + assert_eq!(count_bam_records(&merged), expected); + + // Verify coordinate sort order + let positions = collect_positions(&merged); + for w in positions.windows(2) { + assert!(w[0] <= w[1], "coordinate sort violated with k={k}: {:?} > {:?}", w[0], w[1]); + } + } } diff --git a/src/main.rs b/src/main.rs index 22fdf28d7..05a2f422e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -26,6 +26,7 @@ use commands::extract::Extract; use commands::fastq::Fastq; use commands::filter::Filter; use commands::group::GroupReadsByUmi; +use commands::merge::Merge; use commands::review::Review; use commands::simplex::Simplex; #[cfg(feature = "simulate")] @@ -81,41 +82,43 @@ enum Subcommand { Zipper(Zipper), #[command(display_order = 5)] Sort(Sort), + #[command(display_order = 6)] + Merge(Merge), // Group - #[command(display_order = 6)] + #[command(display_order = 7)] Group(GroupReadsByUmi), // Deduplication - #[command(display_order = 7)] + #[command(display_order = 8)] Dedup(MarkDuplicates), // Consensus Calling - #[command(display_order = 8)] - Simplex(Simplex), #[command(display_order = 9)] - Duplex(Duplex), + Simplex(Simplex), #[command(display_order = 10)] + Duplex(Duplex), + #[command(display_order = 11)] Codec(Codec), // Post-consensus - #[command(display_order = 11)] - Filter(Filter), #[command(display_order = 12)] - Clip(Clip), + Filter(Filter), #[command(display_order = 13)] - DuplexMetrics(DuplexMetrics), + Clip(Clip), #[command(display_order = 14)] + DuplexMetrics(DuplexMetrics), + #[command(display_order = 15)] Review(Review), // Utilities - #[command(display_order = 15)] + #[command(display_order = 16)] Downsample(Downsample), #[cfg(feature = "compare")] - #[command(display_order = 16)] + #[command(display_order = 17)] Compare(Compare), #[cfg(feature = "simulate")] - #[command(display_order = 17)] + #[command(display_order = 18)] Simulate(Simulate), }