dr/src/main.rs

238 lines
9.6 KiB
Rust
Raw Normal View History

2022-11-20 01:34:04 +01:00
mod io;
2023-01-17 12:28:49 +01:00
mod schema;
mod sql;
2022-12-06 15:16:39 +01:00
use clap::{arg, command, ArgAction, Command};
2022-12-26 13:54:17 +01:00
use polars_lazy::prelude::*;
fn main() {
let matches = command!()
2022-12-06 15:16:39 +01:00
.subcommand(
Command::new("csv")
.about("Read csv, output arrow stream")
.arg(arg!([path] "Path to CSV file"))
.arg(arg!(-d --delimiter <String> "Column delimiter. Assume ,").required(false))
2023-01-17 12:50:32 +01:00
.arg(
arg!(-i --stdin ... "Read from stdin")
.required(false)
.action(ArgAction::SetTrue),
)
2022-12-06 15:16:39 +01:00
.arg(arg!(-q --query <String> "Execute query on the file").required(false))
.arg(
arg!(-s --summary ... "Summarize the data")
.required(false)
.action(ArgAction::SetTrue),
)
.arg(
arg!(-t --text ... "Output text instead of binary")
.required(false)
.action(ArgAction::SetTrue),
)
.arg(arg!(-P --parquet <String> "Write output as a parquet file").required(false))
.arg(
arg!(-a --head ... "Print the header of the table")
.required(false)
.action(ArgAction::SetTrue),
),
)
2023-01-17 12:28:49 +01:00
.subcommand(
Command::new("schema")
.about("Several table schema related utilities")
.arg(arg!(-n --name <String> "Table name").required(false))
.arg(arg!(-l --strlen <String> "Default length for string columns").required(false))
.arg(
arg!(-s --summary ... "Summarize the schema")
.required(false)
.action(ArgAction::SetTrue),
)
.arg(
arg!(-p --postgresql ... "Create a postgresql table with schema")
.required(false)
.action(ArgAction::SetTrue),
),
)
.subcommand(
Command::new("sql")
.about("Runs a sql statement on the file")
2023-01-17 15:48:05 +01:00
.arg(arg!(-d --delimiter <String> "Column delimiter. Assume ,").required(false))
.arg(arg!([statement] "SQL statement"))
2022-12-06 15:16:39 +01:00
.arg(
arg!(-t --text ... "Input text instead of binary")
.required(false)
.action(ArgAction::SetTrue),
)
.arg(arg!(-d --delimiter <String> "Column delimiter").required(false)),
)
2022-12-06 15:16:39 +01:00
.subcommand(
2023-01-17 15:48:05 +01:00
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")
.required(false)
.action(ArgAction::SetTrue),
),
2022-12-06 15:16:39 +01:00
)
2022-11-20 23:00:58 +01:00
.subcommand(
Command::new("rpq")
.about("Read parquet file")
2022-12-06 15:16:39 +01:00
.arg(arg!([path] "Path to the parquet file"))
.arg(arg!(-q --query <String> "Execute query on the file").required(false))
.arg(
arg!(-s --summary ... "Summarize the data")
.required(false)
.action(ArgAction::SetTrue),
)
2022-12-26 13:54:17 +01:00
.arg(
arg!(-i --stdin ... "Read from stdin instead than from a file")
.required(false)
.action(ArgAction::SetTrue),
)
2022-12-06 15:16:39 +01:00
.arg(
arg!(-t --text ... "Output text instead of binary")
.required(false)
.action(ArgAction::SetTrue),
)
2022-12-26 13:54:17 +01:00
.arg(
arg!(-P --parquet <String> "Write the result as a parquet file")
.required(false),
)
2022-12-06 15:16:39 +01:00
.arg(
arg!(-a --head ... "Print the header of the table")
.required(false)
.action(ArgAction::SetTrue),
),
2022-11-20 23:00:58 +01:00
)
.subcommand(
Command::new("wpq")
.about("Write to a paquet file")
2023-01-17 15:48:05 +01:00
.arg(arg!(-d --delimiter <String> "Column delimiter. Assume ,").required(false))
2022-12-06 15:16:39 +01:00
.arg(
2023-01-17 12:50:32 +01:00
arg!(-t --text ... "Input text instead of binary")
2022-12-06 15:16:39 +01:00
.required(false)
.action(ArgAction::SetTrue),
)
.arg(arg!([path] "Path to the new parquet file")),
)
.get_matches();
2022-12-06 15:16:39 +01:00
if let Some(_matches) = matches.subcommand_matches("csv") {
2023-01-17 15:48:05 +01:00
let delimiter = match _matches.get_one::<String>("delimiter") {
Some(delimiter) => delimiter.as_bytes()[0],
None => b',',
};
2023-01-17 13:18:08 +01:00
let mut ldf = if _matches.get_flag("stdin") {
2023-01-17 15:48:05 +01:00
io::load_csv_from_stdin(delimiter)
2023-01-17 13:18:08 +01:00
} else {
let path = _matches
.get_one::<String>("path")
.expect("Please, provide a file");
2023-01-17 15:48:05 +01:00
io::read_csv(path.to_string(), delimiter)
2023-01-17 13:18:08 +01:00
};
if let Some(query) = _matches.get_one::<String>("query") {
ldf = sql::execute(ldf, query);
}
if _matches.get_flag("summary") {
let df = ldf.collect().expect("Could not collect");
println!("{}", df.describe(None));
} else if _matches.get_flag("head") {
let df = ldf.fetch(5).expect("Could not fetch");
println!("{}", df)
} else {
if _matches.get_flag("text") {
io::dump_csv_to_stdout(ldf);
2022-12-06 15:16:39 +01:00
} else {
2023-01-17 13:18:08 +01:00
if let Some(path) = _matches.get_one::<String>("parquet") {
io::write_parquet(ldf, path.to_string());
2022-12-06 15:16:39 +01:00
} else {
2023-01-17 13:18:08 +01:00
io::write_ipc(ldf);
2022-12-06 15:16:39 +01:00
}
}
}
} else if let Some(_matches) = matches.subcommand_matches("sql") {
2023-01-17 15:48:05 +01:00
let delimiter = match _matches.get_one::<String>("delimiter") {
Some(delimiter) => delimiter.as_bytes()[0],
None => b',',
};
2022-12-06 15:16:39 +01:00
if let Some(statement) = _matches.get_one::<String>("statement") {
let ldf = if _matches.get_flag("text") {
2023-01-17 15:48:05 +01:00
io::load_csv_from_stdin(delimiter)
2022-12-06 15:16:39 +01:00
} else {
io::read_ipc()
};
let res = sql::execute(ldf, statement);
io::write_ipc(res);
} else {
2022-12-06 15:16:39 +01:00
io::write_ipc(io::read_ipc());
}
2022-11-20 23:00:58 +01:00
} else if let Some(_matches) = matches.subcommand_matches("print") {
2023-01-17 15:48:05 +01:00
let delimiter = match _matches.get_one::<String>("delimiter") {
Some(delimiter) => delimiter.as_bytes()[0],
None => b',',
};
2022-12-06 15:16:39 +01:00
let df = if _matches.get_flag("text") {
2023-01-17 15:48:05 +01:00
io::load_csv_from_stdin(delimiter)
2022-12-06 15:16:39 +01:00
} else {
io::read_ipc()
};
println!("{}", df.collect().expect("Could not collect"));
} else if let Some(_matches) = matches.subcommand_matches("rpq") {
2022-12-26 13:54:17 +01:00
let mut ldf = LazyFrame::default();
if _matches.get_flag("stdin") {
ldf = io::load_parquet_from_stdin();
} else if let Some(path) = _matches.get_one::<String>("path") {
ldf = io::read_parquet(path.to_string());
} else {
eprintln!("File not found or not reading from stdin")
}
if let Some(query) = _matches.get_one::<String>("query") {
ldf = sql::execute(ldf, query);
}
if _matches.get_flag("summary") {
let df = ldf.collect().expect("Could not collect");
println!("{}", df.describe(None));
} else if _matches.get_flag("head") {
let df = ldf.fetch(5).expect("Could not fetch");
println!("{}", df)
} else {
if _matches.get_flag("text") {
io::dump_csv_to_stdout(ldf);
2022-12-06 15:16:39 +01:00
} else {
2022-12-26 13:54:17 +01:00
if let Some(path) = _matches.get_one::<String>("parquet") {
io::write_parquet(ldf, path.to_string());
2022-12-06 15:16:39 +01:00
} else {
2022-12-26 13:54:17 +01:00
io::write_ipc(ldf);
2022-12-06 15:16:39 +01:00
}
}
2022-11-20 23:00:58 +01:00
}
2022-12-06 15:16:39 +01:00
} else if let Some(_matches) = matches.subcommand_matches("wpq") {
2023-01-17 15:48:05 +01:00
let delimiter = match _matches.get_one::<String>("delimiter") {
Some(delimiter) => delimiter.as_bytes()[0],
None => b',',
};
2022-12-06 15:16:39 +01:00
if let Some(path) = _matches.get_one::<String>("path") {
2022-12-26 00:07:50 +01:00
let ldf = if _matches.get_flag("text") {
2023-01-17 15:48:05 +01:00
io::load_csv_from_stdin(delimiter)
2022-12-06 15:16:39 +01:00
} else {
io::read_ipc()
};
2022-12-26 00:07:50 +01:00
io::write_parquet(ldf, path.to_string());
} else {
eprintln!("Could now write to parquet");
}
2023-01-17 12:28:49 +01:00
} else if let Some(_matches) = matches.subcommand_matches("schema") {
if _matches.get_flag("summary") {
schema::print_schema(io::read_ipc());
} else if _matches.get_flag("postgresql") {
let name = _matches
.get_one::<String>("name")
.expect("Please provide a table name");
let strlen: u32 = match _matches.get_one::<String>("strlen") {
Some(strlen) => strlen.parse::<u32>().unwrap(),
None => 128,
};
schema::print_create(io::read_ipc(), name.as_str(), strlen);
}
2022-11-20 23:00:58 +01:00
} else {
println!("No command provided. Please execute dr --help")
2022-11-20 01:34:04 +01:00
}
}