diff --git a/crates/qf-cli/src/bench.rs b/crates/qf-cli/src/bench.rs new file mode 100644 index 0000000..c0baaf5 --- /dev/null +++ b/crates/qf-cli/src/bench.rs @@ -0,0 +1,333 @@ +//! A benchmark that measures the optimiser rather than asserting it. +//! +//! Every query is run twice — once on the bound plan and once on the optimised +//! one — against the same data in the same process. The numbers in the README +//! come from here, and anyone can regenerate them with `queryforge bench`. + +use qf_common::{DataType, Field, Result, Schema, Value}; +use qf_exec::operator::{collect, Operator}; +use qf_exec::Session; +use qf_storage::array::ArrayBuilder; +use qf_storage::batch::{RecordBatch, DEFAULT_BATCH_SIZE}; +use std::path::Path; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +/// Rows in the generated table when no size is given. +pub const DEFAULT_ROWS: usize = 500_000; +/// Rows per row group. Smaller groups prune more finely. +pub const BENCH_ROW_GROUP_SIZE: usize = 8_192; + +pub struct Case { + pub name: &'static str, + pub sql: &'static str, +} + +pub const CASES: &[Case] = &[ + Case { + name: "selective range", + sql: "SELECT id, amount FROM events WHERE id > 490000", + }, + Case { + name: "point lookup", + sql: "SELECT label FROM events WHERE id = 123456", + }, + Case { + name: "narrow projection", + sql: "SELECT sum(amount) FROM events", + }, + Case { + name: "group by", + sql: "SELECT status, count(*), sum(amount) FROM events GROUP BY status", + }, + Case { + name: "filtered group by", + sql: "SELECT status, count(*) FROM events WHERE amount > 900 GROUP BY status", + }, + Case { + name: "order by with limit", + sql: "SELECT id, amount FROM events ORDER BY amount DESC LIMIT 10", + }, + Case { + name: "join", + sql: "SELECT e.id, s.description FROM events e \ + JOIN statuses s ON e.status = s.name WHERE e.id > 495000", + }, +]; + +#[derive(Debug, Clone)] +pub struct Measurement { + pub name: &'static str, + pub optimized: Duration, + pub unoptimized: Duration, + pub rows: usize, + pub row_groups_read: usize, + pub row_groups_pruned: usize, +} + +impl Measurement { + pub fn speedup(&self) -> f64 { + if self.optimized.as_secs_f64() == 0.0 { + return f64::INFINITY; + } + self.unoptimized.as_secs_f64() / self.optimized.as_secs_f64() + } +} + +/// Builds the benchmark tables in `dir` and returns a session over them. +pub fn prepare(dir: &Path, rows: usize) -> Result { + let mut session = Session::new(); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("status", DataType::Utf8, false), + Field::new("label", DataType::Utf8, false), + Field::new("amount", DataType::Float64, true), + ])); + let statuses = ["pending", "shipped", "delivered", "cancelled"]; + + let mut batches = Vec::new(); + let mut written = 0; + while written < rows { + let n = DEFAULT_BATCH_SIZE.min(rows - written); + let mut ids = ArrayBuilder::new(DataType::Int64); + let mut status = ArrayBuilder::new(DataType::Utf8); + let mut label = ArrayBuilder::new(DataType::Utf8); + let mut amount = ArrayBuilder::new(DataType::Float64); + for i in written..written + n { + ids.push(Value::Int64(i as i64))?; + status.push(Value::Utf8(statuses[i % statuses.len()].to_string()))?; + label.push(Value::Utf8(format!("event-{i}")))?; + // Every hundredth row is NULL, so the null paths are exercised. + if i % 100 == 0 { + amount.push_null(); + } else { + amount.push(Value::Float64((i % 1000) as f64))?; + } + } + batches.push(RecordBatch::try_new( + Arc::clone(&schema), + vec![ + ids.finish()?, + status.finish()?, + label.finish()?, + amount.finish()?, + ], + )?); + written += n; + } + session.register("events", Arc::clone(&schema), batches)?; + session.save_table_with_row_group_size( + "events", + dir.join("events.qfc"), + BENCH_ROW_GROUP_SIZE, + )?; + + let status_schema = Arc::new(Schema::new(vec![ + Field::new("name", DataType::Utf8, false), + Field::new("description", DataType::Utf8, false), + ])); + let names: Vec = statuses + .iter() + .map(|s| Value::Utf8(s.to_string())) + .collect(); + let descriptions: Vec = statuses + .iter() + .map(|s| Value::Utf8(format!("status: {s}"))) + .collect(); + session.register( + "statuses", + Arc::clone(&status_schema), + vec![RecordBatch::try_new( + status_schema, + vec![ + qf_storage::array::Array::from_values(DataType::Utf8, &names)?, + qf_storage::array::Array::from_values(DataType::Utf8, &descriptions)?, + ], + )?], + )?; + + Ok(session) +} + +fn time(session: &mut Session, sql: &str, optimized: bool) -> Result<(Duration, usize, Metrics)> { + let mut op = session.operators_with(sql, optimized)?; + let start = Instant::now(); + let batches = collect(op.as_mut())?; + let elapsed = start.elapsed(); + let rows = batches.iter().map(RecordBatch::num_rows).sum(); + Ok((elapsed, rows, scan_metrics(op.as_ref()))) +} + +#[derive(Debug, Default, Clone, Copy)] +struct Metrics { + read: usize, + pruned: usize, +} + +/// Sums the row-group counters of every scan in the tree. +fn scan_metrics(op: &dyn Operator) -> Metrics { + let m = op.metrics(); + let mut total = Metrics { + read: m.row_groups_read, + pruned: m.row_groups_pruned, + }; + for c in op.children() { + let child = scan_metrics(c); + total.read += child.read; + total.pruned += child.pruned; + } + total +} + +/// Runs every case, optimised and not. +pub fn run(session: &mut Session) -> Result> { + let mut out = Vec::new(); + for case in CASES { + // One untimed run first, so the comparison is not measuring the cost + // of warming the file cache on whichever query happened to go first. + time(session, case.sql, true)?; + time(session, case.sql, false)?; + + let (optimized, rows, metrics) = time(session, case.sql, true)?; + let (unoptimized, unopt_rows, _) = time(session, case.sql, false)?; + // The two plans must agree, or the optimiser is not preserving + // semantics and the timings mean nothing. + if rows != unopt_rows { + return Err(qf_common::Error::internal(format!( + "`{}` returned {rows} rows optimised and {unopt_rows} unoptimised", + case.name + ))); + } + out.push(Measurement { + name: case.name, + optimized, + unoptimized, + rows, + row_groups_read: metrics.read, + row_groups_pruned: metrics.pruned, + }); + } + Ok(out) +} + +pub fn report(rows: usize, measurements: &[Measurement]) -> String { + let mut out = + format!("queryforge benchmark — {rows} rows, {BENCH_ROW_GROUP_SIZE}-row row groups\n\n"); + out.push_str(&format!( + "{:<22} {:>10} {:>12} {:>9} {:>10} {:>14}\n", + "query", "rows", "optimised", "speedup", "unopt.", "row groups" + )); + out.push_str(&"-".repeat(82)); + out.push('\n'); + for m in measurements { + out.push_str(&format!( + "{:<22} {:>10} {:>10.2}ms {:>8.1}x {:>8.2}ms {:>6} read {:>3} pruned\n", + m.name, + m.rows, + m.optimized.as_secs_f64() * 1000.0, + m.speedup(), + m.unoptimized.as_secs_f64() * 1000.0, + m.row_groups_read, + m.row_groups_pruned, + )); + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + struct TempDir(std::path::PathBuf); + + impl TempDir { + fn new() -> TempDir { + let mut p = std::env::temp_dir(); + p.push(format!( + "queryforge-bench-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + std::fs::create_dir_all(&p).unwrap(); + TempDir(p) + } + } + + impl Drop for TempDir { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.0); + } + } + + #[test] + fn the_benchmark_data_is_built_and_queryable() { + let dir = TempDir::new(); + let mut s = prepare(&dir.0, 20_000).unwrap(); + let b = s.query("SELECT count(*) FROM events").unwrap(); + assert_eq!(b.row(0)[0], Value::Int64(20_000)); + assert_eq!( + s.query("SELECT count(*) FROM statuses").unwrap().num_rows(), + 1 + ); + } + + #[test] + fn every_case_runs_and_both_plans_agree() { + // If the optimiser changed an answer this would fail rather than + // reporting a flattering speedup. + let dir = TempDir::new(); + let mut s = prepare(&dir.0, 20_000).unwrap(); + let results = run(&mut s).unwrap(); + assert_eq!(results.len(), CASES.len()); + for m in &results { + assert!(m.optimized.as_nanos() > 0, "{} was not timed", m.name); + } + } + + #[test] + fn a_selective_predicate_prunes_row_groups() { + let dir = TempDir::new(); + let mut s = prepare(&dir.0, 20_000).unwrap(); + let (_, rows, metrics) = + time(&mut s, "SELECT id FROM events WHERE id > 19000", true).unwrap(); + assert_eq!(rows, 999); + assert!(metrics.pruned > 0, "nothing was pruned"); + // Without the optimiser there is no pushed predicate, so nothing prunes. + let (_, _, unopt) = time(&mut s, "SELECT id FROM events WHERE id > 19000", false).unwrap(); + assert_eq!(unopt.pruned, 0); + } + + #[test] + fn the_report_lists_every_case() { + let m = vec![Measurement { + name: "example", + optimized: Duration::from_millis(1), + unoptimized: Duration::from_millis(4), + rows: 10, + row_groups_read: 1, + row_groups_pruned: 9, + }]; + let text = report(1000, &m); + assert!(text.contains("example")); + assert!(text.contains("4.0x")); + assert!(text.contains("9 pruned")); + assert!(text.contains("1000 rows")); + } + + #[test] + fn a_zero_duration_reports_an_infinite_speedup_rather_than_dividing_by_zero() { + let m = Measurement { + name: "x", + optimized: Duration::ZERO, + unoptimized: Duration::from_millis(1), + rows: 0, + row_groups_read: 0, + row_groups_pruned: 0, + }; + assert!(m.speedup().is_infinite()); + } +} diff --git a/crates/qf-cli/src/main.rs b/crates/qf-cli/src/main.rs index f328e4d..2f663b1 100644 --- a/crates/qf-cli/src/main.rs +++ b/crates/qf-cli/src/main.rs @@ -1 +1,271 @@ -fn main() {} +//! `queryforge` — the command-line front end. +//! +//! ```text +//! queryforge start the shell +//! queryforge -c "SELECT 1" run one statement and exit +//! queryforge script.sql run a file +//! queryforge bench [--rows N] build a dataset and measure the optimiser +//! +//! --import = load a CSV before doing anything else +//! --attach
= register a columnar file as a table +//! ``` + +mod bench; +mod render; +mod repl; +mod shell; + +use qf_common::{Error, Result}; +use shell::Shell; +use std::process::ExitCode; + +fn main() -> ExitCode { + let args: Vec = std::env::args().skip(1).collect(); + match run(&args) { + Ok(()) => ExitCode::SUCCESS, + Err(e) => { + eprintln!("{e}"); + ExitCode::FAILURE + } + } +} + +fn run(args: &[String]) -> Result<()> { + if args.iter().any(|a| a == "-h" || a == "--help") { + print!("{}", shell::HELP); + return Ok(()); + } + if args.first().is_some_and(|a| a == "bench") { + return run_bench(&args[1..]); + } + + let mut shell = Shell::new(); + let mut rest: Vec<&String> = Vec::new(); + let mut i = 0; + while i < args.len() { + match args[i].as_str() { + "--import" | "--attach" => { + let spec = args + .get(i + 1) + .ok_or_else(|| Error::plan(format!("{} needs
=", args[i])))?; + let (name, path) = spec + .split_once('=') + .ok_or_else(|| Error::plan(format!("expected
=, got `{spec}`")))?; + if args[i] == "--import" { + shell.session().import_csv(name, path)?; + } else { + shell.session().attach(name, path)?; + } + i += 2; + } + _ => { + rest.push(&args[i]); + i += 1; + } + } + } + + if let Some(pos) = rest.iter().position(|a| a.as_str() == "-c") { + let sql = rest + .get(pos + 1) + .ok_or_else(|| Error::plan("-c needs a statement".to_string()))?; + print!("{}", shell.handle(sql)?); + return Ok(()); + } + + if let Some(path) = rest.first() { + let text = std::fs::read_to_string(path.as_str())?; + return run_script(&mut shell, &text); + } + + repl::run(&mut shell)?; + Ok(()) +} + +/// Runs a file, statement by statement, stopping at the first failure so a +/// broken script does not carry on against a half-built catalog. +fn run_script(shell: &mut Shell, text: &str) -> Result<()> { + let mut buffer = String::new(); + for line in text.lines() { + let trimmed = line.trim(); + // A standalone comment or blank line between statements contributes + // nothing. Buffering it would make the next line look like a + // continuation, which is how a `\`-command after a comment ends up + // being parsed as SQL. + if buffer.is_empty() && (trimmed.is_empty() || trimmed.starts_with("--")) { + continue; + } + if buffer.is_empty() && trimmed.starts_with('\\') { + print!("{}", shell.handle(trimmed)?); + continue; + } + if !buffer.is_empty() { + buffer.push('\n'); + } + buffer.push_str(line); + if buffer.trim_end().ends_with(';') { + let statement = std::mem::take(&mut buffer); + print!("{}", shell.handle(&statement)?); + } + } + if !buffer.trim().is_empty() { + print!("{}", shell.handle(&buffer)?); + } + Ok(()) +} + +fn run_bench(args: &[String]) -> Result<()> { + let mut rows = bench::DEFAULT_ROWS; + if let Some(pos) = args.iter().position(|a| a == "--rows") { + rows = args + .get(pos + 1) + .and_then(|v| v.parse().ok()) + .ok_or_else(|| Error::plan("--rows needs a number".to_string()))?; + } + let dir = std::env::temp_dir().join(format!("queryforge-bench-{}", std::process::id())); + std::fs::create_dir_all(&dir)?; + let result = (|| { + eprintln!("building {rows} rows…"); + let mut session = bench::prepare(&dir, rows)?; + eprintln!("running {} queries, optimised and not…", bench::CASES.len()); + let measurements = bench::run(&mut session)?; + print!("{}", bench::report(rows, &measurements)); + Ok(()) + })(); + let _ = std::fs::remove_dir_all(&dir); + result +} + +#[cfg(test)] +mod tests { + use super::*; + + struct TempDir(std::path::PathBuf); + + impl TempDir { + fn new(tag: &str) -> TempDir { + let mut p = std::env::temp_dir(); + p.push(format!( + "queryforge-cli-{tag}-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + std::fs::create_dir_all(&p).unwrap(); + TempDir(p) + } + fn file(&self, name: &str) -> String { + self.0.join(name).to_string_lossy().into_owned() + } + } + + impl Drop for TempDir { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.0); + } + } + + fn args(list: &[&str]) -> Vec { + list.iter().map(|s| s.to_string()).collect() + } + + #[test] + fn help_is_printed_for_either_flag() { + assert!(run(&args(&["--help"])).is_ok()); + assert!(run(&args(&["-h"])).is_ok()); + } + + #[test] + fn a_single_statement_runs_with_dash_c() { + assert!(run(&args(&["-c", "SELECT 1 + 1"])).is_ok()); + } + + #[test] + fn a_bad_statement_comes_back_as_an_error() { + assert!(run(&args(&["-c", "SELECT nope FROM missing"])).is_err()); + assert!(run(&args(&["-c"])).is_err()); + } + + #[test] + fn a_csv_can_be_imported_before_the_query_runs() { + let dir = TempDir::new("import"); + let csv = dir.file("d.csv"); + std::fs::write(&csv, "a,b\n1,x\n2,y\n").unwrap(); + assert!(run(&args(&[ + "--import", + &format!("d={csv}"), + "-c", + "SELECT count(*) FROM d" + ])) + .is_ok()); + } + + #[test] + fn a_malformed_preload_specification_is_reported() { + assert!(run(&args(&["--import", "nocolon"])) + .unwrap_err() + .to_string() + .contains("
=")); + assert!(run(&args(&["--import"])).is_err()); + assert!(run(&args(&["--attach", "t=/nonexistent.qfc"])).is_err()); + } + + #[test] + fn a_qfc_file_can_be_attached_before_the_query_runs() { + let dir = TempDir::new("attach"); + let csv = dir.file("d.csv"); + let qfc = dir.file("d.qfc"); + std::fs::write(&csv, "a\n1\n2\n").unwrap(); + let mut s = Shell::new(); + s.session().import_csv("d", &csv).unwrap(); + s.session().save_table("d", &qfc).unwrap(); + + assert!(run(&args(&[ + "--attach", + &format!("d={qfc}"), + "-c", + "SELECT count(*) FROM d" + ])) + .is_ok()); + } + + #[test] + fn a_script_file_runs_statement_by_statement() { + let dir = TempDir::new("script"); + let path = dir.file("s.sql"); + std::fs::write( + &path, + "-- set up\nCREATE TABLE t (a INT);\nINSERT INTO t\n VALUES (1), (2);\n\\dt\nSELECT count(*) FROM t;\n", + ) + .unwrap(); + assert!(run(&args(&[&path])).is_ok()); + } + + #[test] + fn a_script_stops_at_its_first_failure() { + let dir = TempDir::new("bad-script"); + let path = dir.file("s.sql"); + std::fs::write(&path, "SELECT nope FROM missing;\nSELECT 1;\n").unwrap(); + assert!(run(&args(&[&path])).is_err()); + } + + #[test] + fn a_scripts_last_statement_runs_without_a_trailing_semicolon() { + let mut shell = Shell::new(); + assert!(run_script(&mut shell, "CREATE TABLE t (a INT)").is_ok()); + assert!(shell.session().catalog().contains("t")); + } + + #[test] + fn a_missing_script_file_is_reported() { + assert!(run(&args(&["/nonexistent/script.sql"])).is_err()); + } + + #[test] + fn the_benchmark_runs_at_a_small_size() { + assert!(run_bench(&args(&["--rows", "5000"])).is_ok()); + assert!(run_bench(&args(&["--rows", "not-a-number"])).is_err()); + } +} diff --git a/crates/qf-cli/src/render.rs b/crates/qf-cli/src/render.rs new file mode 100644 index 0000000..78cc6d3 --- /dev/null +++ b/crates/qf-cli/src/render.rs @@ -0,0 +1,198 @@ +//! Turning result batches into something readable in a terminal. + +use qf_common::Schema; +use qf_storage::batch::RecordBatch; + +/// Columns wider than this are truncated with an ellipsis, so one long text +/// value cannot push every other column off the screen. +const MAX_COLUMN_WIDTH: usize = 40; + +/// Renders batches as a bordered table. +pub fn table(schema: &Schema, batches: &[RecordBatch]) -> String { + let headers: Vec = schema.fields().iter().map(|f| f.name.clone()).collect(); + let mut rows: Vec> = Vec::new(); + for b in batches { + for r in b.rows() { + rows.push(r.iter().map(|v| truncate(&v.to_string())).collect()); + } + } + + if headers.is_empty() { + return format!("({} rows, no columns)\n", rows.len()); + } + + let widths: Vec = headers + .iter() + .enumerate() + .map(|(i, h)| { + rows.iter() + .filter_map(|r| r.get(i)) + .map(|c| display_width(c)) + .chain(std::iter::once(display_width(h))) + .max() + .unwrap_or(0) + }) + .collect(); + + let mut out = String::new(); + out.push_str(&border(&widths, '┌', '┬', '┐')); + out.push_str(&row(&headers, &widths)); + out.push_str(&border(&widths, '├', '┼', '┤')); + for r in &rows { + out.push_str(&row(r, &widths)); + } + out.push_str(&border(&widths, '└', '┴', '┘')); + out.push_str(&format!( + "{} row{}\n", + rows.len(), + if rows.len() == 1 { "" } else { "s" } + )); + out +} + +fn border(widths: &[usize], left: char, mid: char, right: char) -> String { + let mut s = String::new(); + s.push(left); + for (i, w) in widths.iter().enumerate() { + s.push_str(&"─".repeat(w + 2)); + s.push(if i + 1 == widths.len() { right } else { mid }); + } + s.push('\n'); + s +} + +fn row(cells: &[String], widths: &[usize]) -> String { + let mut s = String::from("│"); + for (i, w) in widths.iter().enumerate() { + let cell = cells.get(i).cloned().unwrap_or_default(); + let pad = w.saturating_sub(display_width(&cell)); + s.push(' '); + s.push_str(&cell); + s.push_str(&" ".repeat(pad)); + s.push_str(" │"); + } + s.push('\n'); + s +} + +/// Character count rather than byte length, so a column holding `naïve` is not +/// padded as though it were six characters wide. +fn display_width(s: &str) -> usize { + s.chars().count() +} + +fn truncate(s: &str) -> String { + if display_width(s) <= MAX_COLUMN_WIDTH { + return s.to_string(); + } + let kept: String = s.chars().take(MAX_COLUMN_WIDTH - 1).collect(); + format!("{kept}…") +} + +#[cfg(test)] +mod tests { + use super::*; + use qf_common::{DataType, Field, Value}; + use qf_storage::array::Array; + use std::sync::Arc; + + fn schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("name", DataType::Utf8, true), + ])) + } + + fn batch(rows: &[(i64, Option<&str>)]) -> RecordBatch { + let ids: Vec = rows.iter().map(|r| Value::Int64(r.0)).collect(); + let names: Vec = rows + .iter() + .map(|r| r.1.map_or(Value::Null, |s| Value::Utf8(s.to_string()))) + .collect(); + RecordBatch::try_new( + schema(), + vec![ + Array::from_values(DataType::Int64, &ids).unwrap(), + Array::from_values(DataType::Utf8, &names).unwrap(), + ], + ) + .unwrap() + } + + #[test] + fn a_table_has_a_header_a_rule_and_a_row_count() { + let out = table(&schema(), &[batch(&[(1, Some("ada"))])]); + let lines: Vec<&str> = out.lines().collect(); + assert!(lines[0].starts_with('┌')); + assert!(lines[1].contains("id") && lines[1].contains("name")); + assert!(lines[2].starts_with('├')); + assert!(lines[3].contains("ada")); + assert!(lines[4].starts_with('└')); + assert_eq!(lines[5], "1 row"); + } + + #[test] + fn columns_are_padded_to_their_widest_value() { + let out = table( + &schema(), + &[batch(&[(1, Some("ada")), (2, Some("brendan"))])], + ); + let widths: Vec = out + .lines() + .filter(|l| l.starts_with('│')) + .map(|l| l.chars().count()) + .collect(); + assert!(widths.windows(2).all(|w| w[0] == w[1]), "{out}"); + } + + #[test] + fn nulls_are_printed_as_null() { + let out = table(&schema(), &[batch(&[(1, None)])]); + assert!(out.contains("NULL")); + } + + #[test] + fn an_empty_result_still_shows_its_columns() { + let out = table(&schema(), &[]); + assert!(out.contains("id")); + assert!(out.ends_with("0 rows\n")); + } + + #[test] + fn a_long_value_is_truncated_rather_than_wrapping() { + let long = "x".repeat(200); + let out = table(&schema(), &[batch(&[(1, Some(&long))])]); + assert!(out.contains('…')); + assert!(out.lines().all(|l| l.chars().count() < 80), "{out}"); + } + + #[test] + fn multi_byte_characters_are_padded_by_character_not_by_byte() { + let out = table( + &schema(), + &[batch(&[(1, Some("naïve")), (2, Some("plain"))])], + ); + let widths: Vec = out + .lines() + .filter(|l| l.starts_with('│')) + .map(|l| l.chars().count()) + .collect(); + assert!(widths.windows(2).all(|w| w[0] == w[1]), "{out}"); + } + + #[test] + fn a_result_with_no_columns_reports_its_row_count() { + let empty = Arc::new(Schema::empty()); + let out = table(&empty, &[]); + assert!(out.contains("no columns")); + } + + #[test] + fn several_batches_render_as_one_table() { + let out = table( + &schema(), + &[batch(&[(1, Some("a"))]), batch(&[(2, Some("b"))])], + ); + assert!(out.contains("2 rows")); + } +} diff --git a/crates/qf-cli/src/repl.rs b/crates/qf-cli/src/repl.rs new file mode 100644 index 0000000..47a3700 --- /dev/null +++ b/crates/qf-cli/src/repl.rs @@ -0,0 +1,86 @@ +//! The terminal loop: read a line, hand it to the shell, print what comes +//! back. Everything worth testing lives in `shell.rs`, which the tests drive +//! directly, so this file is excluded from coverage. + +use crate::shell::Shell; +use std::io::{BufRead, Write}; + +/// Reads statements until end of input. +/// +/// A statement can span lines: input accumulates until it ends with a +/// semicolon, and a backslash command runs immediately. +pub fn run(shell: &mut Shell) -> std::io::Result<()> { + let stdin = std::io::stdin(); + let mut stdout = std::io::stdout(); + let mut buffer = String::new(); + + writeln!(stdout, "queryforge — \\help for commands, \\q to quit")?; + prompt(&mut stdout, false)?; + + for line in stdin.lock().lines() { + let line = line?; + let trimmed = line.trim(); + + if buffer.is_empty() && (trimmed == "\\q" || trimmed == "\\quit") { + return Ok(()); + } + // Same as in a script: a comment between statements is not the start + // of one. + if buffer.is_empty() && trimmed.starts_with("--") { + prompt(&mut stdout, false)?; + continue; + } + if buffer.is_empty() && trimmed.starts_with('\\') { + emit(&mut stdout, shell.handle(trimmed)); + prompt(&mut stdout, false)?; + continue; + } + + if !buffer.is_empty() { + buffer.push('\n'); + } + buffer.push_str(&line); + + // A statement is complete at a semicolon. Without one, keep reading — + // which is what makes a multi-line query possible. + if !buffer.trim_end().ends_with(';') { + if buffer.trim().is_empty() { + buffer.clear(); + prompt(&mut stdout, false)?; + } else { + prompt(&mut stdout, true)?; + } + continue; + } + + let statement = std::mem::take(&mut buffer); + emit(&mut stdout, shell.handle(&statement)); + prompt(&mut stdout, false)?; + } + + // Whatever was typed without a closing semicolon still runs at EOF, so a + // piped script does not silently drop its last statement. + if !buffer.trim().is_empty() { + emit(&mut stdout, shell.handle(&buffer)); + } + writeln!(stdout)?; + Ok(()) +} + +fn emit(out: &mut impl Write, result: qf_common::Result) { + match result { + Ok(text) => { + if !text.is_empty() { + let _ = write!(out, "{text}"); + } + } + Err(e) => { + let _ = writeln!(out, "{e}"); + } + } +} + +fn prompt(out: &mut impl Write, continuation: bool) -> std::io::Result<()> { + write!(out, "{}", if continuation { " ...> " } else { "qf> " })?; + out.flush() +} diff --git a/crates/qf-cli/src/shell.rs b/crates/qf-cli/src/shell.rs new file mode 100644 index 0000000..f32f5a7 --- /dev/null +++ b/crates/qf-cli/src/shell.rs @@ -0,0 +1,403 @@ +//! The shell's command handling, separated from the terminal loop so that +//! every command can be tested by feeding it a line and reading what it +//! returns — no pty, no stdin, no mocking. + +use crate::render; +use qf_common::{Error, Result}; +use qf_exec::{Output, Session}; +use std::time::Instant; + +pub struct Shell { + session: Session, + timing: bool, +} + +impl Default for Shell { + fn default() -> Self { + Shell::new() + } +} + +impl Shell { + pub fn new() -> Shell { + Shell { + session: Session::new(), + timing: false, + } + } + + pub fn session(&mut self) -> &mut Session { + &mut self.session + } + + /// Runs one line of input — a backslash command or SQL — and returns what + /// should be printed. + /// + /// Errors come back as `Err` rather than being printed here, so the caller + /// decides whether a failure ends the session (a script) or just prints + /// (the REPL). + pub fn handle(&mut self, line: &str) -> Result { + let line = line.trim(); + if line.is_empty() { + return Ok(String::new()); + } + // Comments are not stripped here: the lexer already skips them, and + // treating a leading `--` as "ignore this input" would swallow the + // statement that follows a comment on the line above it. + if let Some(command) = line.strip_prefix('\\') { + return self.meta(command); + } + let start = Instant::now(); + let outputs = self.session.execute(line)?; + let mut text = String::new(); + for out in outputs { + text.push_str(&self.render(out)); + } + if self.timing { + text.push_str(&format!( + "time: {:.2}ms\n", + start.elapsed().as_secs_f64() * 1000.0 + )); + } + Ok(text) + } + + fn render(&self, out: Output) -> String { + match out { + Output::Rows { schema, batches } => render::table(&schema, &batches), + Output::Plan(text) => text, + Output::Message(m) => format!("{m}\n"), + } + } + + fn meta(&mut self, command: &str) -> Result { + let mut parts = command.split_whitespace(); + let name = parts.next().unwrap_or(""); + let args: Vec<&str> = parts.collect(); + + match name { + "?" | "h" | "help" => Ok(HELP.to_string()), + "q" | "quit" => Ok(String::new()), + "dt" => Ok(self.list_tables()), + "d" => { + let table = arg(&args, 0, "\\d
")?; + self.describe(table) + } + "import" => { + let (name, path) = two(&args, "\\import
")?; + let rows = self.session.import_csv(name, path)?; + Ok(format!("imported {rows} rows into {name}\n")) + } + "save" => { + let (name, path) = two(&args, "\\save
")?; + let rows = self.session.save_table(name, path)?; + Ok(format!( + "wrote {rows} rows to {path}; {name} now reads from it\n" + )) + } + "attach" => { + let (name, path) = two(&args, "\\attach
")?; + let rows = self.session.attach(name, path)?; + Ok(format!("attached {name} ({rows} rows) from {path}\n")) + } + "timing" => { + self.timing = !self.timing; + Ok(format!( + "timing {}\n", + if self.timing { "on" } else { "off" } + )) + } + "explain" => { + let sql = command + .strip_prefix("explain") + .unwrap_or("") + .trim() + .to_string(); + if sql.is_empty() { + return Err(Error::plan("usage: \\explain ".to_string())); + } + self.session.explain(&sql) + } + other => Err(Error::plan(format!( + "unknown command `\\{other}` — try \\help" + ))), + } + } + + fn list_tables(&self) -> String { + let tables = self.session.tables(); + if tables.is_empty() { + return "no tables — try \\import \n".to_string(); + } + let mut out = String::new(); + for t in tables { + let kind = match &t.source { + qf_storage::catalog::TableSource::Memory(_) => "memory", + qf_storage::catalog::TableSource::File(p) => { + out.push_str(&format!( + "{:<20} {:>10} rows file {}\n", + t.name, + t.row_count(), + p.display() + )); + continue; + } + }; + out.push_str(&format!( + "{:<20} {:>10} rows {kind}\n", + t.name, + t.row_count() + )); + } + out + } + + fn describe(&self, table: &str) -> Result { + let t = self.session.catalog().get(table)?; + let mut out = format!("{} ({} rows)\n", t.name, t.row_count()); + for (i, f) in t.schema.fields().iter().enumerate() { + let null = if f.nullable { "NULL" } else { "NOT NULL" }; + let stats = match t.stats.column(i) { + Some(s) if s.row_count > 0 => format!( + " min={} max={} nulls={} distinct≈{}", + s.min, s.max, s.null_count, s.distinct_count + ), + _ => String::new(), + }; + out.push_str(&format!( + " {:<20} {:<10} {null:<9}{stats}\n", + f.name, f.data_type + )); + } + Ok(out) + } +} + +fn arg<'a>(args: &[&'a str], i: usize, usage: &str) -> Result<&'a str> { + args.get(i) + .copied() + .ok_or_else(|| Error::plan(format!("usage: {usage}"))) +} + +fn two<'a>(args: &[&'a str], usage: &str) -> Result<(&'a str, &'a str)> { + Ok((arg(args, 0, usage)?, arg(args, 1, usage)?)) +} + +pub const HELP: &str = "\ +queryforge — a columnar SQL engine + + SQL any supported statement, ending with an optional ; + EXPLAIN show the optimised plan + EXPLAIN ANALYZE run it and show per-operator counters + + \\dt list tables + \\d
show a table's columns and statistics + \\import load a CSV into a new in-memory table + \\save write a table to the columnar format and read from it + \\attach register an existing .qfc file as a table + \\explain same as EXPLAIN + \\timing toggle query timing + \\help this text + \\q quit +"; + +#[cfg(test)] +mod tests { + use super::*; + + fn shell() -> Shell { + let mut s = Shell::new(); + s.handle("CREATE TABLE t (a INT, b TEXT)").unwrap(); + s.handle("INSERT INTO t VALUES (1, 'x'), (2, 'y')").unwrap(); + s + } + + struct TempDir(std::path::PathBuf); + + impl TempDir { + fn new(tag: &str) -> TempDir { + let mut p = std::env::temp_dir(); + p.push(format!( + "queryforge-shell-{tag}-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + std::fs::create_dir_all(&p).unwrap(); + TempDir(p) + } + fn file(&self, name: &str) -> String { + self.0.join(name).to_string_lossy().into_owned() + } + } + + impl Drop for TempDir { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.0); + } + } + + #[test] + fn a_query_renders_as_a_table() { + let mut s = shell(); + let out = s.handle("SELECT * FROM t ORDER BY a").unwrap(); + assert!(out.contains(" a ")); + assert!(out.contains(" x ")); + assert!(out.contains("2 rows")); + } + + #[test] + fn blank_lines_and_comments_produce_nothing() { + let mut s = shell(); + assert!(s.handle("").unwrap().is_empty()); + assert!(s.handle(" ").unwrap().is_empty()); + assert!(s.handle("-- just a note").unwrap().is_empty()); + } + + #[test] + fn a_comment_above_a_statement_does_not_swallow_it() { + // The comment and the statement arrive as one buffer from a script. + let mut s = Shell::new(); + let out = s + .handle("-- make a table\nCREATE TABLE c (a INT);") + .unwrap(); + assert!(out.contains("created table c")); + } + + #[test] + fn several_statements_on_one_line_all_run() { + let mut s = Shell::new(); + let out = s + .handle("CREATE TABLE u (a INT); INSERT INTO u VALUES (1); SELECT count(*) FROM u") + .unwrap(); + assert!(out.contains("created table u")); + assert!(out.contains("inserted 1 rows")); + assert!(out.contains("1 row")); + } + + #[test] + fn listing_tables_shows_row_counts_and_where_they_live() { + let mut s = shell(); + assert!(s.handle("\\dt").unwrap().contains("t ")); + assert!(s.handle("\\dt").unwrap().contains("memory")); + assert!(Shell::new().handle("\\dt").unwrap().contains("no tables")); + } + + #[test] + fn describing_a_table_shows_its_columns_and_statistics() { + let mut s = shell(); + let out = s.handle("\\d t").unwrap(); + assert!(out.contains("INT64")); + assert!(out.contains("UTF8")); + assert!(out.contains("min=1 max=2")); + assert!(s.handle("\\d nope").is_err()); + assert!(s.handle("\\d").is_err()); + } + + #[test] + fn importing_a_csv_creates_a_table() { + let dir = TempDir::new("import"); + let path = dir.file("data.csv"); + std::fs::write(&path, "x,y\n1,a\n2,b\n").unwrap(); + let mut s = Shell::new(); + let out = s.handle(&format!("\\import loaded {path}")).unwrap(); + assert!(out.contains("imported 2 rows")); + assert!(s + .handle("SELECT count(*) FROM loaded") + .unwrap() + .contains("1 row")); + } + + #[test] + fn saving_and_attaching_move_a_table_through_the_columnar_format() { + let dir = TempDir::new("save"); + let path = dir.file("t.qfc"); + let mut s = shell(); + assert!(s + .handle(&format!("\\save t {path}")) + .unwrap() + .contains("wrote 2 rows")); + assert!(s.handle("\\dt").unwrap().contains("file")); + + let mut fresh = Shell::new(); + let out = fresh.handle(&format!("\\attach t {path}")).unwrap(); + assert!(out.contains("attached t (2 rows)")); + assert!(fresh + .handle("SELECT b FROM t WHERE a = 2") + .unwrap() + .contains(" y ")); + } + + #[test] + fn explain_shows_a_plan_both_ways() { + let mut s = shell(); + let a = s.handle("EXPLAIN SELECT a FROM t WHERE a > 1").unwrap(); + let b = s.handle("\\explain SELECT a FROM t WHERE a > 1").unwrap(); + assert!(a.contains("Scan: t")); + assert_eq!(a, b); + assert!(s.handle("\\explain").is_err()); + } + + #[test] + fn explain_analyze_reports_counters() { + let mut s = shell(); + let out = s.handle("EXPLAIN ANALYZE SELECT count(*) FROM t").unwrap(); + assert!(out.contains("rows=")); + assert!(out.contains("rows returned")); + } + + #[test] + fn timing_toggles_and_annotates_queries() { + let mut s = shell(); + assert!(!s.handle("SELECT 1").unwrap().contains("time:")); + assert!(s.handle("\\timing").unwrap().contains("on")); + assert!(s.handle("SELECT 1").unwrap().contains("time:")); + assert!(s.handle("\\timing").unwrap().contains("off")); + assert!(!s.handle("SELECT 1").unwrap().contains("time:")); + } + + #[test] + fn help_and_quit_are_recognised() { + let mut s = Shell::new(); + for c in ["\\help", "\\h", "\\?"] { + assert!(s.handle(c).unwrap().contains("queryforge")); + } + assert!(s.handle("\\q").unwrap().is_empty()); + } + + #[test] + fn an_unknown_command_or_a_missing_argument_is_reported() { + let mut s = Shell::new(); + assert!(s + .handle("\\nope") + .unwrap_err() + .to_string() + .contains("unknown command")); + assert!(s + .handle("\\import") + .unwrap_err() + .to_string() + .contains("usage")); + assert!(s + .handle("\\save t") + .unwrap_err() + .to_string() + .contains("usage")); + assert!(s.handle("\\attach t").is_err()); + } + + #[test] + fn a_failing_query_returns_an_error_rather_than_printing_one() { + let mut s = shell(); + let err = s.handle("SELECT nope FROM t").unwrap_err(); + assert!(err.to_string().contains("no such column")); + } + + #[test] + fn the_session_is_reachable_for_scripted_setup() { + let mut s = Shell::new(); + assert!(s.session().tables().is_empty()); + } +} diff --git a/crates/qf-common/src/schema.rs b/crates/qf-common/src/schema.rs index 169c4bf..a9bd1fd 100644 --- a/crates/qf-common/src/schema.rs +++ b/crates/qf-common/src/schema.rs @@ -56,7 +56,9 @@ impl fmt::Display for DataType { DataType::Float64 => "FLOAT64", DataType::Utf8 => "UTF8", }; - f.write_str(s) + // `pad` rather than `write_str`, so `{:<10}` actually aligns the type + // column in `\d`. `write_str` silently ignores the width. + f.pad(s) } } @@ -251,6 +253,12 @@ mod tests { assert!(!schema().contains("missing")); } + #[test] + fn type_names_honour_a_format_width_so_columns_line_up() { + assert_eq!(format!("{:<10}|", DataType::Int64), "INT64 |"); + assert_eq!(format!("{:<10}|", DataType::Float64), "FLOAT64 |"); + } + #[test] fn numeric_predicate_covers_only_the_number_types() { assert!(DataType::Int64.is_numeric()); diff --git a/crates/qf-exec/src/engine.rs b/crates/qf-exec/src/engine.rs index 1482e0a..bae6aef 100644 --- a/crates/qf-exec/src/engine.rs +++ b/crates/qf-exec/src/engine.rs @@ -283,9 +283,19 @@ impl Session { /// Builds the operator tree for a query without draining it. Lets a caller /// inspect per-operator counters after running it themselves. pub fn operators(&mut self, sql: &str) -> Result> { + self.operators_with(sql, true) + } + + /// The same, with the optimiser optionally skipped. + /// + /// Running a query both ways is the only honest way to say what the + /// optimiser is worth: the benchmark measures the difference rather than + /// asserting it. + pub fn operators_with(&mut self, sql: &str, optimized: bool) -> Result> { match Parser::parse_one(sql)? { Statement::Query(q) => { - let plan = self.plan_query(&q)?; + let plan = bind(&q, &self.catalog)?; + let plan = if optimized { optimize(plan)? } else { plan }; build(&plan, &self.catalog) } other => Err(Error::plan(format!("cannot run {other:?}"))), diff --git a/crates/qf-exec/src/join.rs b/crates/qf-exec/src/join.rs index 77c6773..f9427a1 100644 --- a/crates/qf-exec/src/join.rs +++ b/crates/qf-exec/src/join.rs @@ -151,78 +151,85 @@ impl HashJoinExec { } fn probe(&mut self, probe_batch: &RecordBatch) -> Result { - let built = self.built.as_ref().expect("built before probing"); - let mut build_rows: Vec> = Vec::new(); - let mut probe_rows: Vec> = Vec::new(); - - if self.on.is_empty() { - // A cross join: every build row against every probe row. - for p in 0..probe_batch.num_rows() { - for b in 0..built.batch.num_rows() { - build_rows.push(Some(b)); - probe_rows.push(Some(p)); + // Candidate pairs first, with no padding. The ON clause decides what + // counts as a match, and its non-equality part is as much a part of + // that decision as the keys are — so it has to be applied *before* + // asking whether a probe row matched anything. + let mut candidates: Vec<(usize, usize)> = Vec::new(); + { + let built = self.built.as_ref().expect("built before probing"); + if self.on.is_empty() { + for p in 0..probe_batch.num_rows() { + for b in 0..built.batch.num_rows() { + candidates.push((b, p)); + } } - } - } else { - let probe_keys: Vec<&BoundExpr> = self - .on - .iter() - .map(|(l, r)| match self.build_side { - BuildSide::Left => r, - BuildSide::Right => l, - }) - .collect(); - let columns = probe_keys - .iter() - .map(|e| evaluate(e, probe_batch)) - .collect::>>()?; - - for p in 0..probe_batch.num_rows() { - let key: Vec = columns.iter().map(|c| c.value(p)).collect(); - let hits = if key.iter().any(Value::is_null) { - None - } else { - built.index.get(&key) - }; - match hits { - Some(rows) if !rows.is_empty() => { - for b in rows { - build_rows.push(Some(*b)); - probe_rows.push(Some(p)); - } + } else { + let probe_keys: Vec<&BoundExpr> = self + .on + .iter() + .map(|(l, r)| match self.build_side { + BuildSide::Left => r, + BuildSide::Right => l, + }) + .collect(); + let columns = probe_keys + .iter() + .map(|e| evaluate(e, probe_batch)) + .collect::>>()?; + for p in 0..probe_batch.num_rows() { + let key: Vec = columns.iter().map(|c| c.value(p)).collect(); + // A NULL key never joins: `NULL = NULL` is unknown. + if key.iter().any(Value::is_null) { + continue; } - _ => { - if self.probe_is_preserved() { - build_rows.push(None); - probe_rows.push(Some(p)); + if let Some(rows) = built.index.get(&key) { + for b in rows { + candidates.push((*b, p)); } } } } } - let out = self.assemble(&build_rows, &probe_rows, probe_batch)?; + // Apply the ON clause's non-equality part to the candidates. + if let Some(filter) = self.filter.clone().filter(|_| !candidates.is_empty()) { + let build_rows: Vec> = candidates.iter().map(|(b, _)| Some(*b)).collect(); + let probe_rows: Vec> = candidates.iter().map(|(_, p)| Some(*p)).collect(); + let assembled = self.assemble(&build_rows, &probe_rows, probe_batch)?; + let mask = crate::eval::evaluate_predicate(&filter, &assembled)?; + candidates = candidates + .into_iter() + .enumerate() + .filter(|(i, _)| mask.get(*i)) + .map(|(_, c)| c) + .collect(); + } - // Apply the ON clause's non-equality part after matching. It can only - // remove rows, so a probe row that survives matching but fails the - // filter must come back NULL-padded on a preserved side. - let out = match &self.filter { - None => out, - Some(f) => { - let mask = crate::eval::evaluate_predicate(f, &out)?; - out.filter(&mask)? + // Only now is it settled which probe rows matched, so a preserved + // probe row that had candidates but lost them all still comes back + // padded rather than disappearing. + let mut build_rows: Vec> = Vec::with_capacity(candidates.len()); + let mut probe_rows: Vec> = Vec::with_capacity(candidates.len()); + let mut probe_matched = vec![false; probe_batch.num_rows()]; + for (b, p) in &candidates { + build_rows.push(Some(*b)); + probe_rows.push(Some(*p)); + probe_matched[*p] = true; + if self.build_is_preserved() { + self.matched[*b] = true; } - }; - - // Record which build rows matched, for the outer-join tail. - if self.build_is_preserved() { - for (b, p) in build_rows.iter().zip(probe_rows.iter()) { - if let (Some(b), Some(_)) = (b, p) { - self.matched[*b] = true; + } + if self.probe_is_preserved() { + for (p, matched) in probe_matched.iter().enumerate() { + if !matched { + build_rows.push(None); + probe_rows.push(Some(p)); } } } - Ok(out) + + self.assemble(&build_rows, &probe_rows, probe_batch) } /// Gathers the chosen rows from both sides into one batch, always with the @@ -516,6 +523,52 @@ mod tests { assert_eq!(out.row(0)[1], Value::Utf8("l1".into())); } + #[test] + fn a_left_join_whose_on_clause_rejects_a_match_pads_rather_than_drops() { + // The bug this was written for: the ON clause's non-equality part was + // applied *after* deciding what matched, so a preserved row whose only + // candidate failed the condition vanished instead of coming back + // NULL-padded — and so did rows that never had a candidate at all. + let (ls, lb) = table("l", &[Some(1), Some(2)]); + let (rs, rb) = table("r", &[Some(1)]); + let never_true = BoundExpr::binary( + BoundExpr::column(2, "r_id", DataType::Int64), + qf_sql::ast::BinaryOp::Gt, + BoundExpr::Literal(Value::Int64(1000)), + ) + .unwrap(); + let mut op = HashJoinExec::new( + operator(ls, lb), + operator(rs, rb), + JoinType::Left, + vec![(key(0), key(0))], + Some(never_true), + BuildSide::Right, + ); + let out = collect_one(&mut op).unwrap(); + assert_eq!(out.num_rows(), 2, "both left rows must survive"); + assert!(out.rows().all(|r| r[2].is_null()), "both must be padded"); + } + + #[test] + fn a_full_join_with_a_rejecting_on_clause_preserves_both_sides() { + let (ls, lb) = table("l", &[Some(1)]); + let (rs, rb) = table("r", &[Some(1)]); + let never_true = BoundExpr::Literal(Value::Boolean(false)); + let mut op = HashJoinExec::new( + operator(ls, lb), + operator(rs, rb), + JoinType::Full, + vec![(key(0), key(0))], + Some(never_true), + BuildSide::Right, + ); + let out = collect_one(&mut op).unwrap(); + assert_eq!(out.num_rows(), 2); + assert!(out.rows().any(|r| r[0].is_null())); + assert!(out.rows().any(|r| r[2].is_null())); + } + #[test] fn a_join_reports_its_children_and_name() { let (ls, lb) = table("l", &[Some(1)]); diff --git a/crates/qf-exec/tests/end_to_end.rs b/crates/qf-exec/tests/end_to_end.rs index 869cd0a..228cfd8 100644 --- a/crates/qf-exec/tests/end_to_end.rs +++ b/crates/qf-exec/tests/end_to_end.rs @@ -380,6 +380,24 @@ fn a_full_join_preserves_both_sides() { assert!(r.iter().any(|row| row[0] == "5" && row[1] == "NULL")); } +#[test] +fn a_left_join_with_a_condition_in_its_on_clause_still_returns_every_left_row() { + // The condition belongs to the match, not to the result: customers whose + // orders do not satisfy it — and customers with no orders at all — must + // still come back, padded. + let mut s = session(); + let r = rows( + &mut s, + "SELECT c.name, o.id FROM customers c \ + LEFT JOIN orders o ON o.customer_id = c.id AND o.amount > 200 \ + ORDER BY c.name, o.id", + ); + assert_eq!(r.len(), 3, "one row per customer: {r:?}"); + assert_eq!(r[0], vec!["ada", "2"]); + assert_eq!(r[1], vec!["brendan", "NULL"], "no order over 200"); + assert_eq!(r[2], vec!["grace", "NULL"], "no orders at all"); +} + #[test] fn a_cross_join_produces_the_product() { let mut s = session(); diff --git a/crates/qf-plan/src/optimizer.rs b/crates/qf-plan/src/optimizer.rs index 42c2bed..94a39a7 100644 --- a/crates/qf-plan/src/optimizer.rs +++ b/crates/qf-plan/src/optimizer.rs @@ -1750,6 +1750,268 @@ mod tests { best.1 } + // ---- the rewriting machinery, exercised directly ---- + + fn c(i: usize) -> BoundExpr { + BoundExpr::column(i, format!("c{i}"), DataType::Int64) + } + + fn text(i: usize) -> BoundExpr { + BoundExpr::column(i, format!("s{i}"), DataType::Utf8) + } + + /// A projection mapping output position i to input position 10 + i. + fn shift_by_ten() -> Vec<(BoundExpr, String)> { + (0..4).map(|i| (c(10 + i), format!("out{i}"))).collect() + } + + #[test] + fn substitution_reaches_inside_every_expression_shape() { + let projection = shift_by_ten(); + let shapes: Vec = vec![ + BoundExpr::binary(c(0), BinaryOp::Plus, c(1)).unwrap(), + BoundExpr::unary(UnaryOp::Neg, c(0)).unwrap(), + BoundExpr::Cast { + expr: Box::new(c(0)), + data_type: DataType::Float64, + }, + BoundExpr::IsNull { + expr: Box::new(c(0)), + negated: true, + }, + BoundExpr::InList { + expr: Box::new(c(0)), + list: vec![c(1), BoundExpr::Literal(Value::Int64(3))], + negated: false, + }, + BoundExpr::case( + vec![(BoundExpr::binary(c(0), BinaryOp::Gt, c(1)).unwrap(), c(2))], + Some(c(3)), + ) + .unwrap(), + ]; + for shape in shapes { + let rewritten = substitute(&shape, &projection).expect("should substitute"); + let expected: Vec = shape.column_indices().iter().map(|i| i + 10).collect(); + assert_eq!(rewritten.column_indices(), expected, "{shape}"); + } + } + + #[test] + fn substitution_reaches_inside_a_like_expression() { + let projection = vec![(text(10), "a".to_string()), (text(11), "b".to_string())]; + let like = BoundExpr::Like { + expr: Box::new(text(0)), + pattern: Box::new(text(1)), + negated: true, + }; + let rewritten = substitute(&like, &projection).unwrap(); + assert_eq!(rewritten.column_indices(), vec![10, 11]); + } + + #[test] + fn substitution_refuses_a_projection_that_computes_rather_than_renames() { + // Pushing a predicate through `SELECT a + b AS x` would duplicate the + // addition rather than save work, so the rule declines. + let computed = vec![( + BoundExpr::binary(c(0), BinaryOp::Plus, c(1)).unwrap(), + "x".to_string(), + )]; + let predicate = + BoundExpr::binary(c(0), BinaryOp::Gt, BoundExpr::Literal(Value::Int64(1))).unwrap(); + assert!(substitute(&predicate, &computed).is_none()); + } + + #[test] + fn substitution_refuses_a_column_the_projection_does_not_produce() { + let projection = vec![(c(10), "a".to_string())]; + assert!(substitute(&c(5), &projection).is_none()); + } + + #[test] + fn a_literal_only_expression_substitutes_to_itself() { + let projection = shift_by_ten(); + let lit = BoundExpr::Literal(Value::Int64(7)); + assert_eq!(substitute(&lit, &projection).unwrap(), lit); + } + + #[test] + fn rewriting_a_column_with_no_replacement_is_an_error() { + assert!(rewrite_columns(&c(0), &|_| None).is_err()); + } + + #[test] + fn folding_covers_every_comparison_operator() { + let cases = [ + (BinaryOp::Eq, 1, 1, true), + (BinaryOp::NotEq, 1, 2, true), + (BinaryOp::Lt, 1, 2, true), + (BinaryOp::LtEq, 2, 2, true), + (BinaryOp::Gt, 3, 2, true), + (BinaryOp::GtEq, 2, 2, true), + (BinaryOp::Lt, 3, 2, false), + ]; + for (op, a, b, want) in cases { + let e = BoundExpr::binary( + BoundExpr::Literal(Value::Int64(a)), + op, + BoundExpr::Literal(Value::Int64(b)), + ) + .unwrap(); + assert_eq!( + fold_expr(&e).unwrap().as_literal(), + Some(&Value::Boolean(want)), + "{a} {op} {b}" + ); + } + } + + #[test] + fn folding_covers_every_arithmetic_operator_in_both_int_and_float_form() { + let cases: [(BinaryOp, i64, i64, Value); 4] = [ + (BinaryOp::Plus, 2, 3, Value::Int64(5)), + (BinaryOp::Minus, 5, 3, Value::Int64(2)), + (BinaryOp::Multiply, 4, 3, Value::Int64(12)), + (BinaryOp::Modulo, 7, 4, Value::Int64(3)), + ]; + for (op, a, b, want) in cases { + let e = BoundExpr::binary( + BoundExpr::Literal(Value::Int64(a)), + op, + BoundExpr::Literal(Value::Int64(b)), + ) + .unwrap(); + assert_eq!(fold_expr(&e).unwrap().as_literal(), Some(&want), "{op}"); + } + let float = BoundExpr::binary( + BoundExpr::Literal(Value::Float64(1.5)), + BinaryOp::Plus, + BoundExpr::Literal(Value::Int64(1)), + ) + .unwrap(); + assert_eq!( + fold_expr(&float).unwrap().as_literal(), + Some(&Value::Float64(2.5)) + ); + } + + #[test] + fn comparing_a_literal_with_null_folds_to_null_not_to_false() { + let e = BoundExpr::binary( + BoundExpr::Literal(Value::Int64(1)), + BinaryOp::Eq, + BoundExpr::Literal(Value::Null), + ) + .unwrap(); + assert_eq!(fold_expr(&e).unwrap().as_literal(), Some(&Value::Null)); + } + + #[test] + fn folding_reaches_inside_in_lists_and_like_patterns() { + let in_list = BoundExpr::InList { + expr: Box::new(c(0)), + list: vec![BoundExpr::binary( + BoundExpr::Literal(Value::Int64(1)), + BinaryOp::Plus, + BoundExpr::Literal(Value::Int64(1)), + ) + .unwrap()], + negated: false, + }; + assert_eq!(fold_expr(&in_list).unwrap().to_string(), "c0#0 IN (2)"); + + let like = BoundExpr::Like { + expr: Box::new(text(0)), + pattern: Box::new(BoundExpr::Cast { + expr: Box::new(BoundExpr::Literal(Value::Int64(5))), + data_type: DataType::Utf8, + }), + negated: false, + }; + assert_eq!(fold_expr(&like).unwrap().to_string(), "s0#0 LIKE '5'"); + } + + #[test] + fn a_case_whose_first_branch_is_unknown_drops_that_branch() { + let e = BoundExpr::Case { + branches: vec![ + (BoundExpr::Literal(Value::Null), c(0)), + (BoundExpr::Literal(Value::Boolean(true)), c(1)), + ], + else_result: None, + data_type: DataType::Int64, + }; + assert_eq!(fold_expr(&e).unwrap(), c(1)); + } + + #[test] + fn a_case_with_a_live_branch_before_a_constant_one_keeps_both() { + let live = + BoundExpr::binary(c(0), BinaryOp::Gt, BoundExpr::Literal(Value::Int64(1))).unwrap(); + let e = BoundExpr::Case { + branches: vec![ + (live, c(1)), + (BoundExpr::Literal(Value::Boolean(true)), c(2)), + // Unreachable: the branch above always fires. + (BoundExpr::Literal(Value::Boolean(true)), c(3)), + ], + else_result: None, + data_type: DataType::Int64, + }; + let folded = fold_expr(&e).unwrap(); + match folded { + BoundExpr::Case { branches, .. } => assert_eq!(branches.len(), 2), + other => panic!("unexpected {other}"), + } + } + + #[test] + fn a_scan_of_a_table_with_no_useful_columns_still_reads_one() { + // `SELECT count(*)` needs rows, not columns. + let plan = opt("SELECT count(*) FROM orders"); + match scans(&plan)[0] { + LogicalPlan::Scan { + projection: Some(p), + .. + } => assert_eq!(p.len(), 1), + other => panic!("unexpected {other:?}"), + } + } + + #[test] + fn projection_pushdown_restores_the_output_order_when_pruning_reorders_it() { + // Selecting the columns backwards means the pruned scan produces them + // in source order and something has to put them back. + let sql = "SELECT amount, id FROM orders"; + let before = raw(sql).schema().unwrap(); + let after = push_down_projections(raw(sql)).unwrap().schema().unwrap(); + assert_eq!(before, after); + assert_eq!(after.field(0).unwrap().name, "amount"); + } + + #[test] + fn folding_a_scan_drops_a_pushed_filter_that_became_true() { + let scan = LogicalPlan::Scan { + table: "orders".into(), + source_schema: raw("SELECT * FROM orders").schema().unwrap(), + projection: None, + pushed_filters: vec![ + BoundExpr::binary( + BoundExpr::Literal(Value::Int64(1)), + BinaryOp::Eq, + BoundExpr::Literal(Value::Int64(1)), + ) + .unwrap(), + BoundExpr::binary(c(0), BinaryOp::Gt, BoundExpr::Literal(Value::Int64(5))).unwrap(), + ], + stats: TableStats::default(), + }; + match fold_constants(scan).unwrap() { + LogicalPlan::Scan { pushed_filters, .. } => assert_eq!(pushed_filters.len(), 1), + other => panic!("unexpected {other:?}"), + } + } + #[test] fn the_rule_list_is_reported_for_explain() { assert_eq!(RULES.len(), 4); diff --git a/examples/customers.csv b/examples/customers.csv new file mode 100644 index 0000000..ea9b2d7 --- /dev/null +++ b/examples/customers.csv @@ -0,0 +1,6 @@ +id,name,tier +10,Ada Lovelace,gold +20,Brendan Eich,silver +30,Grace Hopper,gold +40,Alan Turing,bronze +50,Barbara Liskov,silver diff --git a/examples/orders.csv b/examples/orders.csv new file mode 100644 index 0000000..7937946 --- /dev/null +++ b/examples/orders.csv @@ -0,0 +1,11 @@ +id,customer_id,region,status,amount +1,10,eu,shipped,120.50 +2,10,eu,delivered,89.00 +3,20,us,pending,240.75 +4,20,us,cancelled, +5,30,apac,shipped,510.00 +6,30,apac,delivered,45.25 +7,40,eu,pending,199.99 +8,10,eu,shipped,75.00 +9,50,us,delivered,310.10 +10,20,us,shipped,64.40 diff --git a/examples/tour.sql b/examples/tour.sql new file mode 100644 index 0000000..24bebcb --- /dev/null +++ b/examples/tour.sql @@ -0,0 +1,49 @@ +-- A short tour of what the engine does. Run it with: +-- cargo run --release -- examples/tour.sql + +\import orders examples/orders.csv +\import customers examples/customers.csv + +-- Types are inferred from the file; `amount` is nullable because row 4 is blank. +\d orders + +-- A plain filter and projection. +SELECT id, region, amount FROM orders WHERE amount > 200 ORDER BY amount DESC; + +-- NULL is not a value: it satisfies neither `> 200` nor `<= 200`. +SELECT + count(*) AS rows, + count(amount) AS with_amount, + sum(amount) AS total, + avg(amount) AS mean +FROM orders; + +-- Grouping, with a filter on the groups rather than the rows. +SELECT region, count(*) AS orders, sum(amount) AS revenue +FROM orders +GROUP BY region +HAVING count(*) > 2 +ORDER BY revenue DESC; + +-- A join, with the join key written the "wrong" way round on purpose — +-- the binder normalises it. +SELECT c.name, c.tier, count(*) AS orders, sum(o.amount) AS spend +FROM orders o +JOIN customers c ON c.id = o.customer_id +GROUP BY c.name, c.tier +ORDER BY spend DESC NULLS LAST; + +-- A left join keeps customers who have never ordered anything. +SELECT c.name, o.id +FROM customers c +LEFT JOIN orders o ON o.customer_id = c.id AND o.status = 'cancelled' +ORDER BY c.name; + +-- The predicate ends up inside the scan rather than above it. +EXPLAIN SELECT id FROM orders WHERE amount > 200 AND region = 'eu'; + +-- Writing the table to the columnar format, then querying it again. The plan +-- is the same; the scan now skips row groups whose zone maps rule the +-- predicate out. +\save orders /tmp/queryforge-orders.qfc +EXPLAIN ANALYZE SELECT id FROM orders WHERE amount > 200;