use std::ffi::OsStr;
use std::io::{self, BufRead, Read, Write};
use std::path::Path;
use std::time::{Duration, Instant};
use polars_core::prelude::*;
use polars_lazy::frame::LazyFrame;
#[cfg(feature = "ipc")]
use polars_lazy::frame::ScanArgsIpc;
#[cfg(feature = "parquet")]
use polars_lazy::frame::ScanArgsParquet;
#[cfg(feature = "csv")]
use polars_lazy::frame::{LazyCsvReader, LazyFileListReader};
use polars_sql::SQLContext;
use rustyline::completion::FilenameCompleter;
use rustyline::error::ReadlineError;
use rustyline::{DefaultEditor, Result};
use sqlparser::ast::{Select, SetExpr, Statement, TableFactor, TableWithJoins};
use sqlparser::dialect::GenericDialect;
use sqlparser::parser::Parser;
const SUPPORTED_FILE_EXTENSIONS: &[&str] = &[
#[cfg(feature = "csv")]
"csv",
#[cfg(feature = "parquet")]
"parquet",
#[cfg(feature = "ipc")]
"ipc",
];
fn print_dataframes(dataframes: &Vec<(String, String)>) {
println!(
"{} dataframes registered{}",
dataframes.len(),
if dataframes.is_empty() { "." } else { ":" }
);
for (name, file) in dataframes.iter() {
println!("{}:\t {}", name, file);
}
}
fn register_dataframe(
context: &mut SQLContext,
dataframes: &mut Vec<(String, String)>,
command: Vec<&str>,
) {
if command.len() < 3 {
println!("Usage: \\rd <name> <file>");
return;
}
let name = command[1];
let source = command[2];
let df = create_dataframe_from_filename(source);
match df {
Ok(frame) => {
context.register(name, frame);
dataframes.push((name.to_owned(), source.to_owned()));
println!("Added dataframe \"{}\" from file {}", name, source)
}
Err(e) => eprintln!("{}", e),
}
}
fn print_help() {
println!("List of all client commands:");
for (name, short, desc) in vec![
("dataframes", "dd", "Show registered frames."),
("help", "?", "Display this help."),
(
"register",
"rd",
"Register new dataframe: \\rd <name> <source>",
),
("quit", "q", "Exit"),
]
.iter()
{
println!("{:20}\\{:10} {}", name, short, desc);
}
}
fn create_dataframe_from_filename(filename: &str) -> PolarsResult<LazyFrame> {
return match get_extension_from_filename(filename) {
#[cfg(feature = "csv")]
Some("csv") => LazyCsvReader::new(filename).finish(),
#[cfg(feature = "parquet")]
Some("parquet") => LazyFrame::scan_parquet(filename, ScanArgsParquet::default()),
#[cfg(feature = "ipc")]
Some("ipc") => LazyFrame::scan_ipc(filename, ScanArgsIpc::default()),
None => polars_bail!(
ComputeError: "unable to infer dataframe format from filename \"{}\"; \"\
either specify a dataframe name registered with \\rd or an absolute path to a file",
filename,
),
Some(ext) => polars_bail!(
ComputeError: "unsupported file extension: \"{}\"; \
supported file extensions are {} and {}",
ext,
SUPPORTED_FILE_EXTENSIONS[0..SUPPORTED_FILE_EXTENSIONS.len() - 1].join(", "),
SUPPORTED_FILE_EXTENSIONS.last().unwrap()
),
};
}
fn create_dataframe_from_tablename(
context: &mut SQLContext,
relation: &TableFactor,
) -> PolarsResult<()> {
if let TableFactor::Table { name, alias, .. } = relation {
let source = name.0.get(0).unwrap().value.as_str();
let name = match alias {
Some(alias) => alias.name.value.to_string(),
None => source.to_string(),
};
if context.table_map.contains_key(&name) {
return Ok(());
}
let frame = create_dataframe_from_filename(source)?;
context.register(&name, frame);
};
Ok(())
}
fn create_dataframes_from_statement(context: &mut SQLContext, stmt: &Select) -> PolarsResult<()> {
let sql_tbl: &TableWithJoins = stmt
.from
.get(0)
.ok_or_else(|| polars_err!(ComputeError: "no table name provided in query"))?;
create_dataframe_from_tablename(context, &sql_tbl.relation)?;
if !sql_tbl.joins.is_empty() {
for tbl in &sql_tbl.joins {
create_dataframe_from_tablename(context, &tbl.relation)?;
}
}
Ok(())
}
fn execute_query(context: &mut SQLContext, query: &str) -> PolarsResult<DataFrame> {
let ast = Parser::parse_sql(&GenericDialect::default(), query)
.map_err(|e| polars_err!(ComputeError: "error parsing SQL: {}", e))?;
polars_ensure!(ast.len() == 1, ComputeError: "one and only one statement at a time please");
let ast = ast.get(0).unwrap();
if let Statement::Query(query) = ast {
if let SetExpr::Select(select_stmt) = &query.body.as_ref() {
create_dataframes_from_statement(context, select_stmt)?;
}
}
context.execute_statement(ast)?.collect()
}
fn get_extension_from_filename(filename: &str) -> Option<&str> {
Path::new(filename).extension().and_then(OsStr::to_str)
}
pub fn run_tty() -> std::io::Result<()> {
let mut stdout = io::stdout();
let mut context = SQLContext::try_new().unwrap();
let mut dataframes = Vec::new();
let mut rl = DefaultEditor::new().unwrap();
println!("Welcome to Polars CLI. Commands end with ; or \\n");
println!("Type help or \\? for help.");
loop {
let input = match rl.readline(">> ") {
Ok(line) => {
rl.add_history_entry(&line);
line.trim().to_owned()
}
Err(ReadlineError::Interrupted) => "exit".to_string(),
Err(e) => {
eprintln!("Error: {:?}", e);
"".to_string()
}
};
let command: Vec<&str> = input.trim().split(' ').collect();
if command[0].is_empty() {
continue;
}
match command[0] {
"\\dd" | "dataframes" => print_dataframes(&dataframes),
"\\rd" | "register" => register_dataframe(&mut context, &mut dataframes, command),
"\\?" | "help" | "?" | "\\h" => print_help(),
"\\q" | "quit" | "exit" => {
println!("Bye");
return Ok(());
}
_ => {
if command[0].starts_with('\\') {
print!("Unknown command: {}\n\n", command[0]);
print_help();
continue;
}
let start = Instant::now();
match execute_query(&mut context, input.trim()) {
Ok(lf) => {
println!("{}", lf);
println!(
"{} rows in set ({:.3} sec)",
lf.shape().0,
start.elapsed().as_secs_f32()
)
}
Err(e) => eprintln!("{}", e),
}
}
}
println!();
}
}
pub fn run() -> io::Result<()> {
if atty::is(atty::Stream::Stdin) {
return run_tty();
}
let mut context = SQLContext::try_new().unwrap();
let mut input: Vec<u8> = Vec::with_capacity(1024);
let mut stdin = std::io::stdin();
loop {
input.clear();
stdin.lock().read_until(b';', &mut input);
let query = std::str::from_utf8(&input).unwrap_or("").trim();
if query.is_empty() {
break;
}
match execute_query(&mut context, query) {
Ok(lf) => println!("{}", lf),
Err(e) => eprintln!("{}", e),
}
}
Ok(())
}