Merge pull request #116 from noamteyssier/115-build-a-join-subcommand

115 build a join subcommand

Author: noamteyssier

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "gia"
-version = "0.2.22"
+version = "0.2.23"
 edition = "2021"
 description = "A tool for set theoretic operations of genomic intervals"
 license = "MIT"

src/cli/commands.rs

@@ -1,7 +1,8 @@
 use super::{
     bam::BamCommand, bcf::BcfCommand, ClosestArgs, ClusterArgs, ComplementArgs, CoverageArgs,
-    ExtendArgs, FlankArgs, GetFastaArgs, IntersectArgs, MergeArgs, RandomArgs, SampleArgs,
-    SegmentArgs, ShiftArgs, SortArgs, SpacingArgs, SubtractArgs, UnionBedGraphArgs, WindowArgs,
+    ExtendArgs, FlankArgs, GetFastaArgs, IntersectArgs, JoinArgs, MergeArgs, RandomArgs,
+    SampleArgs, SegmentArgs, ShiftArgs, SortArgs, SpacingArgs, SubtractArgs, UnionBedGraphArgs,
+    WindowArgs,
 };
 use clap::Subcommand;
 
@@ -49,6 +50,9 @@ pub enum Command {
     /// Intersects two BED files
     Intersect(IntersectArgs),
 
+    /// Joins two BED files
+    Join(JoinArgs),
+
     /// Merges intervals of a BED file with overlapping regions
     Merge(MergeArgs),
 

src/cli/join.rs

@@ -0,0 +1,38 @@
+use super::{DualInput, Output, OverlapPredicates};
+use clap::{Parser, ValueEnum};
+
+#[derive(Parser, Debug)]
+pub struct JoinArgs {
+    #[clap(flatten)]
+    pub inputs: DualInput,
+
+    #[clap(flatten)]
+    pub params: JoinParams,
+
+    #[clap(flatten)]
+    pub output: Output,
+}
+
+#[derive(Parser, Debug)]
+#[clap(next_help_heading = "Parameters")]
+pub struct JoinParams {
+    #[clap(short = 'H', long, default_value = "inner")]
+    pub how: JoinMethod,
+
+    /// Assert the inputs are pre-sorted
+    #[clap(short = 'S', long)]
+    pub sorted: bool,
+
+    #[clap(flatten)]
+    pub overlap_predicates: OverlapPredicates,
+}
+
+#[derive(Parser, Debug, ValueEnum, Clone, Copy)]
+pub enum JoinMethod {
+    /// Return all records in the left input even if no match is found in right
+    Left,
+    /// Return all records in the right input even if no match is found in left
+    Right,
+    /// Return only records that have a match in both inputs
+    Inner,
+}

src/cli/mod.rs

@@ -11,6 +11,7 @@ mod get_fasta;
 mod growth;
 mod inputs;
 mod intersect;
+mod join;
 mod merge;
 mod outputs;
 mod overlap_predicates;
@@ -37,6 +38,7 @@ pub use inputs::{
     DualInput, MixedInputBam, MixedInputVcf, MultiInput, SingleInput, SingleInputBam,
 };
 pub use intersect::{IntersectArgs, IntersectParams, OutputMethod};
+pub use join::{JoinArgs, JoinMethod, JoinParams};
 pub use merge::{MergeArgs, MergeParams};
 pub use outputs::{BamOutput, Output};
 pub use overlap_predicates::OverlapPredicates;

src/commands/join.rs

@@ -0,0 +1,142 @@
+use crate::{
+    cli::{JoinArgs, JoinMethod, JoinParams},
+    dispatch_pair, dispatch_pair_multi,
+    io::{write_pairs_iter_with, WriteNamedIter, WriteNamedIterImpl},
+    types::{IntervalPair, Rename, Renamer, SplitTranslater},
+    utils::sort_pairs,
+};
+use anyhow::Result;
+use bedrs::{traits::IntervalBounds, IntervalContainer};
+use serde::Serialize;
+use std::io::Write;
+
+fn join_sets_inner<'a, Ia, Ib, Na, Nb, W>(
+    mut set_a: IntervalContainer<Ia, usize, usize>,
+    mut set_b: IntervalContainer<Ib, usize, usize>,
+    translater: Option<&'a SplitTranslater>,
+    params: JoinParams,
+    output: W,
+) -> Result<()>
+where
+    Ia: IntervalBounds<usize, usize> + Copy + Serialize,
+    Ib: IntervalBounds<usize, usize> + Copy + Serialize,
+    Na: IntervalBounds<&'a str, usize> + Serialize,
+    Nb: IntervalBounds<&'a str, usize> + Serialize,
+    W: Write,
+    WriteNamedIterImpl: WriteNamedIter<Ia> + WriteNamedIter<Ib>,
+    Renamer: Rename<'a, Ia, Na> + Rename<'a, Ib, Nb>,
+{
+    let query_method = params.overlap_predicates.into();
+    sort_pairs(&mut set_a, &mut set_b, params.sorted);
+    let pairs_iter = set_a.records().iter().flat_map(|iv| {
+        let overlaps = set_b
+            .query_iter(iv, query_method)
+            .expect("Error in finding overlaps")
+            .copied();
+        overlaps.map(|ov| IntervalPair::new(*iv, ov, translater))
+    });
+    write_pairs_iter_with(pairs_iter, output, translater)
+}
+
+fn join_sets_left<'a, Ia, Ib, Na, Nb, W>(
+    mut set_a: IntervalContainer<Ia, usize, usize>,
+    mut set_b: IntervalContainer<Ib, usize, usize>,
+    translater: Option<&'a SplitTranslater>,
+    params: JoinParams,
+    output: W,
+) -> Result<()>
+where
+    Ia: IntervalBounds<usize, usize> + Copy + Serialize,
+    Ib: IntervalBounds<usize, usize> + Copy + Serialize,
+    Na: IntervalBounds<&'a str, usize> + Serialize,
+    Nb: IntervalBounds<&'a str, usize> + Serialize,
+    W: Write,
+    WriteNamedIterImpl: WriteNamedIter<Ia> + WriteNamedIter<Ib>,
+    Renamer: Rename<'a, Ia, Na> + Rename<'a, Ib, Nb>,
+{
+    let query_method = params.overlap_predicates.into();
+    sort_pairs(&mut set_a, &mut set_b, params.sorted);
+    let pairs_iter = set_a.records().iter().flat_map(|iv| {
+        let overlaps = set_b
+            .query_iter(iv, query_method)
+            .expect("Error in finding overlaps")
+            .copied();
+
+        // Peek at the next item to see if there are overlaps or not
+        let mut peek = overlaps.peekable();
+        let adjusted_overlap = if peek.peek().is_none() {
+            let null_pair = IntervalPair::from_option(Some(*iv), None, translater);
+            let new_iter = std::iter::once(null_pair);
+            Box::new(new_iter) as Box<dyn Iterator<Item = IntervalPair<Ia, Ib, Na, Nb>>>
+        } else {
+            let new_iter = peek.map(|ov| IntervalPair::new(*iv, ov, translater));
+            Box::new(new_iter) as Box<dyn Iterator<Item = IntervalPair<Ia, Ib, Na, Nb>>>
+        };
+        adjusted_overlap
+    });
+    write_pairs_iter_with(pairs_iter, output, translater)
+}
+
+fn join_sets_right<'a, Ia, Ib, Na, Nb, W>(
+    mut set_a: IntervalContainer<Ia, usize, usize>,
+    mut set_b: IntervalContainer<Ib, usize, usize>,
+    translater: Option<&'a SplitTranslater>,
+    params: JoinParams,
+    output: W,
+) -> Result<()>
+where
+    Ia: IntervalBounds<usize, usize> + Copy + Serialize,
+    Ib: IntervalBounds<usize, usize> + Copy + Serialize,
+    Na: IntervalBounds<&'a str, usize> + Serialize,
+    Nb: IntervalBounds<&'a str, usize> + Serialize,
+    W: Write,
+    WriteNamedIterImpl: WriteNamedIter<Ia> + WriteNamedIter<Ib>,
+    Renamer: Rename<'a, Ia, Na> + Rename<'a, Ib, Nb>,
+{
+    let query_method = params.overlap_predicates.into();
+    sort_pairs(&mut set_a, &mut set_b, params.sorted);
+    let pairs_iter = set_b.records().iter().flat_map(|iv| {
+        let overlaps = set_a
+            .query_iter(iv, query_method)
+            .expect("Error in finding overlaps")
+            .copied();
+
+        // Peek at the next item to see if there are overlaps or not
+        let mut peek = overlaps.peekable();
+        let adjusted_overlap = if peek.peek().is_none() {
+            let null_pair = IntervalPair::from_option(None, Some(*iv), translater);
+            let new_iter = std::iter::once(null_pair);
+            Box::new(new_iter) as Box<dyn Iterator<Item = IntervalPair<Ia, Ib, Na, Nb>>>
+        } else {
+            let new_iter = peek.map(|ov| IntervalPair::new(ov, *iv, translater));
+            Box::new(new_iter) as Box<dyn Iterator<Item = IntervalPair<Ia, Ib, Na, Nb>>>
+        };
+        adjusted_overlap
+    });
+    write_pairs_iter_with(pairs_iter, output, translater)
+}
+
+pub fn join(args: JoinArgs) -> Result<()> {
+    let writer = args.output.get_writer()?;
+    if args.inputs.is_multi() {
+        let (bed_a, bed_b) = args.inputs.get_multi_readers()?;
+        match args.params.how {
+            JoinMethod::Inner => {
+                dispatch_pair_multi!(bed_a, bed_b, writer, args.params, join_sets_inner)
+            }
+            JoinMethod::Left => {
+                dispatch_pair_multi!(bed_a, bed_b, writer, args.params, join_sets_left)
+            }
+            JoinMethod::Right => {
+                dispatch_pair_multi!(bed_a, bed_b, writer, args.params, join_sets_right)
+            }
+        }
+    } else {
+        let (bed_a, bed_b) = args.inputs.get_readers()?;
+        match args.params.how {
+            JoinMethod::Inner => dispatch_pair!(bed_a, bed_b, writer, args.params, join_sets_inner),
+            JoinMethod::Left => dispatch_pair!(bed_a, bed_b, writer, args.params, join_sets_left),
+            JoinMethod::Right => dispatch_pair!(bed_a, bed_b, writer, args.params, join_sets_right),
+        }
+    }
+}

src/commands/mod.rs

@@ -8,6 +8,7 @@ mod extend;
 mod flank;
 mod get_fasta;
 mod intersect;
+mod join;
 mod merge;
 mod random;
 mod sample;
@@ -27,6 +28,7 @@ pub use extend::extend;
 pub use flank::flank;
 pub use get_fasta::get_fasta;
 pub use intersect::intersect;
+pub use join::join;
 pub use merge::merge;
 pub use random::random;
 pub use sample::sample;

src/io/write/utils.rs

@@ -206,8 +206,8 @@ pub fn write_pairs_iter_with<'a, W, Ia, Ib, Na, Nb, It>(
     translater: Option<&SplitTranslater>,
 ) -> Result<()>
 where
-    Ia: IntervalBounds<usize, usize> + Serialize,
-    Ib: IntervalBounds<usize, usize> + Serialize,
+    Ia: IntervalBounds<usize, usize> + Serialize + Copy,
+    Ib: IntervalBounds<usize, usize> + Serialize + Copy,
     Na: IntervalBounds<&'a str, usize> + Serialize,
     Nb: IntervalBounds<&'a str, usize> + Serialize,
     W: Write,
@@ -224,8 +224,8 @@ where
 
 pub fn write_pairs_iter<'a, W, Ia, Ib, Na, Nb, It>(records: It, writer: W) -> Result<()>
 where
-    Ia: IntervalBounds<usize, usize> + Serialize,
-    Ib: IntervalBounds<usize, usize> + Serialize,
+    Ia: IntervalBounds<usize, usize> + Serialize + Copy,
+    Ib: IntervalBounds<usize, usize> + Serialize + Copy,
     Na: IntervalBounds<&'a str, usize> + Serialize,
     Nb: IntervalBounds<&'a str, usize> + Serialize,
     W: Write,
@@ -241,8 +241,8 @@ where
 
 pub fn write_named_pairs_iter<'a, Ia, Ib, Na, Nb, W, It>(records: It, writer: W) -> Result<()>
 where
-    Ia: IntervalBounds<usize, usize> + Serialize,
-    Ib: IntervalBounds<usize, usize> + Serialize,
+    Ia: IntervalBounds<usize, usize> + Serialize + Copy,
+    Ib: IntervalBounds<usize, usize> + Serialize + Copy,
     Na: IntervalBounds<&'a str, usize> + Serialize,
     Nb: IntervalBounds<&'a str, usize> + Serialize,
     W: Write,

src/main.rs

@@ -9,8 +9,8 @@ use anyhow::Result;
 use clap::Parser;
 use cli::{bam::BamCommand, bcf::BcfCommand, Cli, Command};
 use commands::{
-    bam, bcf, closest, cluster, complement, coverage, extend, flank, get_fasta, intersect, merge,
-    random, sample, segment, shift, sort, spacing, subtract, unionbedgraph, window,
+    bam, bcf, closest, cluster, complement, coverage, extend, flank, get_fasta, intersect, join,
+    merge, random, sample, segment, shift, sort, spacing, subtract, unionbedgraph, window,
 };
 
 fn main() -> Result<()> {
@@ -32,6 +32,7 @@ fn main() -> Result<()> {
         Command::Flank(args) => flank(args)?,
         Command::GetFasta(args) => get_fasta(args)?,
         Command::Intersect(args) => intersect(args)?,
+        Command::Join(args) => join(args)?,
         Command::Merge(args) => merge(args)?,
         Command::Random(args) => random(args)?,
         Command::Sample(args) => sample(args)?,

src/types/formats/genome.rs

@@ -165,14 +165,15 @@ mod testing {
     fn test_genome_named() {
         let mut translater = Translater::new();
         let genome = Genome::from_reader_named(GENOME_NAMED, &mut translater).unwrap();
-        assert_eq!(genome.chr_size_unchecked(0), 1000);
-        assert_eq!(genome.chr_size_unchecked(1), 2000);
-        assert_eq!(genome.chr_size_unchecked(2), 3000);
+        assert_eq!(genome.chr_size_unchecked(1), 1000);
+        assert_eq!(genome.chr_size_unchecked(2), 2000);
+        assert_eq!(genome.chr_size_unchecked(3), 3000);
         assert!(genome.translater().is_some());
         let translater = genome.translater().unwrap();
-        assert_eq!(translater.get_chr_name(0).unwrap(), "chr1");
-        assert_eq!(translater.get_chr_name(1).unwrap(), "chr2");
-        assert_eq!(translater.get_chr_name(2).unwrap(), "chr3");
+        assert_eq!(translater.get_chr_name(0).unwrap(), ".");
+        assert_eq!(translater.get_chr_name(1).unwrap(), "chr1");
+        assert_eq!(translater.get_chr_name(2).unwrap(), "chr2");
+        assert_eq!(translater.get_chr_name(3).unwrap(), "chr3");
     }
 
     #[test]