Add -R Cmdline Option

One given `@RG` SAM header line will replace existing read group
lines in the header. This read group will be added to all reads
in the output file.

Author: jch-13

src/distributed/dispatcher.rs

@@ -8,11 +8,16 @@ use std::{
 };
 
 use bio::data_structures::suffix_array::SuffixArray;
+use bstr::{BStr, BString, ByteSlice};
 use log::{debug, info, warn};
 use mio::{net, Events, Interest, Poll, Token};
 use noodles::{
     bam,
-    sam::{self, alignment::io::Write as AlignmentWrite},
+    sam::{
+        self,
+        alignment::io::Write as AlignmentWrite,
+        header::record::value::map::{Map, ReadGroup},
+    },
 };
 use rayon::prelude::*;
 
@@ -56,6 +61,7 @@ pub struct Dispatcher<'a, 'b> {
     reference_path: &'b Path,
     out_file_path: &'b Path,
     force_overwrite: bool,
+    read_group: Option<(BString, Map<ReadGroup>)>,
     alignment_parameters: &'a AlignmentParameters,
     connections: HashMap<Token, Connection>,
     accept_connections: bool,
@@ -71,6 +77,7 @@ impl<'a, 'b> Dispatcher<'a, 'b> {
         out_file_path: &'b str,
         force_overwrite: bool,
         alignment_parameters: &'a AlignmentParameters,
+        read_group: Option<(BString, Map<ReadGroup>)>,
     ) -> Result<Self> {
         let reads_path = Path::new(reads_path);
         let out_file_path = Path::new(out_file_path);
@@ -87,6 +94,7 @@ impl<'a, 'b> Dispatcher<'a, 'b> {
             out_file_path,
             force_overwrite,
             alignment_parameters,
+            read_group,
             connections: HashMap::new(),
             accept_connections: true,
         })
@@ -133,7 +141,11 @@ impl<'a, 'b> Dispatcher<'a, 'b> {
         );
 
         let mut input_source = InputSource::from_path(self.reads_path)?;
-        let out_header = create_bam_header(input_source.header(), &identifier_position_map)?;
+        let out_header = create_bam_header(
+            input_source.header(),
+            &identifier_position_map,
+            self.read_group.clone(),
+        )?;
         out_file.write_header(&out_header)?;
         let mut task_queue = input_source.task_queue(self.alignment_parameters.chunk_size);
         self.run_inner(
@@ -229,6 +241,7 @@ impl<'a, 'b> Dispatcher<'a, 'b> {
                                     identifier_position_map,
                                     original_symbols,
                                     self.alignment_parameters,
+                                    self.read_group.as_ref().map(|(id, _map)| id.as_bstr()),
                                     out_header,
                                     out_file,
                                 )?;
@@ -331,6 +344,7 @@ impl<'a, 'b> Dispatcher<'a, 'b> {
         identifier_position_map: &FastaIdPositions,
         original_symbols: &OriginalSymbols,
         alignment_parameters: &AlignmentParameters,
+        read_group: Option<&BStr>,
         out_header: &sam::Header,
         out_file: &mut bam::io::Writer<W>,
     ) -> Result<()>
@@ -350,6 +364,7 @@ impl<'a, 'b> Dispatcher<'a, 'b> {
                     original_symbols,
                     Some(&duration),
                     alignment_parameters,
+                    read_group,
                     &mut rng,
                 )
             })

src/main.rs

@@ -1,5 +1,7 @@
+use bstr::BString;
 use clap::{crate_authors, crate_description, value_parser, Arg, ArgAction, ArgMatches, Command};
 use log::{error, info, warn};
+use noodles::sam::{self, header::record::value::map::Map, header::record::value::map::ReadGroup};
 use simple_logger::SimpleLogger;
 use time::UtcOffset;
 
@@ -37,6 +39,21 @@ fn define_cli() -> ArgMatches {
             .ok_or_else(|| clap::Error::new(clap::error::ErrorKind::ValueValidation))
     }
 
+    fn parse_validate_read_group(raw_arg: &str) -> Result<(BString, Map<ReadGroup>), clap::Error> {
+        let mut read_groups = raw_arg
+            .replace(r"\t", "\t")
+            .parse()
+            .map(|header: sam::Header| header.read_groups().to_owned())
+            .map_err(|_e| clap::Error::new(clap::error::ErrorKind::ValueValidation))?;
+        if read_groups.len() != 1 {
+            return Err(clap::Error::new(clap::error::ErrorKind::ValueValidation));
+        }
+        let rg = read_groups
+            .pop()
+            .expect("As per the above check, there is exactly one RG");
+        Ok(rg)
+    }
+
     Command::new(CRATE_NAME)
         .about(crate_description!())
         .version(build_info::get_software_version())
@@ -261,6 +278,15 @@ fn define_cli() -> ArgMatches {
                         it already exists."))
                         .action(ArgAction::SetTrue)
                 )
+                .arg(
+                    Arg::new("read_group")
+                        .short('R')
+                        .long("read_group")
+                        .help("Read group SAM header line. The given read group ID will be \
+                        added to every read in the output file. \
+                        Example: '@RG\tID:identifier1\tSM:sample2'.")
+                        .value_parser(parse_validate_read_group)
+                )
         )
         .subcommand(
             Command::new("worker")
@@ -328,6 +354,7 @@ fn start_mapper(map_matches: &ArgMatches, _seed: u64) {
         .get_one::<String>("output")
         .expect("Presence is ensured by CLI definition");
     let force_overwrite = map_matches.get_flag("force_overwrite");
+    let read_group = map_matches.get_one("read_group").cloned();
 
     let num_threads = *map_matches
         .get_one("num_threads")
@@ -350,6 +377,7 @@ fn start_mapper(map_matches: &ArgMatches, _seed: u64) {
             out_file_path,
             force_overwrite,
             &alignment_parameters,
+            read_group,
         )
         .and_then(|mut dispatcher| dispatcher.run(port))
     } else {
@@ -359,6 +387,7 @@ fn start_mapper(map_matches: &ArgMatches, _seed: u64) {
             out_file_path,
             force_overwrite,
             &alignment_parameters,
+            read_group,
         )
     } {
         error!("{}", e);

src/map/mapping.rs

@@ -12,13 +12,17 @@ use std::{
 };
 
 use bio::{alphabets::dna, data_structures::suffix_array::SuffixArray};
-use bstr::{BString, ByteSlice};
+use bstr::{BStr, BString, ByteSlice};
 use clap::crate_description;
 use log::{debug, info, trace};
 use min_max_heap::MinMaxHeap;
 use noodles::{
     bam,
-    sam::{self, alignment::io::Write as SamWrite},
+    sam::{
+        self,
+        alignment::io::Write as SamWrite,
+        header::record::value::map::{Map, ReadGroup},
+    },
 };
 use rand::RngCore;
 use rayon::prelude::*;
@@ -56,6 +60,7 @@ pub fn run(
     out_file_path: &str,
     force_overwrite: bool,
     alignment_parameters: &AlignmentParameters,
+    read_group: Option<(BString, Map<ReadGroup>)>,
 ) -> Result<()> {
     let reads_path = Path::new(reads_path);
     let out_file_path = Path::new(out_file_path);
@@ -97,7 +102,11 @@ pub fn run(
 
     info!("Map reads");
     let mut input_source = InputSource::from_path(reads_path)?;
-    let out_header = create_bam_header(input_source.header(), &identifier_position_map)?;
+    let out_header = create_bam_header(
+        input_source.header(),
+        &identifier_position_map,
+        read_group.clone(),
+    )?;
     out_file.write_header(&out_header)?;
     run_inner(
         input_source.task_queue(alignment_parameters.chunk_size),
@@ -107,6 +116,7 @@ pub fn run(
         &identifier_position_map,
         &original_symbols,
         &out_header,
+        read_group.as_ref().map(|(id, _map)| id.as_bstr()),
         &mut out_file,
     )?;
 
@@ -125,6 +135,7 @@ fn run_inner<S, T, W>(
     identifier_position_map: &FastaIdPositions,
     original_symbols: &OriginalSymbols,
     out_header: &sam::Header,
+    read_group: Option<&BStr>,
     out_file: &mut bam::io::Writer<W>,
 ) -> Result<()>
 where
@@ -269,6 +280,7 @@ where
                         original_symbols,
                         Some(&duration),
                         alignment_parameters,
+                        read_group,
                         &mut rng,
                     )
                 },
@@ -288,6 +300,7 @@ where
 pub fn create_bam_header(
     src_header: Option<&sam::Header>,
     identifier_position_map: &FastaIdPositions,
+    read_group: Option<(BString, Map<ReadGroup>)>,
 ) -> Result<sam::Header> {
     use sam::header::record::value::map::{self, Map};
 
@@ -343,9 +356,13 @@ pub fn create_bam_header(
             sam_header_builder = sam_header_builder.add_comment(comment.as_bstr());
         }
 
-        for (id, read_group) in src_header.read_groups() {
-            sam_header_builder =
-                sam_header_builder.add_read_group(id.as_bstr(), read_group.clone());
+        if let Some((id, map)) = read_group {
+            sam_header_builder = sam_header_builder.add_read_group(id, map);
+        } else {
+            for (id, read_group) in src_header.read_groups() {
+                sam_header_builder =
+                    sam_header_builder.add_read_group(id.as_bstr(), read_group.clone());
+            }
         }
     }
 
@@ -390,6 +407,7 @@ pub fn intervals_to_bam<R, S>(
     original_symbols: &OriginalSymbols,
     duration: Option<&Duration>,
     alignment_parameters: &AlignmentParameters,
+    read_group: Option<&BStr>,
     rng: &mut R,
 ) -> Result<sam::alignment::RecordBuf>
 where
@@ -517,6 +535,7 @@ where
                     duration,
                     Some(alternative_hits),
                     original_symbols,
+                    read_group,
                 );
             }
             None => {
@@ -544,6 +563,7 @@ where
         duration,
         None,
         original_symbols,
+        read_group,
     )
 }
 
@@ -712,6 +732,7 @@ fn create_bam_record(
     // Contains valid content for the `YA` tag
     alternative_hits: Option<AlternativeAlignments>,
     original_symbols: &OriginalSymbols,
+    read_group: Option<&BStr>,
 ) -> Result<sam::alignment::RecordBuf> {
     let mut bam_builder = sam::alignment::RecordBuf::builder();
 
@@ -822,9 +843,18 @@ fn create_bam_record(
         .into_iter()
         // Remove BWA (+ mapAD) specific auxiliary fields (avoids potential confusion)
         .filter(|(tag, _v)| !tag_filter.contains(&tag))
+        // Don't copy input read group if one is given via cmdline
+        .filter(|(tag, _v)| !(tag == b"RG" && read_group.is_some()))
         .map(|(tag, value)| Ok((tag.into(), value.into())))
         .collect::<Result<Vec<_>>>()?;
 
+    if let Some(read_group) = read_group {
+        aux_data.push((
+            sam::alignment::record::data::field::tag::Tag::READ_GROUP,
+            sam::alignment::record_buf::data::field::Value::String(read_group.to_owned()),
+        ));
+    }
+
     if let Some(hit_interval) = hit_interval {
         aux_data.push((
             sam::alignment::record::data::field::tag::Tag::ALIGNMENT_SCORE,