From 84d37db0c406cfe8c49ad958843631d939426229 Mon Sep 17 00:00:00 2001 From: Sahil Kalgutkar Date: Fri, 28 Aug 2026 14:48:32 -0400 Subject: [PATCH] Add the vectorised execution engine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Operators pull batches from each other, which is what lets a LIMIT stop a scan early instead of running it to completion and discarding the rest. Expression evaluation is vectorised: the operator is resolved once and the loop runs over a contiguous buffer. Comparing a column against a literal has its own path so no constant column is materialised, and same-type column pairs skip boxing through Value. Three-valued logic is implemented properly — `false AND NULL` is false, `true OR NULL` is true — because getting it wrong changes which rows a query returns. The scan is where predicate pushdown stops being a plan rewrite and starts saving I/O: before reading a row group it asks each predicate whether the group's zone map could satisfy it, and skips the column chunks entirely if not. A conjunction prunes when either half rules the group out; a disjunction only when both do. Sorting has three modes. Feeding a LIMIT it keeps only the best k rows, so memory is bounded by the limit rather than the input. Fitting in memory it sorts once. Otherwise it writes sorted runs to spill files and merges them, which is the only reason a database sorts differently from Vec::sort. Spill files are removed even when a query is abandoned partway through. The hash join builds one side and streams the other. For an inner join the smaller estimate is built; for an outer join the choice is forced, because the side that may be NULL-padded has to be the one that can be scanned for non-matches afterwards. Output column order is left-then-right either way. Fixed along the way: a projection producing no columns dropped its row count, which would have made `SELECT count(*)` return zero after column pruning. --- crates/qf-exec/src/aggregate.rs | 505 +++++++++++ crates/qf-exec/src/engine.rs | 294 +++++++ crates/qf-exec/src/eval.rs | 1253 ++++++++++++++++++++++++++++ crates/qf-exec/src/join.rs | 576 +++++++++++++ crates/qf-exec/src/lib.rs | 12 +- crates/qf-exec/src/operator.rs | 209 +++++ crates/qf-exec/src/operators.rs | 485 +++++++++++ crates/qf-exec/src/physical.rs | 432 ++++++++++ crates/qf-exec/src/scan.rs | 492 +++++++++++ crates/qf-exec/src/sort.rs | 658 +++++++++++++++ crates/qf-exec/tests/end_to_end.rs | 741 ++++++++++++++++ crates/qf-plan/src/binder.rs | 10 + 12 files changed, 5666 insertions(+), 1 deletion(-) create mode 100644 crates/qf-exec/src/aggregate.rs create mode 100644 crates/qf-exec/src/engine.rs create mode 100644 crates/qf-exec/src/eval.rs create mode 100644 crates/qf-exec/src/join.rs create mode 100644 crates/qf-exec/src/operator.rs create mode 100644 crates/qf-exec/src/operators.rs create mode 100644 crates/qf-exec/src/physical.rs create mode 100644 crates/qf-exec/src/scan.rs create mode 100644 crates/qf-exec/src/sort.rs create mode 100644 crates/qf-exec/tests/end_to_end.rs diff --git a/crates/qf-exec/src/aggregate.rs b/crates/qf-exec/src/aggregate.rs new file mode 100644 index 0000000..ec29cb1 --- /dev/null +++ b/crates/qf-exec/src/aggregate.rs @@ -0,0 +1,505 @@ +//! Hash aggregation. +//! +//! One pass over the input builds a hash table from group key to a set of +//! accumulators, then the table is turned into output rows. Sorting the input +//! instead would be asymptotically worse and, for the grouped queries people +//! actually write, much slower. +//! +//! The accumulators are the interesting part: `count` and `sum` need one +//! number, `avg` needs two, and any of them under `DISTINCT` needs the set of +//! values seen — which is why `DISTINCT` is opt-in rather than the default. + +use crate::eval::evaluate; +use crate::operator::{Metrics, Operator}; +use qf_common::{DataType, Error, Field, Result, Schema, Value}; +use qf_plan::expr::{AggFunc, BoundAggregate, BoundExpr}; +use qf_storage::array::ArrayBuilder; +use qf_storage::batch::RecordBatch; +use std::collections::HashMap; +use std::collections::HashSet; +use std::sync::Arc; +use std::time::Instant; + +#[derive(Debug, Clone)] +enum Accumulator { + Count { + n: i64, + seen: Option>, + }, + Sum { + total: Option, + seen: Option>, + }, + MinMax { + value: Value, + is_min: bool, + }, + Avg { + total: f64, + n: i64, + seen: Option>, + }, +} + +impl Accumulator { + fn new(agg: &BoundAggregate) -> Accumulator { + let seen = agg.distinct.then(HashSet::new); + match agg.func { + AggFunc::Count => Accumulator::Count { n: 0, seen }, + AggFunc::Sum => Accumulator::Sum { total: None, seen }, + AggFunc::Min => Accumulator::MinMax { + value: Value::Null, + is_min: true, + }, + AggFunc::Max => Accumulator::MinMax { + value: Value::Null, + is_min: false, + }, + AggFunc::Avg => Accumulator::Avg { + total: 0.0, + n: 0, + seen, + }, + } + } + + /// Folds one value in. `None` is the `count(*)` case, where there is no + /// argument and every row counts. + fn update(&mut self, value: Option<&Value>) -> Result<()> { + // Every aggregate ignores NULL inputs. `count(*)` has no input at all, + // so it counts the row regardless. + let v = match value { + None => { + if let Accumulator::Count { n, .. } = self { + *n += 1; + } + return Ok(()); + } + Some(v) if v.is_null() => return Ok(()), + Some(v) => v, + }; + + let fresh = |seen: &mut Option>| match seen { + None => true, + Some(s) => s.insert(v.clone()), + }; + + match self { + Accumulator::Count { n, seen } => { + if fresh(seen) { + *n += 1; + } + } + Accumulator::Sum { total, seen } => { + if !fresh(seen) { + return Ok(()); + } + *total = Some(match total.take() { + None => v.clone(), + Some(t) => add(&t, v)?, + }); + } + Accumulator::MinMax { value: cur, is_min } => { + *cur = if *is_min { + Value::min(cur, v) + } else { + Value::max(cur, v) + }; + } + Accumulator::Avg { total, n, seen } => { + if !fresh(seen) { + return Ok(()); + } + let x = v + .as_f64() + .ok_or_else(|| Error::typ(format!("cannot average {}", v.type_name())))?; + *total += x; + *n += 1; + } + } + Ok(()) + } + + fn finish(&self) -> Value { + match self { + Accumulator::Count { n, .. } => Value::Int64(*n), + // `sum` over no rows is NULL, not zero — SQL distinguishes "no + // rows" from "rows that summed to nothing". + Accumulator::Sum { total, .. } => total.clone().unwrap_or(Value::Null), + Accumulator::MinMax { value, .. } => value.clone(), + Accumulator::Avg { total, n, .. } => { + if *n == 0 { + Value::Null + } else { + Value::Float64(total / *n as f64) + } + } + } + } +} + +fn add(a: &Value, b: &Value) -> Result { + Ok(match (a, b) { + (Value::Int64(x), Value::Int64(y)) => Value::Int64(x.wrapping_add(*y)), + _ => { + let (Some(x), Some(y)) = (a.as_f64(), b.as_f64()) else { + return Err(Error::typ(format!( + "cannot sum {} and {}", + a.type_name(), + b.type_name() + ))); + }; + Value::Float64(x + y) + } + }) +} + +pub struct HashAggregateExec { + input: Box, + group_exprs: Vec, + aggregates: Vec, + schema: Arc, + /// Insertion order is preserved so the output is deterministic run to run, + /// which matters for tests and for anyone diffing two query results. + order: Vec>, + table: HashMap, usize>, + accumulators: Vec>, + done: bool, + metrics: Metrics, +} + +impl HashAggregateExec { + pub fn new( + input: Box, + group_exprs: Vec<(BoundExpr, String)>, + aggregates: Vec, + ) -> HashAggregateExec { + let mut fields: Vec = group_exprs + .iter() + .map(|(e, n)| Field::new(n, e.data_type(), true)) + .collect(); + fields.extend( + aggregates + .iter() + .map(|a| Field::new(&a.output_name, a.data_type, true)), + ); + HashAggregateExec { + input, + group_exprs: group_exprs.into_iter().map(|(e, _)| e).collect(), + aggregates, + schema: Arc::new(Schema::new(fields)), + order: Vec::new(), + table: HashMap::new(), + accumulators: Vec::new(), + done: false, + metrics: Metrics::default(), + } + } + + fn consume(&mut self, batch: &RecordBatch) -> Result<()> { + let keys = self + .group_exprs + .iter() + .map(|e| evaluate(e, batch)) + .collect::>>()?; + let args = self + .aggregates + .iter() + .map(|a| match &a.arg { + Some(e) => evaluate(e, batch).map(Some), + None => Ok(None), + }) + .collect::>>()?; + + for row in 0..batch.num_rows() { + let key: Vec = keys.iter().map(|k| k.value(row)).collect(); + let slot = match self.table.get(&key) { + Some(i) => *i, + None => { + let i = self.accumulators.len(); + self.accumulators + .push(self.aggregates.iter().map(Accumulator::new).collect()); + self.table.insert(key.clone(), i); + self.order.push(key); + i + } + }; + for (a, arg) in self.accumulators[slot].iter_mut().zip(args.iter()) { + let v = arg.as_ref().map(|c| c.value(row)); + a.update(v.as_ref())?; + } + } + Ok(()) + } + + fn build_output(&mut self) -> Result { + // A global aggregate over no rows still returns one row: count 0, + // everything else NULL. + if self.group_exprs.is_empty() && self.accumulators.is_empty() { + self.accumulators + .push(self.aggregates.iter().map(Accumulator::new).collect()); + self.order.push(vec![]); + } + + let mut builders: Vec = self + .schema + .fields() + .iter() + .map(|f| ArrayBuilder::new(f.data_type)) + .collect(); + + for (key, accs) in self.order.iter().zip(self.accumulators.iter()) { + for (i, v) in key.iter().enumerate() { + builders[i].push(v.clone())?; + } + for (i, a) in accs.iter().enumerate() { + let target = self.group_exprs.len() + i; + let value = a.finish(); + // `sum` over integers can overflow into a float column and + // vice versa; the schema is the authority. + let coerced = match (&value, builders[target].data_type()) { + (Value::Null, _) => Value::Null, + (v, DataType::Float64) if v.as_f64().is_some() => { + Value::Float64(v.as_f64().unwrap()) + } + (v, _) => v.clone(), + }; + builders[target].push(coerced)?; + } + } + + let columns = builders + .into_iter() + .map(ArrayBuilder::finish) + .collect::>>()?; + RecordBatch::try_new(Arc::clone(&self.schema), columns) + } +} + +impl Operator for HashAggregateExec { + fn schema(&self) -> Arc { + Arc::clone(&self.schema) + } + + fn next_batch(&mut self) -> Result> { + if self.done { + return Ok(None); + } + let start = Instant::now(); + while let Some(b) = self.input.next_batch()? { + self.consume(&b)?; + } + let out = self.build_output()?; + self.done = true; + self.metrics.elapsed += start.elapsed(); + self.metrics.record(&out); + Ok(Some(out)) + } + + fn metrics(&self) -> &Metrics { + &self.metrics + } + + fn children(&self) -> Vec<&dyn Operator> { + vec![self.input.as_ref()] + } + + fn name(&self) -> &'static str { + "HashAggregate" + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::operator::collect_one; + use crate::operators::ValuesExec; + use qf_plan::expr::AggFunc; + use qf_storage::array::Array; + + fn agg(func: AggFunc, arg: Option, distinct: bool, t: DataType) -> BoundAggregate { + BoundAggregate { + func, + arg, + distinct, + output_name: func.name().to_string(), + data_type: t, + } + } + + fn schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("g", DataType::Utf8, true), + Field::new("n", DataType::Int64, true), + ])) + } + + fn source(rows: &[(Option<&str>, Option)]) -> Box { + let gs: Vec = rows + .iter() + .map(|(g, _)| g.map_or(Value::Null, |s| Value::Utf8(s.to_string()))) + .collect(); + let ns: Vec = rows + .iter() + .map(|(_, n)| n.map_or(Value::Null, Value::Int64)) + .collect(); + let batch = RecordBatch::try_new( + schema(), + vec![ + Array::from_values(DataType::Utf8, &gs).unwrap(), + Array::from_values(DataType::Int64, &ns).unwrap(), + ], + ) + .unwrap(); + Box::new(ValuesExec::new(schema(), vec![batch])) + } + + fn n_col() -> BoundExpr { + BoundExpr::column(1, "n", DataType::Int64) + } + + fn run(rows: &[(Option<&str>, Option)], aggs: Vec) -> RecordBatch { + let mut op = HashAggregateExec::new(source(rows), vec![], aggs); + collect_one(&mut op).unwrap() + } + + #[test] + fn sum_over_no_non_null_values_is_null_not_zero() { + // SQL distinguishes "no rows to add up" from "rows adding to nothing". + let out = run( + &[(None, None), (None, None)], + vec![agg(AggFunc::Sum, Some(n_col()), false, DataType::Int64)], + ); + assert!(out.row(0)[0].is_null()); + } + + #[test] + fn count_star_counts_null_rows_and_count_of_a_column_does_not() { + let out = run( + &[(None, Some(1)), (None, None)], + vec![ + agg(AggFunc::Count, None, false, DataType::Int64), + agg(AggFunc::Count, Some(n_col()), false, DataType::Int64), + ], + ); + assert_eq!(out.row(0)[0], Value::Int64(2)); + assert_eq!(out.row(0)[1], Value::Int64(1)); + } + + #[test] + fn min_and_max_ignore_nulls_and_work_on_text() { + let g = BoundExpr::column(0, "g", DataType::Utf8); + let out = run( + &[(Some("m"), None), (None, None), (Some("a"), None)], + vec![ + agg(AggFunc::Min, Some(g.clone()), false, DataType::Utf8), + agg(AggFunc::Max, Some(g), false, DataType::Utf8), + ], + ); + assert_eq!(out.row(0)[0], Value::Utf8("a".into())); + assert_eq!(out.row(0)[1], Value::Utf8("m".into())); + } + + #[test] + fn avg_divides_by_the_non_null_count_and_is_null_over_nothing() { + let out = run( + &[(None, Some(1)), (None, Some(2)), (None, None)], + vec![agg(AggFunc::Avg, Some(n_col()), false, DataType::Float64)], + ); + assert_eq!(out.row(0)[0], Value::Float64(1.5)); + + let empty = run( + &[(None, None)], + vec![agg(AggFunc::Avg, Some(n_col()), false, DataType::Float64)], + ); + assert!(empty.row(0)[0].is_null()); + } + + #[test] + fn distinct_folds_repeats_out_of_count_sum_and_avg() { + let rows = [(None, Some(2)), (None, Some(2)), (None, Some(4))]; + let out = run( + &rows, + vec![ + agg(AggFunc::Count, Some(n_col()), true, DataType::Int64), + agg(AggFunc::Sum, Some(n_col()), true, DataType::Int64), + agg(AggFunc::Avg, Some(n_col()), true, DataType::Float64), + ], + ); + assert_eq!(out.row(0)[0], Value::Int64(2)); + assert_eq!(out.row(0)[1], Value::Int64(6)); + assert_eq!(out.row(0)[2], Value::Float64(3.0)); + } + + #[test] + fn a_global_aggregate_over_an_empty_input_still_returns_one_row() { + let mut op = HashAggregateExec::new( + Box::new(ValuesExec::new(schema(), vec![])), + vec![], + vec![agg(AggFunc::Count, None, false, DataType::Int64)], + ); + let out = collect_one(&mut op).unwrap(); + assert_eq!(out.num_rows(), 1); + assert_eq!(out.row(0)[0], Value::Int64(0)); + } + + #[test] + fn a_grouped_aggregate_over_an_empty_input_returns_no_rows() { + let mut op = HashAggregateExec::new( + Box::new(ValuesExec::new(schema(), vec![])), + vec![(BoundExpr::column(0, "g", DataType::Utf8), "g".into())], + vec![agg(AggFunc::Count, None, false, DataType::Int64)], + ); + assert_eq!(collect_one(&mut op).unwrap().num_rows(), 0); + } + + #[test] + fn groups_come_back_in_the_order_they_were_first_seen() { + let mut op = HashAggregateExec::new( + source(&[ + (Some("z"), Some(1)), + (Some("a"), Some(1)), + (Some("z"), Some(1)), + ]), + vec![(BoundExpr::column(0, "g", DataType::Utf8), "g".into())], + vec![agg(AggFunc::Count, None, false, DataType::Int64)], + ); + let out = collect_one(&mut op).unwrap(); + assert_eq!(out.row(0)[0], Value::Utf8("z".into())); + assert_eq!(out.row(0)[1], Value::Int64(2)); + assert_eq!(out.row(1)[0], Value::Utf8("a".into())); + assert_eq!(op.name(), "HashAggregate"); + assert_eq!(op.children().len(), 1); + } + + #[test] + fn averaging_text_is_an_error_rather_than_a_wrong_number() { + let g = BoundExpr::column(0, "g", DataType::Utf8); + let mut op = HashAggregateExec::new( + source(&[(Some("a"), None)]), + vec![], + vec![agg(AggFunc::Avg, Some(g), false, DataType::Float64)], + ); + assert!(collect_one(&mut op).is_err()); + } + + #[test] + fn summing_text_is_an_error() { + let g = BoundExpr::column(0, "g", DataType::Utf8); + let mut op = HashAggregateExec::new( + source(&[(Some("a"), None), (Some("b"), None)]), + vec![], + vec![agg(AggFunc::Sum, Some(g), false, DataType::Int64)], + ); + assert!(collect_one(&mut op).is_err()); + } + + #[test] + fn a_sum_declared_as_a_float_column_widens_its_integer_result() { + let out = run( + &[(None, Some(3))], + vec![agg(AggFunc::Sum, Some(n_col()), false, DataType::Float64)], + ); + assert_eq!(out.row(0)[0], Value::Float64(3.0)); + } +} diff --git a/crates/qf-exec/src/engine.rs b/crates/qf-exec/src/engine.rs new file mode 100644 index 0000000..1482e0a --- /dev/null +++ b/crates/qf-exec/src/engine.rs @@ -0,0 +1,294 @@ +//! The session: a catalog plus the ability to run statements against it. +//! +//! Everything the CLI does goes through here, and so does every end-to-end +//! test, which is deliberate — it means the tests exercise the same path a +//! user does rather than a convenient shortcut past it. + +use crate::operator::{collect, Operator}; +use crate::physical::{build, explain_analyze}; +use qf_common::{Error, Result, Schema}; +use qf_plan::binder::{bind, bind_constant_row, schema_from_columns}; +use qf_plan::logical::LogicalPlan; +use qf_plan::optimizer::optimize; +use qf_sql::ast::{Query, Statement}; +use qf_sql::Parser; +use qf_storage::batch::RecordBatch; +use qf_storage::catalog::{Catalog, Table, TableSource}; +use qf_storage::csv::{read_csv_file, read_csv_with_schema}; +use qf_storage::format::{write_file, DEFAULT_ROW_GROUP_SIZE}; +use std::path::Path; +use std::sync::Arc; + +/// What running one statement produced. +#[derive(Debug, Clone, PartialEq)] +pub enum Output { + Rows { + schema: Arc, + batches: Vec, + }, + /// A plan, from `EXPLAIN`. + Plan(String), + /// A confirmation, from a statement that does not return rows. + Message(String), +} + +impl Output { + pub fn row_count(&self) -> usize { + match self { + Output::Rows { batches, .. } => batches.iter().map(RecordBatch::num_rows).sum(), + _ => 0, + } + } +} + +#[derive(Default)] +pub struct Session { + catalog: Catalog, +} + +impl Session { + pub fn new() -> Session { + Session::default() + } + + pub fn catalog(&self) -> &Catalog { + &self.catalog + } + + /// Runs every statement in `sql`. + pub fn execute(&mut self, sql: &str) -> Result> { + Parser::parse_statements(sql)? + .into_iter() + .map(|s| self.run(s)) + .collect() + } + + /// Runs exactly one statement. + pub fn execute_one(&mut self, sql: &str) -> Result { + let mut out = self.execute(sql)?; + match out.len() { + 1 => Ok(out.remove(0)), + n => Err(Error::plan(format!("expected one statement, found {n}"))), + } + } + + /// Runs a query and returns its rows as a single batch. The shape most + /// tests want. + pub fn query(&mut self, sql: &str) -> Result { + match self.execute_one(sql)? { + Output::Rows { schema, batches } => RecordBatch::concat(schema, &batches), + other => Err(Error::plan(format!( + "expected a query, but the statement produced {other:?}" + ))), + } + } + + fn run(&mut self, statement: Statement) -> Result { + match statement { + Statement::Query(q) => self.run_query(&q), + Statement::Explain { analyze, query } => self.run_explain(&query, analyze), + Statement::CreateTable { name, columns } => { + let schema = schema_from_columns(&columns)?; + self.catalog + .register_batches(name.clone(), schema, vec![])?; + Ok(Output::Message(format!("created table {name}"))) + } + Statement::DropTable { name } => { + self.catalog.drop_table(&name)?; + Ok(Output::Message(format!("dropped table {name}"))) + } + Statement::Insert { table, rows } => self.run_insert(&table, rows), + Statement::Copy { table, path } => self.run_copy(&table, &path), + } + } + + fn plan_query(&self, query: &Query) -> Result { + optimize(bind(query, &self.catalog)?) + } + + fn run_query(&mut self, query: &Query) -> Result { + let plan = self.plan_query(query)?; + let mut op = build(&plan, &self.catalog)?; + let schema = op.schema(); + let batches = collect(op.as_mut())?; + Ok(Output::Rows { schema, batches }) + } + + fn run_explain(&mut self, query: &Query, analyze: bool) -> Result { + let plan = self.plan_query(query)?; + if !analyze { + return Ok(Output::Plan(plan.explain())); + } + // ANALYZE has to actually run the query — the counters are the point. + let mut op = build(&plan, &self.catalog)?; + let rows: usize = collect(op.as_mut())? + .iter() + .map(RecordBatch::num_rows) + .sum(); + let mut text = explain_analyze(op.as_ref()); + text.push_str(&format!("\n{rows} rows returned\n")); + Ok(Output::Plan(text)) + } + + fn run_insert(&mut self, table: &str, rows: Vec>) -> Result { + let entry = self.catalog.get(table)?.clone(); + let TableSource::Memory(existing) = &entry.source else { + return Err(Error::plan(format!( + "`{table}` is backed by a file; INSERT only works on in-memory tables" + ))); + }; + + // The values go through the same binder and evaluator as everything + // else — evaluated against one empty row — rather than having their + // own literal-only path that could disagree with the rest of the + // engine about types. + let seed = RecordBatch::empty(Arc::new(Schema::empty()))?.take(&[0usize; 1])?; + let mut builders: Vec = entry + .schema + .fields() + .iter() + .map(|f| qf_storage::array::ArrayBuilder::new(f.data_type)) + .collect(); + + for row in &rows { + if row.len() != entry.schema.len() { + return Err(Error::plan(format!( + "`{table}` has {} columns but a row supplied {}", + entry.schema.len(), + row.len() + ))); + } + let bound = bind_constant_row(row, &self.catalog)?; + for (i, e) in bound.iter().enumerate() { + let column = crate::eval::evaluate(e, &seed)?; + let target = entry.schema.field(i)?.data_type; + builders[i].push(column.value(0).cast_to(target)?)?; + } + } + + let columns = builders + .into_iter() + .map(qf_storage::array::ArrayBuilder::finish) + .collect::>>()?; + let mut batches = existing.clone(); + batches.push(RecordBatch::try_new(Arc::clone(&entry.schema), columns)?); + let inserted = rows.len(); + self.replace_batches(&entry.name, entry.schema.clone(), batches)?; + Ok(Output::Message(format!("inserted {inserted} rows"))) + } + + fn run_copy(&mut self, table: &str, path: &str) -> Result { + let entry = self.catalog.get(table)?.clone(); + let TableSource::Memory(existing) = &entry.source else { + return Err(Error::plan(format!( + "`{table}` is backed by a file; COPY only loads into in-memory tables" + ))); + }; + let text = std::fs::read_to_string(path)?; + let loaded = read_csv_with_schema(&text, Arc::clone(&entry.schema))?; + let rows: usize = loaded.iter().map(RecordBatch::num_rows).sum(); + let mut batches = existing.clone(); + batches.extend(loaded); + self.replace_batches(&entry.name, entry.schema.clone(), batches)?; + Ok(Output::Message(format!("copied {rows} rows into {table}"))) + } + + fn replace_batches( + &mut self, + name: &str, + schema: Arc, + batches: Vec, + ) -> Result<()> { + self.catalog.drop_table(name)?; + self.catalog.register_batches(name, schema, batches) + } + + /// Creates a table from a CSV file, inferring its schema. + pub fn import_csv(&mut self, name: &str, path: impl AsRef) -> Result { + let (schema, batches) = read_csv_file(path, 1000)?; + let rows = batches.iter().map(RecordBatch::num_rows).sum(); + if self.catalog.contains(name) { + self.catalog.drop_table(name)?; + } + self.catalog.register_batches(name, schema, batches)?; + Ok(rows) + } + + /// Writes a table to a `.qfc` file and re-registers it as file-backed, so + /// subsequent queries go through the columnar reader and its zone maps. + pub fn save_table(&mut self, name: &str, path: impl AsRef) -> Result { + self.save_table_with_row_group_size(name, path, DEFAULT_ROW_GROUP_SIZE) + } + + pub fn save_table_with_row_group_size( + &mut self, + name: &str, + path: impl AsRef, + row_group_size: usize, + ) -> Result { + let entry = self.catalog.get(name)?.clone(); + let batches = match &entry.source { + TableSource::Memory(b) => b.clone(), + TableSource::File(p) => { + let mut reader = qf_storage::format::QfcReader::open(p)?; + reader.read_all()? + } + }; + let meta = write_file( + path.as_ref(), + Arc::clone(&entry.schema), + &batches, + row_group_size, + )?; + self.catalog.drop_table(name)?; + self.catalog.register_file(&entry.name, path.as_ref())?; + Ok(meta.num_rows()) + } + + /// Registers an existing `.qfc` file as a table. + pub fn attach(&mut self, name: &str, path: impl AsRef) -> Result { + if self.catalog.contains(name) { + self.catalog.drop_table(name)?; + } + self.catalog.register_file(name, path)?; + Ok(self.catalog.get(name)?.row_count()) + } + + /// Registers batches directly. Used by tests and by anything embedding the + /// engine. + pub fn register( + &mut self, + name: &str, + schema: Arc, + batches: Vec, + ) -> Result<()> { + if self.catalog.contains(name) { + self.catalog.drop_table(name)?; + } + self.catalog.register_batches(name, schema, batches) + } + + pub fn tables(&self) -> Vec<&Table> { + self.catalog.tables().collect() + } + + /// The optimised plan for a query, without running it. + pub fn explain(&mut self, sql: &str) -> Result { + match Parser::parse_one(sql)? { + Statement::Query(q) => Ok(self.plan_query(&q)?.explain()), + other => Err(Error::plan(format!("cannot explain {other:?}"))), + } + } + + /// 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> { + match Parser::parse_one(sql)? { + Statement::Query(q) => { + let plan = self.plan_query(&q)?; + build(&plan, &self.catalog) + } + other => Err(Error::plan(format!("cannot run {other:?}"))), + } + } +} diff --git a/crates/qf-exec/src/eval.rs b/crates/qf-exec/src/eval.rs new file mode 100644 index 0000000..2fce361 --- /dev/null +++ b/crates/qf-exec/src/eval.rs @@ -0,0 +1,1253 @@ +//! Vectorised expression evaluation. +//! +//! Every kernel here takes whole columns and returns a whole column. The +//! operator is resolved once, outside the loop, and the loop then runs over a +//! contiguous buffer — which is the entire reason for the columnar layout. The +//! row-at-a-time alternative would walk the expression tree once per row and +//! spend most of its time on dispatch rather than on the comparison. +//! +//! Two cases get their own paths because they dominate real queries: +//! comparing a column against a literal (no constant column is materialised), +//! and comparing two columns of the same type (no boxing through `Value`). + +use qf_common::{DataType, Error, Result, Value}; +use qf_plan::expr::BoundExpr; +use qf_sql::ast::{BinaryOp, UnaryOp}; +use qf_storage::array::{Array, ArrayBuilder, ArrayData}; +use qf_storage::batch::RecordBatch; +use qf_storage::bitmap::Bitmap; +use std::cmp::Ordering; + +/// Evaluates `expr` over every row of `batch`, producing one column. +pub fn evaluate(expr: &BoundExpr, batch: &RecordBatch) -> Result { + match expr { + BoundExpr::Column { index, .. } => Ok(batch.column(*index)?.clone()), + BoundExpr::Literal(v) => constant(v, batch.num_rows(), expr.data_type()), + BoundExpr::Binary { + left, op, right, .. + } => eval_binary(left, *op, right, batch), + BoundExpr::Unary { op, expr, .. } => { + let input = evaluate(expr, batch)?; + eval_unary(*op, &input) + } + BoundExpr::Cast { expr, data_type } => { + let input = evaluate(expr, batch)?; + cast_array(&input, *data_type) + } + BoundExpr::IsNull { expr, negated } => { + let input = evaluate(expr, batch)?; + let mut bits = Bitmap::with_capacity(input.len()); + for i in 0..input.len() { + bits.push(input.is_valid(i) == *negated); + } + // `IS NULL` is never itself null, so the result carries no mask. + Array::new(ArrayData::Boolean(bits), None, input.len()) + } + BoundExpr::InList { + expr, + list, + negated, + } => eval_in_list(expr, list, *negated, batch), + BoundExpr::Like { + expr, + pattern, + negated, + } => eval_like(expr, pattern, *negated, batch), + BoundExpr::Case { + branches, + else_result, + data_type, + } => eval_case(branches, else_result.as_deref(), *data_type, batch), + } +} + +/// Evaluates a predicate into a selection mask, treating SQL's unknown as +/// "not selected" — which is what `WHERE` does with NULL. +pub fn evaluate_predicate(expr: &BoundExpr, batch: &RecordBatch) -> Result { + let array = evaluate(expr, batch)?; + let ArrayData::Boolean(bits) = array.data() else { + return Err(Error::exec(format!( + "a predicate evaluated to {}, not a boolean", + array.data_type() + ))); + }; + let mut mask = Bitmap::with_capacity(array.len()); + for i in 0..array.len() { + mask.push(array.is_valid(i) && bits.get(i)); + } + Ok(mask) +} + +fn constant(value: &Value, rows: usize, data_type: DataType) -> Result { + let mut b = ArrayBuilder::new(data_type); + b.reserve(rows); + for _ in 0..rows { + b.push(value.clone())?; + } + b.finish() +} + +fn eval_binary( + left: &BoundExpr, + op: BinaryOp, + right: &BoundExpr, + batch: &RecordBatch, +) -> Result { + if op.is_logical() { + let l = evaluate(left, batch)?; + let r = evaluate(right, batch)?; + return eval_logical(op, &l, &r); + } + + // Column against literal, without building a column of the literal. + if let (BoundExpr::Column { index, .. }, Some(v)) = (left, right.as_literal()) { + let column = batch.column(*index)?; + if op.is_comparison() { + return compare_scalar(column, op, v, false); + } + return arithmetic_scalar(column, op, v, false); + } + if let (Some(v), BoundExpr::Column { index, .. }) = (left.as_literal(), right) { + let column = batch.column(*index)?; + if op.is_comparison() { + return compare_scalar(column, op, v, true); + } + return arithmetic_scalar(column, op, v, true); + } + + let l = evaluate(left, batch)?; + let r = evaluate(right, batch)?; + if op.is_comparison() { + compare_arrays(&l, op, &r) + } else { + arithmetic_arrays(&l, op, &r) + } +} + +fn ordering_matches(op: BinaryOp, ord: Ordering) -> bool { + match op { + BinaryOp::Eq => ord.is_eq(), + BinaryOp::NotEq => !ord.is_eq(), + BinaryOp::Lt => ord.is_lt(), + BinaryOp::LtEq => ord.is_le(), + BinaryOp::Gt => ord.is_gt(), + _ => ord.is_ge(), + } +} + +fn boolean_array(bits: Bitmap, validity: Option, len: usize) -> Result { + Array::new(ArrayData::Boolean(bits), validity, len) +} + +/// Column against a scalar. `flipped` means the literal was on the left, so +/// the comparison is mirrored rather than the data being copied. +fn compare_scalar(column: &Array, op: BinaryOp, value: &Value, flipped: bool) -> Result { + let len = column.len(); + let mut bits = Bitmap::with_capacity(len); + let mut validity = Bitmap::with_capacity(len); + let mut any_null = false; + + if value.is_null() { + // Comparing with NULL is unknown for every row. + return boolean_array( + Bitmap::filled(len, false), + Some(Bitmap::filled(len, false)), + len, + ); + } + + // Typed fast paths: the comparison is resolved once, then the loop runs + // over the raw buffer. + match (column.data(), value) { + (ArrayData::Int64(values), Value::Int64(v)) => { + for (i, x) in values.iter().enumerate() { + let valid = column.is_valid(i); + any_null |= !valid; + validity.push(valid); + let ord = if flipped { v.cmp(x) } else { x.cmp(v) }; + bits.push(valid && ordering_matches(op, ord)); + } + } + (ArrayData::Float64(values), _) if value.as_f64().is_some() => { + let v = value.as_f64().unwrap(); + for (i, x) in values.iter().enumerate() { + let valid = column.is_valid(i); + any_null |= !valid; + validity.push(valid); + let ord = if flipped { + v.partial_cmp(x) + } else { + x.partial_cmp(&v) + }; + bits.push(valid && ord.is_some_and(|o| ordering_matches(op, o))); + } + } + (ArrayData::Utf8 { .. }, Value::Utf8(v)) => { + for i in 0..len { + let valid = column.is_valid(i); + any_null |= !valid; + validity.push(valid); + let matched = match column.str_value(i) { + Some(s) => { + let ord = if flipped { + v.as_str().cmp(s) + } else { + s.cmp(v.as_str()) + }; + ordering_matches(op, ord) + } + None => false, + }; + bits.push(valid && matched); + } + } + _ => { + // Mixed or unusual types fall back to the scalar comparison, which + // reports a type error rather than guessing. + for i in 0..len { + let x = column.value(i); + let valid = !x.is_null(); + any_null |= !valid; + validity.push(valid); + let ord = if flipped { + value.sql_compare(&x)? + } else { + x.sql_compare(value)? + }; + bits.push(match ord { + Some(o) => ordering_matches(op, o), + None => false, + }); + } + } + } + + boolean_array(bits, if any_null { Some(validity) } else { None }, len) +} + +fn compare_arrays(left: &Array, op: BinaryOp, right: &Array) -> Result { + let len = left.len(); + if right.len() != len { + return Err(Error::internal(format!( + "comparing a {len}-row column with a {}-row column", + right.len() + ))); + } + let mut bits = Bitmap::with_capacity(len); + let mut validity = Bitmap::with_capacity(len); + let mut any_null = false; + + match (left.data(), right.data()) { + (ArrayData::Int64(a), ArrayData::Int64(b)) => { + for i in 0..len { + let valid = left.is_valid(i) && right.is_valid(i); + any_null |= !valid; + validity.push(valid); + bits.push(valid && ordering_matches(op, a[i].cmp(&b[i]))); + } + } + (ArrayData::Float64(a), ArrayData::Float64(b)) => { + for i in 0..len { + let valid = left.is_valid(i) && right.is_valid(i); + any_null |= !valid; + validity.push(valid); + bits.push( + valid + && a[i] + .partial_cmp(&b[i]) + .is_some_and(|o| ordering_matches(op, o)), + ); + } + } + (ArrayData::Utf8 { .. }, ArrayData::Utf8 { .. }) => { + for i in 0..len { + let valid = left.is_valid(i) && right.is_valid(i); + any_null |= !valid; + validity.push(valid); + let matched = match (left.str_value(i), right.str_value(i)) { + (Some(a), Some(b)) => ordering_matches(op, a.cmp(b)), + _ => false, + }; + bits.push(valid && matched); + } + } + _ => { + for i in 0..len { + let (a, b) = (left.value(i), right.value(i)); + let valid = !a.is_null() && !b.is_null(); + any_null |= !valid; + validity.push(valid); + bits.push(match a.sql_compare(&b)? { + Some(o) => ordering_matches(op, o), + None => false, + }); + } + } + } + boolean_array(bits, if any_null { Some(validity) } else { None }, len) +} + +fn arith_values(op: BinaryOp, a: &Value, b: &Value) -> Result { + if a.is_null() || b.is_null() { + return Ok(Value::Null); + } + Ok(match (a, b) { + (Value::Int64(x), Value::Int64(y)) => match op { + BinaryOp::Plus => Value::Int64(x.wrapping_add(*y)), + BinaryOp::Minus => Value::Int64(x.wrapping_sub(*y)), + BinaryOp::Multiply => Value::Int64(x.wrapping_mul(*y)), + // Dividing by zero yields NULL rather than aborting the query. + // A whole scan should not die because one row held a zero. + BinaryOp::Divide => { + if *y == 0 { + Value::Null + } else { + Value::Float64(*x as f64 / *y as f64) + } + } + BinaryOp::Modulo => { + if *y == 0 { + Value::Null + } else { + Value::Int64(x % y) + } + } + _ => return Err(Error::exec(format!("{op} is not arithmetic"))), + }, + _ => { + let (Some(x), Some(y)) = (a.as_f64(), b.as_f64()) else { + return Err(Error::typ(format!( + "cannot apply {op} to {} and {}", + a.type_name(), + b.type_name() + ))); + }; + match op { + BinaryOp::Plus => Value::Float64(x + y), + BinaryOp::Minus => Value::Float64(x - y), + BinaryOp::Multiply => Value::Float64(x * y), + BinaryOp::Divide | BinaryOp::Modulo => { + if y == 0.0 { + Value::Null + } else if op == BinaryOp::Divide { + Value::Float64(x / y) + } else { + Value::Float64(x % y) + } + } + _ => return Err(Error::exec(format!("{op} is not arithmetic"))), + } + } + }) +} + +fn result_type(op: BinaryOp, a: DataType, b: DataType) -> Result { + if op == BinaryOp::Divide { + return Ok(DataType::Float64); + } + DataType::unify(a, b) +} + +fn arithmetic_scalar(column: &Array, op: BinaryOp, value: &Value, flipped: bool) -> Result { + let value_type = value.data_type().unwrap_or(column.data_type()); + let out_type = result_type(op, column.data_type(), value_type)?; + let mut b = ArrayBuilder::new(out_type); + b.reserve(column.len()); + for i in 0..column.len() { + let x = column.value(i); + let v = if flipped { + arith_values(op, value, &x)? + } else { + arith_values(op, &x, value)? + }; + b.push(v)?; + } + b.finish() +} + +fn arithmetic_arrays(left: &Array, op: BinaryOp, right: &Array) -> Result { + if left.len() != right.len() { + return Err(Error::internal( + "arithmetic over unequal columns".to_string(), + )); + } + let out_type = result_type(op, left.data_type(), right.data_type())?; + let mut b = ArrayBuilder::new(out_type); + b.reserve(left.len()); + for i in 0..left.len() { + b.push(arith_values(op, &left.value(i), &right.value(i))?)?; + } + b.finish() +} + +/// Three-valued AND/OR. +/// +/// `false AND NULL` is `false`, not NULL — the row cannot match whatever the +/// unknown turns out to be. Getting this wrong changes which rows a query +/// returns, so both asymmetric cases have their own test. +fn eval_logical(op: BinaryOp, left: &Array, right: &Array) -> Result { + let len = left.len(); + let (ArrayData::Boolean(a), ArrayData::Boolean(b)) = (left.data(), right.data()) else { + return Err(Error::typ(format!( + "{op} needs booleans, found {} and {}", + left.data_type(), + right.data_type() + ))); + }; + let mut bits = Bitmap::with_capacity(len); + let mut validity = Bitmap::with_capacity(len); + let mut any_null = false; + for i in 0..len { + let l = left.is_valid(i).then(|| a.get(i)); + let r = right.is_valid(i).then(|| b.get(i)); + let value = match op { + BinaryOp::And => match (l, r) { + (Some(false), _) | (_, Some(false)) => Some(false), + (Some(true), Some(true)) => Some(true), + _ => None, + }, + _ => match (l, r) { + (Some(true), _) | (_, Some(true)) => Some(true), + (Some(false), Some(false)) => Some(false), + _ => None, + }, + }; + match value { + Some(v) => { + validity.push(true); + bits.push(v); + } + None => { + any_null = true; + validity.push(false); + bits.push(false); + } + } + } + boolean_array(bits, if any_null { Some(validity) } else { None }, len) +} + +fn eval_unary(op: UnaryOp, input: &Array) -> Result { + match op { + UnaryOp::Not => { + let ArrayData::Boolean(bits) = input.data() else { + return Err(Error::typ(format!( + "NOT needs a boolean, found {}", + input.data_type() + ))); + }; + let mut out = Bitmap::with_capacity(input.len()); + for i in 0..input.len() { + out.push(input.is_valid(i) && !bits.get(i)); + } + boolean_array(out, input.validity().cloned(), input.len()) + } + UnaryOp::Neg => { + let mut b = ArrayBuilder::new(input.data_type()); + b.reserve(input.len()); + for i in 0..input.len() { + b.push(match input.value(i) { + Value::Int64(v) => Value::Int64(-v), + Value::Float64(v) => Value::Float64(-v), + Value::Null => Value::Null, + other => { + return Err(Error::typ(format!("cannot negate {}", other.type_name()))) + } + })?; + } + b.finish() + } + } +} + +fn cast_array(input: &Array, target: DataType) -> Result { + let mut b = ArrayBuilder::new(target); + b.reserve(input.len()); + for i in 0..input.len() { + b.push(input.value(i).cast_to(target)?)?; + } + b.finish() +} + +fn eval_in_list( + expr: &BoundExpr, + list: &[BoundExpr], + negated: bool, + batch: &RecordBatch, +) -> Result { + let column = evaluate(expr, batch)?; + // Literal lists — the overwhelmingly common case — become a set once + // rather than a comparison per row per item. + let literals: Option> = list + .iter() + .map(|e| e.as_literal().cloned()) + .collect::>>(); + + let len = column.len(); + let mut bits = Bitmap::with_capacity(len); + let mut validity = Bitmap::with_capacity(len); + let mut any_null = false; + + let columns: Vec = match &literals { + Some(_) => vec![], + None => list + .iter() + .map(|e| evaluate(e, batch)) + .collect::>>()?, + }; + let set: Option> = literals + .as_ref() + .map(|l| l.iter().filter(|v| !v.is_null()).cloned().collect()); + let list_has_null = literals + .as_ref() + .is_some_and(|l| l.iter().any(Value::is_null)); + + for i in 0..len { + let v = column.value(i); + if v.is_null() { + any_null = true; + validity.push(false); + bits.push(false); + continue; + } + let found = match &set { + Some(s) => s.contains(&v), + None => { + let mut hit = false; + for c in &columns { + if c.value(i) == v { + hit = true; + break; + } + } + hit + } + }; + // `x IN (1, NULL)` is unknown when x is not 1: the NULL might have + // been x. + if !found && list_has_null { + any_null = true; + validity.push(false); + bits.push(false); + continue; + } + validity.push(true); + bits.push(found != negated); + } + boolean_array(bits, if any_null { Some(validity) } else { None }, len) +} + +/// SQL `LIKE`: `%` matches any run of characters, `_` matches exactly one. +/// +/// Backtracking rather than a regex, and iterative rather than recursive so a +/// pattern of many `%`s cannot blow the stack. +pub fn like_matches(text: &str, pattern: &str) -> bool { + let t: Vec = text.chars().collect(); + let p: Vec = pattern.chars().collect(); + let (mut ti, mut pi) = (0usize, 0usize); + // Where to resume if the current `%` guess turns out to be wrong. + let (mut star, mut star_ti) = (None, 0usize); + + while ti < t.len() { + if pi < p.len() && (p[pi] == '_' || p[pi] == t[ti]) { + ti += 1; + pi += 1; + } else if pi < p.len() && p[pi] == '%' { + star = Some(pi); + star_ti = ti; + pi += 1; + } else if let Some(s) = star { + // Let the last `%` swallow one more character and try again. + pi = s + 1; + star_ti += 1; + ti = star_ti; + } else { + return false; + } + } + while pi < p.len() && p[pi] == '%' { + pi += 1; + } + pi == p.len() +} + +fn eval_like( + expr: &BoundExpr, + pattern: &BoundExpr, + negated: bool, + batch: &RecordBatch, +) -> Result { + let column = evaluate(expr, batch)?; + let len = column.len(); + let mut bits = Bitmap::with_capacity(len); + let mut validity = Bitmap::with_capacity(len); + let mut any_null = false; + + let constant_pattern = pattern.as_literal().and_then(|v| match v { + Value::Utf8(s) => Some(s.clone()), + _ => None, + }); + let pattern_column = match &constant_pattern { + Some(_) => None, + None => Some(evaluate(pattern, batch)?), + }; + + for i in 0..len { + let text = column.str_value(i); + let pat = match (&constant_pattern, &pattern_column) { + (Some(p), _) => Some(p.as_str()), + (None, Some(c)) => c.str_value(i), + _ => None, + }; + match (text, pat) { + (Some(t), Some(p)) => { + validity.push(true); + bits.push(like_matches(t, p) != negated); + } + _ => { + any_null = true; + validity.push(false); + bits.push(false); + } + } + } + boolean_array(bits, if any_null { Some(validity) } else { None }, len) +} + +fn eval_case( + branches: &[(BoundExpr, BoundExpr)], + else_result: Option<&BoundExpr>, + data_type: DataType, + batch: &RecordBatch, +) -> Result { + let len = batch.num_rows(); + let conditions = branches + .iter() + .map(|(w, _)| evaluate_predicate(w, batch)) + .collect::>>()?; + let results = branches + .iter() + .map(|(_, t)| evaluate(t, batch)) + .collect::>>()?; + let fallback = match else_result { + Some(e) => Some(evaluate(e, batch)?), + None => None, + }; + + let mut b = ArrayBuilder::new(data_type); + b.reserve(len); + 'row: for i in 0..len { + for (c, r) in conditions.iter().zip(results.iter()) { + if c.get(i) { + b.push(r.value(i).cast_to(data_type)?)?; + continue 'row; + } + } + match &fallback { + Some(f) => b.push(f.value(i).cast_to(data_type)?)?, + None => b.push_null(), + } + } + b.finish() +} + +#[cfg(test)] +mod tests { + use super::*; + use qf_common::{Field, Schema}; + use std::sync::Arc; + + fn batch(columns: Vec<(&str, DataType, Vec)>) -> RecordBatch { + let schema = Arc::new(Schema::new( + columns + .iter() + .map(|(n, t, _)| Field::new(*n, *t, true)) + .collect(), + )); + let arrays = columns + .iter() + .map(|(_, t, v)| Array::from_values(*t, v).unwrap()) + .collect(); + RecordBatch::try_new(schema, arrays).unwrap() + } + + fn ints(name: &str, values: &[Option]) -> (String, DataType, Vec) { + ( + name.to_string(), + DataType::Int64, + values + .iter() + .map(|v| v.map_or(Value::Null, Value::Int64)) + .collect(), + ) + } + + fn sample() -> RecordBatch { + let (n, t, v) = ints("a", &[Some(1), Some(5), None, Some(10)]); + batch(vec![ + (Box::leak(n.into_boxed_str()), t, v), + ( + "s", + DataType::Utf8, + vec![ + Value::Utf8("alpha".into()), + Value::Utf8("beta".into()), + Value::Null, + Value::Utf8("gamma".into()), + ], + ), + ( + "f", + DataType::Float64, + vec![ + Value::Float64(1.5), + Value::Float64(2.5), + Value::Float64(3.5), + Value::Null, + ], + ), + ]) + } + + fn col(i: usize, t: DataType) -> BoundExpr { + BoundExpr::column(i, format!("c{i}"), t) + } + + fn lit(v: Value) -> BoundExpr { + BoundExpr::Literal(v) + } + + fn mask(expr: &BoundExpr, b: &RecordBatch) -> Vec { + evaluate_predicate(expr, b).unwrap().iter().collect() + } + + #[test] + fn a_column_evaluates_to_itself_and_a_literal_to_a_constant_column() { + let b = sample(); + let c = evaluate(&col(0, DataType::Int64), &b).unwrap(); + assert_eq!(c.value(1), Value::Int64(5)); + let k = evaluate(&lit(Value::Int64(7)), &b).unwrap(); + assert_eq!(k.len(), 4); + assert_eq!(k.value(3), Value::Int64(7)); + } + + #[test] + fn comparing_a_column_with_a_literal_selects_the_expected_rows() { + let b = sample(); + let e = + BoundExpr::binary(col(0, DataType::Int64), BinaryOp::Gt, lit(Value::Int64(4))).unwrap(); + assert_eq!(mask(&e, &b), vec![false, true, false, true]); + } + + #[test] + fn a_literal_on_the_left_mirrors_the_comparison_rather_than_copying_the_column() { + let b = sample(); + // `4 < a` must select the same rows as `a > 4`. + let e = + BoundExpr::binary(lit(Value::Int64(4)), BinaryOp::Lt, col(0, DataType::Int64)).unwrap(); + assert_eq!(mask(&e, &b), vec![false, true, false, true]); + } + + #[test] + fn a_null_row_never_satisfies_a_comparison() { + let b = sample(); + for op in [ + BinaryOp::Eq, + BinaryOp::NotEq, + BinaryOp::Lt, + BinaryOp::Gt, + BinaryOp::LtEq, + BinaryOp::GtEq, + ] { + let e = BoundExpr::binary(col(0, DataType::Int64), op, lit(Value::Int64(5))).unwrap(); + assert!(!mask(&e, &b)[2], "{op} selected the null row"); + } + } + + #[test] + fn comparing_against_null_selects_nothing() { + let b = sample(); + let e = BoundExpr::binary(col(0, DataType::Int64), BinaryOp::Eq, lit(Value::Null)).unwrap(); + assert_eq!(mask(&e, &b), vec![false; 4]); + } + + #[test] + fn string_and_float_comparisons_use_their_own_paths() { + let b = sample(); + let s = BoundExpr::binary( + col(1, DataType::Utf8), + BinaryOp::Lt, + lit(Value::Utf8("delta".into())), + ) + .unwrap(); + assert_eq!(mask(&s, &b), vec![true, true, false, false]); + + let f = BoundExpr::binary( + col(2, DataType::Float64), + BinaryOp::GtEq, + lit(Value::Float64(2.5)), + ) + .unwrap(); + assert_eq!(mask(&f, &b), vec![false, true, true, false]); + } + + #[test] + fn an_integer_column_compares_against_a_float_literal() { + let b = sample(); + let e = BoundExpr::binary( + col(0, DataType::Int64), + BinaryOp::Gt, + lit(Value::Float64(4.5)), + ) + .unwrap(); + assert_eq!(mask(&e, &b), vec![false, true, false, true]); + } + + #[test] + fn two_columns_compare_element_by_element() { + let b = batch(vec![ + ("a", DataType::Int64, vec![Value::Int64(1), Value::Int64(5)]), + ("b", DataType::Int64, vec![Value::Int64(3), Value::Int64(2)]), + ]); + let e = BoundExpr::binary( + col(0, DataType::Int64), + BinaryOp::Lt, + col(1, DataType::Int64), + ) + .unwrap(); + assert_eq!(mask(&e, &b), vec![true, false]); + } + + #[test] + fn two_string_columns_and_two_float_columns_compare() { + let b = batch(vec![ + ( + "a", + DataType::Utf8, + vec![Value::Utf8("x".into()), Value::Null], + ), + ( + "b", + DataType::Utf8, + vec![Value::Utf8("y".into()), Value::Utf8("z".into())], + ), + ( + "c", + DataType::Float64, + vec![Value::Float64(1.0), Value::Float64(2.0)], + ), + ( + "d", + DataType::Float64, + vec![Value::Float64(2.0), Value::Float64(1.0)], + ), + ]); + let s = BoundExpr::binary(col(0, DataType::Utf8), BinaryOp::Lt, col(1, DataType::Utf8)) + .unwrap(); + assert_eq!(mask(&s, &b), vec![true, false]); + let f = BoundExpr::binary( + col(2, DataType::Float64), + BinaryOp::Lt, + col(3, DataType::Float64), + ) + .unwrap(); + assert_eq!(mask(&f, &b), vec![true, false]); + } + + #[test] + fn mixed_type_columns_fall_back_to_the_scalar_path() { + let b = batch(vec![ + ("a", DataType::Int64, vec![Value::Int64(2)]), + ("b", DataType::Float64, vec![Value::Float64(2.0)]), + ]); + let e = BoundExpr::binary( + col(0, DataType::Int64), + BinaryOp::Eq, + col(1, DataType::Float64), + ) + .unwrap(); + assert_eq!(mask(&e, &b), vec![true]); + } + + #[test] + fn arithmetic_runs_over_columns_and_scalars_alike() { + let b = sample(); + let plus = BoundExpr::binary( + col(0, DataType::Int64), + BinaryOp::Plus, + lit(Value::Int64(1)), + ) + .unwrap(); + let out = evaluate(&plus, &b).unwrap(); + assert_eq!(out.value(0), Value::Int64(2)); + assert!(out.value(2).is_null()); + + let times = BoundExpr::binary( + col(0, DataType::Int64), + BinaryOp::Multiply, + col(0, DataType::Int64), + ) + .unwrap(); + assert_eq!(evaluate(×, &b).unwrap().value(1), Value::Int64(25)); + } + + #[test] + fn a_literal_on_the_left_of_a_subtraction_keeps_the_operands_in_order() { + let b = sample(); + let e = BoundExpr::binary( + lit(Value::Int64(10)), + BinaryOp::Minus, + col(0, DataType::Int64), + ) + .unwrap(); + let out = evaluate(&e, &b).unwrap(); + assert_eq!(out.value(0), Value::Int64(9)); + assert_eq!(out.value(1), Value::Int64(5)); + } + + #[test] + fn dividing_by_zero_yields_null_rather_than_killing_the_query() { + // One bad row should not abort a scan of a million. + let b = batch(vec![ + ("a", DataType::Int64, vec![Value::Int64(6), Value::Int64(6)]), + ("b", DataType::Int64, vec![Value::Int64(3), Value::Int64(0)]), + ]); + let e = BoundExpr::binary( + col(0, DataType::Int64), + BinaryOp::Divide, + col(1, DataType::Int64), + ) + .unwrap(); + let out = evaluate(&e, &b).unwrap(); + assert_eq!(out.value(0), Value::Float64(2.0)); + assert!(out.value(1).is_null()); + + let m = BoundExpr::binary( + col(0, DataType::Int64), + BinaryOp::Modulo, + col(1, DataType::Int64), + ) + .unwrap(); + assert!(evaluate(&m, &b).unwrap().value(1).is_null()); + } + + #[test] + fn integer_division_produces_a_float_column() { + let b = batch(vec![ + ("a", DataType::Int64, vec![Value::Int64(7)]), + ("b", DataType::Int64, vec![Value::Int64(2)]), + ]); + let e = BoundExpr::binary( + col(0, DataType::Int64), + BinaryOp::Divide, + col(1, DataType::Int64), + ) + .unwrap(); + let out = evaluate(&e, &b).unwrap(); + assert_eq!(out.data_type(), DataType::Float64); + assert_eq!(out.value(0), Value::Float64(3.5)); + } + + #[test] + fn three_valued_and_treats_false_as_decisive() { + // `false AND NULL` is false: the row cannot match whatever the unknown + // turns out to be. + let b = batch(vec![ + ( + "p", + DataType::Boolean, + vec![Value::Boolean(false), Value::Boolean(true), Value::Null], + ), + ("q", DataType::Boolean, vec![Value::Null; 3]), + ]); + let and = BoundExpr::binary( + col(0, DataType::Boolean), + BinaryOp::And, + col(1, DataType::Boolean), + ) + .unwrap(); + let out = evaluate(&and, &b).unwrap(); + assert_eq!(out.value(0), Value::Boolean(false)); + assert!(out.value(1).is_null()); + assert!(out.value(2).is_null()); + } + + #[test] + fn three_valued_or_treats_true_as_decisive() { + let b = batch(vec![ + ( + "p", + DataType::Boolean, + vec![Value::Boolean(true), Value::Boolean(false), Value::Null], + ), + ("q", DataType::Boolean, vec![Value::Null; 3]), + ]); + let or = BoundExpr::binary( + col(0, DataType::Boolean), + BinaryOp::Or, + col(1, DataType::Boolean), + ) + .unwrap(); + let out = evaluate(&or, &b).unwrap(); + assert_eq!(out.value(0), Value::Boolean(true)); + assert!(out.value(1).is_null()); + assert!(out.value(2).is_null()); + } + + #[test] + fn a_predicate_that_is_unknown_does_not_select_its_row() { + let b = batch(vec![("p", DataType::Boolean, vec![Value::Null])]); + assert_eq!(mask(&col(0, DataType::Boolean), &b), vec![false]); + } + + #[test] + fn not_inverts_and_leaves_nulls_null() { + let b = batch(vec![( + "p", + DataType::Boolean, + vec![Value::Boolean(true), Value::Boolean(false), Value::Null], + )]); + let e = BoundExpr::unary(UnaryOp::Not, col(0, DataType::Boolean)).unwrap(); + let out = evaluate(&e, &b).unwrap(); + assert_eq!(out.value(0), Value::Boolean(false)); + assert_eq!(out.value(1), Value::Boolean(true)); + assert!(out.value(2).is_null()); + } + + #[test] + fn negation_and_casting_run_over_whole_columns() { + let b = sample(); + let neg = BoundExpr::unary(UnaryOp::Neg, col(0, DataType::Int64)).unwrap(); + assert_eq!(evaluate(&neg, &b).unwrap().value(1), Value::Int64(-5)); + + let cast = BoundExpr::Cast { + expr: Box::new(col(0, DataType::Int64)), + data_type: DataType::Utf8, + }; + let out = evaluate(&cast, &b).unwrap(); + assert_eq!(out.value(0), Value::Utf8("1".into())); + assert!(out.value(2).is_null()); + } + + #[test] + fn is_null_and_is_not_null_are_never_themselves_null() { + let b = sample(); + let is_null = BoundExpr::IsNull { + expr: Box::new(col(0, DataType::Int64)), + negated: false, + }; + assert_eq!(mask(&is_null, &b), vec![false, false, true, false]); + let not_null = BoundExpr::IsNull { + expr: Box::new(col(0, DataType::Int64)), + negated: true, + }; + assert_eq!(mask(¬_null, &b), vec![true, true, false, true]); + assert!(evaluate(&is_null, &b).unwrap().validity().is_none()); + } + + #[test] + fn in_list_matches_against_a_set_built_once() { + let b = sample(); + let e = BoundExpr::InList { + expr: Box::new(col(0, DataType::Int64)), + list: vec![lit(Value::Int64(1)), lit(Value::Int64(10))], + negated: false, + }; + assert_eq!(mask(&e, &b), vec![true, false, false, true]); + let n = BoundExpr::InList { + expr: Box::new(col(0, DataType::Int64)), + list: vec![lit(Value::Int64(1))], + negated: true, + }; + assert_eq!(mask(&n, &b), vec![false, true, false, true]); + } + + #[test] + fn a_null_in_the_list_makes_a_non_match_unknown() { + // `5 IN (1, NULL)` is unknown, not false — the NULL might have been 5. + let b = sample(); + let e = BoundExpr::InList { + expr: Box::new(col(0, DataType::Int64)), + list: vec![lit(Value::Int64(1)), lit(Value::Null)], + negated: false, + }; + let out = evaluate(&e, &b).unwrap(); + assert_eq!(out.value(0), Value::Boolean(true)); + assert!(out.value(1).is_null()); + } + + #[test] + fn an_in_list_of_expressions_is_evaluated_per_row() { + let b = batch(vec![ + ("a", DataType::Int64, vec![Value::Int64(3), Value::Int64(4)]), + ("b", DataType::Int64, vec![Value::Int64(3), Value::Int64(9)]), + ]); + let e = BoundExpr::InList { + expr: Box::new(col(0, DataType::Int64)), + list: vec![BoundExpr::binary( + col(1, DataType::Int64), + BinaryOp::Plus, + lit(Value::Int64(0)), + ) + .unwrap()], + negated: false, + }; + assert_eq!(mask(&e, &b), vec![true, false]); + } + + #[test] + fn like_handles_the_wildcards_and_the_awkward_patterns() { + assert!(like_matches("alpha", "a%")); + assert!(like_matches("alpha", "%a")); + assert!(like_matches("alpha", "%lph%")); + assert!(like_matches("alpha", "_lpha")); + assert!(like_matches("alpha", "alpha")); + assert!(like_matches("alpha", "%")); + assert!(like_matches("", "%")); + assert!(like_matches("", "")); + assert!(!like_matches("alpha", "b%")); + assert!(!like_matches("alpha", "_lph")); + assert!(!like_matches("", "_")); + } + + #[test] + fn like_backtracks_correctly_on_a_pattern_that_needs_it() { + // A greedy `%` has to give characters back for this to match. + assert!(like_matches("aaa", "%a")); + assert!(like_matches("abcabc", "%abc")); + assert!(!like_matches("abcabd", "%abc")); + assert!(like_matches("xaybz", "x%y%z")); + assert!(!like_matches("xaybz", "x%y%q")); + } + + #[test] + fn like_over_a_column_uses_the_constant_pattern_path() { + let b = sample(); + let e = BoundExpr::Like { + expr: Box::new(col(1, DataType::Utf8)), + pattern: Box::new(lit(Value::Utf8("%a".into()))), + negated: false, + }; + assert_eq!(mask(&e, &b), vec![true, true, false, true]); + let n = BoundExpr::Like { + expr: Box::new(col(1, DataType::Utf8)), + pattern: Box::new(lit(Value::Utf8("a%".into()))), + negated: true, + }; + assert_eq!(mask(&n, &b), vec![false, true, false, true]); + } + + #[test] + fn like_with_a_per_row_pattern_column_works_too() { + let b = batch(vec![ + ( + "s", + DataType::Utf8, + vec![Value::Utf8("abc".into()), Value::Utf8("abc".into())], + ), + ( + "p", + DataType::Utf8, + vec![Value::Utf8("a%".into()), Value::Utf8("z%".into())], + ), + ]); + let e = BoundExpr::Like { + expr: Box::new(col(0, DataType::Utf8)), + pattern: Box::new(col(1, DataType::Utf8)), + negated: false, + }; + assert_eq!(mask(&e, &b), vec![true, false]); + } + + #[test] + fn case_picks_the_first_matching_branch_and_falls_back_to_else() { + let b = sample(); + let e = BoundExpr::case( + vec![ + ( + BoundExpr::binary(col(0, DataType::Int64), BinaryOp::Lt, lit(Value::Int64(3))) + .unwrap(), + lit(Value::Utf8("small".into())), + ), + ( + BoundExpr::binary(col(0, DataType::Int64), BinaryOp::Lt, lit(Value::Int64(8))) + .unwrap(), + lit(Value::Utf8("medium".into())), + ), + ], + Some(lit(Value::Utf8("large".into()))), + ) + .unwrap(); + let out = evaluate(&e, &b).unwrap(); + assert_eq!(out.value(0), Value::Utf8("small".into())); + assert_eq!(out.value(1), Value::Utf8("medium".into())); + // The NULL row matches no condition, so it takes the ELSE. + assert_eq!(out.value(2), Value::Utf8("large".into())); + assert_eq!(out.value(3), Value::Utf8("large".into())); + } + + #[test] + fn a_case_without_else_yields_null_where_nothing_matches() { + let b = sample(); + let e = BoundExpr::case( + vec![( + BoundExpr::binary( + col(0, DataType::Int64), + BinaryOp::Gt, + lit(Value::Int64(100)), + ) + .unwrap(), + lit(Value::Int64(1)), + )], + None, + ) + .unwrap(); + let out = evaluate(&e, &b).unwrap(); + assert!(out.iter().all(|v| v.is_null())); + } + + #[test] + fn a_non_boolean_predicate_is_an_execution_error() { + let b = sample(); + let err = evaluate_predicate(&col(0, DataType::Int64), &b).unwrap_err(); + assert!(err.to_string().contains("not a boolean")); + } + + #[test] + fn logical_operators_over_non_booleans_are_rejected() { + let b = sample(); + let e = BoundExpr::Binary { + left: Box::new(col(0, DataType::Int64)), + op: BinaryOp::And, + right: Box::new(col(0, DataType::Int64)), + data_type: DataType::Boolean, + }; + assert!(evaluate(&e, &b).is_err()); + let n = BoundExpr::Unary { + op: UnaryOp::Not, + expr: Box::new(col(0, DataType::Int64)), + data_type: DataType::Boolean, + }; + assert!(evaluate(&n, &b).is_err()); + } + + #[test] + fn negating_or_adding_text_is_rejected_at_run_time_too() { + let b = sample(); + let neg = BoundExpr::Unary { + op: UnaryOp::Neg, + expr: Box::new(col(1, DataType::Utf8)), + data_type: DataType::Utf8, + }; + assert!(evaluate(&neg, &b).is_err()); + } + + #[test] + fn an_empty_batch_evaluates_to_an_empty_column() { + let b = batch(vec![("a", DataType::Int64, vec![])]); + let e = + BoundExpr::binary(col(0, DataType::Int64), BinaryOp::Gt, lit(Value::Int64(0))).unwrap(); + assert_eq!(evaluate(&e, &b).unwrap().len(), 0); + assert!(evaluate_predicate(&e, &b).unwrap().is_empty()); + } +} diff --git a/crates/qf-exec/src/join.rs b/crates/qf-exec/src/join.rs new file mode 100644 index 0000000..77c6773 --- /dev/null +++ b/crates/qf-exec/src/join.rs @@ -0,0 +1,576 @@ +//! Hash join. +//! +//! One side is collected and hashed; the other streams past it. Which side is +//! built is a real decision — the build side is what has to fit in memory — +//! and for an inner join the physical planner picks the smaller estimate. For +//! an outer join the choice is forced: the side that may be NULL-padded has to +//! be the one being built, so that unmatched rows can be found at the end. +//! +//! Output column order is always left-then-right, whichever side was built. + +use crate::eval::evaluate; +use crate::operator::{collect_one, Metrics, Operator}; +use qf_common::{Result, Schema, Value}; +use qf_plan::expr::BoundExpr; +use qf_sql::ast::JoinType; +use qf_storage::array::Array; +use qf_storage::batch::{RecordBatch, DEFAULT_BATCH_SIZE}; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Instant; + +/// Which input is loaded into the hash table. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BuildSide { + Left, + Right, +} + +pub struct HashJoinExec { + left: Box, + right: Box, + join_type: JoinType, + on: Vec<(BoundExpr, BoundExpr)>, + filter: Option, + build_side: BuildSide, + schema: Arc, + + built: Option, + /// Which build rows found a partner, for the outer-join tail. + matched: Vec, + tail_emitted: bool, + pending: Vec, + metrics: Metrics, +} + +struct BuiltSide { + batch: RecordBatch, + /// Key to the rows holding it. A vector because keys repeat. + index: HashMap, Vec>, +} + +impl HashJoinExec { + pub fn new( + left: Box, + right: Box, + join_type: JoinType, + on: Vec<(BoundExpr, BoundExpr)>, + filter: Option, + build_side: BuildSide, + ) -> HashJoinExec { + // An outer join has to be able to find the unmatched rows of the side + // it preserves the *other* way round, so the choice is not free. + let build_side = match join_type { + JoinType::Left => BuildSide::Right, + JoinType::Right => BuildSide::Left, + JoinType::Full => BuildSide::Right, + _ => build_side, + }; + let schema = Arc::new(Schema::join(&left.schema(), &right.schema())); + HashJoinExec { + left, + right, + join_type, + on, + filter, + build_side, + schema, + built: None, + matched: Vec::new(), + tail_emitted: false, + pending: Vec::new(), + metrics: Metrics::default(), + } + } + + pub fn build_side(&self) -> BuildSide { + self.build_side + } + + fn build(&mut self) -> Result<()> { + if self.built.is_some() { + return Ok(()); + } + let batch = match self.build_side { + BuildSide::Left => collect_one(self.left.as_mut())?, + BuildSide::Right => collect_one(self.right.as_mut())?, + }; + let keys: Vec<&BoundExpr> = self + .on + .iter() + .map(|(l, r)| match self.build_side { + BuildSide::Left => l, + BuildSide::Right => r, + }) + .collect(); + + let mut index: HashMap, Vec> = HashMap::new(); + if !keys.is_empty() { + let columns = keys + .iter() + .map(|e| evaluate(e, &batch)) + .collect::>>()?; + for row in 0..batch.num_rows() { + let key: Vec = columns.iter().map(|c| c.value(row)).collect(); + // A NULL key never joins — `NULL = NULL` is unknown, not true. + if key.iter().any(Value::is_null) { + continue; + } + index.entry(key).or_default().push(row); + } + } + self.matched = vec![false; batch.num_rows()]; + self.built = Some(BuiltSide { batch, index }); + Ok(()) + } + + fn probe_operator(&mut self) -> &mut dyn Operator { + match self.build_side { + BuildSide::Left => self.right.as_mut(), + BuildSide::Right => self.left.as_mut(), + } + } + + /// Does the probe side hold rows that must survive without a match? + fn probe_is_preserved(&self) -> bool { + match self.join_type { + JoinType::Left => self.build_side == BuildSide::Right, + JoinType::Right => self.build_side == BuildSide::Left, + JoinType::Full => true, + _ => false, + } + } + + fn build_is_preserved(&self) -> bool { + match self.join_type { + JoinType::Left => self.build_side == BuildSide::Left, + JoinType::Right => self.build_side == BuildSide::Right, + JoinType::Full => true, + _ => false, + } + } + + 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)); + } + } + } 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)); + } + } + _ => { + if self.probe_is_preserved() { + build_rows.push(None); + probe_rows.push(Some(p)); + } + } + } + } + } + + let out = self.assemble(&build_rows, &probe_rows, probe_batch)?; + + // 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)? + } + }; + + // 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; + } + } + } + Ok(out) + } + + /// Gathers the chosen rows from both sides into one batch, always with the + /// left input's columns first. + fn assemble( + &self, + build_rows: &[Option], + probe_rows: &[Option], + probe_batch: &RecordBatch, + ) -> Result { + let built = self.built.as_ref().expect("built before assembling"); + let build_part = gather(&built.batch, build_rows)?; + let probe_part = gather(probe_batch, probe_rows)?; + match self.build_side { + BuildSide::Left => RecordBatch::hstack(&build_part, &probe_part), + BuildSide::Right => RecordBatch::hstack(&probe_part, &build_part), + } + } + + /// The build-side rows that never matched, padded with NULLs. Emitted once + /// the probe side is exhausted. + fn tail(&mut self) -> Result> { + if self.tail_emitted || !self.build_is_preserved() { + return Ok(None); + } + self.tail_emitted = true; + let built = self.built.as_ref().expect("built before the tail"); + let rows: Vec> = (0..built.batch.num_rows()) + .filter(|i| !self.matched[*i]) + .map(Some) + .collect(); + if rows.is_empty() { + return Ok(None); + } + let build_part = gather(&built.batch, &rows)?; + let other_schema = match self.build_side { + BuildSide::Left => self.right.schema(), + BuildSide::Right => self.left.schema(), + }; + let nulls = null_batch(&other_schema, rows.len())?; + let out = match self.build_side { + BuildSide::Left => RecordBatch::hstack(&build_part, &nulls), + BuildSide::Right => RecordBatch::hstack(&nulls, &build_part), + }?; + Ok(Some(out)) + } +} + +/// Takes rows by index, turning `None` into a row of NULLs. +fn gather(batch: &RecordBatch, rows: &[Option]) -> Result { + if rows.iter().all(Option::is_some) { + let indices: Vec = rows.iter().map(|r| r.unwrap()).collect(); + return batch.take(&indices); + } + let schema = Arc::clone(batch.schema()); + let mut columns = Vec::with_capacity(batch.num_columns()); + for c in 0..batch.num_columns() { + let source = batch.column(c)?; + let mut b = qf_storage::array::ArrayBuilder::new(source.data_type()); + b.reserve(rows.len()); + for r in rows { + match r { + Some(i) => b.push_from(source, *i), + None => b.push_null(), + } + } + columns.push(b.finish()?); + } + if columns.is_empty() { + return RecordBatch::empty(schema)?.take(&vec![0usize; rows.len()]); + } + RecordBatch::try_new(schema, columns) +} + +fn null_batch(schema: &Arc, rows: usize) -> Result { + let columns = schema + .fields() + .iter() + .map(|f| Array::nulls(f.data_type, rows)) + .collect::>>()?; + if columns.is_empty() { + return RecordBatch::empty(Arc::clone(schema))?.take(&vec![0usize; rows]); + } + RecordBatch::try_new(Arc::clone(schema), columns) +} + +impl Operator for HashJoinExec { + fn schema(&self) -> Arc { + Arc::clone(&self.schema) + } + + fn next_batch(&mut self) -> Result> { + let start = Instant::now(); + self.build()?; + + if let Some(b) = self.pending.pop() { + self.metrics.elapsed += start.elapsed(); + self.metrics.record(&b); + return Ok(Some(b)); + } + + while let Some(probe_batch) = self.probe_operator().next_batch()? { + if probe_batch.num_rows() == 0 { + continue; + } + let out = self.probe(&probe_batch)?; + if out.num_rows() == 0 { + continue; + } + // A join can multiply rows, so cut the result back to batch size. + let mut chunks = out.chunks(DEFAULT_BATCH_SIZE)?; + let first = chunks.remove(0); + chunks.reverse(); + self.pending = chunks; + self.metrics.elapsed += start.elapsed(); + self.metrics.record(&first); + return Ok(Some(first)); + } + + let tail = self.tail()?; + self.metrics.elapsed += start.elapsed(); + if let Some(b) = &tail { + self.metrics.record(b); + } + Ok(tail) + } + + fn metrics(&self) -> &Metrics { + &self.metrics + } + + fn children(&self) -> Vec<&dyn Operator> { + vec![self.left.as_ref(), self.right.as_ref()] + } + + fn name(&self) -> &'static str { + "HashJoin" + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::operator::collect_one; + use crate::operators::ValuesExec; + use qf_common::{DataType, Field}; + use qf_storage::array::Array; + + fn table(name: &str, ids: &[Option]) -> (Arc, RecordBatch) { + let schema = Arc::new(Schema::new(vec![ + Field::new(format!("{name}_id"), DataType::Int64, true), + Field::new(format!("{name}_tag"), DataType::Utf8, false), + ])); + let id_values: Vec = ids + .iter() + .map(|i| i.map_or(Value::Null, Value::Int64)) + .collect(); + let tags: Vec = ids + .iter() + .enumerate() + .map(|(i, _)| Value::Utf8(format!("{name}{i}"))) + .collect(); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Array::from_values(DataType::Int64, &id_values).unwrap(), + Array::from_values(DataType::Utf8, &tags).unwrap(), + ], + ) + .unwrap(); + (schema, batch) + } + + fn operator(schema: Arc, batch: RecordBatch) -> Box { + Box::new(ValuesExec::new(schema, vec![batch])) + } + + fn key(index: usize) -> BoundExpr { + BoundExpr::column(index, "id", DataType::Int64) + } + + fn run(join_type: JoinType, build_side: BuildSide, on: bool) -> RecordBatch { + let (ls, lb) = table("l", &[Some(1), Some(2), None]); + let (rs, rb) = table("r", &[Some(2), Some(3)]); + let mut op = HashJoinExec::new( + operator(ls, lb), + operator(rs, rb), + join_type, + if on { vec![(key(0), key(0))] } else { vec![] }, + None, + build_side, + ); + collect_one(&mut op).unwrap() + } + + fn pairs(batch: &RecordBatch) -> Vec<(String, String)> { + batch + .rows() + .map(|r| (r[1].to_string(), r[3].to_string())) + .collect() + } + + #[test] + fn an_inner_join_gives_the_same_answer_whichever_side_is_built() { + let a = run(JoinType::Inner, BuildSide::Left, true); + let b = run(JoinType::Inner, BuildSide::Right, true); + assert_eq!(pairs(&a), vec![("l1".into(), "r0".into())]); + assert_eq!(pairs(&a), pairs(&b)); + // Column order is left-then-right regardless of the build side. + assert_eq!(a.schema().field(0).unwrap().name, "l_id"); + assert_eq!(b.schema().field(0).unwrap().name, "l_id"); + } + + #[test] + fn a_null_key_row_never_matches() { + let out = run(JoinType::Inner, BuildSide::Left, true); + assert!(!pairs(&out).iter().any(|(l, _)| l == "l2")); + } + + #[test] + fn an_outer_join_forces_the_build_side_regardless_of_what_was_asked_for() { + // A LEFT join has to be able to pad its right side, so it builds right + // even when the planner suggested otherwise. + let (ls, lb) = table("l", &[Some(1)]); + let (rs, rb) = table("r", &[Some(1)]); + let op = HashJoinExec::new( + operator(ls, lb), + operator(rs, rb), + JoinType::Left, + vec![(key(0), key(0))], + None, + BuildSide::Left, + ); + assert_eq!(op.build_side(), BuildSide::Right); + } + + #[test] + fn a_left_join_pads_unmatched_left_rows() { + let out = run(JoinType::Left, BuildSide::Right, true); + let p = pairs(&out); + assert_eq!(p.len(), 3); + assert!(p.contains(&("l0".into(), "NULL".into()))); + assert!(p.contains(&("l2".into(), "NULL".into()))); + } + + #[test] + fn a_right_join_pads_unmatched_right_rows() { + let out = run(JoinType::Right, BuildSide::Left, true); + let p = pairs(&out); + assert_eq!(p.len(), 2); + assert!(p.contains(&("NULL".into(), "r1".into()))); + } + + #[test] + fn a_full_join_pads_both_sides() { + let out = run(JoinType::Full, BuildSide::Right, true); + let p = pairs(&out); + assert_eq!(p.len(), 4); + assert!(p.contains(&("l0".into(), "NULL".into()))); + assert!(p.contains(&("NULL".into(), "r1".into()))); + assert!(p.contains(&("l1".into(), "r0".into()))); + } + + #[test] + fn a_cross_join_pairs_every_row_with_every_other() { + let out = run(JoinType::Cross, BuildSide::Left, false); + assert_eq!(out.num_rows(), 6); + assert_eq!(out.num_columns(), 4); + } + + #[test] + fn a_residual_filter_removes_rows_after_matching() { + let (ls, lb) = table("l", &[Some(1), Some(2)]); + let (rs, rb) = table("r", &[Some(1), Some(2)]); + let filter = BoundExpr::binary( + BoundExpr::column(0, "l_id", DataType::Int64), + qf_sql::ast::BinaryOp::Gt, + BoundExpr::Literal(Value::Int64(1)), + ) + .unwrap(); + let mut op = HashJoinExec::new( + operator(ls, lb), + operator(rs, rb), + JoinType::Inner, + vec![(key(0), key(0))], + Some(filter), + BuildSide::Right, + ); + let out = collect_one(&mut op).unwrap(); + assert_eq!(out.num_rows(), 1); + assert_eq!(out.row(0)[1], Value::Utf8("l1".into())); + } + + #[test] + fn a_join_reports_its_children_and_name() { + let (ls, lb) = table("l", &[Some(1)]); + let (rs, rb) = table("r", &[Some(1)]); + let op = HashJoinExec::new( + operator(ls, lb), + operator(rs, rb), + JoinType::Inner, + vec![], + None, + BuildSide::Left, + ); + assert_eq!(op.name(), "HashJoin"); + assert_eq!(op.children().len(), 2); + assert_eq!(op.schema().len(), 4); + } + + #[test] + fn joining_against_an_empty_side_yields_nothing_or_padding() { + let (ls, lb) = table("l", &[Some(1)]); + let (rs, _) = table("r", &[]); + let empty = RecordBatch::empty(Arc::clone(&rs)).unwrap(); + let mut inner = HashJoinExec::new( + operator(Arc::clone(&ls), lb.clone()), + operator(Arc::clone(&rs), empty.clone()), + JoinType::Inner, + vec![(key(0), key(0))], + None, + BuildSide::Right, + ); + assert_eq!(collect_one(&mut inner).unwrap().num_rows(), 0); + + let mut left = HashJoinExec::new( + operator(ls, lb), + operator(rs, empty), + JoinType::Left, + vec![(key(0), key(0))], + None, + BuildSide::Right, + ); + assert_eq!(collect_one(&mut left).unwrap().num_rows(), 1); + } + + #[test] + fn a_repeated_key_on_the_build_side_produces_a_row_per_match() { + let (ls, lb) = table("l", &[Some(7)]); + let (rs, rb) = table("r", &[Some(7), Some(7), Some(8)]); + let mut op = HashJoinExec::new( + operator(ls, lb), + operator(rs, rb), + JoinType::Inner, + vec![(key(0), key(0))], + None, + BuildSide::Right, + ); + assert_eq!(collect_one(&mut op).unwrap().num_rows(), 2); + } +} diff --git a/crates/qf-exec/src/lib.rs b/crates/qf-exec/src/lib.rs index ff7bd09..739cba9 100644 --- a/crates/qf-exec/src/lib.rs +++ b/crates/qf-exec/src/lib.rs @@ -1 +1,11 @@ -// placeholder +pub mod aggregate; +pub mod engine; +pub mod eval; +pub mod join; +pub mod operator; +pub mod operators; +pub mod physical; +pub mod scan; +pub mod sort; + +pub use engine::{Output, Session}; diff --git a/crates/qf-exec/src/operator.rs b/crates/qf-exec/src/operator.rs new file mode 100644 index 0000000..27c2dee --- /dev/null +++ b/crates/qf-exec/src/operator.rs @@ -0,0 +1,209 @@ +//! The operator interface and the counters `EXPLAIN ANALYZE` reports. +//! +//! Operators pull: each asks its input for the next batch and returns a batch +//! of its own, or `None` when it is finished. Pulling is what lets `LIMIT` +//! stop a scan early — a push model would have to run the scan to completion +//! and throw the extra rows away. +//! +//! Some operators cannot stream. A sort has to see every row before it can +//! emit the first one, and a hash join has to finish building before it can +//! probe. Those buffer internally and stream their *output*, so the operator +//! above them never has to know the difference. + +use qf_common::{Result, Schema}; +use qf_storage::batch::RecordBatch; +use std::sync::Arc; +use std::time::Duration; + +/// What one operator did, for `EXPLAIN ANALYZE`. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct Metrics { + pub rows_out: usize, + pub batches_out: usize, + pub elapsed: Duration, + /// Row groups actually read. Scans only. + pub row_groups_read: usize, + /// Row groups skipped because a zone map ruled the predicate out. This is + /// the number that shows whether pushdown did anything. + pub row_groups_pruned: usize, + /// Runs written to disk by an external sort. + pub spilled_runs: usize, + pub spilled_rows: usize, +} + +impl Metrics { + pub fn record(&mut self, batch: &RecordBatch) { + self.rows_out += batch.num_rows(); + self.batches_out += 1; + } + + /// The parenthesised suffix `EXPLAIN ANALYZE` appends to a plan line. + pub fn describe(&self) -> String { + let mut parts = vec![ + format!("rows={}", self.rows_out), + format!("time={:.2}ms", self.elapsed.as_secs_f64() * 1000.0), + ]; + if self.row_groups_read + self.row_groups_pruned > 0 { + parts.push(format!( + "row groups={} read, {} pruned", + self.row_groups_read, self.row_groups_pruned + )); + } + if self.spilled_runs > 0 { + parts.push(format!( + "spilled={} runs, {} rows", + self.spilled_runs, self.spilled_rows + )); + } + parts.join(", ") + } +} + +/// A node in the running query. +pub trait Operator { + fn schema(&self) -> Arc; + + /// The next batch, or `None` when the operator is exhausted. Returning an + /// empty batch is allowed and means "nothing this time, ask again" — a + /// filter that rejected everything in one batch is not finished. + fn next_batch(&mut self) -> Result>; + + fn metrics(&self) -> &Metrics; + + /// This operator's children, so the driver can walk the tree for + /// `EXPLAIN ANALYZE`. + fn children(&self) -> Vec<&dyn Operator> { + vec![] + } + + /// A short label for `EXPLAIN ANALYZE`. + fn name(&self) -> &'static str; +} + +/// Drains an operator into batches. Used by the CLI and by the tests; the +/// operators themselves never call it. +pub fn collect(op: &mut dyn Operator) -> Result> { + let mut out = Vec::new(); + while let Some(b) = op.next_batch()? { + if b.num_rows() > 0 { + out.push(b); + } + } + Ok(out) +} + +/// Drains an operator into a single batch. +pub fn collect_one(op: &mut dyn Operator) -> Result { + let schema = op.schema(); + let batches = collect(op)?; + RecordBatch::concat(schema, &batches) +} + +#[cfg(test)] +mod tests { + use super::*; + use qf_common::{DataType, Field, Value}; + use qf_storage::array::Array; + + struct Fixed { + schema: Arc, + batches: Vec, + metrics: Metrics, + } + + impl Operator for Fixed { + fn schema(&self) -> Arc { + Arc::clone(&self.schema) + } + fn next_batch(&mut self) -> Result> { + if self.batches.is_empty() { + return Ok(None); + } + let b = self.batches.remove(0); + self.metrics.record(&b); + Ok(Some(b)) + } + fn metrics(&self) -> &Metrics { + &self.metrics + } + fn name(&self) -> &'static str { + "Fixed" + } + } + + fn fixed(values: &[&[i64]]) -> Fixed { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, true)])); + let batches = values + .iter() + .map(|vs| { + let vals: Vec = vs.iter().map(|v| Value::Int64(*v)).collect(); + RecordBatch::try_new( + Arc::clone(&schema), + vec![Array::from_values(DataType::Int64, &vals).unwrap()], + ) + .unwrap() + }) + .collect(); + Fixed { + schema, + batches, + metrics: Metrics::default(), + } + } + + #[test] + fn collecting_drains_every_batch_and_drops_empty_ones() { + let mut op = fixed(&[&[1, 2], &[], &[3]]); + let batches = collect(&mut op).unwrap(); + assert_eq!(batches.len(), 2); + assert_eq!(op.metrics().rows_out, 3); + assert_eq!(op.metrics().batches_out, 3); + } + + #[test] + fn collecting_into_one_batch_concatenates_in_order() { + let mut op = fixed(&[&[1, 2], &[3]]); + let b = collect_one(&mut op).unwrap(); + assert_eq!(b.num_rows(), 3); + assert_eq!(b.column(0).unwrap().value(2), Value::Int64(3)); + } + + #[test] + fn an_exhausted_operator_keeps_returning_none() { + let mut op = fixed(&[&[1]]); + assert!(op.next_batch().unwrap().is_some()); + assert!(op.next_batch().unwrap().is_none()); + assert!(op.next_batch().unwrap().is_none()); + assert!(op.children().is_empty()); + assert_eq!(op.name(), "Fixed"); + } + + #[test] + fn metrics_render_only_the_counters_that_apply() { + let plain = Metrics { + rows_out: 10, + batches_out: 1, + ..Metrics::default() + }; + let text = plain.describe(); + assert!(text.contains("rows=10")); + assert!(text.contains("time=")); + assert!(!text.contains("row groups")); + assert!(!text.contains("spilled")); + + let scan = Metrics { + rows_out: 5, + row_groups_read: 2, + row_groups_pruned: 8, + ..Metrics::default() + }; + assert!(scan.describe().contains("2 read, 8 pruned")); + + let sort = Metrics { + spilled_runs: 3, + spilled_rows: 300, + ..Metrics::default() + }; + assert!(sort.describe().contains("3 runs, 300 rows")); + } +} diff --git a/crates/qf-exec/src/operators.rs b/crates/qf-exec/src/operators.rs new file mode 100644 index 0000000..b0867ab --- /dev/null +++ b/crates/qf-exec/src/operators.rs @@ -0,0 +1,485 @@ +//! The streaming operators: filter, projection, limit, distinct and a literal +//! source. None of them buffer, so a `LIMIT 10` over a billion-row table stops +//! the scan after the first batch or two. + +use crate::eval::{evaluate, evaluate_predicate}; +use crate::operator::{Metrics, Operator}; +use qf_common::{Field, Result, Schema, Value}; +use qf_plan::expr::BoundExpr; +use qf_storage::batch::RecordBatch; +use std::collections::HashSet; +use std::sync::Arc; +use std::time::Instant; + +pub struct FilterExec { + input: Box, + predicate: BoundExpr, + metrics: Metrics, +} + +impl FilterExec { + pub fn new(input: Box, predicate: BoundExpr) -> FilterExec { + FilterExec { + input, + predicate, + metrics: Metrics::default(), + } + } +} + +impl Operator for FilterExec { + fn schema(&self) -> Arc { + self.input.schema() + } + + fn next_batch(&mut self) -> Result> { + let Some(batch) = self.input.next_batch()? else { + return Ok(None); + }; + let start = Instant::now(); + let mask = evaluate_predicate(&self.predicate, &batch)?; + let out = batch.filter(&mask)?; + self.metrics.elapsed += start.elapsed(); + self.metrics.record(&out); + Ok(Some(out)) + } + + fn metrics(&self) -> &Metrics { + &self.metrics + } + + fn children(&self) -> Vec<&dyn Operator> { + vec![self.input.as_ref()] + } + + fn name(&self) -> &'static str { + "Filter" + } +} + +pub struct ProjectionExec { + input: Box, + exprs: Vec<(BoundExpr, String)>, + schema: Arc, + metrics: Metrics, +} + +impl ProjectionExec { + pub fn new(input: Box, exprs: Vec<(BoundExpr, String)>) -> ProjectionExec { + let schema = Arc::new(Schema::new( + exprs + .iter() + .map(|(e, n)| Field::new(n, e.data_type(), true)) + .collect(), + )); + ProjectionExec { + input, + exprs, + schema, + metrics: Metrics::default(), + } + } +} + +impl Operator for ProjectionExec { + fn schema(&self) -> Arc { + Arc::clone(&self.schema) + } + + fn next_batch(&mut self) -> Result> { + let Some(batch) = self.input.next_batch()? else { + return Ok(None); + }; + let start = Instant::now(); + let columns = self + .exprs + .iter() + .map(|(e, _)| evaluate(e, &batch)) + .collect::>>()?; + // A projection that produces no columns still has to carry the row + // count, or `SELECT count(*)` above it counts nothing. + let out = if columns.is_empty() { + RecordBatch::empty(Arc::clone(&self.schema))?.take(&vec![0usize; batch.num_rows()])? + } else { + RecordBatch::try_new(Arc::clone(&self.schema), columns)? + }; + self.metrics.elapsed += start.elapsed(); + self.metrics.record(&out); + Ok(Some(out)) + } + + fn metrics(&self) -> &Metrics { + &self.metrics + } + + fn children(&self) -> Vec<&dyn Operator> { + vec![self.input.as_ref()] + } + + fn name(&self) -> &'static str { + "Projection" + } +} + +/// `LIMIT n OFFSET k`, applied while streaming. +pub struct LimitExec { + input: Box, + limit: Option, + offset: usize, + skipped: usize, + emitted: usize, + metrics: Metrics, +} + +impl LimitExec { + pub fn new(input: Box, limit: Option, offset: usize) -> LimitExec { + LimitExec { + input, + limit, + offset, + skipped: 0, + emitted: 0, + metrics: Metrics::default(), + } + } + + fn finished(&self) -> bool { + self.limit.is_some_and(|l| self.emitted >= l) + } +} + +impl Operator for LimitExec { + fn schema(&self) -> Arc { + self.input.schema() + } + + fn next_batch(&mut self) -> Result> { + loop { + if self.finished() { + // Stop pulling. This is the whole point of a pull model: the + // scan below never reads the rest of the table. + return Ok(None); + } + let Some(batch) = self.input.next_batch()? else { + return Ok(None); + }; + let rows = batch.num_rows(); + if rows == 0 { + continue; + } + let start = Instant::now(); + + let skip = self.offset.saturating_sub(self.skipped).min(rows); + self.skipped += skip; + let available = rows - skip; + if available == 0 { + self.metrics.elapsed += start.elapsed(); + continue; + } + let take = match self.limit { + Some(l) => available.min(l - self.emitted), + None => available, + }; + let indices: Vec = (skip..skip + take).collect(); + let out = batch.take(&indices)?; + self.emitted += take; + self.metrics.elapsed += start.elapsed(); + self.metrics.record(&out); + return Ok(Some(out)); + } + } + + fn metrics(&self) -> &Metrics { + &self.metrics + } + + fn children(&self) -> Vec<&dyn Operator> { + vec![self.input.as_ref()] + } + + fn name(&self) -> &'static str { + "Limit" + } +} + +/// `SELECT DISTINCT`, streaming and order-preserving. +/// +/// Keeping the first occurrence rather than sorting means a `DISTINCT` under +/// an `ORDER BY` does not undo the sort, and rows come out as soon as they are +/// seen instead of after the input is exhausted. +pub struct DistinctExec { + input: Box, + seen: HashSet>, + metrics: Metrics, +} + +impl DistinctExec { + pub fn new(input: Box) -> DistinctExec { + DistinctExec { + input, + seen: HashSet::new(), + metrics: Metrics::default(), + } + } +} + +impl Operator for DistinctExec { + fn schema(&self) -> Arc { + self.input.schema() + } + + fn next_batch(&mut self) -> Result> { + let Some(batch) = self.input.next_batch()? else { + return Ok(None); + }; + let start = Instant::now(); + let mut keep = Vec::new(); + for i in 0..batch.num_rows() { + let key = batch.row(i); + if self.seen.insert(key) { + keep.push(i); + } + } + let out = batch.take(&keep)?; + self.metrics.elapsed += start.elapsed(); + self.metrics.record(&out); + Ok(Some(out)) + } + + fn metrics(&self) -> &Metrics { + &self.metrics + } + + fn children(&self) -> Vec<&dyn Operator> { + vec![self.input.as_ref()] + } + + fn name(&self) -> &'static str { + "Distinct" + } +} + +/// Literal rows: `INSERT ... VALUES`, and the single empty row a `SELECT` with +/// no `FROM` runs over. +pub struct ValuesExec { + schema: Arc, + batches: Vec, + next: usize, + metrics: Metrics, +} + +impl ValuesExec { + pub fn new(schema: Arc, batches: Vec) -> ValuesExec { + ValuesExec { + schema, + batches, + next: 0, + metrics: Metrics::default(), + } + } + + /// One row with no columns — what `SELECT 1 + 1` evaluates against. + pub fn single_row() -> Result { + let schema = Arc::new(Schema::empty()); + let batch = RecordBatch::empty(Arc::clone(&schema))?.take(&[0usize; 1])?; + Ok(ValuesExec::new(schema, vec![batch])) + } +} + +impl Operator for ValuesExec { + fn schema(&self) -> Arc { + Arc::clone(&self.schema) + } + + fn next_batch(&mut self) -> Result> { + if self.next >= self.batches.len() { + return Ok(None); + } + let b = self.batches[self.next].clone(); + self.next += 1; + self.metrics.record(&b); + Ok(Some(b)) + } + + fn metrics(&self) -> &Metrics { + &self.metrics + } + + fn name(&self) -> &'static str { + "Values" + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::operator::{collect_one, Operator}; + use qf_common::DataType; + use qf_sql::ast::BinaryOp; + use qf_storage::array::Array; + + fn schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Utf8, false), + ])) + } + + fn source(values: &[&[i64]]) -> Box { + let batches = values + .iter() + .map(|vs| { + let a: Vec = vs.iter().map(|v| Value::Int64(*v)).collect(); + let b: Vec = vs.iter().map(|v| Value::Utf8(format!("v{v}"))).collect(); + RecordBatch::try_new( + schema(), + vec![ + Array::from_values(DataType::Int64, &a).unwrap(), + Array::from_values(DataType::Utf8, &b).unwrap(), + ], + ) + .unwrap() + }) + .collect(); + Box::new(ValuesExec::new(schema(), batches)) + } + + fn column(op: &mut dyn Operator, index: usize) -> Vec { + collect_one(op) + .unwrap() + .column(index) + .unwrap() + .iter() + .map(|v| v.to_string()) + .collect() + } + + fn col(i: usize, t: DataType) -> BoundExpr { + BoundExpr::column(i, format!("c{i}"), t) + } + + #[test] + fn a_filter_keeps_matching_rows_across_batches() { + let mut op = FilterExec::new( + source(&[&[1, 2, 3], &[4, 5]]), + BoundExpr::binary( + col(0, DataType::Int64), + BinaryOp::Gt, + BoundExpr::Literal(Value::Int64(2)), + ) + .unwrap(), + ); + assert_eq!(column(&mut op, 0), vec!["3", "4", "5"]); + assert_eq!(op.name(), "Filter"); + assert_eq!(op.children().len(), 1); + assert_eq!(op.metrics().rows_out, 3); + } + + #[test] + fn a_filter_that_rejects_a_whole_batch_keeps_going() { + // An empty batch is not the end of the input. + let mut op = FilterExec::new( + source(&[&[1], &[9]]), + BoundExpr::binary( + col(0, DataType::Int64), + BinaryOp::Gt, + BoundExpr::Literal(Value::Int64(5)), + ) + .unwrap(), + ); + assert_eq!(column(&mut op, 0), vec!["9"]); + } + + #[test] + fn a_projection_computes_its_expressions_and_names_its_columns() { + let mut op = ProjectionExec::new( + source(&[&[1, 2]]), + vec![( + BoundExpr::binary( + col(0, DataType::Int64), + BinaryOp::Multiply, + BoundExpr::Literal(Value::Int64(10)), + ) + .unwrap(), + "tens".into(), + )], + ); + assert_eq!(op.schema().field(0).unwrap().name, "tens"); + assert_eq!(column(&mut op, 0), vec!["10", "20"]); + assert_eq!(op.name(), "Projection"); + assert_eq!(op.children().len(), 1); + } + + #[test] + fn a_projection_with_no_columns_still_carries_the_row_count() { + let mut op = ProjectionExec::new(source(&[&[1, 2, 3]]), vec![]); + let out = collect_one(&mut op).unwrap(); + assert_eq!(out.num_columns(), 0); + assert_eq!(out.num_rows(), 3); + } + + #[test] + fn a_limit_takes_the_first_rows_and_stops_pulling() { + let mut op = LimitExec::new(source(&[&[1, 2, 3], &[4, 5, 6]]), Some(2), 0); + assert_eq!(column(&mut op, 0), vec!["1", "2"]); + assert_eq!(op.name(), "Limit"); + assert_eq!(op.children().len(), 1); + } + + #[test] + fn an_offset_can_span_several_batches() { + let mut op = LimitExec::new(source(&[&[1, 2], &[3, 4], &[5, 6]]), Some(2), 3); + assert_eq!(column(&mut op, 0), vec!["4", "5"]); + } + + #[test] + fn an_offset_past_the_end_returns_nothing() { + let mut op = LimitExec::new(source(&[&[1, 2]]), None, 10); + assert_eq!(collect_one(&mut op).unwrap().num_rows(), 0); + } + + #[test] + fn a_limit_of_zero_returns_nothing_and_a_limit_without_one_returns_everything() { + let mut zero = LimitExec::new(source(&[&[1, 2]]), Some(0), 0); + assert_eq!(collect_one(&mut zero).unwrap().num_rows(), 0); + let mut all = LimitExec::new(source(&[&[1, 2], &[3]]), None, 0); + assert_eq!(column(&mut all, 0), vec!["1", "2", "3"]); + } + + #[test] + fn a_limit_larger_than_the_input_returns_everything() { + let mut op = LimitExec::new(source(&[&[1, 2]]), Some(50), 0); + assert_eq!(column(&mut op, 0), vec!["1", "2"]); + } + + #[test] + fn distinct_removes_repeats_across_batch_boundaries_and_keeps_first_order() { + let mut op = DistinctExec::new(source(&[&[3, 1, 3], &[1, 2]])); + assert_eq!(column(&mut op, 0), vec!["3", "1", "2"]); + assert_eq!(op.name(), "Distinct"); + assert_eq!(op.children().len(), 1); + } + + #[test] + fn values_replays_its_batches_once() { + let mut op = ValuesExec::new(schema(), vec![]); + assert!(op.next_batch().unwrap().is_none()); + assert_eq!(op.name(), "Values"); + assert!(op.children().is_empty()); + assert_eq!(op.schema().len(), 2); + } + + #[test] + fn the_single_empty_row_has_no_columns_but_one_row() { + let mut op = ValuesExec::single_row().unwrap(); + let b = collect_one(&mut op).unwrap(); + assert_eq!(b.num_rows(), 1); + assert_eq!(b.num_columns(), 0); + } + + #[test] + fn empty_input_batches_are_skipped_by_the_limit() { + let mut op = LimitExec::new(source(&[&[], &[1, 2]]), Some(1), 0); + assert_eq!(column(&mut op, 0), vec!["1"]); + } +} diff --git a/crates/qf-exec/src/physical.rs b/crates/qf-exec/src/physical.rs new file mode 100644 index 0000000..9609b7f --- /dev/null +++ b/crates/qf-exec/src/physical.rs @@ -0,0 +1,432 @@ +//! Turning a logical plan into running operators. +//! +//! The logical plan says *what*; this decides *how*. Two choices are made here +//! rather than in the optimiser, because both depend on physical facts the +//! logical plan does not model: +//! +//! * **Which side of a join to build.** The build side is what has to fit in +//! memory, so the smaller estimate wins — except on an outer join, where the +//! preserved side has to be the one that can be scanned for non-matches. +//! * **Whether a sort can be a top-k.** A sort feeding a `LIMIT` never needs to +//! hold more than `limit + offset` rows, so the limit is pushed into the sort +//! as a fetch hint. Nothing about the logical plan changes; the sort just +//! stops keeping rows it can prove it will throw away. + +use crate::aggregate::HashAggregateExec; +use crate::eval::evaluate; +use crate::join::{BuildSide, HashJoinExec}; +use crate::operator::Operator; +use crate::operators::{DistinctExec, FilterExec, LimitExec, ProjectionExec, ValuesExec}; +use crate::scan::ScanExec; +use crate::sort::SortExec; +use qf_common::{Error, Result, Schema}; +use qf_plan::cost::estimate_rows; +use qf_plan::logical::LogicalPlan; +use qf_sql::ast::JoinType; +use qf_storage::array::ArrayBuilder; +use qf_storage::batch::RecordBatch; +use qf_storage::catalog::{Catalog, TableSource}; +use std::sync::Arc; + +/// Builds the operator tree for a plan. +pub fn build(plan: &LogicalPlan, catalog: &Catalog) -> Result> { + build_with_fetch(plan, catalog, None) +} + +fn build_with_fetch( + plan: &LogicalPlan, + catalog: &Catalog, + fetch: Option, +) -> Result> { + Ok(match plan { + LogicalPlan::Scan { + table, + source_schema, + projection, + pushed_filters, + .. + } => { + let entry = catalog.get(table)?; + let indices = projection + .clone() + .unwrap_or_else(|| (0..source_schema.len()).collect()); + let schema = Arc::new(source_schema.project(&indices)?); + match &entry.source { + TableSource::Memory(batches) => Box::new(ScanExec::from_batches( + batches.clone(), + indices, + pushed_filters.clone(), + schema, + )), + TableSource::File(path) => Box::new(ScanExec::from_file( + path, + indices, + pushed_filters.clone(), + schema, + )?), + } + } + + LogicalPlan::Filter { input, predicate } => Box::new(FilterExec::new( + build_with_fetch(input, catalog, None)?, + predicate.clone(), + )), + + // A projection neither adds nor removes rows, so a fetch hint from a + // limit above it passes straight through to whatever is below. + LogicalPlan::Projection { input, exprs } => Box::new(ProjectionExec::new( + build_with_fetch(input, catalog, fetch)?, + exprs.clone(), + )), + + LogicalPlan::Aggregate { + input, + group_exprs, + aggregates, + } => Box::new(HashAggregateExec::new( + build_with_fetch(input, catalog, None)?, + group_exprs.clone(), + aggregates.clone(), + )), + + LogicalPlan::Join { + left, + right, + join_type, + on, + filter, + } => { + let build_side = choose_build_side(left, right, *join_type); + Box::new(HashJoinExec::new( + build_with_fetch(left, catalog, None)?, + build_with_fetch(right, catalog, None)?, + *join_type, + on.clone(), + filter.clone(), + build_side, + )) + } + + LogicalPlan::Sort { input, exprs } => Box::new(SortExec::new( + build_with_fetch(input, catalog, None)?, + exprs.clone(), + fetch, + )), + + LogicalPlan::Limit { + input, + limit, + offset, + } => { + // A sort below only ever needs the rows this limit can reach. + let hint = limit.map(|l| l.saturating_add(*offset)); + Box::new(LimitExec::new( + build_with_fetch(input, catalog, hint)?, + *limit, + *offset, + )) + } + + LogicalPlan::Distinct { input } => { + Box::new(DistinctExec::new(build_with_fetch(input, catalog, None)?)) + } + + LogicalPlan::Values { schema, rows } => Box::new(ValuesExec::new( + Arc::clone(schema), + values_batches(schema, rows)?, + )), + }) +} + +/// The smaller input becomes the hash table. On an outer join the choice is +/// forced by which side has to survive without a match, so this only decides +/// the inner and cross cases. +fn choose_build_side(left: &LogicalPlan, right: &LogicalPlan, join_type: JoinType) -> BuildSide { + match join_type { + JoinType::Left | JoinType::Full => BuildSide::Right, + JoinType::Right => BuildSide::Left, + _ => { + if estimate_rows(left) <= estimate_rows(right) { + BuildSide::Left + } else { + BuildSide::Right + } + } + } +} + +/// Evaluates a `VALUES` list into a batch. The expressions are constant, so +/// they are evaluated against a single empty row. +fn values_batches( + schema: &Arc, + rows: &[Vec], +) -> Result> { + if schema.is_empty() { + // The single empty row a `SELECT` with no `FROM` runs against. + let batch = RecordBatch::empty(Arc::clone(schema))?.take(&vec![0usize; rows.len()])?; + return Ok(vec![batch]); + } + let seed = RecordBatch::empty(Arc::new(Schema::empty()))?.take(&[0usize; 1])?; + let mut builders: Vec = schema + .fields() + .iter() + .map(|f| ArrayBuilder::new(f.data_type)) + .collect(); + for row in rows { + if row.len() != schema.len() { + return Err(Error::plan(format!( + "a row has {} values but the table has {} columns", + row.len(), + schema.len() + ))); + } + for (i, e) in row.iter().enumerate() { + let column = evaluate(e, &seed)?; + builders[i].push(column.value(0))?; + } + } + let columns = builders + .into_iter() + .map(ArrayBuilder::finish) + .collect::>>()?; + Ok(vec![RecordBatch::try_new(Arc::clone(schema), columns)?]) +} + +/// Renders the running operator tree with its counters, for `EXPLAIN ANALYZE`. +pub fn explain_analyze(op: &dyn Operator) -> String { + fn walk(op: &dyn Operator, depth: usize, out: &mut String) { + for _ in 0..depth { + out.push_str(" "); + } + out.push_str(op.name()); + out.push_str(" ("); + out.push_str(&op.metrics().describe()); + out.push_str(")\n"); + for c in op.children() { + walk(c, depth + 1, out); + } + } + let mut out = String::new(); + walk(op, 0, &mut out); + out +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::operator::collect_one; + use qf_common::{DataType, Field, Value}; + use qf_plan::binder::bind; + use qf_plan::expr::BoundExpr; + use qf_plan::optimizer::optimize; + use qf_storage::array::Array; + use qf_storage::stats::{ColumnStats, TableStats}; + + fn catalog() -> Catalog { + let mut c = Catalog::new(); + add(&mut c, "big", 100_000); + add(&mut c, "small", 10); + c + } + + fn add(c: &mut Catalog, name: &str, rows: usize) { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("v", DataType::Int64, true), + ])); + let values: Vec = (0..rows.min(10) as i64).map(Value::Int64).collect(); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Array::from_values(DataType::Int64, &values).unwrap(), + Array::from_values(DataType::Int64, &values).unwrap(), + ], + ) + .unwrap(); + let stats = TableStats { + row_count: rows, + columns: vec![ + ColumnStats { + min: Value::Int64(0), + max: Value::Int64(rows as i64), + null_count: 0, + row_count: rows, + distinct_count: rows, + }; + 2 + ], + }; + c.replace(qf_storage::catalog::Table { + name: name.to_string(), + schema, + source: qf_storage::catalog::TableSource::Memory(vec![batch]), + stats, + }); + } + + fn plan(sql: &str, c: &Catalog) -> LogicalPlan { + match qf_sql::parse(sql).unwrap() { + qf_sql::ast::Statement::Query(q) => optimize(bind(&q, c).unwrap()).unwrap(), + other => panic!("expected a query, got {other:?}"), + } + } + + fn find_join(plan: &LogicalPlan) -> (&LogicalPlan, &LogicalPlan, JoinType) { + fn walk(p: &LogicalPlan) -> Option<(&LogicalPlan, &LogicalPlan, JoinType)> { + if let LogicalPlan::Join { + left, + right, + join_type, + .. + } = p + { + return Some((left, right, *join_type)); + } + p.children().into_iter().find_map(walk) + } + walk(plan).expect("a join") + } + + #[test] + fn the_smaller_input_becomes_the_build_side() { + let c = catalog(); + let p = plan("SELECT b.id FROM big b JOIN small s ON b.id = s.id", &c); + let (l, r, t) = find_join(&p); + assert_eq!( + choose_build_side(l, r, t), + BuildSide::Right, + "small is right" + ); + + let p = plan("SELECT b.id FROM small s JOIN big b ON s.id = b.id", &c); + let (l, r, t) = find_join(&p); + assert_eq!(choose_build_side(l, r, t), BuildSide::Left, "small is left"); + } + + #[test] + fn an_outer_join_s_build_side_is_forced_by_the_side_it_preserves() { + let c = catalog(); + for (sql, want) in [ + ( + "SELECT b.id FROM big b LEFT JOIN small s ON b.id = s.id", + BuildSide::Right, + ), + ( + "SELECT b.id FROM big b RIGHT JOIN small s ON b.id = s.id", + BuildSide::Left, + ), + ( + "SELECT b.id FROM big b FULL JOIN small s ON b.id = s.id", + BuildSide::Right, + ), + ] { + let p = plan(sql, &c); + let (l, r, t) = find_join(&p); + assert_eq!(choose_build_side(l, r, t), want, "{sql}"); + } + } + + #[test] + fn a_limit_pushes_a_fetch_hint_through_a_projection_into_the_sort() { + // The sort must see the limit, or `ORDER BY ... LIMIT 5` would hold + // the whole table. + let c = catalog(); + let p = plan("SELECT id FROM big ORDER BY id LIMIT 5 OFFSET 2", &c); + let mut op = build(&p, &c).unwrap(); + let out = collect_one(op.as_mut()).unwrap(); + assert!(out.num_rows() <= 5); + let text = explain_analyze(op.as_ref()); + assert!(text.contains("Sort"), "{text}"); + // The sort emitted at most limit + offset rows. + let sort_line = text + .lines() + .find(|l| l.trim_start().starts_with("Sort")) + .unwrap(); + assert!( + sort_line.contains("rows=7") || sort_line.contains("rows=8"), + "{sort_line}" + ); + } + + #[test] + fn a_values_plan_evaluates_its_constant_expressions() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Utf8, false), + ])); + let rows = vec![vec![ + BoundExpr::Literal(Value::Int64(1)), + BoundExpr::Literal(Value::Utf8("x".into())), + ]]; + let batches = values_batches(&schema, &rows).unwrap(); + assert_eq!(batches[0].num_rows(), 1); + assert_eq!(batches[0].row(0)[1], Value::Utf8("x".into())); + } + + #[test] + fn a_values_row_of_the_wrong_width_is_refused() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let rows = vec![vec![ + BoundExpr::Literal(Value::Int64(1)), + BoundExpr::Literal(Value::Int64(2)), + ]]; + assert!(values_batches(&schema, &rows) + .unwrap_err() + .to_string() + .contains("but the table has")); + } + + #[test] + fn a_schemaless_values_plan_produces_bare_rows() { + let schema = Arc::new(Schema::empty()); + let batches = values_batches(&schema, &[vec![]]).unwrap(); + assert_eq!(batches[0].num_rows(), 1); + assert_eq!(batches[0].num_columns(), 0); + } + + #[test] + fn every_logical_node_has_a_physical_counterpart() { + let c = catalog(); + let sqls = [ + "SELECT id FROM big", + "SELECT id FROM big WHERE id > 1", + "SELECT DISTINCT id FROM big", + "SELECT id FROM big ORDER BY id", + "SELECT id FROM big LIMIT 1", + "SELECT count(*) FROM big", + "SELECT b.id FROM big b JOIN small s ON b.id = s.id", + "SELECT 1", + ]; + for sql in sqls { + let p = plan(sql, &c); + let mut op = build(&p, &c).unwrap(); + assert!(collect_one(op.as_mut()).is_ok(), "{sql}"); + assert!(!explain_analyze(op.as_ref()).is_empty(), "{sql}"); + } + } + + #[test] + fn explain_analyze_indents_children_under_their_parent() { + let c = catalog(); + let p = plan("SELECT count(*) FROM big WHERE id > 1", &c); + let mut op = build(&p, &c).unwrap(); + collect_one(op.as_mut()).unwrap(); + let text = explain_analyze(op.as_ref()); + let lines: Vec<&str> = text.lines().collect(); + assert!(lines[0].starts_with("Projection (")); + assert!(lines.last().unwrap().starts_with(" Scan (")); + } + + #[test] + fn scanning_a_table_that_is_not_registered_is_an_error() { + let c = Catalog::new(); + let p = LogicalPlan::Scan { + table: "ghost".into(), + source_schema: Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])), + projection: None, + pushed_filters: vec![], + stats: TableStats::default(), + }; + assert!(build(&p, &c).is_err()); + } +} diff --git a/crates/qf-exec/src/scan.rs b/crates/qf-exec/src/scan.rs new file mode 100644 index 0000000..9dd1874 --- /dev/null +++ b/crates/qf-exec/src/scan.rs @@ -0,0 +1,492 @@ +//! The scan, and the one place where predicate pushdown stops being a plan +//! rewrite and starts saving I/O. +//! +//! For a file-backed table the scan walks row groups. Before reading one it +//! asks each pushed predicate whether the group's zone map could possibly +//! satisfy it; if not, the group's column chunks are never read. What is left +//! is decoded — only the projected columns — and the predicates are applied to +//! the rows that survive. + +use crate::eval::evaluate_predicate; +use crate::operator::{Metrics, Operator}; +use qf_common::{Result, Schema, Value}; +use qf_plan::expr::BoundExpr; +use qf_sql::ast::BinaryOp; +use qf_storage::batch::{RecordBatch, DEFAULT_BATCH_SIZE}; +use qf_storage::format::QfcReader; +use qf_storage::stats::PruneOp; +use std::path::Path; +use std::sync::Arc; +use std::time::Instant; + +enum Source { + /// Batches already in memory, still in the table's full schema. + Memory { + batches: Vec, + next: usize, + }, + File { + reader: Box, + next_group: usize, + /// Rows left over from a group larger than one batch. + pending: Vec, + }, +} + +pub struct ScanExec { + source: Source, + /// Source column positions this scan reads, in output order. + projection: Vec, + /// Predicates over the *projected* schema. + filters: Vec, + schema: Arc, + metrics: Metrics, +} + +impl ScanExec { + pub fn from_batches( + batches: Vec, + projection: Vec, + filters: Vec, + schema: Arc, + ) -> ScanExec { + ScanExec { + source: Source::Memory { batches, next: 0 }, + projection, + filters, + schema, + metrics: Metrics::default(), + } + } + + pub fn from_file( + path: impl AsRef, + projection: Vec, + filters: Vec, + schema: Arc, + ) -> Result { + Ok(ScanExec { + source: Source::File { + reader: Box::new(QfcReader::open(path)?), + next_group: 0, + pending: Vec::new(), + }, + projection, + filters, + schema, + metrics: Metrics::default(), + }) + } + + fn apply_filters(&self, batch: RecordBatch) -> Result { + let mut out = batch; + for f in &self.filters { + let mask = evaluate_predicate(f, &out)?; + out = out.filter(&mask)?; + if out.num_rows() == 0 { + break; + } + } + Ok(out) + } +} + +/// Can any row of this row group satisfy every pushed predicate? +/// +/// The answer is only ever used to *skip* work, so it errs toward reading: +/// anything the zone maps cannot decide comes back `true`. +fn group_may_match( + filters: &[BoundExpr], + projection: &[usize], + group: &qf_storage::format::RowGroupMeta, +) -> bool { + filters + .iter() + .all(|f| conjunct_may_match(f, projection, group)) +} + +fn conjunct_may_match( + expr: &BoundExpr, + projection: &[usize], + group: &qf_storage::format::RowGroupMeta, +) -> bool { + match expr { + // Both halves of an AND must be satisfiable. + BoundExpr::Binary { + left, + op: BinaryOp::And, + right, + .. + } => { + conjunct_may_match(left, projection, group) + && conjunct_may_match(right, projection, group) + } + // Either half of an OR is enough. + BoundExpr::Binary { + left, + op: BinaryOp::Or, + right, + .. + } => { + conjunct_may_match(left, projection, group) + || conjunct_may_match(right, projection, group) + } + BoundExpr::Binary { + left, op, right, .. + } if op.is_comparison() => { + let Some((index, literal, op)) = column_literal(left, *op, right) else { + return true; + }; + let Some(source) = projection.get(index) else { + return true; + }; + let Ok(chunk) = group.column(*source) else { + return true; + }; + let Some(prune_op) = prune_op(op) else { + return true; + }; + chunk.stats.may_match(prune_op, literal) + } + _ => true, + } +} + +/// Recognises `column literal`, mirroring the operator when the literal +/// is on the left. +fn column_literal<'a>( + left: &'a BoundExpr, + op: BinaryOp, + right: &'a BoundExpr, +) -> Option<(usize, &'a Value, BinaryOp)> { + match (left, right) { + (BoundExpr::Column { index, .. }, BoundExpr::Literal(v)) => Some((*index, v, op)), + (BoundExpr::Literal(v), BoundExpr::Column { index, .. }) => { + Some((*index, v, op.swap_operands()?)) + } + _ => None, + } +} + +fn prune_op(op: BinaryOp) -> Option { + Some(match op { + BinaryOp::Eq => PruneOp::Eq, + BinaryOp::NotEq => PruneOp::NotEq, + BinaryOp::Lt => PruneOp::Lt, + BinaryOp::LtEq => PruneOp::LtEq, + BinaryOp::Gt => PruneOp::Gt, + BinaryOp::GtEq => PruneOp::GtEq, + _ => return None, + }) +} + +impl Operator for ScanExec { + fn schema(&self) -> Arc { + Arc::clone(&self.schema) + } + + fn next_batch(&mut self) -> Result> { + let start = Instant::now(); + let result = match &mut self.source { + Source::Memory { batches, next } => { + if *next >= batches.len() { + None + } else { + let b = batches[*next].project(&self.projection)?; + *next += 1; + Some(b) + } + } + Source::File { + reader, + next_group, + pending, + } => { + if !pending.is_empty() { + Some(pending.remove(0)) + } else { + let mut found = None; + while *next_group < reader.num_row_groups() { + let index = *next_group; + *next_group += 1; + let meta = reader.row_group(index)?.clone(); + if !group_may_match(&self.filters, &self.projection, &meta) { + self.metrics.row_groups_pruned += 1; + continue; + } + self.metrics.row_groups_read += 1; + let batch = reader.read_row_group(index, &self.projection)?; + let mut chunks = batch.chunks(DEFAULT_BATCH_SIZE)?; + let first = chunks.remove(0); + pending.extend(chunks); + found = Some(first); + break; + } + found + } + } + }; + + let Some(batch) = result else { + self.metrics.elapsed += start.elapsed(); + return Ok(None); + }; + let out = self.apply_filters(batch)?; + self.metrics.elapsed += start.elapsed(); + self.metrics.record(&out); + Ok(Some(out)) + } + + fn metrics(&self) -> &Metrics { + &self.metrics + } + + fn name(&self) -> &'static str { + "Scan" + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::operator::collect_one; + use qf_common::{DataType, Field}; + use qf_storage::array::Array; + use qf_storage::format::write_file; + use std::path::PathBuf; + + struct TempFile(PathBuf); + + impl TempFile { + fn new(tag: &str) -> TempFile { + let mut p = std::env::temp_dir(); + p.push(format!( + "queryforge-scan-{tag}-{}-{}.qfc", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + TempFile(p) + } + } + + impl Drop for TempFile { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.0); + } + } + + fn schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("tag", DataType::Utf8, false), + ])) + } + + fn batch(range: std::ops::Range) -> RecordBatch { + let ids: Vec = range.clone().map(Value::Int64).collect(); + let tags: Vec = range.map(|i| Value::Utf8(format!("t{i}"))).collect(); + RecordBatch::try_new( + schema(), + vec![ + Array::from_values(DataType::Int64, &ids).unwrap(), + Array::from_values(DataType::Utf8, &tags).unwrap(), + ], + ) + .unwrap() + } + + fn id_col() -> BoundExpr { + BoundExpr::column(0, "id", DataType::Int64) + } + + fn cmp(op: BinaryOp, v: i64) -> BoundExpr { + BoundExpr::binary(id_col(), op, BoundExpr::Literal(Value::Int64(v))).unwrap() + } + + fn file_scan(path: &PathBuf, filters: Vec) -> ScanExec { + ScanExec::from_file(path, vec![0, 1], filters, schema()).unwrap() + } + + fn write(path: &PathBuf, rows: std::ops::Range, group: usize) { + write_file(path, schema(), std::slice::from_ref(&batch(rows)), group).unwrap(); + } + + #[test] + fn a_memory_scan_projects_and_filters() { + let mut op = ScanExec::from_batches( + vec![batch(0..10)], + vec![0], + vec![cmp(BinaryOp::Gt, 7)], + Arc::new(schema().project(&[0]).unwrap()), + ); + let out = collect_one(&mut op).unwrap(); + assert_eq!(out.num_columns(), 1); + assert_eq!(out.num_rows(), 2); + assert_eq!(op.name(), "Scan"); + } + + #[test] + fn a_range_predicate_prunes_the_groups_it_cannot_reach() { + let f = TempFile::new("range"); + write(&f.0, 0..500, 100); + let mut op = file_scan(&f.0, vec![cmp(BinaryOp::GtEq, 400)]); + let out = collect_one(&mut op).unwrap(); + assert_eq!(out.num_rows(), 100); + assert_eq!(op.metrics().row_groups_read, 1); + assert_eq!(op.metrics().row_groups_pruned, 4); + } + + #[test] + fn a_literal_on_the_left_prunes_just_as_well() { + let f = TempFile::new("flipped"); + write(&f.0, 0..500, 100); + // `400 <= id` must prune the same groups as `id >= 400`. + let expr = BoundExpr::binary( + BoundExpr::Literal(Value::Int64(400)), + BinaryOp::LtEq, + id_col(), + ) + .unwrap(); + let mut op = file_scan(&f.0, vec![expr]); + assert_eq!(collect_one(&mut op).unwrap().num_rows(), 100); + assert_eq!(op.metrics().row_groups_pruned, 4); + } + + #[test] + fn a_conjunction_prunes_a_group_that_either_half_rules_out() { + let f = TempFile::new("and"); + write(&f.0, 0..500, 100); + let both = BoundExpr::binary( + cmp(BinaryOp::GtEq, 150), + BinaryOp::And, + cmp(BinaryOp::Lt, 180), + ) + .unwrap(); + let mut op = file_scan(&f.0, vec![both]); + assert_eq!(collect_one(&mut op).unwrap().num_rows(), 30); + assert_eq!(op.metrics().row_groups_read, 1); + assert_eq!(op.metrics().row_groups_pruned, 4); + } + + #[test] + fn a_disjunction_only_prunes_a_group_that_neither_half_can_reach() { + let f = TempFile::new("or"); + write(&f.0, 0..500, 100); + let either = BoundExpr::binary( + cmp(BinaryOp::Lt, 50), + BinaryOp::Or, + cmp(BinaryOp::GtEq, 450), + ) + .unwrap(); + let mut op = file_scan(&f.0, vec![either]); + assert_eq!(collect_one(&mut op).unwrap().num_rows(), 100); + // The first and last groups survive; the middle three are ruled out by + // both halves. + assert_eq!(op.metrics().row_groups_read, 2); + assert_eq!(op.metrics().row_groups_pruned, 3); + } + + #[test] + fn a_predicate_the_zone_maps_cannot_decide_leaves_every_group_readable() { + let f = TempFile::new("undecidable"); + write(&f.0, 0..300, 100); + // Neither side is a plain column against a literal. + let expr = BoundExpr::binary( + BoundExpr::binary( + id_col(), + BinaryOp::Plus, + BoundExpr::Literal(Value::Int64(1)), + ) + .unwrap(), + BinaryOp::Gt, + BoundExpr::Literal(Value::Int64(250)), + ) + .unwrap(); + let mut op = file_scan(&f.0, vec![expr]); + assert_eq!(collect_one(&mut op).unwrap().num_rows(), 50); + assert_eq!(op.metrics().row_groups_pruned, 0); + assert_eq!(op.metrics().row_groups_read, 3); + } + + #[test] + fn a_non_comparison_predicate_is_applied_but_never_prunes() { + let f = TempFile::new("isnull"); + write(&f.0, 0..200, 100); + let expr = BoundExpr::IsNull { + expr: Box::new(id_col()), + negated: true, + }; + let mut op = file_scan(&f.0, vec![expr]); + assert_eq!(collect_one(&mut op).unwrap().num_rows(), 200); + assert_eq!(op.metrics().row_groups_pruned, 0); + } + + #[test] + fn pruning_everything_returns_no_rows_and_reads_nothing() { + let f = TempFile::new("all-pruned"); + write(&f.0, 0..300, 100); + let mut op = file_scan(&f.0, vec![cmp(BinaryOp::Gt, 10_000)]); + assert_eq!(collect_one(&mut op).unwrap().num_rows(), 0); + assert_eq!(op.metrics().row_groups_read, 0); + assert_eq!(op.metrics().row_groups_pruned, 3); + } + + #[test] + fn a_scan_with_no_predicate_reads_every_group() { + let f = TempFile::new("plain"); + write(&f.0, 0..250, 100); + let mut op = file_scan(&f.0, vec![]); + assert_eq!(collect_one(&mut op).unwrap().num_rows(), 250); + assert_eq!(op.metrics().row_groups_read, 3); + assert_eq!(op.metrics().row_groups_pruned, 0); + } + + #[test] + fn a_projected_scan_maps_predicate_positions_back_to_source_columns() { + // Reading only `tag` puts it at position 0, but its statistics live at + // source position 1. + let f = TempFile::new("projected"); + write(&f.0, 0..300, 100); + let tag = BoundExpr::column(0, "tag", DataType::Utf8); + let expr = BoundExpr::binary( + tag, + BinaryOp::Eq, + BoundExpr::Literal(Value::Utf8("t150".into())), + ) + .unwrap(); + let mut op = ScanExec::from_file( + &f.0, + vec![1], + vec![expr], + Arc::new(schema().project(&[1]).unwrap()), + ) + .unwrap(); + let out = collect_one(&mut op).unwrap(); + assert_eq!(out.num_rows(), 1); + assert_eq!(out.num_columns(), 1); + } + + #[test] + fn a_row_group_larger_than_a_batch_is_split() { + let f = TempFile::new("chunked"); + write(&f.0, 0..20_000, 20_000); + let mut op = file_scan(&f.0, vec![]); + let mut batches = 0; + let mut rows = 0; + while let Some(b) = op.next_batch().unwrap() { + batches += 1; + rows += b.num_rows(); + } + assert_eq!(rows, 20_000); + assert!(batches > 1, "one row group came back as a single batch"); + assert_eq!(op.metrics().row_groups_read, 1); + } + + #[test] + fn a_missing_file_is_reported_rather_than_panicking() { + assert!(ScanExec::from_file("/nonexistent/nope.qfc", vec![0], vec![], schema()).is_err()); + } +} diff --git a/crates/qf-exec/src/sort.rs b/crates/qf-exec/src/sort.rs new file mode 100644 index 0000000..f2c4227 --- /dev/null +++ b/crates/qf-exec/src/sort.rs @@ -0,0 +1,658 @@ +//! Sorting, in three modes depending on what the query asks for. +//! +//! * **Top-k** — when the sort feeds a `LIMIT`, only the best `k` rows are ever +//! held. Memory is bounded by `k` no matter how large the input, which is the +//! difference between `ORDER BY ts DESC LIMIT 10` being instant and being +//! impossible. +//! * **In memory** — when everything fits, sort once and emit. +//! * **External** — when it does not, sorted runs are written to spill files +//! and merged. This is the path that lets a sort finish on an input larger +//! than memory, which is the only reason a database sorts differently from +//! `Vec::sort`. + +use crate::eval::evaluate; +use crate::operator::{Metrics, Operator}; +use qf_common::{Result, Schema, Value}; +use qf_plan::logical::SortExpr; +use qf_storage::batch::{RecordBatch, DEFAULT_BATCH_SIZE}; +use qf_storage::format::{write_file, QfcReader}; +use std::cmp::Ordering; +use std::path::PathBuf; +use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering}; +use std::sync::Arc; +use std::time::Instant; + +/// Rows held before spilling. Deliberately low so the external path is +/// exercised by real queries and by the tests, not just in theory. +pub const DEFAULT_SPILL_THRESHOLD: usize = 100_000; + +static SPILL_COUNTER: AtomicU64 = AtomicU64::new(0); + +/// A sorted run on disk. Removes its file when dropped, including when a query +/// fails partway through. +struct SpillFile { + path: PathBuf, +} + +impl SpillFile { + fn create(schema: Arc, batch: &RecordBatch) -> Result { + let n = SPILL_COUNTER.fetch_add(1, AtomicOrdering::Relaxed); + let mut path = std::env::temp_dir(); + path.push(format!("queryforge-sort-{}-{n}.qfc", std::process::id())); + write_file( + &path, + schema, + std::slice::from_ref(batch), + DEFAULT_BATCH_SIZE, + )?; + Ok(SpillFile { path }) + } +} + +impl Drop for SpillFile { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.path); + } +} + +/// Reads one spill run in order, a batch at a time. +struct RunCursor { + reader: QfcReader, + columns: Vec, + group: usize, + batch: Option, + keys: Vec>, + row: usize, + _file: SpillFile, +} + +impl RunCursor { + fn open(file: SpillFile, exprs: &[SortExpr]) -> Result { + let reader = QfcReader::open(&file.path)?; + let columns = (0..reader.schema().len()).collect(); + let mut c = RunCursor { + reader, + columns, + group: 0, + batch: None, + keys: Vec::new(), + row: 0, + _file: file, + }; + c.advance_group(exprs)?; + Ok(c) + } + + fn advance_group(&mut self, exprs: &[SortExpr]) -> Result<()> { + if self.group >= self.reader.num_row_groups() { + self.batch = None; + return Ok(()); + } + let batch = self.reader.read_row_group(self.group, &self.columns)?; + self.group += 1; + self.keys = sort_keys(exprs, &batch)?; + self.row = 0; + self.batch = Some(batch); + Ok(()) + } + + fn current(&self) -> Option<(&RecordBatch, usize)> { + self.batch.as_ref().map(|b| (b, self.row)) + } + + fn key(&self) -> Option<&Vec> { + self.keys.get(self.row) + } + + fn step(&mut self, exprs: &[SortExpr]) -> Result<()> { + self.row += 1; + if let Some(b) = &self.batch { + if self.row >= b.num_rows() { + self.advance_group(exprs)?; + } + } + Ok(()) + } +} + +/// Materialises the sort key of every row. +fn sort_keys(exprs: &[SortExpr], batch: &RecordBatch) -> Result>> { + let columns = exprs + .iter() + .map(|s| evaluate(&s.expr, batch)) + .collect::>>()?; + Ok((0..batch.num_rows()) + .map(|r| columns.iter().map(|c| c.value(r)).collect()) + .collect()) +} + +/// Compares two keys under the sort specification. +/// +/// NULL ordering is explicit rather than falling out of the value ordering, +/// because SQL's default differs by direction: ascending puts NULLs last, +/// descending puts them first. +pub fn compare_keys(a: &[Value], b: &[Value], exprs: &[SortExpr]) -> Ordering { + for ((x, y), spec) in a.iter().zip(b.iter()).zip(exprs.iter()) { + let ord = match (x.is_null(), y.is_null()) { + (true, true) => Ordering::Equal, + (true, false) => { + if spec.nulls_first { + Ordering::Less + } else { + Ordering::Greater + } + } + (false, true) => { + if spec.nulls_first { + Ordering::Greater + } else { + Ordering::Less + } + } + (false, false) => { + let o = x.total_cmp(y); + if spec.ascending { + o + } else { + o.reverse() + } + } + }; + if ord != Ordering::Equal { + return ord; + } + } + Ordering::Equal +} + +pub struct SortExec { + input: Box, + exprs: Vec, + schema: Arc, + /// Rows the consumer will actually take, when known. Enables top-k. + fetch: Option, + spill_threshold: usize, + + buffered: Vec, + buffered_rows: usize, + runs: Vec, + /// Paths of every run this sort wrote, kept after the runs themselves are + /// consumed so a caller can confirm they were cleaned up. + spill_paths: Vec, + + output: Vec, + next: usize, + done: bool, + metrics: Metrics, +} + +impl SortExec { + pub fn new(input: Box, exprs: Vec, fetch: Option) -> SortExec { + SortExec::with_spill_threshold(input, exprs, fetch, DEFAULT_SPILL_THRESHOLD) + } + + /// Every spill file this sort wrote, whether or not it still exists. + pub fn spill_paths(&self) -> &[PathBuf] { + &self.spill_paths + } + + pub fn with_spill_threshold( + input: Box, + exprs: Vec, + fetch: Option, + spill_threshold: usize, + ) -> SortExec { + let schema = input.schema(); + SortExec { + input, + exprs, + schema, + fetch, + spill_threshold: spill_threshold.max(1), + buffered: Vec::new(), + buffered_rows: 0, + runs: Vec::new(), + spill_paths: Vec::new(), + output: Vec::new(), + next: 0, + done: false, + metrics: Metrics::default(), + } + } + + fn sort_batch(&self, batch: &RecordBatch) -> Result { + let keys = sort_keys(&self.exprs, batch)?; + let mut order: Vec = (0..batch.num_rows()).collect(); + // A stable sort keeps rows that tie in the order they arrived, which + // makes results reproducible run to run. + order.sort_by(|a, b| compare_keys(&keys[*a], &keys[*b], &self.exprs)); + batch.take(&order) + } + + fn drain_buffer(&mut self) -> Result { + let combined = RecordBatch::concat(Arc::clone(&self.schema), &self.buffered)?; + self.buffered.clear(); + self.buffered_rows = 0; + self.sort_batch(&combined) + } + + fn run(&mut self) -> Result<()> { + while let Some(b) = self.input.next_batch()? { + if b.num_rows() == 0 { + continue; + } + self.buffered_rows += b.num_rows(); + self.buffered.push(b); + + match self.fetch { + // Top-k: trim back to the best `k` whenever the buffer grows + // past twice that, so memory stays O(k). + Some(k) if self.buffered_rows > (k.max(1)).saturating_mul(2) => { + let sorted = self.drain_buffer()?; + let keep: Vec = (0..sorted.num_rows().min(k)).collect(); + let trimmed = sorted.take(&keep)?; + self.buffered_rows = trimmed.num_rows(); + self.buffered.push(trimmed); + } + None if self.buffered_rows >= self.spill_threshold => { + let sorted = self.drain_buffer()?; + self.metrics.spilled_rows += sorted.num_rows(); + self.metrics.spilled_runs += 1; + let file = SpillFile::create(Arc::clone(&self.schema), &sorted)?; + self.spill_paths.push(file.path.clone()); + self.runs.push(file); + } + _ => {} + } + } + + if self.runs.is_empty() { + let sorted = if self.buffered.is_empty() { + RecordBatch::empty(Arc::clone(&self.schema))? + } else { + self.drain_buffer()? + }; + let sorted = match self.fetch { + Some(k) if sorted.num_rows() > k => sorted.take(&(0..k).collect::>())?, + _ => sorted, + }; + self.output = sorted.chunks(DEFAULT_BATCH_SIZE)?; + return Ok(()); + } + + // Whatever is still buffered becomes one last run, so the merge has a + // uniform set of inputs. + if !self.buffered.is_empty() { + let sorted = self.drain_buffer()?; + self.metrics.spilled_rows += sorted.num_rows(); + self.metrics.spilled_runs += 1; + let file = SpillFile::create(Arc::clone(&self.schema), &sorted)?; + self.spill_paths.push(file.path.clone()); + self.runs.push(file); + } + self.merge() + } + + /// Merges the sorted runs. + /// + /// The smallest head is found by scanning the cursors rather than with a + /// heap. Runs are few — one per spill threshold's worth of rows — so the + /// linear scan costs less than maintaining the heap would, and it is far + /// easier to be sure it is correct. + fn merge(&mut self) -> Result<()> { + let files: Vec = std::mem::take(&mut self.runs); + let mut cursors = files + .into_iter() + .map(|f| RunCursor::open(f, &self.exprs)) + .collect::>>()?; + + let mut emitted = 0usize; + let mut chunk: Vec<(usize, usize)> = Vec::new(); // (cursor, row) + let mut sources: Vec = Vec::new(); + let mut out = Vec::new(); + + loop { + if self.fetch.is_some_and(|k| emitted >= k) { + break; + } + let mut best: Option = None; + for (i, c) in cursors.iter().enumerate() { + let Some(k) = c.key() else { continue }; + match best { + None => best = Some(i), + Some(b) => { + let bk = cursors[b].key().expect("a cursor with a key"); + if compare_keys(k, bk, &self.exprs) == Ordering::Less { + best = Some(i); + } + } + } + } + let Some(i) = best else { break }; + + let (batch, row) = cursors[i].current().expect("a cursor with a row"); + // Rows are copied out one at a time; keeping the source batch + // alongside lets the gather below stay a single take per run. + if sources.len() <= i { + sources.resize(i + 1, batch.clone()); + } + sources[i] = batch.clone(); + chunk.push((i, row)); + emitted += 1; + cursors[i].step(&self.exprs)?; + + if chunk.len() >= DEFAULT_BATCH_SIZE { + out.push(assemble(&self.schema, &sources, &chunk)?); + chunk.clear(); + } + } + if !chunk.is_empty() { + out.push(assemble(&self.schema, &sources, &chunk)?); + } + self.output = out; + Ok(()) + } +} + +/// Builds one output batch from rows picked out of several run batches. +fn assemble( + schema: &Arc, + sources: &[RecordBatch], + rows: &[(usize, usize)], +) -> Result { + let mut builders: Vec = schema + .fields() + .iter() + .map(|f| qf_storage::array::ArrayBuilder::new(f.data_type)) + .collect(); + for (cursor, row) in rows { + let batch = &sources[*cursor]; + for (c, b) in builders.iter_mut().enumerate() { + b.push_from(batch.column(c)?, *row); + } + } + let columns = builders + .into_iter() + .map(qf_storage::array::ArrayBuilder::finish) + .collect::>>()?; + if columns.is_empty() { + return RecordBatch::empty(Arc::clone(schema))?.take(&vec![0usize; rows.len()]); + } + RecordBatch::try_new(Arc::clone(schema), columns) +} + +impl Operator for SortExec { + fn schema(&self) -> Arc { + Arc::clone(&self.schema) + } + + fn next_batch(&mut self) -> Result> { + let start = Instant::now(); + if !self.done { + self.run()?; + self.done = true; + } + if self.next >= self.output.len() { + self.metrics.elapsed += start.elapsed(); + return Ok(None); + } + let b = self.output[self.next].clone(); + self.next += 1; + self.metrics.elapsed += start.elapsed(); + self.metrics.record(&b); + Ok(Some(b)) + } + + fn metrics(&self) -> &Metrics { + &self.metrics + } + + fn children(&self) -> Vec<&dyn Operator> { + vec![self.input.as_ref()] + } + + fn name(&self) -> &'static str { + "Sort" + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::operator::collect_one; + use crate::operators::ValuesExec; + use qf_common::{DataType, Field}; + use qf_plan::expr::BoundExpr; + use qf_storage::array::Array; + + fn schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int64, true), + Field::new("tag", DataType::Utf8, false), + ])) + } + + fn source(keys: &[Option], batch_size: usize) -> Box { + let ks: Vec = keys + .iter() + .map(|k| k.map_or(Value::Null, Value::Int64)) + .collect(); + let tags: Vec = keys + .iter() + .enumerate() + .map(|(i, _)| Value::Utf8(format!("r{i}"))) + .collect(); + let whole = RecordBatch::try_new( + schema(), + vec![ + Array::from_values(DataType::Int64, &ks).unwrap(), + Array::from_values(DataType::Utf8, &tags).unwrap(), + ], + ) + .unwrap(); + let batches = whole.chunks(batch_size.max(1)).unwrap(); + Box::new(ValuesExec::new(schema(), batches)) + } + + fn spec(ascending: bool, nulls_first: bool) -> Vec { + vec![SortExpr { + expr: BoundExpr::column(0, "k", DataType::Int64), + ascending, + nulls_first, + }] + } + + fn keys_of(batch: &RecordBatch) -> Vec { + batch + .column(0) + .unwrap() + .iter() + .map(|v| v.to_string()) + .collect() + } + + #[test] + fn a_small_input_sorts_in_memory_without_spilling() { + let mut op = SortExec::new( + source(&[Some(3), Some(1), Some(2)], 8), + spec(true, false), + None, + ); + let out = collect_one(&mut op).unwrap(); + assert_eq!(keys_of(&out), vec!["1", "2", "3"]); + assert_eq!(op.metrics().spilled_runs, 0); + } + + #[test] + fn nulls_go_last_ascending_and_first_descending_by_default() { + let data = [Some(2), None, Some(1)]; + let mut asc = SortExec::new(source(&data, 8), spec(true, false), None); + assert_eq!( + keys_of(&collect_one(&mut asc).unwrap()), + vec!["1", "2", "NULL"] + ); + + let mut desc = SortExec::new(source(&data, 8), spec(false, true), None); + assert_eq!( + keys_of(&collect_one(&mut desc).unwrap()), + vec!["NULL", "2", "1"] + ); + } + + #[test] + fn an_input_larger_than_the_spill_threshold_still_comes_back_fully_sorted() { + // The whole point of the external path: the answer must not depend on + // whether the input fitted in memory. + let data: Vec> = (0..500).map(|i| Some((i * 37) % 500)).collect(); + let mut spilled = + SortExec::with_spill_threshold(source(&data, 32), spec(true, false), None, 50); + let out = collect_one(&mut spilled).unwrap(); + let keys = keys_of(&out); + assert_eq!(keys.len(), 500); + assert!(spilled.metrics().spilled_runs > 1, "nothing spilled"); + assert_eq!(spilled.metrics().spilled_rows, 500); + + let expected: Vec = (0..500).map(|i| i.to_string()).collect(); + assert_eq!(keys, expected); + } + + #[test] + fn spilling_and_not_spilling_produce_the_same_rows_in_the_same_order() { + let data: Vec> = (0..200) + .map(|i| { + if i % 17 == 0 { + None + } else { + Some((i * 13) % 200) + } + }) + .collect(); + let mut memory = SortExec::new(source(&data, 16), spec(true, false), None); + let mut disk = + SortExec::with_spill_threshold(source(&data, 16), spec(true, false), None, 25); + let a = collect_one(&mut memory).unwrap(); + let b = collect_one(&mut disk).unwrap(); + assert!(disk.metrics().spilled_runs > 0); + assert_eq!(a, b); + } + + #[test] + fn spill_files_are_removed_once_the_merge_has_consumed_them() { + let data: Vec> = (0..300).map(Some).collect(); + let mut op = SortExec::with_spill_threshold(source(&data, 16), spec(true, false), None, 40); + collect_one(&mut op).unwrap(); + assert!(op.metrics().spilled_runs > 0); + let paths: Vec = op.spill_paths().to_vec(); + assert!(!paths.is_empty()); + // The merge takes ownership of each run and drops it as it finishes, + // so by the time the output is drained nothing is left on disk. + for p in &paths { + assert!(!p.exists(), "{} was left behind", p.display()); + } + } + + #[test] + fn spill_files_are_removed_even_when_the_sort_is_abandoned_mid_query() { + let data: Vec> = (0..300).map(Some).collect(); + let paths; + { + let mut op = + SortExec::with_spill_threshold(source(&data, 16), spec(true, false), None, 40); + // Pull one batch, then drop the operator without draining it. + op.next_batch().unwrap(); + paths = op.spill_paths().to_vec(); + assert!(!paths.is_empty()); + } + for p in &paths { + assert!(!p.exists(), "{} survived the drop", p.display()); + } + } + + #[test] + fn a_top_k_sort_returns_the_right_rows_without_holding_the_input() { + let data: Vec> = (0..1000).map(|i| Some((i * 91) % 1000)).collect(); + let mut op = SortExec::new(source(&data, 64), spec(true, false), Some(5)); + let out = collect_one(&mut op).unwrap(); + assert_eq!(keys_of(&out), vec!["0", "1", "2", "3", "4"]); + // Top-k never spills, however large the input. + assert_eq!(op.metrics().spilled_runs, 0); + } + + #[test] + fn top_k_agrees_with_a_full_sort_then_truncate() { + let data: Vec> = (0..300) + .map(|i| { + if i % 23 == 0 { + None + } else { + Some((i * 7) % 300) + } + }) + .collect(); + let mut full = SortExec::new(source(&data, 32), spec(false, true), None); + let all = collect_one(&mut full).unwrap(); + let expected = all.take(&(0..10).collect::>()).unwrap(); + + let mut top = SortExec::new(source(&data, 32), spec(false, true), Some(10)); + assert_eq!(collect_one(&mut top).unwrap(), expected); + } + + #[test] + fn a_top_k_larger_than_the_input_returns_everything() { + let mut op = SortExec::new(source(&[Some(2), Some(1)], 8), spec(true, false), Some(50)); + assert_eq!(keys_of(&collect_one(&mut op).unwrap()), vec!["1", "2"]); + } + + #[test] + fn an_empty_input_sorts_to_nothing() { + let mut op = SortExec::new(source(&[], 8), spec(true, false), None); + assert_eq!(collect_one(&mut op).unwrap().num_rows(), 0); + assert_eq!(op.name(), "Sort"); + assert_eq!(op.children().len(), 1); + assert_eq!(op.schema().len(), 2); + } + + #[test] + fn ties_keep_their_input_order() { + // Every key is 1, so a stable sort must leave the tags as they came. + let data = vec![Some(1); 6]; + let mut op = SortExec::new(source(&data, 2), spec(true, false), None); + let out = collect_one(&mut op).unwrap(); + let tags: Vec = out + .column(1) + .unwrap() + .iter() + .map(|v| v.to_string()) + .collect(); + assert_eq!(tags, vec!["r0", "r1", "r2", "r3", "r4", "r5"]); + } + + #[test] + fn comparing_keys_follows_the_specification_field_by_field() { + let asc = spec(true, false); + assert_eq!( + compare_keys(&[Value::Int64(1)], &[Value::Int64(2)], &asc), + Ordering::Less + ); + let desc = spec(false, false); + assert_eq!( + compare_keys(&[Value::Int64(1)], &[Value::Int64(2)], &desc), + Ordering::Greater + ); + assert_eq!( + compare_keys(&[Value::Null], &[Value::Null], &asc), + Ordering::Equal + ); + assert_eq!( + compare_keys(&[Value::Null], &[Value::Int64(1)], &asc), + Ordering::Greater, + "nulls last" + ); + assert_eq!( + compare_keys(&[Value::Int64(1)], &[Value::Null], &spec(true, true)), + Ordering::Greater, + "nulls first" + ); + assert_eq!(compare_keys(&[], &[], &asc), Ordering::Equal); + } +} diff --git a/crates/qf-exec/tests/end_to_end.rs b/crates/qf-exec/tests/end_to_end.rs new file mode 100644 index 0000000..869cd0a --- /dev/null +++ b/crates/qf-exec/tests/end_to_end.rs @@ -0,0 +1,741 @@ +//! End-to-end tests: SQL text in, rows out, through the same path the CLI +//! uses. Anything that only passes because a test reached past the session is +//! not evidence the engine works. + +use qf_common::{DataType, Field, Schema, Value}; +use qf_exec::{Output, Session}; +use qf_storage::array::Array; +use qf_storage::batch::RecordBatch; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +struct TempDir(PathBuf); + +impl TempDir { + fn new(tag: &str) -> TempDir { + let mut p = std::env::temp_dir(); + p.push(format!( + "queryforge-test-{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 join(&self, name: &str) -> PathBuf { + self.0.join(name) + } + fn path(&self) -> &Path { + &self.0 + } +} + +impl Drop for TempDir { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.0); + } +} + +/// orders(id, customer_id, region, amount) with a NULL amount and a NULL +/// customer_id, so the null-handling tests have something to work with. +fn session() -> Session { + let mut s = Session::new(); + let orders = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("customer_id", DataType::Int64, true), + Field::new("region", DataType::Utf8, false), + Field::new("amount", DataType::Float64, true), + ])); + let rows: Vec<(i64, Option, &str, Option)> = vec![ + (1, Some(10), "eu", Some(100.0)), + (2, Some(10), "eu", Some(250.0)), + (3, Some(20), "us", Some(75.0)), + (4, Some(20), "us", None), + (5, None, "apac", Some(500.0)), + (6, Some(30), "eu", Some(30.0)), + ]; + let batch = RecordBatch::try_new( + Arc::clone(&orders), + vec![ + col(DataType::Int64, rows.iter().map(|r| Value::Int64(r.0))), + col( + DataType::Int64, + rows.iter().map(|r| r.1.map_or(Value::Null, Value::Int64)), + ), + col( + DataType::Utf8, + rows.iter().map(|r| Value::Utf8(r.2.to_string())), + ), + col( + DataType::Float64, + rows.iter().map(|r| r.3.map_or(Value::Null, Value::Float64)), + ), + ], + ) + .unwrap(); + s.register("orders", orders, vec![batch]).unwrap(); + + let customers = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("name", DataType::Utf8, false), + ])); + let cust: Vec<(i64, &str)> = vec![(10, "ada"), (20, "brendan"), (40, "grace")]; + let batch = RecordBatch::try_new( + Arc::clone(&customers), + vec![ + col(DataType::Int64, cust.iter().map(|r| Value::Int64(r.0))), + col( + DataType::Utf8, + cust.iter().map(|r| Value::Utf8(r.1.to_string())), + ), + ], + ) + .unwrap(); + s.register("customers", customers, vec![batch]).unwrap(); + s +} + +fn col(t: DataType, values: impl Iterator) -> Array { + let v: Vec = values.collect(); + Array::from_values(t, &v).unwrap() +} + +/// Every row rendered as strings, which makes assertions readable. +fn rows(s: &mut Session, sql: &str) -> Vec> { + let batch = s.query(sql).unwrap(); + batch + .rows() + .map(|r| r.iter().map(ToString::to_string).collect()) + .collect() +} + +fn one(s: &mut Session, sql: &str) -> String { + let r = rows(s, sql); + assert_eq!(r.len(), 1, "expected one row from `{sql}`, got {r:?}"); + assert_eq!(r[0].len(), 1, "expected one column from `{sql}`"); + r[0][0].clone() +} + +#[test] +fn a_projection_returns_the_columns_it_names() { + let mut s = session(); + assert_eq!( + rows(&mut s, "SELECT id, region FROM orders WHERE id <= 2"), + vec![vec!["1", "eu"], vec!["2", "eu"]] + ); +} + +#[test] +fn a_wildcard_returns_every_column() { + let mut s = session(); + let r = rows(&mut s, "SELECT * FROM orders WHERE id = 1"); + assert_eq!(r, vec![vec!["1", "10", "eu", "100.0"]]); +} + +#[test] +fn a_null_never_satisfies_a_comparison_but_is_found_by_is_null() { + let mut s = session(); + // Row 4 has a NULL amount, so neither `> 0` nor `<= 0` selects it. + assert_eq!( + one(&mut s, "SELECT count(*) FROM orders WHERE amount > 0"), + "5" + ); + assert_eq!( + one(&mut s, "SELECT count(*) FROM orders WHERE amount <= 0"), + "0" + ); + assert_eq!( + one(&mut s, "SELECT count(*) FROM orders WHERE amount IS NULL"), + "1" + ); +} + +#[test] +fn arithmetic_and_expressions_evaluate_per_row() { + let mut s = session(); + assert_eq!( + rows( + &mut s, + "SELECT id, amount * 2 AS d FROM orders WHERE id = 2" + ), + vec![vec!["2", "500.0"]] + ); + assert_eq!( + one(&mut s, "SELECT amount / 4 FROM orders WHERE id = 1"), + "25.0" + ); +} + +#[test] +fn in_between_and_like_filter_as_written() { + let mut s = session(); + assert_eq!( + one( + &mut s, + "SELECT count(*) FROM orders WHERE region IN ('eu', 'us')" + ), + "5" + ); + assert_eq!( + one( + &mut s, + "SELECT count(*) FROM orders WHERE amount BETWEEN 50 AND 260" + ), + "3" + ); + assert_eq!( + one( + &mut s, + "SELECT count(*) FROM orders WHERE region LIKE '%u%'" + ), + "5" + ); + assert_eq!( + one( + &mut s, + "SELECT count(*) FROM orders WHERE region NOT LIKE 'e%'" + ), + "3" + ); +} + +#[test] +fn case_expressions_pick_a_branch_per_row() { + let mut s = session(); + let r = rows( + &mut s, + "SELECT CASE WHEN amount > 200 THEN 'big' WHEN amount IS NULL THEN 'unknown' \ + ELSE 'small' END AS size FROM orders ORDER BY id", + ); + let flat: Vec<&str> = r.iter().map(|x| x[0].as_str()).collect(); + assert_eq!( + flat, + vec!["small", "big", "small", "unknown", "big", "small"] + ); +} + +#[test] +fn aggregates_ignore_nulls_and_count_star_does_not() { + let mut s = session(); + let r = rows( + &mut s, + "SELECT count(*), count(amount), sum(amount), min(amount), max(amount) FROM orders", + ); + assert_eq!(r[0][0], "6", "count(*) counts every row"); + assert_eq!(r[0][1], "5", "count(amount) skips the NULL"); + assert_eq!(r[0][2], "955.0"); + assert_eq!(r[0][3], "30.0"); + assert_eq!(r[0][4], "500.0"); +} + +#[test] +fn avg_divides_by_the_non_null_count() { + let mut s = session(); + // 955 over 5 non-null rows, not over 6. + assert_eq!(one(&mut s, "SELECT avg(amount) FROM orders"), "191.0"); +} + +#[test] +fn a_global_aggregate_over_no_rows_returns_one_row() { + let mut s = session(); + let r = rows( + &mut s, + "SELECT count(*), sum(amount) FROM orders WHERE id = 99", + ); + assert_eq!(r.len(), 1); + assert_eq!(r[0][0], "0"); + // sum over nothing is NULL, not zero. + assert_eq!(r[0][1], "NULL"); +} + +#[test] +fn grouping_splits_rows_by_key() { + let mut s = session(); + let r = rows( + &mut s, + "SELECT region, count(*) AS n, sum(amount) AS total FROM orders \ + GROUP BY region ORDER BY region", + ); + assert_eq!( + r, + vec![ + vec!["apac", "1", "500.0"], + vec!["eu", "3", "380.0"], + vec!["us", "2", "75.0"], + ] + ); +} + +#[test] +fn a_null_group_key_forms_its_own_group() { + let mut s = session(); + let r = rows( + &mut s, + "SELECT customer_id, count(*) FROM orders GROUP BY customer_id ORDER BY customer_id", + ); + assert_eq!(r.len(), 4); + // Ascending order puts NULL last. + assert_eq!(r[3][0], "NULL"); + assert_eq!(r[3][1], "1"); +} + +#[test] +fn having_filters_groups_not_rows() { + let mut s = session(); + let r = rows( + &mut s, + "SELECT region, count(*) FROM orders GROUP BY region HAVING count(*) > 1 ORDER BY region", + ); + assert_eq!(r, vec![vec!["eu", "3"], vec!["us", "2"]]); +} + +#[test] +fn count_distinct_counts_each_value_once() { + let mut s = session(); + assert_eq!( + one(&mut s, "SELECT count(DISTINCT region) FROM orders"), + "3" + ); + assert_eq!( + one(&mut s, "SELECT count(DISTINCT customer_id) FROM orders"), + "3" + ); + assert_eq!( + one(&mut s, "SELECT sum(DISTINCT amount) FROM orders"), + "955.0" + ); +} + +#[test] +fn an_inner_join_keeps_only_matching_rows() { + let mut s = session(); + let r = rows( + &mut s, + "SELECT o.id, c.name FROM orders o JOIN customers c ON o.customer_id = c.id ORDER BY o.id", + ); + assert_eq!( + r, + vec![ + vec!["1", "ada"], + vec!["2", "ada"], + vec!["3", "brendan"], + vec!["4", "brendan"], + ] + ); +} + +#[test] +fn a_null_join_key_never_matches() { + // Order 5 has a NULL customer_id; `NULL = anything` is unknown, so it + // must not join even though customers also has no NULL id. + let mut s = session(); + assert_eq!( + one( + &mut s, + "SELECT count(*) FROM orders o JOIN customers c ON o.customer_id = c.id" + ), + "4" + ); +} + +#[test] +fn a_left_join_pads_the_unmatched_side_with_nulls() { + let mut s = session(); + let r = rows( + &mut s, + "SELECT o.id, c.name FROM orders o LEFT JOIN customers c ON o.customer_id = c.id \ + ORDER BY o.id", + ); + assert_eq!(r.len(), 6); + assert_eq!(r[4], vec!["5", "NULL"], "the NULL key row survives padded"); + assert_eq!(r[5], vec!["6", "NULL"], "customer 30 does not exist"); +} + +#[test] +fn a_right_join_preserves_the_right_side() { + let mut s = session(); + let r = rows( + &mut s, + "SELECT o.id, c.name FROM orders o RIGHT JOIN customers c ON o.customer_id = c.id \ + ORDER BY c.name, o.id", + ); + // grace (id 40) has no orders and must still appear. + assert!(r.iter().any(|row| row == &vec!["NULL", "grace"])); + assert_eq!(r.len(), 5); +} + +#[test] +fn a_full_join_preserves_both_sides() { + let mut s = session(); + let r = rows( + &mut s, + "SELECT o.id, c.name FROM orders o FULL JOIN customers c ON o.customer_id = c.id", + ); + // 4 matched + orders 5 and 6 unmatched + customer grace unmatched. + assert_eq!(r.len(), 7); + assert!(r.iter().any(|row| row[1] == "grace" && row[0] == "NULL")); + assert!(r.iter().any(|row| row[0] == "5" && row[1] == "NULL")); +} + +#[test] +fn a_cross_join_produces_the_product() { + let mut s = session(); + assert_eq!( + one(&mut s, "SELECT count(*) FROM orders CROSS JOIN customers"), + "18" + ); +} + +#[test] +fn a_comma_join_with_an_equality_gives_the_same_answer_as_an_inner_join() { + // The optimiser rewrites this into a hash join; the answer must not change. + let mut s = session(); + let comma = rows( + &mut s, + "SELECT o.id, c.name FROM orders o, customers c WHERE o.customer_id = c.id ORDER BY o.id", + ); + let explicit = rows( + &mut s, + "SELECT o.id, c.name FROM orders o JOIN customers c ON o.customer_id = c.id ORDER BY o.id", + ); + assert_eq!(comma, explicit); +} + +#[test] +fn a_join_with_a_non_equality_condition_still_works() { + let mut s = session(); + let r = rows( + &mut s, + "SELECT count(*) FROM orders o JOIN customers c \ + ON o.customer_id = c.id AND o.amount > 100", + ); + assert_eq!(r[0][0], "1"); +} + +#[test] +fn ordering_is_ascending_by_default_with_nulls_last() { + let mut s = session(); + let r = rows(&mut s, "SELECT id, amount FROM orders ORDER BY amount"); + let amounts: Vec<&str> = r.iter().map(|x| x[1].as_str()).collect(); + assert_eq!( + amounts, + vec!["30.0", "75.0", "100.0", "250.0", "500.0", "NULL"] + ); +} + +#[test] +fn descending_order_puts_nulls_first_unless_told_otherwise() { + let mut s = session(); + let r = rows(&mut s, "SELECT amount FROM orders ORDER BY amount DESC"); + assert_eq!(r[0][0], "NULL"); + let r = rows( + &mut s, + "SELECT amount FROM orders ORDER BY amount DESC NULLS LAST", + ); + assert_eq!(r[0][0], "500.0"); + assert_eq!(r[5][0], "NULL"); +} + +#[test] +fn ordering_by_several_keys_breaks_ties_in_order() { + let mut s = session(); + let r = rows( + &mut s, + "SELECT region, id FROM orders ORDER BY region ASC, id DESC", + ); + assert_eq!(r[1], vec!["eu", "6"]); + assert_eq!(r[2], vec!["eu", "2"]); + assert_eq!(r[3], vec!["eu", "1"]); +} + +#[test] +fn ordering_by_a_column_that_is_not_selected_works() { + let mut s = session(); + let r = rows( + &mut s, + "SELECT id FROM orders ORDER BY amount DESC NULLS LAST", + ); + assert_eq!(r[0][0], "5"); +} + +#[test] +fn limit_and_offset_take_a_window_of_the_result() { + let mut s = session(); + let r = rows(&mut s, "SELECT id FROM orders ORDER BY id LIMIT 2 OFFSET 3"); + assert_eq!(r, vec![vec!["4"], vec!["5"]]); + assert!(rows(&mut s, "SELECT id FROM orders LIMIT 0").is_empty()); + assert_eq!(rows(&mut s, "SELECT id FROM orders OFFSET 5").len(), 1); +} + +#[test] +fn distinct_removes_duplicate_rows_but_keeps_their_first_order() { + let mut s = session(); + let r = rows(&mut s, "SELECT DISTINCT region FROM orders"); + assert_eq!(r, vec![vec!["eu"], vec!["us"], vec!["apac"]]); +} + +#[test] +fn a_select_without_from_evaluates_constants() { + let mut s = session(); + assert_eq!(one(&mut s, "SELECT 1 + 1"), "2"); + assert_eq!(one(&mut s, "SELECT CAST('7' AS INT) * 2"), "14"); +} + +#[test] +fn create_insert_and_drop_work_end_to_end() { + let mut s = Session::new(); + assert!(matches!( + s.execute_one("CREATE TABLE t (a INT NOT NULL, b TEXT)") + .unwrap(), + Output::Message(_) + )); + s.execute_one("INSERT INTO t VALUES (1, 'x'), (2, 'y')") + .unwrap(); + assert_eq!(one(&mut s, "SELECT count(*) FROM t"), "2"); + assert_eq!(one(&mut s, "SELECT b FROM t WHERE a = 2"), "y"); + + s.execute_one("INSERT INTO t VALUES (3, 'z')").unwrap(); + assert_eq!(one(&mut s, "SELECT count(*) FROM t"), "3"); + + s.execute_one("DROP TABLE t").unwrap(); + assert!(s.query("SELECT * FROM t").is_err()); +} + +#[test] +fn an_insert_of_the_wrong_width_or_type_is_refused() { + let mut s = Session::new(); + s.execute_one("CREATE TABLE t (a INT, b TEXT)").unwrap(); + assert!(s.execute_one("INSERT INTO t VALUES (1)").is_err()); + assert!(s + .execute_one("INSERT INTO t VALUES ('not a number', 'x')") + .is_err()); + // An integer into a float column is fine, though. + s.execute_one("CREATE TABLE f (x DOUBLE)").unwrap(); + s.execute_one("INSERT INTO f VALUES (3)").unwrap(); + assert_eq!(one(&mut s, "SELECT x FROM f"), "3.0"); +} + +#[test] +fn several_statements_run_in_order() { + let mut s = Session::new(); + let out = s + .execute("CREATE TABLE t (a INT); INSERT INTO t VALUES (1); SELECT count(*) FROM t") + .unwrap(); + assert_eq!(out.len(), 3); + assert_eq!(out[2].row_count(), 1); +} + +#[test] +fn csv_is_imported_with_inferred_types() { + let dir = TempDir::new("csv"); + let path = dir.join("sales.csv"); + std::fs::write(&path, "id,region,amount\n1,eu,10.5\n2,us,20\n3,eu,\n").unwrap(); + + let mut s = Session::new(); + assert_eq!(s.import_csv("sales", &path).unwrap(), 3); + assert_eq!(one(&mut s, "SELECT count(*) FROM sales"), "3"); + assert_eq!(one(&mut s, "SELECT sum(amount) FROM sales"), "30.5"); + assert_eq!( + one(&mut s, "SELECT count(*) FROM sales WHERE amount IS NULL"), + "1" + ); +} + +#[test] +fn copy_loads_a_csv_into_an_existing_table() { + let dir = TempDir::new("copy"); + let path = dir.join("more.csv"); + std::fs::write(&path, "a,b\n7,seven\n8,eight\n").unwrap(); + + let mut s = Session::new(); + s.execute_one("CREATE TABLE t (a INT, b TEXT)").unwrap(); + let out = s + .execute_one(&format!("COPY t FROM '{}'", path.display())) + .unwrap(); + assert!(matches!(out, Output::Message(m) if m.contains("2 rows"))); + assert_eq!(one(&mut s, "SELECT b FROM t WHERE a = 8"), "eight"); +} + +#[test] +fn a_table_survives_a_round_trip_through_the_columnar_format() { + let dir = TempDir::new("qfc"); + let path = dir.join("orders.qfc"); + + let mut s = session(); + let before = rows(&mut s, "SELECT id, region, amount FROM orders ORDER BY id"); + assert_eq!(s.save_table("orders", &path).unwrap(), 6); + // The table is now file-backed, so this reads through the columnar reader. + let after = rows(&mut s, "SELECT id, region, amount FROM orders ORDER BY id"); + assert_eq!(before, after); + + let mut fresh = Session::new(); + fresh.attach("orders", &path).unwrap(); + assert_eq!( + rows( + &mut fresh, + "SELECT id, region, amount FROM orders ORDER BY id" + ), + before + ); +} + +#[test] +fn a_predicate_prunes_row_groups_it_cannot_match() { + let dir = TempDir::new("prune"); + let path = dir.join("wide.qfc"); + + // 10 row groups of 100 sorted rows each. + let mut s = Session::new(); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("label", DataType::Utf8, false), + ])); + let ids: Vec = (0..1000).map(Value::Int64).collect(); + let labels: Vec = (0..1000).map(|i| Value::Utf8(format!("row{i}"))).collect(); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Array::from_values(DataType::Int64, &ids).unwrap(), + Array::from_values(DataType::Utf8, &labels).unwrap(), + ], + ) + .unwrap(); + s.register("wide", schema, vec![batch]).unwrap(); + s.save_table_with_row_group_size("wide", &path, 100) + .unwrap(); + + // Only the group holding 950..999 can satisfy this. + let mut op = s + .operators("SELECT label FROM wide WHERE id > 980") + .unwrap(); + let out = qf_exec::operator::collect(op.as_mut()).unwrap(); + let returned: usize = out.iter().map(RecordBatch::num_rows).sum(); + assert_eq!(returned, 19); + + let text = qf_exec::physical::explain_analyze(op.as_ref()); + assert!(text.contains("1 read, 9 pruned"), "{text}"); +} + +#[test] +fn pruning_does_not_change_the_answer() { + let dir = TempDir::new("prune-answer"); + let path = dir.join("t.qfc"); + let mut s = Session::new(); + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let ids: Vec = (0..500).map(|i| Value::Int64((i * 7) % 500)).collect(); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Array::from_values(DataType::Int64, &ids).unwrap()], + ) + .unwrap(); + s.register("t", Arc::clone(&schema), vec![batch.clone()]) + .unwrap(); + let in_memory = rows( + &mut s, + "SELECT count(*) FROM t WHERE id BETWEEN 100 AND 200", + ); + + s.save_table_with_row_group_size("t", &path, 50).unwrap(); + let on_disk = rows( + &mut s, + "SELECT count(*) FROM t WHERE id BETWEEN 100 AND 200", + ); + assert_eq!(in_memory, on_disk); + assert_eq!(on_disk[0][0], "101"); +} + +#[test] +fn explain_shows_the_optimised_plan() { + let mut s = session(); + let text = s + .explain("SELECT id FROM orders WHERE amount > 100") + .unwrap(); + assert!(text.contains("Scan: orders"), "{text}"); + assert!(text.contains("pushed:"), "{text}"); + assert!(!text.contains("Filter"), "{text}"); +} + +#[test] +fn explain_analyze_reports_rows_per_operator() { + let mut s = session(); + let out = s + .execute_one("EXPLAIN ANALYZE SELECT region, count(*) FROM orders GROUP BY region") + .unwrap(); + let Output::Plan(text) = out else { + panic!("expected a plan"); + }; + assert!(text.contains("HashAggregate"), "{text}"); + assert!(text.contains("rows="), "{text}"); + assert!(text.contains("3 rows returned"), "{text}"); +} + +#[test] +fn a_limit_stops_the_scan_early() { + let dir = TempDir::new("early-stop"); + let path = dir.join("big.qfc"); + let mut s = Session::new(); + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); + let ids: Vec = (0..1000).map(Value::Int64).collect(); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Array::from_values(DataType::Int64, &ids).unwrap()], + ) + .unwrap(); + s.register("big", schema, vec![batch]).unwrap(); + s.save_table_with_row_group_size("big", &path, 100).unwrap(); + + let mut op = s.operators("SELECT id FROM big LIMIT 5").unwrap(); + let out = qf_exec::operator::collect(op.as_mut()).unwrap(); + assert_eq!(out.iter().map(RecordBatch::num_rows).sum::(), 5); + // Only the first row group was ever read. + let text = qf_exec::physical::explain_analyze(op.as_ref()); + assert!(text.contains("1 read"), "{text}"); +} + +#[test] +fn errors_reach_the_caller_rather_than_panicking() { + let mut s = session(); + for sql in [ + "SELECT nope FROM orders", + "SELECT * FROM ghosts", + "SELECT id FROM orders WHERE region > 3", + "SELECT region FROM orders GROUP BY amount", + "SELECT sum(region) FROM orders", + "SELECT FROM orders", + ] { + assert!(s.query(sql).is_err(), "`{sql}` should have failed"); + } +} + +#[test] +fn the_catalog_lists_what_has_been_registered() { + let s = session(); + let names: Vec<&str> = s.tables().iter().map(|t| t.name.as_str()).collect(); + assert_eq!(names, vec!["customers", "orders"]); + assert_eq!(s.catalog().get("orders").unwrap().row_count(), 6); +} + +#[test] +fn saving_a_file_backed_table_again_works() { + let dir = TempDir::new("resave"); + let mut s = session(); + s.save_table("orders", dir.join("a.qfc")).unwrap(); + assert_eq!(s.save_table("orders", dir.join("b.qfc")).unwrap(), 6); + assert_eq!(one(&mut s, "SELECT count(*) FROM orders"), "6"); + assert!(dir.path().join("b.qfc").exists()); +} + +#[test] +fn insert_and_copy_are_refused_on_a_file_backed_table() { + let dir = TempDir::new("readonly"); + let mut s = session(); + s.save_table("orders", dir.join("o.qfc")).unwrap(); + assert!(s + .execute_one("INSERT INTO orders VALUES (9, 9, 'x', 1.0)") + .unwrap_err() + .to_string() + .contains("backed by a file")); +} diff --git a/crates/qf-plan/src/binder.rs b/crates/qf-plan/src/binder.rs index 87cb2fe..165d86b 100644 --- a/crates/qf-plan/src/binder.rs +++ b/crates/qf-plan/src/binder.rs @@ -843,6 +843,16 @@ pub fn bind(query: &Query, catalog: &Catalog) -> Result { Binder::new(catalog).bind_query(query) } +/// Binds a row of constant expressions — the values of an `INSERT`. +/// +/// The scope is empty, so a column reference in a `VALUES` row is reported as +/// an unknown column rather than silently resolving against some table. +pub fn bind_constant_row(exprs: &[Expr], catalog: &Catalog) -> Result> { + let binder = Binder::new(catalog); + let scope = Scope::default(); + exprs.iter().map(|e| binder.bind_expr(e, &scope)).collect() +} + #[cfg(test)] mod tests { use super::*;