feat: refactor limit handling and improve column validation in FASTA …

…and FASTQ processing
dwpeng · Nov 8, 2024 · 086d599 · 086d599
1 parent f2ff0ef
commit 086d599
Show file tree

Hide file tree

Showing 8 changed files with 69 additions and 30 deletions.
diff --git a/src/cmd/fasta.rs b/src/cmd/fasta.rs
@@ -58,7 +58,10 @@ pub fn filterx_fasta(cmd: FastaCommand) -> FilterxResult<()> {
     vm.set_scope(VmSourceType::Fasta);
     vm.set_writer(output);
     'stop_parse: loop {
-        chunk_size = usize::min(chunk_size, vm.status.limit - vm.status.cosumer_rows);
+        if vm.status.consume_rows >= vm.status.limit_rows {
+            break;
+        }
+        chunk_size = usize::min(chunk_size, vm.status.limit_rows - vm.status.consume_rows);
         let df = source.into_dataframe(chunk_size)?;
         if df.is_none() {
             break;
@@ -73,11 +76,43 @@ pub fn filterx_fasta(cmd: FastaCommand) -> FilterxResult<()> {
             let writer = vm.writer.as_mut().unwrap();
             let df = vm.source.into_df()?;
             let cols = df.get_columns();
+            let seq_col = cols.iter().position(|x| x.name() == "seq");
+            if seq_col.is_none() {
+                let h = &mut vm.hint;
+                h.white("Lost ")
+                    .cyan("'seq'")
+                    .white(" column.")
+                    .print_and_exit();
+            }
+            let name_col = cols.iter().position(|x| x.name() == "name");
+            if name_col.is_none() {
+                let h = &mut vm.hint;
+                h.white("Lost ")
+                    .cyan("'name'")
+                    .white(" column.")
+                    .print_and_exit();
+            }
+
+            let comm_col = cols.iter().position(|x| x.name() == "comm");
+            let cols = df.get_columns();
+            let seq_col = seq_col.unwrap();
+            let name_col = name_col.unwrap();
+
+            let valid_cols;
+            if comm_col.is_some() {
+                valid_cols = vec![name_col, comm_col.unwrap(), seq_col]
+            } else {
+                valid_cols = vec![name_col, seq_col]
+            }
+
             let rows = df.height();
             for i in 0..rows {
-                // TODO: handle the case where the column is not found
-                // TODO: handle the case where the column's rank is not right
-                for col in cols.iter() {
+                if vm.status.consume_rows >= vm.status.limit_rows {
+                    break 'stop_parse;
+                }
+                vm.status.consume_rows += 1;
+                for col_index in &valid_cols {
+                    let col = &cols[*col_index];
                     match col.name().as_str() {
                         "seq" => {
                             let seq = col.get(i).unwrap();
@@ -99,10 +134,6 @@ pub fn filterx_fasta(cmd: FastaCommand) -> FilterxResult<()> {
                         }
                     }
                 }
-                vm.status.cosumer_rows += 1;
-                if vm.status.cosumer_rows >= vm.status.limit {
-                    break 'stop_parse;
-                }
             }
         }
     }

diff --git a/src/cmd/fastq.rs b/src/cmd/fastq.rs
@@ -36,7 +36,7 @@ pub fn filterx_fastq(cmd: FastqCommand) -> FilterxResult<()> {
         None => None,
     };
 
-    let mut names = vec!["name", "seq", "comm", "qual"];
+    let mut names = vec!["name", "comm", "seq", "qual"];
 
     match no_comment {
         Some(true) => {
@@ -71,7 +71,10 @@ pub fn filterx_fastq(cmd: FastqCommand) -> FilterxResult<()> {
     vm.set_scope(VmSourceType::Fastq);
     vm.set_writer(output);
     'stop_parse: loop {
-        chunk_size = usize::min(chunk_size, vm.status.limit - vm.status.cosumer_rows);
+        if vm.status.consume_rows >= vm.status.limit_rows {
+            break;
+        }
+        chunk_size = usize::min(chunk_size, vm.status.limit_rows - vm.status.consume_rows);
         let df = source.into_dataframe(chunk_size)?;
         if df.is_none() {
             break;
@@ -89,6 +92,10 @@ pub fn filterx_fastq(cmd: FastqCommand) -> FilterxResult<()> {
             let rows = df.height();
             for i in 0..rows {
                 for col in cols.iter() {
+                    if vm.status.consume_rows >= vm.status.limit_rows {
+                        break 'stop_parse;
+                    }
+                    vm.status.consume_rows += 1;
                     match col.name().as_str() {
                         "name" => {
                             let name = col.get(i).unwrap();
@@ -116,10 +123,6 @@ pub fn filterx_fastq(cmd: FastqCommand) -> FilterxResult<()> {
                     }
                 }
             }
-            vm.status.cosumer_rows += 1;
-            if vm.status.cosumer_rows >= vm.status.limit {
-                break 'stop_parse;
-            }
         }
     }
     Ok(())

diff --git a/src/engine/eval/call/builtin/head.rs b/src/engine/eval/call/builtin/head.rs
@@ -28,6 +28,6 @@ pub fn head(vm: &mut Vm, args: &Vec<ast::Expr>) -> FilterxResult<value::Value> {
     };
 
     vm.source.slice(0, nrow);
-    vm.status.limit = nrow;
+    vm.status.limit_rows = nrow;
     Ok(value::Value::None)
 }
diff --git a/src/engine/eval/call/builtin/limit.rs b/src/engine/eval/call/builtin/limit.rs
@@ -32,7 +32,7 @@ pub fn limit<'a>(vm: &'a mut Vm, args: &Vec<ast::Expr>) -> FilterxResult<value::
 
     match vm.source_type {
         VmSourceType::Fasta | VmSourceType::Fastq => {
-            vm.status.limit = nrow;
+            vm.status.limit_rows = nrow;
             return Ok(value::Value::None);
         }
         _ => {}

diff --git a/src/engine/eval/call/builtin/mod.rs b/src/engine/eval/call/builtin/mod.rs
@@ -77,5 +77,7 @@ builtin_function! {
     is_null,
     is_na,
     nr,
-    drop_null
+    drop_null,
+    to_fasta,
+    to_fastq,
 }
diff --git a/src/engine/eval/call/builtin/print.rs b/src/engine/eval/call/builtin/print.rs
@@ -128,8 +128,8 @@ pub fn print<'a>(vm: &'a mut Vm, args: &Vec<ast::Expr>) -> FilterxResult<value::
     let fmt = df.column("fmt").unwrap();
     let writer = vm.writer.as_mut().unwrap();
     let writer = writer.as_mut();
-    let mut consmer_rows = vm.status.cosumer_rows;
-    if vm.status.cosumer_rows >= vm.status.limit {
+    println!("limit: {}, consume: {}", vm.status.limit_rows, vm.status.consume_rows);
+    if vm.status.consume_rows >= vm.status.limit_rows {
         return Ok(value::Value::None);
     }
     for i in 0..fmt.len() {
@@ -140,9 +140,8 @@ pub fn print<'a>(vm: &'a mut Vm, args: &Vec<ast::Expr>) -> FilterxResult<value::
         }
         let s = s.unwrap();
         writeln!(writer, "{}", s)?;
-        consmer_rows += 1;
-        if consmer_rows >= vm.status.limit {
-            vm.status.stop = true;
+        vm.status.consume_rows += 1;
+        if vm.status.consume_rows >= vm.status.limit_rows {
             break;
         }
     }

diff --git a/src/engine/eval/call/call.rs b/src/engine/eval/call/call.rs
@@ -109,6 +109,8 @@ impl<'a> Eval<'a> for ast::ExprCall {
             "is_not_na" => call::is_na(vm, &self.args, true),
             "drop_null" => call::drop_null(vm, &self.args, false),
             "drop_null_" => call::drop_null(vm, &self.args, true),
+            "to_fa" | "to_fasta" => call::to_fasta(vm),
+            "to_fq" | "to_fastq" => call::to_fastq(vm),
             _ => {
                 let simi = compute_similarity(&function_name, vec![]);
                 let h = &mut vm.hint;

diff --git a/src/engine/vm.rs b/src/engine/vm.rs
@@ -19,28 +19,24 @@ pub enum VmMode {
 #[derive(Debug)]
 pub struct VmStatus {
     pub apply_lazy: bool,
-    pub skip: bool,
     pub stop: bool,
     pub count: usize,
-    pub limit: usize,
+    pub limit_rows: usize,
     pub offset: usize,
     pub printed: bool,
-    pub cosumer_rows: usize,
-    pub nrows: usize,
+    pub consume_rows: usize,
 }
 
 impl VmStatus {
     pub fn new() -> Self {
         Self {
             apply_lazy: true,
-            skip: false,
             stop: false,
             count: 0,
-            limit: usize::MAX,
+            limit_rows: usize::MAX,
             offset: 0,
             printed: false,
-            cosumer_rows: 0,
-            nrows: 0,
+            consume_rows: 0,
         }
     }
 }
@@ -212,6 +208,12 @@ impl Vm {
                 let expr = eval_expr.as_interactive().unwrap();
                 expr.eval(self)?;
             }
+            // if self.status.stop {
+            //     std::process::exit(0);
+            // }
+            // if self.status.printed {
+            //     return Ok(());
+            // }
         }
         Ok(())
     }