From 657f1f83478c0d2a5490e9d56b05f1f0f2782a5a Mon Sep 17 00:00:00 2001 From: Nils Homer Date: Fri, 3 Apr 2026 23:31:42 -0700 Subject: [PATCH] perf(zipper): eliminate RecordBuf re-encode bottleneck with raw-byte merge Replace the noodles RecordBuf encode path in the zipper with direct raw Vec BAM record manipulation, eliminating the serialization bottleneck visible in profiling. The mapped template is converted to raw bytes once (via the vendored BAM encoder), then all 7 merge operations (fix mate info, tag removal, tag copying with reverse/revcomp, QC flag transfer, AS/XS normalization, PA tags) are performed directly on the raw bytes using fgumi_raw_bam primitives. Output is written via RawBamWriter::write_raw_record, bypassing the noodles re-encode entirely. New raw BAM primitives: field setters (set_mapq, set_mate_ref_id, set_mate_pos, set_template_length, set_ref_id, set_pos), cigar_to_string_from_raw, append_i32_array_tag, append_signed_int_tag, normalize_int_tag_to_smallest_signed, copy_aux_tags. New Template method: fix_mate_info_raw operating on raw_records. Adds tests for AS/XS normalization, negative strand tag copying, and fix_mate_info value verification. --- crates/fgumi-raw-bam/src/cigar.rs | 51 ++ crates/fgumi-raw-bam/src/fields.rs | 88 ++ crates/fgumi-raw-bam/src/tags.rs | 233 +++++ src/commands/zipper.rs | 1334 ++++++++++++++++++++++++---- src/lib/template.rs | 329 +++++++ 5 files changed, 1882 insertions(+), 153 deletions(-) diff --git a/crates/fgumi-raw-bam/src/cigar.rs b/crates/fgumi-raw-bam/src/cigar.rs index d1713dfb1..a558c29c4 100644 --- a/crates/fgumi-raw-bam/src/cigar.rs +++ b/crates/fgumi-raw-bam/src/cigar.rs @@ -49,6 +49,34 @@ pub fn get_cigar_ops(bam: &[u8]) -> Vec { cigar_bytes.chunks_exact(4).map(|c| u32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect() } +/// CIGAR operation type to SAM character mapping. +const CIGAR_OP_CHARS: [u8; 9] = [b'M', b'I', b'D', b'N', b'S', b'H', b'P', b'=', b'X']; + +/// Formats CIGAR operations from a raw BAM record as a SAM-style CIGAR string. +/// +/// Returns an empty string if the record has no CIGAR operations, or `"*"` is not used +/// (callers should check `n_cigar_op == 0` separately if they need `"*"`). +/// +/// # Panics +/// +/// Panics if a CIGAR operation has an invalid type (>= 9). +#[must_use] +pub fn cigar_to_string_from_raw(bam: &[u8]) -> String { + let ops = get_cigar_ops(bam); + if ops.is_empty() { + return String::new(); + } + let mut result = String::with_capacity(ops.len() * 4); + for &op in &ops { + let op_len = op >> 4; + let op_type = (op & 0xF) as usize; + assert!(op_type < CIGAR_OP_CHARS.len(), "invalid CIGAR op type: {op_type}"); + result.push_str(&op_len.to_string()); + result.push(CIGAR_OP_CHARS[op_type] as char); + } + result +} + /// Calculate reference-consuming length from CIGAR operations. /// /// This is the sum of M/D/N/=/X operations, which represents how many @@ -1934,4 +1962,27 @@ mod tests { rec.truncate(36); assert_eq!(unclipped_5prime_raw(&rec, 100, false), 100); } + + // ======================================================================== + // cigar_to_string_from_raw tests + // ======================================================================== + + #[test] + fn test_cigar_to_string_simple() { + let rec = make_bam_bytes(0, 0, 0, b"r1", &[encode_op(0, 100)], 100, -1, -1, &[]); + assert_eq!(cigar_to_string_from_raw(&rec), "100M"); + } + + #[test] + fn test_cigar_to_string_complex() { + let ops = [encode_op(4, 5), encode_op(0, 50), encode_op(1, 3), encode_op(0, 42)]; + let rec = make_bam_bytes(0, 0, 0, b"r1", &ops, 100, -1, -1, &[]); + assert_eq!(cigar_to_string_from_raw(&rec), "5S50M3I42M"); + } + + #[test] + fn test_cigar_to_string_no_cigar() { + let rec = make_bam_bytes(0, 0, 0, b"r1", &[], 4, -1, -1, &[]); + assert_eq!(cigar_to_string_from_raw(&rec), ""); + } } diff --git a/crates/fgumi-raw-bam/src/fields.rs b/crates/fgumi-raw-bam/src/fields.rs index dfa4387f9..6f86d7733 100644 --- a/crates/fgumi-raw-bam/src/fields.rs +++ b/crates/fgumi-raw-bam/src/fields.rs @@ -193,6 +193,42 @@ pub fn set_flags(bam: &mut [u8], new_flags: u16) { bam[14..16].copy_from_slice(&new_flags.to_le_bytes()); } +/// Set the mapping quality in a BAM record. +#[inline] +pub fn set_mapq(bam: &mut [u8], new_mapq: u8) { + bam[9] = new_mapq; +} + +/// Set the mate reference sequence ID in a BAM record. +#[inline] +pub fn set_mate_ref_id(bam: &mut [u8], val: i32) { + bam[20..24].copy_from_slice(&val.to_le_bytes()); +} + +/// Set the mate 0-based position in a BAM record. +#[inline] +pub fn set_mate_pos(bam: &mut [u8], val: i32) { + bam[24..28].copy_from_slice(&val.to_le_bytes()); +} + +/// Set the template length (TLEN) in a BAM record. +#[inline] +pub fn set_template_length(bam: &mut [u8], val: i32) { + bam[28..32].copy_from_slice(&val.to_le_bytes()); +} + +/// Set the reference sequence ID in a BAM record. +#[inline] +pub fn set_ref_id(bam: &mut [u8], val: i32) { + bam[0..4].copy_from_slice(&val.to_le_bytes()); +} + +/// Set the 0-based leftmost position in a BAM record. +#[inline] +pub fn set_pos(bam: &mut [u8], val: i32) { + bam[4..8].copy_from_slice(&val.to_le_bytes()); +} + /// Calculate the offset to auxiliary data in a BAM record. /// /// `aux_offset = 32 + l_read_name + n_cigar_op*4 + (l_seq+1)/2 + l_seq` @@ -841,4 +877,56 @@ mod tests { assert!(fields.flags.paired()); assert_eq!(fields.name, b"rd"); } + + // ======================================================================== + // Field setter tests + // ======================================================================== + + #[test] + fn test_set_mapq_roundtrip() { + let mut rec = make_bam_bytes(0, 0, 0, b"r1", &[], 4, -1, -1, &[]); + assert_eq!(mapq(&rec), 0); + set_mapq(&mut rec, 42); + assert_eq!(mapq(&rec), 42); + } + + #[test] + fn test_set_mate_ref_id_roundtrip() { + let mut rec = make_bam_bytes(0, 0, 0, b"r1", &[], 4, -1, -1, &[]); + assert_eq!(mate_ref_id(&rec), -1); + set_mate_ref_id(&mut rec, 5); + assert_eq!(mate_ref_id(&rec), 5); + } + + #[test] + fn test_set_mate_pos_roundtrip() { + let mut rec = make_bam_bytes(0, 0, 0, b"r1", &[], 4, -1, -1, &[]); + assert_eq!(mate_pos(&rec), -1); + set_mate_pos(&mut rec, 12345); + assert_eq!(mate_pos(&rec), 12345); + } + + #[test] + fn test_set_template_length_roundtrip() { + let mut rec = make_bam_bytes(0, 0, 0, b"r1", &[], 4, -1, -1, &[]); + assert_eq!(template_length(&rec), 0); + set_template_length(&mut rec, -300); + assert_eq!(template_length(&rec), -300); + } + + #[test] + fn test_set_ref_id_roundtrip() { + let mut rec = make_bam_bytes(0, 0, 0, b"r1", &[], 4, -1, -1, &[]); + assert_eq!(ref_id(&rec), 0); + set_ref_id(&mut rec, 3); + assert_eq!(ref_id(&rec), 3); + } + + #[test] + fn test_set_pos_roundtrip() { + let mut rec = make_bam_bytes(0, 0, 0, b"r1", &[], 4, -1, -1, &[]); + assert_eq!(pos(&rec), 0); + set_pos(&mut rec, 9999); + assert_eq!(pos(&rec), 9999); + } } diff --git a/crates/fgumi-raw-bam/src/tags.rs b/crates/fgumi-raw-bam/src/tags.rs index 657865c20..6fc5c2031 100644 --- a/crates/fgumi-raw-bam/src/tags.rs +++ b/crates/fgumi-raw-bam/src/tags.rs @@ -688,6 +688,99 @@ pub fn reverse_complement_string_tag_in_place(record: &mut [u8], aux_offset: usi } } +/// Append an `i32` array (`B:i`-type) tag to a BAM record. +/// +/// Format: `[tag0, tag1, 'B', 'i', count_u32_le, values_i32_le...]` +/// +/// # Panics +/// +/// Panics if `values.len()` exceeds `u32::MAX`. +pub fn append_i32_array_tag(record: &mut Vec, tag: &[u8; 2], values: &[i32]) { + record.push(tag[0]); + record.push(tag[1]); + record.push(b'B'); + record.push(b'i'); + record.extend_from_slice( + &u32::try_from(values.len()).expect("array length exceeds u32").to_le_bytes(), + ); + for &v in values { + record.extend_from_slice(&v.to_le_bytes()); + } +} + +/// Normalize an integer tag to the smallest signed type that fits its value. +/// +/// If the tag exists and is an integer type, it is re-encoded using +/// [`append_signed_int_tag`] semantics (signed types only: i8 → i16 → i32). +/// No-op if the tag is not found or is not an integer type. +pub fn normalize_int_tag_to_smallest_signed(record: &mut Vec, tag: &[u8; 2]) { + let Some(aux_start) = aux_data_offset_from_record(record) else { + return; + }; + if aux_start >= record.len() { + return; + } + let aux_data = &record[aux_start..]; + let Some(value) = find_int_tag(aux_data, tag) else { + return; + }; + let Ok(value_i32) = i32::try_from(value) else { + return; + }; + // Remove and re-append with smallest *signed* encoding + // (i8 -> i16 -> i32, matching fgbio's to_smallest_signed_int) + remove_tag(record, tag); + append_signed_int_tag(record, tag, value_i32); +} + +/// Append an integer tag using the smallest *signed* type that fits. +/// +/// Unlike [`append_int_tag`] (which prefers unsigned types), this uses only +/// signed types: `i8` (type `'c'`) -> `i16` (type `'s'`) -> `i32` (type `'i'`). +/// This matches fgbio's `to_smallest_signed_int` encoding. +pub fn append_signed_int_tag(record: &mut Vec, tag: &[u8; 2], value: i32) { + record.push(tag[0]); + record.push(tag[1]); + if let Ok(v) = i8::try_from(value) { + record.push(b'c'); + record.push(v.cast_unsigned()); + } else if let Ok(v) = i16::try_from(value) { + record.push(b's'); + record.extend_from_slice(&v.to_le_bytes()); + } else { + record.push(b'i'); + record.extend_from_slice(&value.to_le_bytes()); + } +} + +/// Copy auxiliary tags from source aux data to a destination record, skipping specified tags. +/// +/// Iterates all tags in `src_aux` and appends each tag entry (tag + type + value bytes) +/// to `dest`, unless the tag's two-byte key is in `skip_tags`. +pub fn copy_aux_tags(src_aux: &[u8], dest: &mut Vec, skip_tags: &[&[u8; 2]]) { + let mut offset = 0; + while offset + 3 <= src_aux.len() { + let tag_key = [src_aux[offset], src_aux[offset + 1]]; + let val_type = src_aux[offset + 2]; + let value_start = offset + 3; + + let Some(value_size) = tag_value_size(val_type, &src_aux[value_start..]) else { + break; + }; + let entry_end = value_start + value_size; + if entry_end > src_aux.len() { + break; + } + + // Copy unless this tag should be skipped + if !skip_tags.iter().any(|t| **t == tag_key) { + dest.extend_from_slice(&src_aux[offset..entry_end]); + } + + offset = entry_end; + } +} + #[cfg(test)] #[allow(clippy::identity_op)] mod tests { @@ -2171,4 +2264,144 @@ mod tests { // But find_tag_type correctly finds it assert!(find_tag_type(&aux, &pa_tag_bytes).is_some()); } + + // ======================================================================== + // append_i32_array_tag tests + // ======================================================================== + + #[test] + fn test_append_i32_array_tag() { + let mut rec = make_bam_bytes(0, 0, 0, b"r1", &[], 4, -1, -1, &[]); + let values = [1i32, -200, 300_000]; + append_i32_array_tag(&mut rec, b"pa", &values); + + let aux = aux_data_slice(&rec); + let tag_type = find_tag_type(aux, b"pa"); + assert_eq!(tag_type, Some(b'B')); + + // Verify the array contents + let arr = find_array_tag(aux, b"pa").unwrap(); + assert_eq!(arr.count, 3); + assert_eq!(arr.elem_type, b'i'); + } + + #[test] + fn test_append_i32_array_tag_empty() { + let mut rec = make_bam_bytes(0, 0, 0, b"r1", &[], 4, -1, -1, &[]); + append_i32_array_tag(&mut rec, b"pa", &[]); + + let aux = aux_data_slice(&rec); + let arr = find_array_tag(aux, b"pa").unwrap(); + assert_eq!(arr.count, 0); + } + + // ======================================================================== + // normalize_int_tag_to_smallest_signed tests + // ======================================================================== + + #[test] + fn test_normalize_int_tag_i32_to_i8() { + // AS:i:77 (stored as i32) should normalize to AS:c:77 (i8) + let mut aux = Vec::new(); + aux.extend_from_slice(b"ASi"); + aux.extend_from_slice(&77i32.to_le_bytes()); + let mut rec = make_bam_bytes(0, 0, 0, b"r1", &[], 4, -1, -1, &aux); + + normalize_int_tag_to_smallest_signed(&mut rec, b"AS"); + + let aux_data = aux_data_slice(&rec); + let (_, val_type) = find_tag_position(aux_data, *b"AS").unwrap(); + assert_eq!(val_type, b'c', "77 should fit in i8 (type 'c')"); + assert_eq!(find_int_tag(aux_data, b"AS"), Some(77)); + } + + #[test] + fn test_normalize_int_tag_i32_to_i16() { + // AS:i:200 (stored as i32) should normalize to AS:s:200 (i16, signed-only encoding) + let mut aux = Vec::new(); + aux.extend_from_slice(b"ASi"); + aux.extend_from_slice(&200i32.to_le_bytes()); + let mut rec = make_bam_bytes(0, 0, 0, b"r1", &[], 4, -1, -1, &aux); + + normalize_int_tag_to_smallest_signed(&mut rec, b"AS"); + + let aux_data = aux_data_slice(&rec); + let (_, val_type) = find_tag_position(aux_data, *b"AS").unwrap(); + assert_eq!(val_type, b's', "200 should be i16 (type 's') with signed-only encoding"); + assert_eq!(find_int_tag(aux_data, b"AS"), Some(200)); + } + + #[test] + fn test_normalize_int_tag_preserves_large_value() { + // AS:i:100000 should remain i32 + let mut aux = Vec::new(); + aux.extend_from_slice(b"ASi"); + aux.extend_from_slice(&100_000i32.to_le_bytes()); + let mut rec = make_bam_bytes(0, 0, 0, b"r1", &[], 4, -1, -1, &aux); + + normalize_int_tag_to_smallest_signed(&mut rec, b"AS"); + + let aux_data = aux_data_slice(&rec); + let (_, val_type) = find_tag_position(aux_data, *b"AS").unwrap(); + assert_eq!(val_type, b'i', "100000 requires i32 (type 'i')"); + assert_eq!(find_int_tag(aux_data, b"AS"), Some(100_000)); + } + + #[test] + fn test_normalize_int_tag_missing_is_noop() { + let mut rec = make_bam_bytes(0, 0, 0, b"r1", &[], 4, -1, -1, &[]); + let original = rec.clone(); + normalize_int_tag_to_smallest_signed(&mut rec, b"AS"); + assert_eq!(rec, original); + } + + // ======================================================================== + // copy_aux_tags tests + // ======================================================================== + + #[test] + fn test_copy_aux_tags_all() { + // Source has RX:Z:ACGT and NM:C:5 + let mut src_aux = Vec::new(); + src_aux.extend_from_slice(b"RXZ\x41\x43\x47\x54\x00"); // RX:Z:ACGT + src_aux.extend_from_slice(&[b'N', b'M', b'C', 5]); // NM:C:5 + + let mut dest = Vec::new(); + copy_aux_tags(&src_aux, &mut dest, &[]); + + assert_eq!(dest, src_aux, "All tags should be copied when no skip list"); + } + + #[test] + fn test_copy_aux_tags_with_skip() { + // Source has RX:Z:ACGT and NM:C:5 + let mut src_aux = Vec::new(); + src_aux.extend_from_slice(b"RXZ\x41\x43\x47\x54\x00"); // RX:Z:ACGT + src_aux.extend_from_slice(&[b'N', b'M', b'C', 5]); // NM:C:5 + + let mut dest = Vec::new(); + copy_aux_tags(&src_aux, &mut dest, &[b"NM"]); + + // Only RX should be copied + assert_eq!(dest, b"RXZ\x41\x43\x47\x54\x00"); + } + + #[test] + fn test_copy_aux_tags_skip_all() { + let mut src_aux = Vec::new(); + src_aux.extend_from_slice(b"RXZ\x41\x43\x47\x54\x00"); + src_aux.extend_from_slice(&[b'N', b'M', b'C', 5]); + + let mut dest = Vec::new(); + copy_aux_tags(&src_aux, &mut dest, &[b"RX", b"NM"]); + + assert!(dest.is_empty(), "All tags skipped should produce empty dest"); + } + + #[test] + fn test_copy_aux_tags_empty_source() { + let mut dest = Vec::new(); + copy_aux_tags(&[], &mut dest, &[]); + assert!(dest.is_empty()); + } } diff --git a/src/commands/zipper.rs b/src/commands/zipper.rs index 30959badb..e27a69814 100644 --- a/src/commands/zipper.rs +++ b/src/commands/zipper.rs @@ -49,10 +49,9 @@ use crate::commands::command::Command; use crate::commands::common::CompressionOptions; use anyhow::{Context, Result}; +use bstr::ByteSlice; use clap::Parser; -use fgumi_lib::bam_io::{ - BamReaderAuto, create_bam_reader, create_bam_writer, is_stdin_path, is_stdout_path, -}; +use fgumi_lib::bam_io::{BamReaderAuto, create_bam_reader, create_raw_bam_writer, is_stdin_path}; use fgumi_lib::batched_sam_reader::BatchedSamReader; use fgumi_lib::logging::OperationTimer; use fgumi_lib::progress::ProgressTracker; @@ -61,18 +60,19 @@ use fgumi_lib::sam::{ buf_value_to_smallest_signed_int, check_sort, revcomp_buf_value, reverse_buf_value, unclipped_five_prime_position, }; -use fgumi_lib::sort::{PA_TAG, PrimaryAlignmentInfo}; +use fgumi_lib::sort::{PA_TAG, PrimaryAlignmentInfo, bam_fields}; use fgumi_lib::template::{Template, TemplateIterator}; use fgumi_lib::umi::TagInfo; use fgumi_lib::validation::validate_file_exists; +use fgumi_lib::vendored::encode_record_buf; use log::{debug, info, warn}; use noodles::sam::Header; -use noodles::sam::alignment::io::Write as AlignmentWrite; use noodles::sam::alignment::record::data::field::Tag; use noodles::sam::alignment::record_buf::Data; use noodles::sam::alignment::record_buf::RecordBuf; +use noodles::sam::alignment::record_buf::data::field::Value as BufValue; use std::collections::HashSet; -use std::io::{BufRead, BufReader, Write}; +use std::io::{BufRead, BufReader}; use std::path::{Path, PathBuf}; /// Command-line arguments for `zipper` @@ -521,109 +521,478 @@ pub fn merge( Ok(()) } -impl Zipper { - /// Process templates in single-threaded mode - fn process_singlethreaded( - &self, - unmapped_iter: U, - mut mapped_iter: M, - output_header: &Header, - writer: &mut W, - tag_info: &TagInfo, - ) -> Result - where - W: AlignmentWrite, - U: Iterator>, - M: Iterator>, - { - let progress = ProgressTracker::new("Processed records").with_interval(1_000_000); - let mut mapped_peek: Option