Added pretty printing and parquet

This commit is contained in:
Guillem Borrell 2022-11-20 00:34:04 +00:00
parent e1690a49ed
commit 57bb60ba27
4 changed files with 87 additions and 32 deletions

View file

@ -1,5 +1,5 @@
[package] [package]
name = "csvgr" name = "dr"
version = "0.1.0" version = "0.1.0"
edition = "2021" edition = "2021"
@ -7,5 +7,7 @@ edition = "2021"
[dependencies] [dependencies]
clap = {version = "4.0", features = ["cargo"]} clap = {version = "4.0", features = ["cargo"]}
polars = "0.25.1" polars = "0.25"
polars-sql = "0.2.1" polars-sql = "0.2.1"
polars-lazy = "0.25"
polars-io = {"version" = "0.25", features = ["parquet"]}

45
src/io.rs Normal file
View file

@ -0,0 +1,45 @@
use polars::frame::DataFrame;
use polars::prelude::*;
use std::io;
use std::io::Read;
/// Read CSV format from stdin and return a Polars DataFrame
pub fn load_csv_from_stdin() -> DataFrame {
let mut buffer = String::new();
let _res: () = match io::stdin().read_to_string(&mut buffer) {
Ok(_ok) => (),
Err(_e) => (),
};
let cursor = io::Cursor::new(buffer.as_bytes());
let df = match CsvReader::new(cursor).finish() {
Ok(df) => df,
Err(_e) => DataFrame::default(),
};
df
}
/// Take a Polars Dataframe and write it as CSV to stdout
pub fn dump_csv_to_stdout(df: &mut DataFrame) {
let _res: () = match CsvWriter::new(io::stdout().lock()).finish(df) {
Ok(_ok) => (),
Err(_e) => (),
};
}
/// Read parquet format from stdin and return a Polars DataFrame
pub fn load_parquet_from_stdin() -> DataFrame {
let mut buffer: String = String::new();
let _res: () = match io::stdin().read_to_string(&mut buffer) {
Ok(_ok) => (),
Err(_e) => (),
};
let cursor = io::Cursor::new(buffer.as_bytes());
let df = match ParquetReader::new(cursor).finish() {
Ok(df) => df,
Err(e) => {
eprintln!("{e}");
DataFrame::default()
}
};
df
}

View file

@ -1,6 +1,6 @@
mod io;
mod sql; mod sql;
use clap::{arg, command, Command}; use clap::{arg, command, Command};
use polars::prelude::*;
fn main() { fn main() {
let matches = command!() let matches = command!()
@ -10,25 +10,31 @@ fn main() {
.arg(arg!([statement] "SQL statement")) .arg(arg!([statement] "SQL statement"))
.arg(arg!(-d --delimiter <String> "Column delimiter").required(false)), .arg(arg!(-d --delimiter <String> "Column delimiter").required(false)),
) )
.subcommand(Command::new("print").about("Pretty prints the table"))
.subcommand(Command::new("rpq").about("Read parquet file"))
.get_matches(); .get_matches();
if let Some(matches) = matches.subcommand_matches("sql") { if let Some(matches) = matches.subcommand_matches("sql") {
if let Some(delimiter) = matches.get_one::<String>("delimiter") { //if let Some(delimiter) = matches.get_one::<String>("delimiter") {
println!("Delimiter: {delimiter}") // println!("DEBUG: Delimiter: {delimiter}")
} else { //} else {
println!("No delimiter") // println!("DEBUG: No delimiter")
} //}
if let Some(statement) = matches.get_one::<String>("statement") { if let Some(statement) = matches.get_one::<String>("statement") {
println!("Statement: {statement}"); sql::execute(statement);
} else { } else {
let mut df = match sql::load_csv_from_stdin() { let mut df = io::load_csv_from_stdin();
Ok(df) => df, io::dump_csv_to_stdout(&mut df);
Err(_e) => DataFrame::default(),
};
let _res = match sql::dump_csv_to_stdout(&mut df) {
Ok(_df) => (),
Err(_e) => (),
};
} }
} }
if let Some(_matches) = matches.subcommand_matches("print") {
let df = io::load_csv_from_stdin();
println!("{}", df)
}
if let Some(_matches) = matches.subcommand_matches("rpq") {
let mut df = io::load_parquet_from_stdin();
io::dump_csv_to_stdout(&mut df);
}
} }

View file

@ -1,17 +1,19 @@
use polars::frame::DataFrame; use crate::io::dump_csv_to_stdout;
use polars::prelude::*; use crate::io::load_csv_from_stdin;
use std::io; use polars_lazy::frame::IntoLazy;
use std::io::Read; use polars_sql::SQLContext;
/// Read from stdin from CSV format and return a Polars DataFrame pub fn execute(statement: &String) {
pub fn load_csv_from_stdin() -> PolarsResult<DataFrame> { if let Ok(mut context) = SQLContext::try_new() {
let mut buffer = String::new(); let df = load_csv_from_stdin();
io::stdin().read_to_string(&mut buffer)?; context.register("this", df.lazy());
let cursor = io::Cursor::new(buffer.as_bytes()); if let Ok(res) = context.execute(statement) {
CsvReader::new(cursor).finish() if let Ok(mut res) = res.collect() {
} dump_csv_to_stdout(&mut res);
};
/// Take a Polars Dataframe and write it as CSV to stdout };
pub fn dump_csv_to_stdout(df: &mut DataFrame) -> Result<(), PolarsError> { if let Err(e) = context.execute(statement) {
CsvWriter::new(io::stdout().lock()).finish(df) eprintln!("Query execution error {e}")
};
};
} }