diff --git a/src/io.rs b/src/io.rs index 9eba78c..2afa724 100644 --- a/src/io.rs +++ b/src/io.rs @@ -5,8 +5,9 @@ use std::io::Read; use std::path::PathBuf; /// Read CSV file -pub fn read_csv(path: String) -> LazyFrame { +pub fn read_csv(path: String, delimiter: u8) -> LazyFrame { LazyCsvReader::new(path) + .with_delimiter(delimiter) .finish() .expect("Could not load file") } @@ -31,14 +32,14 @@ pub fn read_ipc() -> LazyFrame { } /// Read CSV format from stdin and return a Polars DataFrame -pub fn load_csv_from_stdin() -> LazyFrame { +pub fn load_csv_from_stdin(delimiter: u8) -> LazyFrame { let mut buffer = Vec::new(); let _res: () = match io::stdin().lock().read_to_end(&mut buffer) { Ok(_ok) => (), Err(_e) => (), }; let cursor = io::Cursor::new(buffer); - match CsvReader::new(cursor).finish() { + match CsvReader::new(cursor).with_delimiter(delimiter).finish() { Ok(df) => df.lazy(), Err(_e) => LazyFrame::default(), } diff --git a/src/main.rs b/src/main.rs index 57c4af9..54d155c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -53,6 +53,7 @@ fn main() { .subcommand( Command::new("sql") .about("Runs a sql statement on the file") + .arg(arg!(-d --delimiter "Column delimiter. Assume ,").required(false)) .arg(arg!([statement] "SQL statement")) .arg( arg!(-t --text ... "Input text instead of binary") @@ -62,11 +63,14 @@ fn main() { .arg(arg!(-d --delimiter "Column delimiter").required(false)), ) .subcommand( - Command::new("print").about("Pretty prints the table").arg( - arg!(-t --text ... "Inputs csv instead of binary") - .required(false) - .action(ArgAction::SetTrue), - ), + Command::new("print") + .about("Pretty prints the table") + .arg(arg!(-d --delimiter "Column delimiter. Assume ,").required(false)) + .arg( + arg!(-t --text ... "Inputs csv instead of binary") + .required(false) + .action(ArgAction::SetTrue), + ), ) .subcommand( Command::new("rpq") @@ -101,6 +105,7 @@ fn main() { .subcommand( Command::new("wpq") .about("Write to a paquet file") + .arg(arg!(-d --delimiter "Column delimiter. Assume ,").required(false)) .arg( arg!(-t --text ... "Input text instead of binary") .required(false) @@ -110,13 +115,17 @@ fn main() { ) .get_matches(); if let Some(_matches) = matches.subcommand_matches("csv") { + let delimiter = match _matches.get_one::("delimiter") { + Some(delimiter) => delimiter.as_bytes()[0], + None => b',', + }; let mut ldf = if _matches.get_flag("stdin") { - io::load_csv_from_stdin() + io::load_csv_from_stdin(delimiter) } else { let path = _matches .get_one::("path") .expect("Please, provide a file"); - io::read_csv(path.to_string()) + io::read_csv(path.to_string(), delimiter) }; if let Some(query) = _matches.get_one::("query") { ldf = sql::execute(ldf, query); @@ -139,9 +148,13 @@ fn main() { } } } else if let Some(_matches) = matches.subcommand_matches("sql") { + let delimiter = match _matches.get_one::("delimiter") { + Some(delimiter) => delimiter.as_bytes()[0], + None => b',', + }; if let Some(statement) = _matches.get_one::("statement") { let ldf = if _matches.get_flag("text") { - io::load_csv_from_stdin() + io::load_csv_from_stdin(delimiter) } else { io::read_ipc() }; @@ -151,8 +164,12 @@ fn main() { io::write_ipc(io::read_ipc()); } } else if let Some(_matches) = matches.subcommand_matches("print") { + let delimiter = match _matches.get_one::("delimiter") { + Some(delimiter) => delimiter.as_bytes()[0], + None => b',', + }; let df = if _matches.get_flag("text") { - io::load_csv_from_stdin() + io::load_csv_from_stdin(delimiter) } else { io::read_ipc() }; @@ -187,9 +204,13 @@ fn main() { } } } else if let Some(_matches) = matches.subcommand_matches("wpq") { + let delimiter = match _matches.get_one::("delimiter") { + Some(delimiter) => delimiter.as_bytes()[0], + None => b',', + }; if let Some(path) = _matches.get_one::("path") { let ldf = if _matches.get_flag("text") { - io::load_csv_from_stdin() + io::load_csv_from_stdin(delimiter) } else { io::read_ipc() }; diff --git a/src/schema.rs b/src/schema.rs index 65bd905..5447adb 100644 --- a/src/schema.rs +++ b/src/schema.rs @@ -51,7 +51,7 @@ pub fn print_create(ldf: LazyFrame, table_name: &str, default_strlen: u32) { statements.push(table.to_string(PostgresQueryBuilder)); } - // Finall print all statements + // Finallyls print all statements for statement in statements { println!("{};", statement); }