Support other delimiters
All checks were successful
ci/woodpecker/push/woodpecker Pipeline was successful

This commit is contained in:
Guillem Borrell 2023-01-17 14:48:05 +00:00
parent 1858777c69
commit facae6af40
3 changed files with 36 additions and 14 deletions

View file

@ -5,8 +5,9 @@ use std::io::Read;
use std::path::PathBuf; use std::path::PathBuf;
/// Read CSV file /// Read CSV file
pub fn read_csv(path: String) -> LazyFrame { pub fn read_csv(path: String, delimiter: u8) -> LazyFrame {
LazyCsvReader::new(path) LazyCsvReader::new(path)
.with_delimiter(delimiter)
.finish() .finish()
.expect("Could not load file") .expect("Could not load file")
} }
@ -31,14 +32,14 @@ pub fn read_ipc() -> LazyFrame {
} }
/// Read CSV format from stdin and return a Polars DataFrame /// 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 mut buffer = Vec::new();
let _res: () = match io::stdin().lock().read_to_end(&mut buffer) { let _res: () = match io::stdin().lock().read_to_end(&mut buffer) {
Ok(_ok) => (), Ok(_ok) => (),
Err(_e) => (), Err(_e) => (),
}; };
let cursor = io::Cursor::new(buffer); let cursor = io::Cursor::new(buffer);
match CsvReader::new(cursor).finish() { match CsvReader::new(cursor).with_delimiter(delimiter).finish() {
Ok(df) => df.lazy(), Ok(df) => df.lazy(),
Err(_e) => LazyFrame::default(), Err(_e) => LazyFrame::default(),
} }

View file

@ -53,6 +53,7 @@ fn main() {
.subcommand( .subcommand(
Command::new("sql") Command::new("sql")
.about("Runs a sql statement on the file") .about("Runs a sql statement on the file")
.arg(arg!(-d --delimiter <String> "Column delimiter. Assume ,").required(false))
.arg(arg!([statement] "SQL statement")) .arg(arg!([statement] "SQL statement"))
.arg( .arg(
arg!(-t --text ... "Input text instead of binary") arg!(-t --text ... "Input text instead of binary")
@ -62,7 +63,10 @@ fn main() {
.arg(arg!(-d --delimiter <String> "Column delimiter").required(false)), .arg(arg!(-d --delimiter <String> "Column delimiter").required(false)),
) )
.subcommand( .subcommand(
Command::new("print").about("Pretty prints the table").arg( Command::new("print")
.about("Pretty prints the table")
.arg(arg!(-d --delimiter <String> "Column delimiter. Assume ,").required(false))
.arg(
arg!(-t --text ... "Inputs csv instead of binary") arg!(-t --text ... "Inputs csv instead of binary")
.required(false) .required(false)
.action(ArgAction::SetTrue), .action(ArgAction::SetTrue),
@ -101,6 +105,7 @@ fn main() {
.subcommand( .subcommand(
Command::new("wpq") Command::new("wpq")
.about("Write to a paquet file") .about("Write to a paquet file")
.arg(arg!(-d --delimiter <String> "Column delimiter. Assume ,").required(false))
.arg( .arg(
arg!(-t --text ... "Input text instead of binary") arg!(-t --text ... "Input text instead of binary")
.required(false) .required(false)
@ -110,13 +115,17 @@ fn main() {
) )
.get_matches(); .get_matches();
if let Some(_matches) = matches.subcommand_matches("csv") { if let Some(_matches) = matches.subcommand_matches("csv") {
let delimiter = match _matches.get_one::<String>("delimiter") {
Some(delimiter) => delimiter.as_bytes()[0],
None => b',',
};
let mut ldf = if _matches.get_flag("stdin") { let mut ldf = if _matches.get_flag("stdin") {
io::load_csv_from_stdin() io::load_csv_from_stdin(delimiter)
} else { } else {
let path = _matches let path = _matches
.get_one::<String>("path") .get_one::<String>("path")
.expect("Please, provide a file"); .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::<String>("query") { if let Some(query) = _matches.get_one::<String>("query") {
ldf = sql::execute(ldf, query); ldf = sql::execute(ldf, query);
@ -139,9 +148,13 @@ fn main() {
} }
} }
} else if let Some(_matches) = matches.subcommand_matches("sql") { } else if let Some(_matches) = matches.subcommand_matches("sql") {
let delimiter = match _matches.get_one::<String>("delimiter") {
Some(delimiter) => delimiter.as_bytes()[0],
None => b',',
};
if let Some(statement) = _matches.get_one::<String>("statement") { if let Some(statement) = _matches.get_one::<String>("statement") {
let ldf = if _matches.get_flag("text") { let ldf = if _matches.get_flag("text") {
io::load_csv_from_stdin() io::load_csv_from_stdin(delimiter)
} else { } else {
io::read_ipc() io::read_ipc()
}; };
@ -151,8 +164,12 @@ fn main() {
io::write_ipc(io::read_ipc()); io::write_ipc(io::read_ipc());
} }
} else if let Some(_matches) = matches.subcommand_matches("print") { } else if let Some(_matches) = matches.subcommand_matches("print") {
let delimiter = match _matches.get_one::<String>("delimiter") {
Some(delimiter) => delimiter.as_bytes()[0],
None => b',',
};
let df = if _matches.get_flag("text") { let df = if _matches.get_flag("text") {
io::load_csv_from_stdin() io::load_csv_from_stdin(delimiter)
} else { } else {
io::read_ipc() io::read_ipc()
}; };
@ -187,9 +204,13 @@ fn main() {
} }
} }
} else if let Some(_matches) = matches.subcommand_matches("wpq") { } else if let Some(_matches) = matches.subcommand_matches("wpq") {
let delimiter = match _matches.get_one::<String>("delimiter") {
Some(delimiter) => delimiter.as_bytes()[0],
None => b',',
};
if let Some(path) = _matches.get_one::<String>("path") { if let Some(path) = _matches.get_one::<String>("path") {
let ldf = if _matches.get_flag("text") { let ldf = if _matches.get_flag("text") {
io::load_csv_from_stdin() io::load_csv_from_stdin(delimiter)
} else { } else {
io::read_ipc() io::read_ipc()
}; };

View file

@ -51,7 +51,7 @@ pub fn print_create(ldf: LazyFrame, table_name: &str, default_strlen: u32) {
statements.push(table.to_string(PostgresQueryBuilder)); statements.push(table.to_string(PostgresQueryBuilder));
} }
// Finall print all statements // Finallyls print all statements
for statement in statements { for statement in statements {
println!("{};", statement); println!("{};", statement);
} }