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::*;