rewrite cmd subfq

Author: sharkLoc

src/cli/stats.rs

@@ -236,7 +236,6 @@ impl ParallelProcessor for Info {
     }
 
     fn on_batch_complete(&mut self) -> Result<(), ProcessError> {
-
         *self.total_num_a.lock() += self.num_a;
         *self.total_num_t.lock() += self.num_t;
         *self.total_num_g.lock() += self.num_g;

src/cli/subfq.rs

@@ -5,139 +5,101 @@ use paraseq::{fastq, fastx::Record};
 use rand::{Rng, prelude::*};
 use rand_pcg::Pcg64;
 
-// reduce much memory but cost more time
-fn select_fastq(
+pub fn sample_fastq(
     file: Option<&String>,
     n: usize,
     seed: u64,
+    two_pass: bool,
     out: Option<&String>,
     compression_level: u32,
     stdout_type: char,
 ) -> Result<(), FqkitError> {
-    let mut rng: rand_pcg::Lcg128Xsl64 = Pcg64::seed_from_u64(seed);
-    let mut get: Vec<usize> = Vec::with_capacity(n);
-
     if n == 0 {
         error!("n must be greater than 0");
         std::process::exit(1);
     }
+    info!("subseq number: {}", n);
 
     let mut fq_reader = fastq::Reader::new(file_reader(file)?);
-    info!("rand seed: {}", seed);
-    info!("subseq number: {}", n);
-    info!("reduce much memory but cost more time");
     let mut rset = fastq::RecordSet::default();
+    info!("rand seed: {}", seed);
+    let mut fq_writer = file_writer(out, compression_level, stdout_type)?;
+
+    let mut rng: rand_pcg::Lcg128Xsl64 = Pcg64::seed_from_u64(seed);
     let mut order: usize = 0;
-    while rset.fill(&mut fq_reader)? {
-        for _ in rset.iter().map_while(Result::ok) {
-            if order < n {
-                get.push(order);
-            } else {
-                let ret = rng.random_range(0..=order);
-                if ret < n {
-                    get[ret] = order;
+    if two_pass {
+        info!("two pass mode enabled");
+        let mut get: Vec<usize> = Vec::with_capacity(n);
+        while rset.fill(&mut fq_reader)? {
+            for _ in rset.iter().map_while(Result::ok) {
+                if order < n {
+                    get.push(order);
+                } else {
+                    let ret = rng.random_range(0..=order);
+                    if ret < n {
+                        get[ret] = order;
+                    }
                 }
+                order += 1;
             }
-            order += 1;
         }
-    }
 
-    let mut fq_writer = file_writer(out, compression_level, stdout_type)?;
-    let mut fq_reader2 = fastq::Reader::new(file_reader(file)?);
-    let mut rset2 = fastq::RecordSet::default();
-    let mut order2: usize = 0;
-
-    while rset2.fill(&mut fq_reader2)? {
-        for rec in rset2.iter().map_while(Result::ok) {
-            if get.contains(&order2) {
-                write_record(&mut fq_writer, rec.id(), rec.seq(), rec.qual())?;
+        let mut fq_reader2 = fastq::Reader::new(file_reader(file)?);
+        let mut rset2 = fastq::RecordSet::default();
+        order = 0;
+        get.sort_unstable(); // keep the order
+        let mut idx = 0usize;
+        info!("all records has been readed into memory, start write to output ...");
+        while rset2.fill(&mut fq_reader2)? {
+            for rec in rset2.iter().map_while(Result::ok) {
+                if idx < get.len() && order == get[idx] {
+                    write_record(&mut fq_writer, rec.id(), rec.seq(), rec.qual())?;
+                    idx += 1;
+                }
+                if idx >= get.len() {
+                    break;
+                }
+                order += 1;
             }
-            order2 += 1;
         }
-    }
-    fq_writer.flush()?;
-
-    Ok(())
-}
-
-// fast mode but cost more memory
-fn select_fastq2(
-    file: Option<&String>,
-    n: usize,
-    seed: u64,
-    out: Option<&String>,
-    compression_level: u32,
-    stdout_type: char,
-) -> Result<(), FqkitError> {
-    info!("rand seed: {}", seed);
-    info!("subseq num: {}", n);
-    info!("fast mode but cost more memory");
-
-    let mut rng = Pcg64::seed_from_u64(seed);
-    let mut get = Vec::with_capacity(n);
-    if n == 0 {
-        error!("n must be greater than 0");
-        std::process::exit(1);
-    }
-
-    let mut fq_reader = fastq::Reader::new(file_reader(file)?);
-    let mut rset = fastq::RecordSet::default();
-    let mut order: usize = 0;
-    while rset.fill(&mut fq_reader)? {
-        for rec in rset.iter().map_while(Result::ok) {
-            if order < n {
-                let rec_t = vec![
-                    rec.id_str().to_owned(),
-                    rec.seq_str().to_owned(),
-                    rec.qual_str().to_owned(),
-                ];
-                get.push(rec_t);
-            } else {
-                let ret = rng.random_range(0..=order);
-                if ret < n {
-                    get[ret] = vec![
+    } else {
+        let mut get = Vec::with_capacity(n);
+        while rset.fill(&mut fq_reader)? {
+            for rec in rset.iter().map_while(Result::ok) {
+                if order < n {
+                    get.push((
+                        order,
                         rec.id_str().to_owned(),
                         rec.seq_str().to_owned(),
                         rec.qual_str().to_owned(),
-                    ];
+                    ));
+                } else {
+                    let ret = rng.random_range(0..=order);
+                    if ret < n {
+                        get[ret] = (
+                            order,
+                            rec.id_str().to_owned(),
+                            rec.seq_str().to_owned(),
+                            rec.qual_str().to_owned(),
+                        );
+                    }
                 }
+                order += 1;
             }
-            order += 1;
+        }
+        info!("all records has been readed into memory, start write to output ...");
+        get.sort_unstable_by_key(|x| x.0); // sort by order to keep the raw order
+        for (_, id, seq, qual) in get {
+            write_record(
+                &mut fq_writer,
+                id.as_bytes(),
+                seq.as_bytes(),
+                qual.as_bytes(),
+            )?;
         }
     }
 
-    let mut fq_writer = file_writer(out, compression_level, stdout_type)?;
-    for rec in get.iter() {
-        write_record(
-            &mut fq_writer,
-            rec[0].as_bytes(),
-            rec[1].as_bytes(),
-            rec[2].as_bytes(),
-        )?
-    }
     fq_writer.flush()?;
 
     Ok(())
 }
-
-pub fn subset_fastq(
-    rdc: bool,
-    file: Option<&String>,
-    n: usize,
-    seed: u64,
-    out: Option<&String>,
-    compression_level: u32,
-    stdout_type: char,
-) -> Result<(), FqkitError> {
-    if rdc {
-        if file.is_none() {
-            error!("opt -r used, fastq data can't from stdin.");
-            std::process::exit(1);
-        }
-        select_fastq(file, n, seed, out, compression_level, stdout_type)?;
-    } else {
-        select_fastq2(file, n, seed, out, compression_level, stdout_type)?;
-    }
-
-    Ok(())
-}

src/command.rs

@@ -139,9 +139,9 @@ pub enum Subcli {
         /// subseq number
         #[arg(short = 'n', long = "num", value_name = "INT")]
         num: usize,
-        /// read files twice to reduce much memory but cost more time
-        #[arg(short = 'r', long = "rdc", help_heading = Some("FLAGS"))]
-        rdc: bool,
+        /// if specified, use two-pass mode to reduce memory usage
+        #[arg(short = '2', long = "two-pass", help_heading = Some("FLAGS"))]
+        two_pass: bool,
         /// fastq output file name or write to stdout, files ending in .gz/.bz2/.xz will be compressed automatically
         #[arg(short = 'o', long = "out", value_name = "FILE")]
         out: Option<String>,

src/main.rs

@@ -62,14 +62,14 @@ fn run_main() -> Result<(), Error> {
             input,
             seed,
             num,
-            rdc,
+            two_pass,
             out,
         } => {
-            subset_fastq(
-                rdc,
+            sample_fastq(
                 input.as_ref(),
                 num,
                 seed,
+                two_pass,
                 out.as_ref(),
                 arg.compression_level,
                 arg.stdout_type,

src/utils.rs

@@ -166,14 +166,13 @@ where
                 2 => 3,
                 3 => 7,
                 4 => 11,
-                _ => 3
+                _ => 3,
             };
             Ok(Box::new(BufWriter::with_capacity(
                 BUFF_SIZE,
                 zstd::stream::write::Encoder::new(fp, level)?,
             )))
-        }
-        else {
+        } else {
             Ok(Box::new(BufWriter::with_capacity(BUFF_SIZE, fp)))
         }
     } else if stdout_format == 'g' {
@@ -196,15 +195,18 @@ where
         )))
     } else if stdout_format == 'z' {
         Ok(Box::new(BufWriter::with_capacity(
-                BUFF_SIZE,
-                zstd::stream::write::Encoder::new(io::stdout(), match compression_level {
-                1 => 1,
-                2 => 3,
-                3 => 7,
-                4 => 11,
-                _ => 3
-            })?,
-            )))
+            BUFF_SIZE,
+            zstd::stream::write::Encoder::new(
+                io::stdout(),
+                match compression_level {
+                    1 => 1,
+                    2 => 3,
+                    3 => 7,
+                    4 => 11,
+                    _ => 3,
+                },
+            )?,
+        )))
     } else if stdout_format == 'u' {
         Ok(Box::new(BufWriter::new(io::stdout())))
     } else {
@@ -245,18 +247,22 @@ where
             BUFF_SIZE,
             xz2::write::XzEncoder::new(fp, compression_level),
         )))
-    } else if file_out.as_ref().extension().is_some_and(|ext| ext == "zst") {
+    } else if file_out
+        .as_ref()
+        .extension()
+        .is_some_and(|ext| ext == "zst")
+    {
         let level = match compression_level {
-                1 => 1,
-                2 => 3,
-                3 => 7,
-                4 => 11,
-                _ => 3
-            };
-            Ok(Box::new(BufWriter::with_capacity(
-                BUFF_SIZE,
-                zstd::stream::write::Encoder::new(fp, level)?,
-            )))
+            1 => 1,
+            2 => 3,
+            3 => 7,
+            4 => 11,
+            _ => 3,
+        };
+        Ok(Box::new(BufWriter::with_capacity(
+            BUFF_SIZE,
+            zstd::stream::write::Encoder::new(fp, level)?,
+        )))
     } else {
         Ok(Box::new(BufWriter::with_capacity(BUFF_SIZE, fp)))
     }