diff --git a/crates/qf-plan/src/binder.rs b/crates/qf-plan/src/binder.rs new file mode 100644 index 0000000..87cb2fe --- /dev/null +++ b/crates/qf-plan/src/binder.rs @@ -0,0 +1,1258 @@ +//! Name resolution and type checking: AST in, logical plan out. +//! +//! This is the only place in the engine that answers "does this column exist", +//! "which table did it come from" and "is this expression well typed". Once a +//! plan leaves the binder, every column is a position and every node knows its +//! type, so nothing downstream has to carry a catalog around or produce a +//! user-facing name error. + +use crate::expr::{AggFunc, BoundAggregate, BoundExpr}; +use crate::logical::{LogicalPlan, SortExpr}; +use qf_common::{DataType, Error, Field, Result, Schema, Value}; +use qf_sql::ast::{ + is_aggregate_name, BinaryOp, Expr, OrderByExpr, Query, SelectItem, TableRef, UnaryOp, +}; +use qf_storage::catalog::Catalog; +use std::sync::Arc; + +#[derive(Debug, Clone)] +struct ScopeColumn { + /// The table name or alias a reference may be qualified with. + qualifier: Option, + name: String, + index: usize, + data_type: DataType, +} + +/// The columns visible to an expression, and where each one sits in the input. +#[derive(Debug, Clone, Default)] +struct Scope { + columns: Vec, +} + +impl Scope { + fn from_schema(schema: &Schema, qualifier: Option, offset: usize) -> Scope { + Scope { + columns: schema + .fields() + .iter() + .enumerate() + .map(|(i, f)| ScopeColumn { + qualifier: qualifier.clone(), + name: f.name.clone(), + index: i + offset, + data_type: f.data_type, + }) + .collect(), + } + } + + fn concat(mut left: Scope, right: Scope) -> Scope { + left.columns.extend(right.columns); + left + } + + fn width(&self) -> usize { + self.columns.len() + } + + /// Resolves `[table.]name`. + /// + /// An unqualified name matching two tables is an error rather than a + /// first-match win: silently picking one is how a join query comes back + /// with the wrong column and no complaint. + fn resolve(&self, qualifier: Option<&str>, name: &str) -> Result { + let mut found: Option<&ScopeColumn> = None; + for c in &self.columns { + let qualifier_matches = match qualifier { + None => true, + Some(q) => c + .qualifier + .as_deref() + .is_some_and(|cq| cq.eq_ignore_ascii_case(q)), + }; + if qualifier_matches && c.name.eq_ignore_ascii_case(name) { + if found.is_some() { + return Err(Error::plan(format!( + "column `{name}` is ambiguous — qualify it with a table name" + ))); + } + found = Some(c); + } + } + let c = found.ok_or_else(|| match qualifier { + Some(q) => Error::plan(format!("no column `{name}` in `{q}`")), + None => Error::plan(format!("no such column `{name}`")), + })?; + Ok(BoundExpr::column(c.index, c.name.clone(), c.data_type)) + } +} + +/// Collects the aggregates a query mentions while its expressions are bound. +#[derive(Debug, Default)] +struct AggregateState { + group_asts: Vec, + group_exprs: Vec<(BoundExpr, String)>, + aggregates: Vec, +} + +impl AggregateState { + /// Registers an aggregate, reusing an identical one already seen so that + /// `SELECT count(*) ... HAVING count(*) > 2` computes it once. + fn add(&mut self, agg: BoundAggregate) -> usize { + if let Some(i) = self + .aggregates + .iter() + .position(|a| a.func == agg.func && a.arg == agg.arg && a.distinct == agg.distinct) + { + return i; + } + self.aggregates.push(agg); + self.aggregates.len() - 1 + } + + fn group_index(&self, ast: &Expr) -> Option { + self.group_asts.iter().position(|g| g == ast) + } +} + +pub struct Binder<'a> { + catalog: &'a Catalog, +} + +impl<'a> Binder<'a> { + pub fn new(catalog: &'a Catalog) -> Binder<'a> { + Binder { catalog } + } + + pub fn bind_query(&self, query: &Query) -> Result { + let (mut plan, scope) = match &query.from { + Some(from) => self.bind_from(from)?, + // `SELECT 1 + 1` with no FROM runs over a single empty row. + None => ( + LogicalPlan::Values { + schema: Arc::new(Schema::empty()), + rows: vec![vec![]], + }, + Scope::default(), + ), + }; + + if let Some(predicate) = &query.selection { + // WHERE runs before grouping, so an aggregate there has nothing to + // aggregate over. HAVING is the clause that filters groups. + if predicate.contains_aggregate() { + return Err(Error::plan( + "WHERE cannot contain an aggregate — use HAVING".to_string(), + )); + } + let bound = self.bind_expr(predicate, &scope)?; + require_boolean(&bound, "WHERE")?; + plan = LogicalPlan::Filter { + input: Box::new(plan), + predicate: bound, + }; + } + + let uses_aggregates = !query.group_by.is_empty() + || query.projection.iter().any(|item| match item { + SelectItem::Expr { expr, .. } => expr.contains_aggregate(), + _ => false, + }) + || query.having.as_ref().is_some_and(Expr::contains_aggregate) + || query.order_by.iter().any(|o| o.expr.contains_aggregate()); + + if uses_aggregates { + self.bind_aggregate_query(query, plan, &scope) + } else { + self.bind_simple_query(query, plan, &scope) + } + } + + // ---- FROM ---- + + fn bind_from(&self, from: &TableRef) -> Result<(LogicalPlan, Scope)> { + match from { + TableRef::Table { name, alias } => { + let table = self.catalog.get(name)?; + let qualifier = Some(alias.clone().unwrap_or_else(|| table.name.clone())); + let scope = Scope::from_schema(&table.schema, qualifier, 0); + let plan = LogicalPlan::Scan { + table: table.name.clone(), + source_schema: Arc::clone(&table.schema), + projection: None, + pushed_filters: vec![], + stats: table.stats.clone(), + }; + Ok((plan, scope)) + } + TableRef::Join { + left, + right, + join_type, + on, + } => { + let (left_plan, left_scope) = self.bind_from(left)?; + let (right_plan, right_scope) = self.bind_from(right)?; + let offset = left_scope.width(); + let right_scope = Scope { + columns: right_scope + .columns + .into_iter() + .map(|mut c| { + c.index += offset; + c + }) + .collect(), + }; + let combined = Scope::concat(left_scope, right_scope); + + let (keys, residual) = match on { + None => (vec![], None), + Some(expr) => { + let bound = self.bind_expr(expr, &combined)?; + require_boolean(&bound, "ON")?; + split_join_condition(bound, offset)? + } + }; + + Ok(( + LogicalPlan::Join { + left: Box::new(left_plan), + right: Box::new(right_plan), + join_type: *join_type, + on: keys, + filter: residual, + }, + combined, + )) + } + } + } + + // ---- non-aggregate queries ---- + + fn bind_simple_query( + &self, + query: &Query, + input: LogicalPlan, + scope: &Scope, + ) -> Result { + let projection = self.bind_projection(&query.projection, scope, None)?; + let sorts = self.bind_order_by(&query.order_by, scope, &query.projection, None)?; + + let mut plan = input; + if !sorts.is_empty() { + plan = LogicalPlan::Sort { + input: Box::new(plan), + exprs: sorts, + }; + } + plan = LogicalPlan::Projection { + input: Box::new(plan), + exprs: projection, + }; + Ok(self.finish(plan, query)) + } + + // ---- aggregate queries ---- + + fn bind_aggregate_query( + &self, + query: &Query, + input: LogicalPlan, + scope: &Scope, + ) -> Result { + let mut state = AggregateState::default(); + for g in &query.group_by { + // `GROUP BY 1` refers to the first output column. + let ast = resolve_ordinal(g, &query.projection)?; + if ast.contains_aggregate() { + return Err(Error::plan( + "GROUP BY cannot contain an aggregate".to_string(), + )); + } + let bound = self.bind_expr(&ast, scope)?; + let name = bound.output_name(); + state.group_asts.push(ast); + state.group_exprs.push((bound, name)); + } + + let projection = self.bind_projection(&query.projection, scope, Some(&mut state))?; + let having = match &query.having { + Some(h) => { + let bound = self.bind_post_aggregate(h, scope, &mut state)?; + require_boolean(&bound, "HAVING")?; + Some(bound) + } + None => None, + }; + let sorts = + self.bind_order_by(&query.order_by, scope, &query.projection, Some(&mut state))?; + + let mut plan = LogicalPlan::Aggregate { + input: Box::new(input), + group_exprs: state.group_exprs, + aggregates: state.aggregates, + }; + if let Some(h) = having { + plan = LogicalPlan::Filter { + input: Box::new(plan), + predicate: h, + }; + } + if !sorts.is_empty() { + plan = LogicalPlan::Sort { + input: Box::new(plan), + exprs: sorts, + }; + } + plan = LogicalPlan::Projection { + input: Box::new(plan), + exprs: projection, + }; + Ok(self.finish(plan, query)) + } + + fn finish(&self, mut plan: LogicalPlan, query: &Query) -> LogicalPlan { + if query.distinct { + plan = LogicalPlan::Distinct { + input: Box::new(plan), + }; + } + if query.limit.is_some() || query.offset.unwrap_or(0) > 0 { + plan = LogicalPlan::Limit { + input: Box::new(plan), + limit: query.limit, + offset: query.offset.unwrap_or(0), + }; + } + plan + } + + // ---- projections ---- + + fn bind_projection( + &self, + items: &[SelectItem], + scope: &Scope, + mut state: Option<&mut AggregateState>, + ) -> Result> { + let mut out = Vec::new(); + for item in items { + match item { + SelectItem::Wildcard | SelectItem::QualifiedWildcard(_) => { + if state.is_some() { + return Err(Error::plan( + "`*` cannot be mixed with aggregates — list the grouped columns" + .to_string(), + )); + } + let qualifier = match item { + SelectItem::QualifiedWildcard(q) => Some(q.as_str()), + _ => None, + }; + let mut matched = 0; + for c in &scope.columns { + let ok = qualifier.is_none_or(|q| { + c.qualifier + .as_deref() + .is_some_and(|cq| cq.eq_ignore_ascii_case(q)) + }); + if ok { + matched += 1; + out.push(( + BoundExpr::column(c.index, c.name.clone(), c.data_type), + c.name.clone(), + )); + } + } + if matched == 0 { + return Err(Error::plan(match qualifier { + Some(q) => format!("no table named `{q}` in this query"), + None => "there are no columns to select".to_string(), + })); + } + } + SelectItem::Expr { expr, alias } => { + let bound = match state.as_deref_mut() { + Some(s) => self.bind_post_aggregate(expr, scope, s)?, + None => self.bind_expr(expr, scope)?, + }; + let name = alias.clone().unwrap_or_else(|| bound.output_name()); + out.push((bound, name)); + } + } + } + if out.is_empty() { + return Err(Error::plan("a query must select something".to_string())); + } + Ok(out) + } + + fn bind_order_by( + &self, + items: &[OrderByExpr], + scope: &Scope, + projection: &[SelectItem], + mut state: Option<&mut AggregateState>, + ) -> Result> { + let mut out = Vec::new(); + for item in items { + // `ORDER BY 2` and `ORDER BY alias` both mean an output column, so + // resolve them back to the expression that produced it. + let ast = resolve_ordinal(&item.expr, projection)?; + let ast = resolve_alias(&ast, projection); + let expr = match state.as_deref_mut() { + Some(s) => self.bind_post_aggregate(&ast, scope, s)?, + None => self.bind_expr(&ast, scope)?, + }; + out.push(SortExpr { + expr, + ascending: item.ascending, + nulls_first: item.nulls_first, + }); + } + Ok(out) + } + + // ---- expressions ---- + + /// Binds an expression that may reference aggregates, producing positions + /// in the aggregate node's output rather than in its input. + fn bind_post_aggregate( + &self, + expr: &Expr, + scope: &Scope, + state: &mut AggregateState, + ) -> Result { + // A whole grouped expression resolves to its group column. + if let Some(i) = state.group_index(expr) { + let (e, name) = &state.group_exprs[i]; + return Ok(BoundExpr::column(i, name.clone(), e.data_type())); + } + + if let Expr::Function { + name, + args, + distinct, + } = expr + { + if is_aggregate_name(name) { + let agg = self.bind_aggregate_call(name, args, *distinct, scope)?; + let data_type = agg.data_type; + let output_name = agg.output_name.clone(); + let i = state.add(agg); + return Ok(BoundExpr::column( + state.group_exprs.len() + i, + output_name, + data_type, + )); + } + } + + match expr { + // A bare column that is neither grouped nor aggregated has no + // single value per group. Rejecting it is the difference between + // an error and a silently arbitrary row. + Expr::Column { table, name } => { + let qualified = match table { + Some(t) => format!("{t}.{name}"), + None => name.clone(), + }; + // Resolve first so a genuinely missing column reports that + // rather than the grouping error. + scope.resolve(table.as_deref(), name)?; + Err(Error::plan(format!( + "`{qualified}` must appear in GROUP BY or inside an aggregate" + ))) + } + Expr::Literal(v) => Ok(BoundExpr::Literal(v.clone())), + Expr::Binary { left, op, right } => BoundExpr::binary( + self.bind_post_aggregate(left, scope, state)?, + *op, + self.bind_post_aggregate(right, scope, state)?, + ), + Expr::Unary { op, expr } => { + BoundExpr::unary(*op, self.bind_post_aggregate(expr, scope, state)?) + } + Expr::Cast { expr, data_type } => Ok(BoundExpr::Cast { + expr: Box::new(self.bind_post_aggregate(expr, scope, state)?), + data_type: *data_type, + }), + Expr::IsNull { expr, negated } => Ok(BoundExpr::IsNull { + expr: Box::new(self.bind_post_aggregate(expr, scope, state)?), + negated: *negated, + }), + other => { + // The remaining shapes desugar or bind identically whether or + // not aggregates are in play, so route them through the same + // code with a post-aggregate binder for their children. + self.bind_expr_with(other, scope, &mut |e, sc| { + self.bind_post_aggregate(e, sc, state) + }) + } + } + } + + fn bind_aggregate_call( + &self, + name: &str, + args: &[Expr], + distinct: bool, + scope: &Scope, + ) -> Result { + let func = AggFunc::from_name(name) + .ok_or_else(|| Error::plan(format!("unknown aggregate `{name}`")))?; + + if args.len() != 1 { + return Err(Error::plan(format!( + "{}() takes exactly one argument, found {}", + func.name(), + args.len() + ))); + } + + let (arg, input_type, rendered) = match &args[0] { + Expr::Wildcard => { + if !func.accepts_wildcard() { + return Err(Error::plan(format!( + "{}(*) is not meaningful — name a column", + func.name() + ))); + } + (None, None, "*".to_string()) + } + other => { + if other.contains_aggregate() { + return Err(Error::plan("aggregates cannot be nested".to_string())); + } + let bound = self.bind_expr(other, scope)?; + let t = bound.data_type(); + let rendered = other.to_string(); + (Some(bound), Some(t), rendered) + } + }; + + let data_type = func.output_type(input_type)?; + let d = if distinct { "DISTINCT " } else { "" }; + Ok(BoundAggregate { + func, + arg, + distinct, + output_name: format!("{}({d}{rendered})", func.name()), + data_type, + }) + } + + /// Binds an expression against a scope, with no aggregates allowed. + fn bind_expr(&self, expr: &Expr, scope: &Scope) -> Result { + match expr { + Expr::Column { table, name } => scope.resolve(table.as_deref(), name), + Expr::Literal(v) => Ok(BoundExpr::Literal(v.clone())), + Expr::Function { name, .. } if is_aggregate_name(name) => Err(Error::plan(format!( + "{name}() is only allowed in SELECT, HAVING or ORDER BY" + ))), + Expr::Function { name, .. } => Err(Error::plan(format!("unknown function `{name}`"))), + Expr::Wildcard => Err(Error::plan( + "`*` is only meaningful in SELECT or count(*)".to_string(), + )), + other => self.bind_expr_with(other, scope, &mut |e, sc| self.bind_expr(e, sc)), + } + } + + /// The shapes that bind the same way in both contexts, parameterised by + /// how their children are bound. + fn bind_expr_with( + &self, + expr: &Expr, + scope: &Scope, + bind: &mut dyn FnMut(&Expr, &Scope) -> Result, + ) -> Result { + Ok(match expr { + Expr::Column { .. } | Expr::Literal(_) | Expr::Function { .. } | Expr::Wildcard => { + bind(expr, scope)? + } + Expr::Binary { left, op, right } => { + BoundExpr::binary(bind(left, scope)?, *op, bind(right, scope)?)? + } + Expr::Unary { op, expr } => BoundExpr::unary(*op, bind(expr, scope)?)?, + Expr::Cast { expr, data_type } => BoundExpr::Cast { + expr: Box::new(bind(expr, scope)?), + data_type: *data_type, + }, + Expr::IsNull { expr, negated } => BoundExpr::IsNull { + expr: Box::new(bind(expr, scope)?), + negated: *negated, + }, + // BETWEEN desugars rather than becoming its own node: two + // comparisons are something predicate pushdown already knows how + // to split, prune with and evaluate. + Expr::Between { + expr, + low, + high, + negated, + } => { + let e = bind(expr, scope)?; + let lower = BoundExpr::binary(e.clone(), BinaryOp::GtEq, bind(low, scope)?)?; + let upper = BoundExpr::binary(e, BinaryOp::LtEq, bind(high, scope)?)?; + let both = BoundExpr::binary(lower, BinaryOp::And, upper)?; + if *negated { + BoundExpr::unary(UnaryOp::Not, both)? + } else { + both + } + } + Expr::InList { + expr, + list, + negated, + } => { + let e = bind(expr, scope)?; + let items = list + .iter() + .map(|i| bind(i, scope)) + .collect::>>()?; + for i in &items { + if i.is_null_literal() { + continue; + } + DataType::unify(e.data_type(), i.data_type()).map_err(|_| { + Error::typ(format!( + "IN list holds {} but the column is {}", + i.data_type(), + e.data_type() + )) + })?; + } + BoundExpr::InList { + expr: Box::new(e), + list: items, + negated: *negated, + } + } + Expr::Like { + expr, + pattern, + negated, + } => { + let e = bind(expr, scope)?; + let p = bind(pattern, scope)?; + for (side, t) in [("value", e.data_type()), ("pattern", p.data_type())] { + if t != DataType::Utf8 { + return Err(Error::typ(format!( + "LIKE needs text, but its {side} is {t}" + ))); + } + } + BoundExpr::Like { + expr: Box::new(e), + pattern: Box::new(p), + negated: *negated, + } + } + Expr::Case { + operand, + branches, + else_result, + } => { + let mut bound_branches = Vec::with_capacity(branches.len()); + for (when, then) in branches { + // `CASE x WHEN 1 ...` is `CASE WHEN x = 1 ...`. + let condition = match operand { + Some(o) => { + BoundExpr::binary(bind(o, scope)?, BinaryOp::Eq, bind(when, scope)?)? + } + None => bind(when, scope)?, + }; + bound_branches.push((condition, bind(then, scope)?)); + } + let else_bound = match else_result { + Some(e) => Some(bind(e, scope)?), + None => None, + }; + BoundExpr::case(bound_branches, else_bound)? + } + }) + } +} + +/// Turns `ORDER BY 2` / `GROUP BY 1` into the expression that output column +/// was built from. +fn resolve_ordinal(expr: &Expr, projection: &[SelectItem]) -> Result { + let Expr::Literal(Value::Int64(n)) = expr else { + return Ok(expr.clone()); + }; + let n = *n; + if n < 1 || n as usize > projection.len() { + return Err(Error::plan(format!( + "position {n} is not in the select list, which has {} items", + projection.len() + ))); + } + match &projection[(n - 1) as usize] { + SelectItem::Expr { expr, .. } => Ok(expr.clone()), + _ => Err(Error::plan(format!( + "position {n} refers to `*`, which is not a single column" + ))), + } +} + +/// Turns `ORDER BY alias` into the expression that alias names. +fn resolve_alias(expr: &Expr, projection: &[SelectItem]) -> Expr { + let Expr::Column { table: None, name } = expr else { + return expr.clone(); + }; + for item in projection { + if let SelectItem::Expr { + expr: inner, + alias: Some(a), + } = item + { + if a.eq_ignore_ascii_case(name) { + return inner.clone(); + } + } + } + expr.clone() +} + +fn require_boolean(expr: &BoundExpr, clause: &str) -> Result<()> { + if expr.data_type() != DataType::Boolean { + return Err(Error::typ(format!( + "{clause} needs a boolean, found {}", + expr.data_type() + ))); + } + Ok(()) +} + +/// Splits a bound ON condition into equality key pairs plus whatever is left. +/// +/// A conjunct is a join key only if it is an equality with one side reading +/// only left columns and the other only right. Everything else — including an +/// equality between two columns of the same side — stays as a filter, because +/// a hash join cannot build on it. +/// +/// Returns the key pairs and whatever could not become a key. +fn split_join_condition(condition: BoundExpr, right_offset: usize) -> Result { + let mut keys = Vec::new(); + let mut residual = Vec::new(); + + for conjunct in split_and(condition) { + if let BoundExpr::Binary { + left, + op: BinaryOp::Eq, + right, + .. + } = &conjunct + { + let l_side = side_of(left, right_offset); + let r_side = side_of(right, right_offset); + match (l_side, r_side) { + (Side::Left, Side::Right) => { + keys.push((left.as_ref().clone(), shift_right(right, right_offset)?)); + continue; + } + (Side::Right, Side::Left) => { + keys.push((right.as_ref().clone(), shift_right(left, right_offset)?)); + continue; + } + _ => {} + } + } + residual.push(conjunct); + } + + let filter = residual.into_iter().reduce(|a, b| BoundExpr::Binary { + left: Box::new(a), + op: BinaryOp::And, + right: Box::new(b), + data_type: DataType::Boolean, + }); + Ok((keys, filter)) +} + +/// Join keys, plus the part of the condition that is not a key. +type SplitCondition = (Vec<(BoundExpr, BoundExpr)>, Option); + +#[derive(Debug, PartialEq, Eq, Clone, Copy)] +enum Side { + Left, + Right, + Both, + Neither, +} + +fn side_of(expr: &BoundExpr, right_offset: usize) -> Side { + let indices = expr.column_indices(); + if indices.is_empty() { + return Side::Neither; + } + let left = indices.iter().any(|i| *i < right_offset); + let right = indices.iter().any(|i| *i >= right_offset); + match (left, right) { + (true, false) => Side::Left, + (false, true) => Side::Right, + _ => Side::Both, + } +} + +/// Rebases a right-side expression so its column positions are relative to the +/// right input rather than to the joined pair. +fn shift_right(expr: &BoundExpr, offset: usize) -> Result { + expr.remap_columns(&|i| i.checked_sub(offset)) +} + +fn split_and(expr: BoundExpr) -> Vec { + match expr { + BoundExpr::Binary { + left, + op: BinaryOp::And, + right, + .. + } => { + let mut out = split_and(*left); + out.extend(split_and(*right)); + out + } + other => vec![other], + } +} + +/// Binds a `CREATE TABLE` column list into a schema. +pub fn schema_from_columns(columns: &[qf_sql::ast::ColumnDef]) -> Result> { + let mut seen: Vec<&str> = Vec::new(); + for c in columns { + if seen.iter().any(|s| s.eq_ignore_ascii_case(&c.name)) { + return Err(Error::plan(format!("duplicate column `{}`", c.name))); + } + seen.push(&c.name); + } + Ok(Arc::new(Schema::new( + columns + .iter() + .map(|c| Field::new(&c.name, c.data_type, c.nullable)) + .collect(), + ))) +} + +/// Binds a query against a catalog. +pub fn bind(query: &Query, catalog: &Catalog) -> Result { + Binder::new(catalog).bind_query(query) +} + +#[cfg(test)] +mod tests { + use super::*; + use qf_storage::array::Array; + use qf_storage::batch::RecordBatch; + + fn catalog() -> Catalog { + let mut c = Catalog::new(); + let orders = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("customer_id", DataType::Int64, false), + Field::new("region", DataType::Utf8, false), + Field::new("amount", DataType::Float64, true), + ])); + let customers = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("name", DataType::Utf8, false), + ])); + c.register_batches("orders", Arc::clone(&orders), vec![rows(&orders)]) + .unwrap(); + c.register_batches("customers", Arc::clone(&customers), vec![rows(&customers)]) + .unwrap(); + c + } + + /// One row of type-appropriate values, enough to give the catalog stats. + fn rows(schema: &Arc) -> RecordBatch { + let columns = schema + .fields() + .iter() + .map(|f| { + let v = match f.data_type { + DataType::Int64 => Value::Int64(1), + DataType::Float64 => Value::Float64(1.0), + DataType::Utf8 => Value::Utf8("x".into()), + DataType::Boolean => Value::Boolean(true), + }; + Array::from_values(f.data_type, &[v]).unwrap() + }) + .collect(); + RecordBatch::try_new(Arc::clone(schema), columns).unwrap() + } + + fn plan(sql: &str) -> Result { + let stmt = qf_sql::parse(sql)?; + match stmt { + qf_sql::ast::Statement::Query(q) => bind(&q, &catalog()), + other => panic!("expected a query, got {other:?}"), + } + } + + fn explain(sql: &str) -> String { + plan(sql).unwrap().explain() + } + + fn error(sql: &str) -> String { + plan(sql).unwrap_err().to_string() + } + + #[test] + fn a_projection_sits_above_a_scan() { + let text = explain("SELECT id FROM orders"); + assert_eq!(text, "Projection: id#0\n Scan: orders [*]\n"); + } + + #[test] + fn a_wildcard_expands_to_every_column_in_order() { + let p = plan("SELECT * FROM orders").unwrap(); + let s = p.schema().unwrap(); + assert_eq!(s.len(), 4); + assert_eq!(s.field(3).unwrap().name, "amount"); + } + + #[test] + fn a_qualified_wildcard_expands_to_one_side_of_a_join() { + let p = plan("SELECT c.* FROM orders o JOIN customers c ON o.customer_id = c.id").unwrap(); + let s = p.schema().unwrap(); + assert_eq!(s.len(), 2); + assert_eq!(s.field(1).unwrap().name, "name"); + assert!(error("SELECT z.* FROM orders").contains("no table named `z`")); + } + + #[test] + fn columns_resolve_to_positions_and_missing_ones_are_reported() { + let p = plan("SELECT amount FROM orders").unwrap(); + match &p { + LogicalPlan::Projection { exprs, .. } => { + assert_eq!(exprs[0].0.column_indices(), vec![3]); + } + other => panic!("unexpected {other:?}"), + } + assert!(error("SELECT nope FROM orders").contains("no such column `nope`")); + } + + #[test] + fn an_ambiguous_column_across_a_join_is_refused() { + let err = error("SELECT id FROM orders o JOIN customers c ON o.id = c.id"); + assert!(err.contains("ambiguous")); + // ...but qualifying it resolves. + assert!(plan("SELECT o.id FROM orders o JOIN customers c ON o.id = c.id").is_ok()); + } + + #[test] + fn an_alias_replaces_the_table_name_as_the_qualifier() { + assert!(plan("SELECT o.id FROM orders o").is_ok()); + assert!(error("SELECT orders.id FROM orders o").contains("no column `id` in `orders`")); + } + + #[test] + fn a_where_clause_becomes_a_filter_under_the_projection() { + let text = explain("SELECT id FROM orders WHERE amount > 10"); + assert!(text.contains("Filter: (amount#3 > 10)")); + assert!(text.lines().nth(1).unwrap().starts_with(" Filter")); + } + + #[test] + fn a_non_boolean_where_clause_is_refused() { + assert!(error("SELECT id FROM orders WHERE amount").contains("WHERE needs a boolean")); + } + + #[test] + fn comparing_a_text_column_with_a_number_is_refused_at_plan_time() { + assert!(error("SELECT id FROM orders WHERE region > 3").contains("cannot compare")); + } + + #[test] + fn an_aggregate_in_where_is_redirected_to_having() { + assert!(error("SELECT id FROM orders WHERE count(*) > 1").contains("use HAVING")); + } + + #[test] + fn between_desugars_into_two_comparisons() { + let text = explain("SELECT id FROM orders WHERE amount BETWEEN 1 AND 9"); + assert!(text.contains("(amount#3 >= 1)")); + assert!(text.contains("(amount#3 <= 9)")); + assert!(text.contains("AND")); + } + + #[test] + fn not_between_wraps_the_desugared_form_in_a_negation() { + let text = explain("SELECT id FROM orders WHERE amount NOT BETWEEN 1 AND 9"); + assert!(text.contains("NOT (")); + } + + #[test] + fn a_case_with_an_operand_desugars_into_equality_conditions() { + let text = explain("SELECT CASE region WHEN 'eu' THEN 1 ELSE 0 END FROM orders"); + assert!(text.contains("WHEN (region#2 = 'eu')")); + } + + #[test] + fn in_lists_and_like_patterns_are_type_checked() { + assert!(plan("SELECT id FROM orders WHERE region IN ('a', 'b')").is_ok()); + assert!(plan("SELECT id FROM orders WHERE region IN ('a', NULL)").is_ok()); + assert!(error("SELECT id FROM orders WHERE region IN (1)").contains("IN list holds")); + assert!(plan("SELECT id FROM orders WHERE region LIKE 'a%'").is_ok()); + assert!(error("SELECT id FROM orders WHERE amount LIKE 'a%'").contains("LIKE needs text")); + } + + #[test] + fn an_equijoin_condition_becomes_join_keys_rebased_to_each_side() { + let p = plan("SELECT o.id FROM orders o JOIN customers c ON o.customer_id = c.id").unwrap(); + let join = find_join(&p).expect("a join node"); + match join { + LogicalPlan::Join { on, filter, .. } => { + assert_eq!(on.len(), 1); + // Left key is position 1 in orders; right key is position 0 in + // customers, *not* position 4 of the joined pair. + assert_eq!(on[0].0.column_indices(), vec![1]); + assert_eq!(on[0].1.column_indices(), vec![0]); + assert!(filter.is_none()); + } + other => panic!("unexpected {other:?}"), + } + } + + #[test] + fn a_non_equality_in_on_stays_a_filter_rather_than_a_join_key() { + let p = plan("SELECT o.id FROM orders o JOIN customers c ON o.customer_id > c.id").unwrap(); + match find_join(&p).unwrap() { + LogicalPlan::Join { on, filter, .. } => { + assert!(on.is_empty()); + assert!(filter.is_some()); + } + other => panic!("unexpected {other:?}"), + } + } + + #[test] + fn an_equality_within_one_side_is_not_a_join_key() { + let p = plan("SELECT o.id FROM orders o JOIN customers c ON o.id = o.customer_id").unwrap(); + match find_join(&p).unwrap() { + LogicalPlan::Join { on, filter, .. } => { + assert!(on.is_empty()); + assert!(filter.is_some()); + } + other => panic!("unexpected {other:?}"), + } + } + + #[test] + fn a_mixed_on_clause_splits_into_keys_and_a_residual_filter() { + let p = plan( + "SELECT o.id FROM orders o JOIN customers c \ + ON o.customer_id = c.id AND o.amount > 5", + ) + .unwrap(); + match find_join(&p).unwrap() { + LogicalPlan::Join { on, filter, .. } => { + assert_eq!(on.len(), 1); + assert!(filter.is_some()); + } + other => panic!("unexpected {other:?}"), + } + } + + #[test] + fn a_reversed_equality_still_becomes_a_join_key_with_the_sides_in_order() { + let p = plan("SELECT o.id FROM orders o JOIN customers c ON c.id = o.customer_id").unwrap(); + match find_join(&p).unwrap() { + LogicalPlan::Join { on, .. } => { + assert_eq!(on[0].0.column_indices(), vec![1], "left key"); + assert_eq!(on[0].1.column_indices(), vec![0], "right key"); + } + other => panic!("unexpected {other:?}"), + } + } + + #[test] + fn a_cross_join_has_no_keys_and_needs_no_on_clause() { + let p = plan("SELECT o.id FROM orders o CROSS JOIN customers c").unwrap(); + match find_join(&p).unwrap() { + LogicalPlan::Join { on, filter, .. } => { + assert!(on.is_empty()); + assert!(filter.is_none()); + } + other => panic!("unexpected {other:?}"), + } + } + + #[test] + fn a_non_boolean_on_clause_is_refused() { + assert!(error("SELECT o.id FROM orders o JOIN customers c ON c.id") + .contains("ON needs a boolean")); + } + + #[test] + fn grouping_produces_an_aggregate_node_with_groups_then_aggregates() { + let text = explain("SELECT region, count(*) FROM orders GROUP BY region"); + assert!(text.contains("Aggregate: group by [region#2], count(*)")); + let s = plan("SELECT region, count(*) FROM orders GROUP BY region") + .unwrap() + .schema() + .unwrap(); + assert_eq!(s.len(), 2); + assert_eq!(s.field(1).unwrap().data_type, DataType::Int64); + } + + #[test] + fn a_bare_column_outside_group_by_is_refused() { + let err = error("SELECT region, amount FROM orders GROUP BY region"); + assert!(err.contains("must appear in GROUP BY")); + // A genuinely missing column reports that instead. + assert!(error("SELECT nope FROM orders GROUP BY region").contains("no such column")); + } + + #[test] + fn an_aggregate_without_group_by_collapses_to_a_single_row() { + let text = explain("SELECT count(*), sum(amount) FROM orders"); + assert!(text.contains("Aggregate: count(*), sum(amount#3)")); + assert!(!text.contains("group by")); + } + + #[test] + fn identical_aggregates_are_computed_once() { + let p = plan("SELECT count(*) FROM orders GROUP BY region HAVING count(*) > 2").unwrap(); + let agg = find_aggregate(&p).unwrap(); + match agg { + LogicalPlan::Aggregate { aggregates, .. } => assert_eq!(aggregates.len(), 1), + other => panic!("unexpected {other:?}"), + } + } + + #[test] + fn having_becomes_a_filter_above_the_aggregate() { + let text = explain("SELECT region FROM orders GROUP BY region HAVING count(*) > 2"); + let lines: Vec<&str> = text.lines().collect(); + assert!(lines[0].starts_with("Projection")); + assert!(lines[1].trim_start().starts_with("Filter")); + assert!(lines[2].trim_start().starts_with("Aggregate")); + } + + #[test] + fn aggregate_argument_and_nesting_rules_are_enforced() { + assert!(error("SELECT sum(*) FROM orders").contains("not meaningful")); + assert!(error("SELECT sum() FROM orders").contains("exactly one argument")); + assert!(error("SELECT sum(amount, id) FROM orders").contains("exactly one argument")); + assert!(error("SELECT sum(count(id)) FROM orders").contains("cannot be nested")); + assert!(error("SELECT sum(region) FROM orders").contains("cannot sum")); + assert!(error("SELECT median(amount) FROM orders").contains("unknown")); + assert!(error("SELECT upper(region) FROM orders").contains("unknown function")); + } + + #[test] + fn grouping_by_an_aggregate_or_mixing_a_wildcard_in_is_refused() { + assert!(error("SELECT count(*) FROM orders GROUP BY count(*)") + .contains("GROUP BY cannot contain an aggregate")); + assert!(error("SELECT *, count(*) FROM orders GROUP BY region") + .contains("cannot be mixed with aggregates")); + } + + #[test] + fn avg_widens_to_a_float_even_over_integers() { + let s = plan("SELECT avg(id) FROM orders") + .unwrap() + .schema() + .unwrap(); + assert_eq!(s.field(0).unwrap().data_type, DataType::Float64); + } + + #[test] + fn order_by_sorts_before_the_projection_so_unselected_columns_still_work() { + let text = explain("SELECT id FROM orders ORDER BY amount DESC"); + let lines: Vec<&str> = text.lines().collect(); + assert!(lines[0].starts_with("Projection")); + assert!(lines[1].trim_start().starts_with("Sort: amount#3 DESC")); + } + + #[test] + fn order_by_accepts_an_output_position_or_an_alias() { + assert!(explain("SELECT amount AS a FROM orders ORDER BY a").contains("Sort: amount#3")); + assert!(explain("SELECT amount FROM orders ORDER BY 1").contains("Sort: amount#3")); + assert!(error("SELECT amount FROM orders ORDER BY 4").contains("position 4 is not")); + assert!(error("SELECT * FROM orders ORDER BY 1").contains("refers to `*`")); + } + + #[test] + fn group_by_accepts_an_output_position() { + assert!(explain("SELECT region, count(*) FROM orders GROUP BY 1") + .contains("group by [region#2]")); + } + + #[test] + fn distinct_and_limit_wrap_the_query_outermost() { + let text = explain("SELECT DISTINCT region FROM orders LIMIT 5 OFFSET 2"); + let lines: Vec<&str> = text.lines().collect(); + assert_eq!(lines[0], "Limit: 5, offset 2"); + assert_eq!(lines[1].trim(), "Distinct"); + assert!(lines[2].trim().starts_with("Projection")); + } + + #[test] + fn a_query_with_no_limit_and_no_offset_gets_neither_node() { + assert!(!explain("SELECT id FROM orders").contains("Limit")); + assert!(explain("SELECT id FROM orders OFFSET 3").contains("Offset: 3")); + } + + #[test] + fn a_select_without_from_runs_over_one_empty_row() { + let text = explain("SELECT 1 + 1"); + assert!(text.contains("Values: 1 rows")); + assert!(text.starts_with("Projection: (1 + 1)")); + } + + #[test] + fn a_missing_table_is_reported() { + assert!(error("SELECT id FROM ghosts").contains("no such table")); + } + + #[test] + fn create_table_columns_become_a_schema_and_duplicates_are_refused() { + use qf_sql::ast::ColumnDef; + let cols = vec![ + ColumnDef { + name: "a".into(), + data_type: DataType::Int64, + nullable: false, + }, + ColumnDef { + name: "b".into(), + data_type: DataType::Utf8, + nullable: true, + }, + ]; + let s = schema_from_columns(&cols).unwrap(); + assert_eq!(s.len(), 2); + assert!(!s.field(0).unwrap().nullable); + + let dup = vec![cols[0].clone(), cols[0].clone()]; + assert!(schema_from_columns(&dup) + .unwrap_err() + .to_string() + .contains("duplicate column")); + } + + fn find_join(plan: &LogicalPlan) -> Option<&LogicalPlan> { + find(plan, &|p| matches!(p, LogicalPlan::Join { .. })) + } + + fn find_aggregate(plan: &LogicalPlan) -> Option<&LogicalPlan> { + find(plan, &|p| matches!(p, LogicalPlan::Aggregate { .. })) + } + + fn find<'p>( + plan: &'p LogicalPlan, + pred: &dyn Fn(&LogicalPlan) -> bool, + ) -> Option<&'p LogicalPlan> { + if pred(plan) { + return Some(plan); + } + plan.children().into_iter().find_map(|c| find(c, pred)) + } +} diff --git a/crates/qf-plan/src/cost.rs b/crates/qf-plan/src/cost.rs new file mode 100644 index 0000000..73c9530 --- /dev/null +++ b/crates/qf-plan/src/cost.rs @@ -0,0 +1,589 @@ +//! Cardinality estimation. +//! +//! Everything here is an estimate, and the numbers it produces are only ever +//! used to *choose between* plans, never to decide what a query returns. A bad +//! estimate costs a slower plan; it cannot cost a wrong answer. That is why +//! the defaults below are allowed to be crude. +//! +//! Where a column's zone map is available — which is exactly at a scan — the +//! estimate uses real statistics. Above a scan it falls back to the standard +//! textbook constants, and says so. + +use crate::expr::BoundExpr; +use crate::logical::LogicalPlan; +use qf_sql::ast::{BinaryOp, UnaryOp}; +use qf_storage::stats::{ColumnStats, TableStats}; + +/// Fraction of rows a range comparison is assumed to keep when nothing better +/// is known. The classic value, and about as defensible as any constant. +const RANGE_SELECTIVITY: f64 = 1.0 / 3.0; +/// Fallback for a predicate whose shape the estimator does not model. +const UNKNOWN_SELECTIVITY: f64 = 0.25; +/// Fallback join selectivity when neither key has usable statistics. +const UNKNOWN_JOIN_SELECTIVITY: f64 = 0.1; + +/// Estimated rows out of a plan node. +pub fn estimate_rows(plan: &LogicalPlan) -> f64 { + match plan { + LogicalPlan::Scan { + pushed_filters, + stats, + .. + } => { + let base = stats.row_count as f64; + let s: f64 = pushed_filters + .iter() + .map(|f| selectivity(f, Some(stats))) + .product(); + (base * s).max(1.0) + } + LogicalPlan::Filter { input, predicate } => { + let stats = source_stats(input); + (estimate_rows(input) * selectivity(predicate, stats)).max(1.0) + } + LogicalPlan::Projection { input, .. } => estimate_rows(input), + LogicalPlan::Aggregate { + input, group_exprs, .. + } => { + if group_exprs.is_empty() { + return 1.0; + } + // Grouping cannot produce more rows than it consumes; without a + // distinct count for the group keys, assume a tenth. + let rows = estimate_rows(input); + (rows * 0.1).clamp(1.0, rows) + } + LogicalPlan::Join { + left, + right, + join_type, + on, + .. + } => { + let (l, r) = (estimate_rows(left), estimate_rows(right)); + let product = l * r; + if on.is_empty() { + // A cross join really does produce the product. + return product.max(1.0); + } + // The textbook equi-join estimate: dividing by the larger distinct + // count assumes the smaller key set is contained in the larger. + let divisor: f64 = on + .iter() + .map(|(lk, rk)| { + let ld = key_distinct(left, lk); + let rd = key_distinct(right, rk); + match (ld, rd) { + (Some(a), Some(b)) => a.max(b).max(1.0), + (Some(a), None) | (None, Some(a)) => a.max(1.0), + (None, None) => 1.0 / UNKNOWN_JOIN_SELECTIVITY, + } + }) + .product(); + let estimate = (product / divisor.max(1.0)).max(1.0); + match join_type { + // An outer join emits at least every row of its preserved side. + qf_sql::ast::JoinType::Left => estimate.max(l), + qf_sql::ast::JoinType::Right => estimate.max(r), + qf_sql::ast::JoinType::Full => estimate.max(l).max(r), + _ => estimate, + } + } + LogicalPlan::Sort { input, .. } => estimate_rows(input), + LogicalPlan::Distinct { input } => { + let rows = estimate_rows(input); + (rows * 0.5).clamp(1.0, rows) + } + LogicalPlan::Limit { + input, + limit, + offset, + } => { + let rows = (estimate_rows(input) - *offset as f64).max(0.0); + match limit { + Some(l) => rows.min(*l as f64).max(1.0), + None => rows.max(1.0), + } + } + LogicalPlan::Values { rows, .. } => rows.len().max(1) as f64, + } +} + +/// The table statistics behind a subtree, if it is a scan (possibly with +/// filters or a projection stacked on it). +fn source_stats(plan: &LogicalPlan) -> Option<&TableStats> { + match plan { + LogicalPlan::Scan { stats, .. } => Some(stats), + LogicalPlan::Filter { input, .. } | LogicalPlan::Sort { input, .. } => source_stats(input), + _ => None, + } +} + +/// Distinct values of a join key, when the key is a plain column of a scan. +fn key_distinct(plan: &LogicalPlan, key: &BoundExpr) -> Option { + let BoundExpr::Column { index, .. } = key else { + return None; + }; + let stats = source_stats(plan)?; + let column = column_stats_for(plan, stats, *index)?; + if column.distinct_count == 0 { + return None; + } + Some(column.distinct_count as f64) +} + +/// Maps an output column position back to the statistics of the source column, +/// accounting for a scan that has already been projected. +fn column_stats_for<'a>( + plan: &LogicalPlan, + stats: &'a TableStats, + index: usize, +) -> Option<&'a ColumnStats> { + let source_index = match scan_of(plan) { + Some(LogicalPlan::Scan { + projection: Some(p), + .. + }) => *p.get(index)?, + _ => index, + }; + stats.column(source_index) +} + +fn scan_of(plan: &LogicalPlan) -> Option<&LogicalPlan> { + match plan { + LogicalPlan::Scan { .. } => Some(plan), + LogicalPlan::Filter { input, .. } | LogicalPlan::Sort { input, .. } => scan_of(input), + _ => None, + } +} + +/// The fraction of rows a predicate is expected to keep, in `[0, 1]`. +pub fn selectivity(expr: &BoundExpr, stats: Option<&TableStats>) -> f64 { + match expr { + BoundExpr::Literal(qf_common::Value::Boolean(true)) => 1.0, + BoundExpr::Literal(qf_common::Value::Boolean(false)) => 0.0, + BoundExpr::Binary { + left, op, right, .. + } => match op { + BinaryOp::And => selectivity(left, stats) * selectivity(right, stats), + // Assumes independence, which is wrong whenever the two sides + // correlate — the usual failure mode of this whole approach. + BinaryOp::Or => { + let (a, b) = (selectivity(left, stats), selectivity(right, stats)); + a + b - a * b + } + BinaryOp::Eq => column_literal(left, right) + .and_then(|i| stats?.column(i)) + .map_or(0.1, ColumnStats::selectivity_eq), + BinaryOp::NotEq => { + 1.0 - column_literal(left, right) + .and_then(|i| stats?.column(i)) + .map_or(0.1, ColumnStats::selectivity_eq) + } + BinaryOp::Lt | BinaryOp::LtEq | BinaryOp::Gt | BinaryOp::GtEq => RANGE_SELECTIVITY, + _ => UNKNOWN_SELECTIVITY, + }, + BoundExpr::Unary { + op: UnaryOp::Not, + expr, + .. + } => 1.0 - selectivity(expr, stats), + BoundExpr::IsNull { expr, negated } => { + let fraction = match (expr.as_ref(), stats) { + (BoundExpr::Column { index, .. }, Some(s)) => match s.column(*index) { + Some(c) if c.row_count > 0 => c.null_count as f64 / c.row_count as f64, + _ => 0.1, + }, + _ => 0.1, + }; + if *negated { + 1.0 - fraction + } else { + fraction + } + } + BoundExpr::InList { + expr, + list, + negated, + } => { + let one = match (expr.as_ref(), stats) { + (BoundExpr::Column { index, .. }, Some(s)) => { + s.column(*index).map_or(0.1, ColumnStats::selectivity_eq) + } + _ => 0.1, + }; + let s = (one * list.len() as f64).min(1.0); + if *negated { + 1.0 - s + } else { + s + } + } + // A LIKE pattern's selectivity depends on the pattern in ways this + // estimator does not model; a leading-wildcard pattern in particular + // may match everything. + BoundExpr::Like { negated, .. } => { + if *negated { + 1.0 - UNKNOWN_SELECTIVITY + } else { + UNKNOWN_SELECTIVITY + } + } + _ => UNKNOWN_SELECTIVITY, + } +} + +/// Recognises `column literal` in either order, returning the column +/// position. +fn column_literal(left: &BoundExpr, right: &BoundExpr) -> Option { + match (left, right) { + (BoundExpr::Column { index, .. }, r) if r.is_literal() => Some(*index), + (l, BoundExpr::Column { index, .. }) if l.is_literal() => Some(*index), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use qf_common::{DataType, Field, Schema, Value}; + use qf_sql::ast::JoinType; + use std::sync::Arc; + + fn stats(row_count: usize, distinct: &[usize], nulls: &[usize]) -> TableStats { + TableStats { + row_count, + columns: distinct + .iter() + .zip(nulls) + .map(|(d, n)| ColumnStats { + min: Value::Int64(0), + max: Value::Int64(100), + null_count: *n, + row_count, + distinct_count: *d, + }) + .collect(), + } + } + + fn schema(n: usize) -> Arc { + Arc::new(Schema::new( + (0..n) + .map(|i| Field::new(format!("c{i}"), DataType::Int64, true)) + .collect(), + )) + } + + fn scan(rows: usize, distinct: &[usize]) -> LogicalPlan { + let nulls = vec![0; distinct.len()]; + LogicalPlan::Scan { + table: "t".into(), + source_schema: schema(distinct.len()), + projection: None, + pushed_filters: vec![], + stats: stats(rows, distinct, &nulls), + } + } + + fn col(i: usize) -> BoundExpr { + BoundExpr::column(i, format!("c{i}"), DataType::Int64) + } + + fn eq(i: usize, v: i64) -> BoundExpr { + BoundExpr::binary(col(i), BinaryOp::Eq, BoundExpr::Literal(Value::Int64(v))).unwrap() + } + + #[test] + fn a_scan_estimates_its_own_row_count() { + assert_eq!(estimate_rows(&scan(1000, &[10])), 1000.0); + } + + #[test] + fn equality_on_a_ten_value_column_keeps_a_tenth() { + let plan = LogicalPlan::Filter { + input: Box::new(scan(1000, &[10])), + predicate: eq(0, 5), + }; + assert!((estimate_rows(&plan) - 100.0).abs() < 1e-6); + } + + #[test] + fn a_pushed_filter_is_reflected_in_the_scan_s_own_estimate() { + let mut s = scan(1000, &[10]); + if let LogicalPlan::Scan { pushed_filters, .. } = &mut s { + pushed_filters.push(eq(0, 5)); + } + assert!((estimate_rows(&s) - 100.0).abs() < 1e-6); + } + + #[test] + fn conjunctions_multiply_and_disjunctions_do_not() { + let and = BoundExpr::binary(eq(0, 1), BinaryOp::And, eq(1, 2)).unwrap(); + let or = BoundExpr::binary(eq(0, 1), BinaryOp::Or, eq(1, 2)).unwrap(); + let s = stats(1000, &[10, 10], &[0, 0]); + let sa = selectivity(&and, Some(&s)); + let so = selectivity(&or, Some(&s)); + assert!((sa - 0.01).abs() < 1e-9); + assert!((so - 0.19).abs() < 1e-9); + assert!(so > sa); + } + + #[test] + fn a_range_predicate_uses_the_standard_third() { + let lt = + BoundExpr::binary(col(0), BinaryOp::Lt, BoundExpr::Literal(Value::Int64(5))).unwrap(); + assert!((selectivity(<, None) - RANGE_SELECTIVITY).abs() < 1e-9); + } + + #[test] + fn negation_inverts_a_selectivity() { + let e = eq(0, 1); + let s = stats(100, &[4], &[0]); + let not = BoundExpr::unary(UnaryOp::Not, e.clone()).unwrap(); + assert!((selectivity(¬, Some(&s)) + selectivity(&e, Some(&s)) - 1.0).abs() < 1e-9); + let ne = BoundExpr::binary(col(0), BinaryOp::NotEq, BoundExpr::Literal(Value::Int64(1))) + .unwrap(); + assert!((selectivity(&ne, Some(&s)) - 0.75).abs() < 1e-9); + } + + #[test] + fn a_literal_predicate_is_all_or_nothing() { + assert_eq!( + selectivity(&BoundExpr::Literal(Value::Boolean(true)), None), + 1.0 + ); + assert_eq!( + selectivity(&BoundExpr::Literal(Value::Boolean(false)), None), + 0.0 + ); + } + + #[test] + fn is_null_uses_the_column_s_recorded_null_count() { + let s = stats(100, &[10], &[25]); + let is_null = BoundExpr::IsNull { + expr: Box::new(col(0)), + negated: false, + }; + let not_null = BoundExpr::IsNull { + expr: Box::new(col(0)), + negated: true, + }; + assert!((selectivity(&is_null, Some(&s)) - 0.25).abs() < 1e-9); + assert!((selectivity(¬_null, Some(&s)) - 0.75).abs() < 1e-9); + // With no statistics it falls back rather than guessing zero. + assert!(selectivity(&is_null, None) > 0.0); + } + + #[test] + fn an_in_list_scales_with_its_length_and_saturates_at_one() { + let s = stats(100, &[10], &[0]); + let short = BoundExpr::InList { + expr: Box::new(col(0)), + list: vec![BoundExpr::Literal(Value::Int64(1))], + negated: false, + }; + let long = BoundExpr::InList { + expr: Box::new(col(0)), + list: (0..50) + .map(|i| BoundExpr::Literal(Value::Int64(i))) + .collect(), + negated: false, + }; + assert!((selectivity(&short, Some(&s)) - 0.1).abs() < 1e-9); + assert_eq!(selectivity(&long, Some(&s)), 1.0); + let negated = BoundExpr::InList { + expr: Box::new(col(0)), + list: vec![BoundExpr::Literal(Value::Int64(1))], + negated: true, + }; + assert!((selectivity(&negated, Some(&s)) - 0.9).abs() < 1e-9); + } + + #[test] + fn like_is_treated_as_unknown_in_both_polarities() { + let like = BoundExpr::Like { + expr: Box::new(col(0)), + pattern: Box::new(BoundExpr::Literal(Value::Utf8("a%".into()))), + negated: false, + }; + assert!((selectivity(&like, None) - UNKNOWN_SELECTIVITY).abs() < 1e-9); + let not_like = BoundExpr::Like { + expr: Box::new(col(0)), + pattern: Box::new(BoundExpr::Literal(Value::Utf8("a%".into()))), + negated: true, + }; + assert!((selectivity(¬_like, None) - 0.75).abs() < 1e-9); + } + + #[test] + fn an_equijoin_divides_by_the_larger_distinct_count() { + // 1000 x 100 rows joined on a key with 100 distinct values on the + // right and 10 on the left → 1000 * 100 / 100. + let join = LogicalPlan::Join { + left: Box::new(scan(1000, &[10])), + right: Box::new(scan(100, &[100])), + join_type: JoinType::Inner, + on: vec![(col(0), col(0))], + filter: None, + }; + assert!((estimate_rows(&join) - 1000.0).abs() < 1e-6); + } + + #[test] + fn a_cross_join_estimates_the_full_product() { + let join = LogicalPlan::Join { + left: Box::new(scan(50, &[5])), + right: Box::new(scan(20, &[5])), + join_type: JoinType::Cross, + on: vec![], + filter: None, + }; + assert!((estimate_rows(&join) - 1000.0).abs() < 1e-6); + } + + #[test] + fn an_outer_join_never_estimates_fewer_rows_than_its_preserved_side() { + let mk = |t| LogicalPlan::Join { + left: Box::new(scan(1000, &[1000])), + right: Box::new(scan(2, &[2])), + join_type: t, + on: vec![(col(0), col(0))], + filter: None, + }; + assert!(estimate_rows(&mk(JoinType::Left)) >= 1000.0); + assert!(estimate_rows(&mk(JoinType::Right)) >= 2.0); + assert!(estimate_rows(&mk(JoinType::Full)) >= 1000.0); + assert!(estimate_rows(&mk(JoinType::Inner)) < 1000.0); + } + + #[test] + fn a_join_key_that_is_not_a_plain_column_falls_back_to_a_constant() { + let expr = + BoundExpr::binary(col(0), BinaryOp::Plus, BoundExpr::Literal(Value::Int64(1))).unwrap(); + let join = LogicalPlan::Join { + left: Box::new(scan(100, &[10])), + right: Box::new(scan(100, &[10])), + join_type: JoinType::Inner, + on: vec![(expr.clone(), expr)], + filter: None, + }; + // 100 * 100 * 0.1 + assert!((estimate_rows(&join) - 1000.0).abs() < 1e-6); + } + + #[test] + fn aggregates_collapse_and_a_global_aggregate_yields_one_row() { + let grouped = LogicalPlan::Aggregate { + input: Box::new(scan(1000, &[10])), + group_exprs: vec![(col(0), "c0".into())], + aggregates: vec![], + }; + assert!(estimate_rows(&grouped) < 1000.0); + let global = LogicalPlan::Aggregate { + input: Box::new(scan(1000, &[10])), + group_exprs: vec![], + aggregates: vec![], + }; + assert_eq!(estimate_rows(&global), 1.0); + } + + #[test] + fn limit_caps_the_estimate_and_offset_reduces_it() { + let input = Box::new(scan(1000, &[10])); + assert_eq!( + estimate_rows(&LogicalPlan::Limit { + input: input.clone(), + limit: Some(10), + offset: 0 + }), + 10.0 + ); + assert_eq!( + estimate_rows(&LogicalPlan::Limit { + input: input.clone(), + limit: None, + offset: 400 + }), + 600.0 + ); + assert_eq!( + estimate_rows(&LogicalPlan::Limit { + input, + limit: Some(10), + offset: 5000 + }), + 1.0 + ); + } + + #[test] + fn projection_sort_and_distinct_pass_through_or_shrink() { + let s = scan(100, &[10]); + assert_eq!( + estimate_rows(&LogicalPlan::Projection { + input: Box::new(s.clone()), + exprs: vec![(col(0), "c0".into())] + }), + 100.0 + ); + assert_eq!( + estimate_rows(&LogicalPlan::Sort { + input: Box::new(s.clone()), + exprs: vec![] + }), + 100.0 + ); + assert!(estimate_rows(&LogicalPlan::Distinct { input: Box::new(s) }) < 100.0); + } + + #[test] + fn values_estimates_its_own_row_count() { + assert_eq!( + estimate_rows(&LogicalPlan::Values { + schema: schema(1), + rows: vec![vec![], vec![], vec![]] + }), + 3.0 + ); + } + + #[test] + fn an_estimate_never_drops_below_one_row() { + // Stacking selective filters must not talk the planner into believing + // a branch is empty and choosing a plan on that basis. + let mut plan = scan(10, &[1000]); + for i in 0..8 { + plan = LogicalPlan::Filter { + input: Box::new(plan), + predicate: eq(0, i), + }; + } + assert!(estimate_rows(&plan) >= 1.0); + } + + #[test] + fn statistics_are_read_through_a_projected_scan() { + // After projection pushdown the join key is position 0, but its + // statistics still live at the source position. + let s = LogicalPlan::Scan { + table: "t".into(), + source_schema: schema(3), + projection: Some(vec![2]), + pushed_filters: vec![], + stats: stats(1000, &[1, 1, 500], &[0, 0, 0]), + }; + let join = LogicalPlan::Join { + left: Box::new(s), + right: Box::new(scan(10, &[10])), + join_type: JoinType::Inner, + on: vec![(col(0), col(0))], + filter: None, + }; + // Divides by 500 (the real distinct count of the projected column), + // not by 1 (the count of source column 0). + assert!((estimate_rows(&join) - 20.0).abs() < 1e-6); + } +} diff --git a/crates/qf-plan/src/expr.rs b/crates/qf-plan/src/expr.rs new file mode 100644 index 0000000..8e1a386 --- /dev/null +++ b/crates/qf-plan/src/expr.rs @@ -0,0 +1,843 @@ +//! Expressions after binding: every column is a position, every node knows its +//! type, and nothing is left to resolve at execution time. +//! +//! The split from the AST matters. `qf_sql::Expr` is what the user typed; +//! `BoundExpr` is what the engine will run. Once an expression is bound, the +//! evaluator never looks up a name, never asks whether a type is legal, and +//! never has to report a user-facing error about either — all of that +//! happened here, once, at plan time. + +use qf_common::{DataType, Error, Result, Value}; +use qf_sql::ast::{BinaryOp, UnaryOp}; +use std::fmt; + +#[derive(Debug, Clone, PartialEq)] +pub enum BoundExpr { + /// A position in the input batch, not a name. + Column { + index: usize, + name: String, + data_type: DataType, + }, + Literal(Value), + Binary { + left: Box, + op: BinaryOp, + right: Box, + data_type: DataType, + }, + Unary { + op: UnaryOp, + expr: Box, + data_type: DataType, + }, + Cast { + expr: Box, + data_type: DataType, + }, + IsNull { + expr: Box, + negated: bool, + }, + InList { + expr: Box, + list: Vec, + negated: bool, + }, + Like { + expr: Box, + pattern: Box, + negated: bool, + }, + Case { + branches: Vec<(BoundExpr, BoundExpr)>, + else_result: Option>, + data_type: DataType, + }, +} + +impl BoundExpr { + pub fn data_type(&self) -> DataType { + match self { + BoundExpr::Column { data_type, .. } + | BoundExpr::Binary { data_type, .. } + | BoundExpr::Unary { data_type, .. } + | BoundExpr::Cast { data_type, .. } + | BoundExpr::Case { data_type, .. } => *data_type, + BoundExpr::IsNull { .. } | BoundExpr::Like { .. } | BoundExpr::InList { .. } => { + DataType::Boolean + } + // A bare NULL literal has no type of its own; it takes the type of + // whatever it is compared against, and Int64 is the harmless + // default for the cases where nothing constrains it. + BoundExpr::Literal(v) => v.data_type().unwrap_or(DataType::Int64), + } + } + + pub fn column(index: usize, name: impl Into, data_type: DataType) -> BoundExpr { + BoundExpr::Column { + index, + name: name.into(), + data_type, + } + } + + /// Builds a binary node, inferring and checking its result type. + /// + /// This is where a query like `WHERE name > 3` is rejected — at plan time, + /// before a single row is read. + pub fn binary(left: BoundExpr, op: BinaryOp, right: BoundExpr) -> Result { + let (lt, rt) = (left.data_type(), right.data_type()); + let data_type = if op.is_logical() { + for (side, t) in [("left", lt), ("right", rt)] { + if t != DataType::Boolean { + return Err(Error::typ(format!( + "{op} needs a boolean on its {side}, found {t}" + ))); + } + } + DataType::Boolean + } else if op.is_comparison() { + // NULL compares against anything; otherwise the types must reconcile. + if !left.is_null_literal() && !right.is_null_literal() { + DataType::unify(lt, rt) + .map_err(|_| Error::typ(format!("cannot compare {lt} with {rt} using {op}")))?; + } + DataType::Boolean + } else { + if !lt.is_numeric() || !rt.is_numeric() { + return Err(Error::typ(format!( + "{op} needs numbers, found {lt} and {rt}" + ))); + } + // Division always widens: integer division silently truncating + // `total / count` is a classic source of quietly wrong reports. + if op == BinaryOp::Divide { + DataType::Float64 + } else { + DataType::unify(lt, rt)? + } + }; + Ok(BoundExpr::Binary { + left: Box::new(left), + op, + right: Box::new(right), + data_type, + }) + } + + pub fn unary(op: UnaryOp, expr: BoundExpr) -> Result { + let t = expr.data_type(); + let data_type = match op { + UnaryOp::Neg => { + if !t.is_numeric() { + return Err(Error::typ(format!("cannot negate {t}"))); + } + t + } + UnaryOp::Not => { + if t != DataType::Boolean { + return Err(Error::typ(format!("NOT needs a boolean, found {t}"))); + } + DataType::Boolean + } + }; + Ok(BoundExpr::Unary { + op, + expr: Box::new(expr), + data_type, + }) + } + + /// All branches of a CASE must agree on a type, so the column it produces + /// has one. + pub fn case( + branches: Vec<(BoundExpr, BoundExpr)>, + else_result: Option, + ) -> Result { + if branches.is_empty() { + return Err(Error::plan( + "CASE needs at least one WHEN branch".to_string(), + )); + } + for (when, _) in &branches { + if when.data_type() != DataType::Boolean { + return Err(Error::typ(format!( + "a CASE condition must be boolean, found {}", + when.data_type() + ))); + } + } + let mut data_type = None; + for result in branches.iter().map(|(_, t)| t).chain(else_result.iter()) { + if result.is_null_literal() { + continue; + } + data_type = Some(match data_type { + None => result.data_type(), + Some(t) => DataType::unify(t, result.data_type()).map_err(|_| { + Error::typ(format!( + "CASE branches disagree on type: {t} and {}", + result.data_type() + )) + })?, + }); + } + Ok(BoundExpr::Case { + branches, + else_result: else_result.map(Box::new), + data_type: data_type.unwrap_or(DataType::Int64), + }) + } + + pub fn is_null_literal(&self) -> bool { + matches!(self, BoundExpr::Literal(Value::Null)) + } + + pub fn is_literal(&self) -> bool { + matches!(self, BoundExpr::Literal(_)) + } + + pub fn as_literal(&self) -> Option<&Value> { + match self { + BoundExpr::Literal(v) => Some(v), + _ => None, + } + } + + /// Every column position this expression reads. + pub fn column_indices(&self) -> Vec { + let mut out = Vec::new(); + self.walk(&mut |e| { + if let BoundExpr::Column { index, .. } = e { + out.push(*index); + } + }); + out.sort_unstable(); + out.dedup(); + out + } + + pub fn walk(&self, f: &mut impl FnMut(&BoundExpr)) { + f(self); + match self { + BoundExpr::Column { .. } | BoundExpr::Literal(_) => {} + BoundExpr::Binary { left, right, .. } => { + left.walk(f); + right.walk(f); + } + BoundExpr::Unary { expr, .. } + | BoundExpr::Cast { expr, .. } + | BoundExpr::IsNull { expr, .. } => expr.walk(f), + BoundExpr::InList { expr, list, .. } => { + expr.walk(f); + list.iter().for_each(|e| e.walk(f)); + } + BoundExpr::Like { expr, pattern, .. } => { + expr.walk(f); + pattern.walk(f); + } + BoundExpr::Case { + branches, + else_result, + .. + } => { + for (w, t) in branches { + w.walk(f); + t.walk(f); + } + if let Some(e) = else_result { + e.walk(f); + } + } + } + } + + /// Rebuilds the expression with every column index remapped. + /// + /// Needed whenever a rule changes the shape of an operator's input — a + /// scan that stops reading columns it does not need, or a join whose + /// inputs are reordered. Getting this wrong is silent: the query still + /// runs and reads the wrong column. + pub fn remap_columns(&self, map: &dyn Fn(usize) -> Option) -> Result { + Ok(match self { + BoundExpr::Column { + index, + name, + data_type, + } => { + let new = map(*index).ok_or_else(|| { + Error::internal(format!( + "column `{name}` (index {index}) has no place in the new input" + )) + })?; + BoundExpr::Column { + index: new, + name: name.clone(), + data_type: *data_type, + } + } + BoundExpr::Literal(v) => BoundExpr::Literal(v.clone()), + BoundExpr::Binary { + left, + op, + right, + data_type, + } => BoundExpr::Binary { + left: Box::new(left.remap_columns(map)?), + op: *op, + right: Box::new(right.remap_columns(map)?), + data_type: *data_type, + }, + BoundExpr::Unary { + op, + expr, + data_type, + } => BoundExpr::Unary { + op: *op, + expr: Box::new(expr.remap_columns(map)?), + data_type: *data_type, + }, + BoundExpr::Cast { expr, data_type } => BoundExpr::Cast { + expr: Box::new(expr.remap_columns(map)?), + data_type: *data_type, + }, + BoundExpr::IsNull { expr, negated } => BoundExpr::IsNull { + expr: Box::new(expr.remap_columns(map)?), + negated: *negated, + }, + BoundExpr::InList { + expr, + list, + negated, + } => BoundExpr::InList { + expr: Box::new(expr.remap_columns(map)?), + list: list + .iter() + .map(|e| e.remap_columns(map)) + .collect::>>()?, + negated: *negated, + }, + BoundExpr::Like { + expr, + pattern, + negated, + } => BoundExpr::Like { + expr: Box::new(expr.remap_columns(map)?), + pattern: Box::new(pattern.remap_columns(map)?), + negated: *negated, + }, + BoundExpr::Case { + branches, + else_result, + data_type, + } => BoundExpr::Case { + branches: branches + .iter() + .map(|(w, t)| Ok((w.remap_columns(map)?, t.remap_columns(map)?))) + .collect::>>()?, + else_result: match else_result { + Some(e) => Some(Box::new(e.remap_columns(map)?)), + None => None, + }, + data_type: *data_type, + }, + }) + } + + /// A readable name for the column this expression produces, used when the + /// user did not supply an alias. + pub fn output_name(&self) -> String { + match self { + BoundExpr::Column { name, .. } => name.clone(), + other => other.to_string(), + } + } +} + +impl fmt::Display for BoundExpr { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + BoundExpr::Column { name, index, .. } => write!(f, "{name}#{index}"), + BoundExpr::Literal(Value::Utf8(s)) => write!(f, "'{s}'"), + BoundExpr::Literal(v) => write!(f, "{v}"), + BoundExpr::Binary { + left, op, right, .. + } => write!(f, "({left} {op} {right})"), + BoundExpr::Unary { op, expr, .. } => write!(f, "{op}{expr}"), + BoundExpr::Cast { expr, data_type } => write!(f, "CAST({expr} AS {data_type})"), + BoundExpr::IsNull { expr, negated } => { + let not = if *negated { "NOT " } else { "" }; + write!(f, "{expr} IS {not}NULL") + } + BoundExpr::InList { + expr, + list, + negated, + } => { + let not = if *negated { "NOT " } else { "" }; + let items: Vec = list.iter().map(BoundExpr::to_string).collect(); + write!(f, "{expr} {not}IN ({})", items.join(", ")) + } + BoundExpr::Like { + expr, + pattern, + negated, + } => { + let not = if *negated { "NOT " } else { "" }; + write!(f, "{expr} {not}LIKE {pattern}") + } + BoundExpr::Case { + branches, + else_result, + .. + } => { + write!(f, "CASE")?; + for (w, t) in branches { + write!(f, " WHEN {w} THEN {t}")?; + } + if let Some(e) = else_result { + write!(f, " ELSE {e}")?; + } + f.write_str(" END") + } + } + } +} + +/// The aggregate functions the engine implements. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum AggFunc { + Count, + Sum, + Min, + Max, + Avg, +} + +impl AggFunc { + pub fn from_name(name: &str) -> Option { + Some(match name.to_ascii_lowercase().as_str() { + "count" => AggFunc::Count, + "sum" => AggFunc::Sum, + "min" => AggFunc::Min, + "max" => AggFunc::Max, + "avg" => AggFunc::Avg, + _ => return None, + }) + } + + pub fn name(self) -> &'static str { + match self { + AggFunc::Count => "count", + AggFunc::Sum => "sum", + AggFunc::Min => "min", + AggFunc::Max => "max", + AggFunc::Avg => "avg", + } + } + + /// The type this aggregate produces over an input of `input` type. + pub fn output_type(self, input: Option) -> Result { + Ok(match self { + AggFunc::Count => DataType::Int64, + // avg of integers is a float — the same reasoning as division. + AggFunc::Avg => DataType::Float64, + AggFunc::Sum => match input { + Some(t) if t.is_numeric() => t, + Some(t) => return Err(Error::typ(format!("cannot sum {t}"))), + None => return Err(Error::plan("sum() needs an argument".to_string())), + }, + AggFunc::Min | AggFunc::Max => { + input.ok_or_else(|| Error::plan(format!("{}() needs an argument", self.name())))? + } + }) + } + + /// Whether `f(*)` is meaningful. Only `count` is. + pub fn accepts_wildcard(self) -> bool { + self == AggFunc::Count + } +} + +#[derive(Debug, Clone, PartialEq)] +pub struct BoundAggregate { + pub func: AggFunc, + /// `None` only for `count(*)`. + pub arg: Option, + pub distinct: bool, + pub output_name: String, + pub data_type: DataType, +} + +impl fmt::Display for BoundAggregate { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let d = if self.distinct { "DISTINCT " } else { "" }; + match &self.arg { + Some(a) => write!(f, "{}({d}{a})", self.func.name()), + None => write!(f, "{}(*)", self.func.name()), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn col(i: usize, t: DataType) -> BoundExpr { + BoundExpr::column(i, format!("c{i}"), t) + } + + #[test] + fn comparisons_produce_a_boolean_whatever_they_compare() { + let e = BoundExpr::binary( + col(0, DataType::Int64), + BinaryOp::Lt, + BoundExpr::Literal(Value::Float64(1.5)), + ) + .unwrap(); + assert_eq!(e.data_type(), DataType::Boolean); + } + + #[test] + fn comparing_text_with_a_number_is_rejected_at_plan_time() { + let err = BoundExpr::binary( + col(0, DataType::Utf8), + BinaryOp::Gt, + BoundExpr::Literal(Value::Int64(3)), + ) + .unwrap_err(); + assert!(matches!(err, Error::Type(_))); + assert!(err.to_string().contains("cannot compare")); + } + + #[test] + fn null_may_be_compared_with_anything() { + assert!(BoundExpr::binary( + col(0, DataType::Utf8), + BinaryOp::Eq, + BoundExpr::Literal(Value::Null) + ) + .is_ok()); + } + + #[test] + fn arithmetic_widens_int_and_float_and_refuses_text() { + let e = BoundExpr::binary( + col(0, DataType::Int64), + BinaryOp::Plus, + col(1, DataType::Float64), + ) + .unwrap(); + assert_eq!(e.data_type(), DataType::Float64); + assert!(BoundExpr::binary( + col(0, DataType::Utf8), + BinaryOp::Plus, + col(1, DataType::Int64) + ) + .is_err()); + } + + #[test] + fn dividing_two_integers_yields_a_float() { + // Integer division truncating `total / count` is a quiet way to report + // the wrong number, so division always widens. + let e = BoundExpr::binary( + col(0, DataType::Int64), + BinaryOp::Divide, + col(1, DataType::Int64), + ) + .unwrap(); + assert_eq!(e.data_type(), DataType::Float64); + } + + #[test] + fn logical_operators_demand_booleans_on_both_sides() { + assert!(BoundExpr::binary( + col(0, DataType::Boolean), + BinaryOp::And, + col(1, DataType::Boolean) + ) + .is_ok()); + let err = BoundExpr::binary( + col(0, DataType::Int64), + BinaryOp::And, + col(1, DataType::Boolean), + ) + .unwrap_err(); + assert!(err.to_string().contains("on its left")); + assert!(BoundExpr::binary( + col(0, DataType::Boolean), + BinaryOp::Or, + col(1, DataType::Int64) + ) + .unwrap_err() + .to_string() + .contains("on its right")); + } + + #[test] + fn unary_operators_check_their_operand() { + assert!(BoundExpr::unary(UnaryOp::Neg, col(0, DataType::Int64)).is_ok()); + assert!(BoundExpr::unary(UnaryOp::Neg, col(0, DataType::Utf8)).is_err()); + assert!(BoundExpr::unary(UnaryOp::Not, col(0, DataType::Boolean)).is_ok()); + assert!(BoundExpr::unary(UnaryOp::Not, col(0, DataType::Int64)).is_err()); + } + + #[test] + fn case_branches_must_agree_on_a_type() { + let ok = BoundExpr::case( + vec![(col(0, DataType::Boolean), col(1, DataType::Int64))], + Some(col(2, DataType::Float64)), + ) + .unwrap(); + assert_eq!(ok.data_type(), DataType::Float64); + + assert!(BoundExpr::case( + vec![(col(0, DataType::Boolean), col(1, DataType::Int64))], + Some(col(2, DataType::Utf8)), + ) + .unwrap_err() + .to_string() + .contains("disagree on type")); + } + + #[test] + fn a_case_condition_must_be_boolean_and_a_null_branch_does_not_decide_the_type() { + assert!(BoundExpr::case( + vec![(col(0, DataType::Int64), col(1, DataType::Int64))], + None + ) + .is_err()); + let e = BoundExpr::case( + vec![(col(0, DataType::Boolean), BoundExpr::Literal(Value::Null))], + Some(col(2, DataType::Utf8)), + ) + .unwrap(); + assert_eq!(e.data_type(), DataType::Utf8); + assert!(BoundExpr::case(vec![], None).is_err()); + } + + #[test] + fn column_indices_are_collected_deduplicated_and_sorted() { + let e = BoundExpr::binary( + col(3, DataType::Int64), + BinaryOp::Plus, + BoundExpr::binary( + col(1, DataType::Int64), + BinaryOp::Plus, + col(3, DataType::Int64), + ) + .unwrap(), + ) + .unwrap(); + assert_eq!(e.column_indices(), vec![1, 3]); + } + + #[test] + fn remapping_rewrites_every_column_position() { + let e = BoundExpr::binary( + col(5, DataType::Int64), + BinaryOp::Lt, + col(9, DataType::Int64), + ) + .unwrap(); + let mapped = e + .remap_columns(&|i| match i { + 5 => Some(0), + 9 => Some(1), + _ => None, + }) + .unwrap(); + assert_eq!(mapped.column_indices(), vec![0, 1]); + } + + #[test] + fn remapping_a_column_with_nowhere_to_go_is_an_error_not_a_silent_shift() { + let e = col(7, DataType::Int64); + assert!(e.remap_columns(&|_| None).is_err()); + } + + #[test] + fn remapping_reaches_inside_every_expression_shape() { + let shapes = vec![ + BoundExpr::Cast { + expr: Box::new(col(4, DataType::Int64)), + data_type: DataType::Float64, + }, + BoundExpr::IsNull { + expr: Box::new(col(4, DataType::Int64)), + negated: false, + }, + BoundExpr::InList { + expr: Box::new(col(4, DataType::Int64)), + list: vec![BoundExpr::Literal(Value::Int64(1))], + negated: false, + }, + BoundExpr::Like { + expr: Box::new(col(4, DataType::Utf8)), + pattern: Box::new(BoundExpr::Literal(Value::Utf8("a%".into()))), + negated: false, + }, + BoundExpr::unary(UnaryOp::Neg, col(4, DataType::Int64)).unwrap(), + BoundExpr::case( + vec![( + BoundExpr::IsNull { + expr: Box::new(col(4, DataType::Int64)), + negated: false, + }, + col(4, DataType::Int64), + )], + Some(col(4, DataType::Int64)), + ) + .unwrap(), + ]; + for s in shapes { + assert_eq!(s.column_indices(), vec![4], "{s}"); + let m = s.remap_columns(&|i| Some(i + 10)).unwrap(); + assert_eq!(m.column_indices(), vec![14], "{s}"); + } + } + + #[test] + fn bound_expressions_render_with_their_column_positions() { + assert_eq!(col(2, DataType::Int64).to_string(), "c2#2"); + assert_eq!( + BoundExpr::Literal(Value::Utf8("x".into())).to_string(), + "'x'" + ); + assert_eq!( + BoundExpr::binary( + col(0, DataType::Int64), + BinaryOp::Plus, + col(1, DataType::Int64) + ) + .unwrap() + .to_string(), + "(c0#0 + c1#1)" + ); + assert_eq!( + BoundExpr::IsNull { + expr: Box::new(col(0, DataType::Int64)), + negated: true + } + .to_string(), + "c0#0 IS NOT NULL" + ); + } + + #[test] + fn the_remaining_display_shapes_render() { + assert!(BoundExpr::InList { + expr: Box::new(col(0, DataType::Int64)), + list: vec![BoundExpr::Literal(Value::Int64(1))], + negated: true, + } + .to_string() + .contains("NOT IN")); + assert!(BoundExpr::Like { + expr: Box::new(col(0, DataType::Utf8)), + pattern: Box::new(BoundExpr::Literal(Value::Utf8("a%".into()))), + negated: false, + } + .to_string() + .contains("LIKE")); + assert!(BoundExpr::Cast { + expr: Box::new(col(0, DataType::Int64)), + data_type: DataType::Utf8, + } + .to_string() + .starts_with("CAST")); + assert!(BoundExpr::case( + vec![(col(0, DataType::Boolean), col(1, DataType::Int64))], + Some(col(2, DataType::Int64)) + ) + .unwrap() + .to_string() + .contains("ELSE")); + assert!(BoundExpr::unary(UnaryOp::Neg, col(0, DataType::Int64)) + .unwrap() + .to_string() + .starts_with('-')); + } + + #[test] + fn literals_and_output_names_are_reported() { + let lit = BoundExpr::Literal(Value::Int64(3)); + assert!(lit.is_literal()); + assert_eq!(lit.as_literal(), Some(&Value::Int64(3))); + assert!(!col(0, DataType::Int64).is_literal()); + assert!(col(0, DataType::Int64).as_literal().is_none()); + assert_eq!(col(0, DataType::Int64).output_name(), "c0"); + assert_eq!(lit.output_name(), "3"); + assert!(BoundExpr::Literal(Value::Null).is_null_literal()); + assert_eq!(BoundExpr::Literal(Value::Null).data_type(), DataType::Int64); + } + + #[test] + fn aggregate_output_types_follow_their_input() { + assert_eq!(AggFunc::Count.output_type(None).unwrap(), DataType::Int64); + assert_eq!( + AggFunc::Avg.output_type(Some(DataType::Int64)).unwrap(), + DataType::Float64 + ); + assert_eq!( + AggFunc::Sum.output_type(Some(DataType::Int64)).unwrap(), + DataType::Int64 + ); + assert_eq!( + AggFunc::Max.output_type(Some(DataType::Utf8)).unwrap(), + DataType::Utf8 + ); + } + + #[test] + fn summing_text_or_aggregating_nothing_is_rejected() { + assert!(AggFunc::Sum.output_type(Some(DataType::Utf8)).is_err()); + assert!(AggFunc::Sum.output_type(None).is_err()); + assert!(AggFunc::Min.output_type(None).is_err()); + } + + #[test] + fn only_count_accepts_a_wildcard() { + assert!(AggFunc::Count.accepts_wildcard()); + for f in [AggFunc::Sum, AggFunc::Min, AggFunc::Max, AggFunc::Avg] { + assert!(!f.accepts_wildcard()); + } + } + + #[test] + fn aggregate_names_round_trip() { + for f in [ + AggFunc::Count, + AggFunc::Sum, + AggFunc::Min, + AggFunc::Max, + AggFunc::Avg, + ] { + assert_eq!(AggFunc::from_name(f.name()), Some(f)); + } + assert_eq!(AggFunc::from_name("COUNT"), Some(AggFunc::Count)); + assert_eq!(AggFunc::from_name("median"), None); + } + + #[test] + fn aggregates_render_for_plan_output() { + let a = BoundAggregate { + func: AggFunc::Sum, + arg: Some(col(1, DataType::Int64)), + distinct: true, + output_name: "s".into(), + data_type: DataType::Int64, + }; + assert_eq!(a.to_string(), "sum(DISTINCT c1#1)"); + let c = BoundAggregate { + func: AggFunc::Count, + arg: None, + distinct: false, + output_name: "n".into(), + data_type: DataType::Int64, + }; + assert_eq!(c.to_string(), "count(*)"); + } +} diff --git a/crates/qf-plan/src/lib.rs b/crates/qf-plan/src/lib.rs index ff7bd09..9d57b55 100644 --- a/crates/qf-plan/src/lib.rs +++ b/crates/qf-plan/src/lib.rs @@ -1 +1,5 @@ -// placeholder +pub mod binder; +pub mod cost; +pub mod expr; +pub mod logical; +pub mod optimizer; diff --git a/crates/qf-plan/src/logical.rs b/crates/qf-plan/src/logical.rs new file mode 100644 index 0000000..2c88de8 --- /dev/null +++ b/crates/qf-plan/src/logical.rs @@ -0,0 +1,701 @@ +//! The logical plan: what the query means, with no commitment yet to how it +//! runs. +//! +//! Every node computes its own output schema, which is what lets the optimiser +//! rewrite the tree and immediately check that the result still type-checks — +//! a rule that drops a column or reorders a join is caught by the schema, not +//! by a wrong answer at the end of a run. + +use crate::expr::{BoundAggregate, BoundExpr}; +use qf_common::{Error, Field, Result, Schema}; +use qf_sql::ast::JoinType; +use qf_storage::stats::TableStats; +use std::fmt; +use std::sync::Arc; + +#[derive(Debug, Clone, PartialEq)] +pub struct SortExpr { + pub expr: BoundExpr, + pub ascending: bool, + pub nulls_first: bool, +} + +#[derive(Debug, Clone, PartialEq)] +pub enum LogicalPlan { + /// A table read. + /// + /// `projection` and `pushed_filters` start empty and are filled in by the + /// optimiser. They are part of the scan rather than separate nodes because + /// that is exactly what the storage layer can act on: read fewer column + /// chunks, and skip row groups whose zone map rules the predicate out. + Scan { + table: String, + /// The table's full schema, before projection. + source_schema: Arc, + projection: Option>, + pushed_filters: Vec, + stats: TableStats, + }, + Filter { + input: Box, + predicate: BoundExpr, + }, + Projection { + input: Box, + exprs: Vec<(BoundExpr, String)>, + }, + Aggregate { + input: Box, + group_exprs: Vec<(BoundExpr, String)>, + aggregates: Vec, + }, + Join { + left: Box, + right: Box, + join_type: JoinType, + /// Equality pairs, left expression against right expression. These are + /// what the hash join builds on. + on: Vec<(BoundExpr, BoundExpr)>, + /// Anything in the ON clause that is not an equality. + filter: Option, + }, + Sort { + input: Box, + exprs: Vec, + }, + Limit { + input: Box, + limit: Option, + offset: usize, + }, + Distinct { + input: Box, + }, + /// Literal rows, from `INSERT ... VALUES` or a `SELECT` with no `FROM`. + Values { + schema: Arc, + rows: Vec>, + }, +} + +impl LogicalPlan { + /// The schema this node produces. + pub fn schema(&self) -> Result> { + Ok(match self { + LogicalPlan::Scan { + source_schema, + projection, + .. + } => match projection { + None => Arc::clone(source_schema), + Some(indices) => Arc::new(source_schema.project(indices)?), + }, + LogicalPlan::Filter { input, .. } + | LogicalPlan::Sort { input, .. } + | LogicalPlan::Limit { input, .. } + | LogicalPlan::Distinct { input } => input.schema()?, + LogicalPlan::Projection { exprs, .. } => Arc::new(Schema::new( + exprs + .iter() + .map(|(e, name)| Field::new(name, e.data_type(), true)) + .collect(), + )), + LogicalPlan::Aggregate { + group_exprs, + aggregates, + .. + } => { + let mut fields: Vec = group_exprs + .iter() + .map(|(e, name)| Field::new(name, e.data_type(), true)) + .collect(); + fields.extend( + aggregates + .iter() + .map(|a| Field::new(&a.output_name, a.data_type, true)), + ); + Arc::new(Schema::new(fields)) + } + LogicalPlan::Join { + left, + right, + join_type, + .. + } => { + // An outer join can produce NULLs on the side that did not + // match, so those columns become nullable whatever the source + // said. + let (l, r) = (left.schema()?, right.schema()?); + let widen_left = matches!(join_type, JoinType::Right | JoinType::Full); + let widen_right = matches!(join_type, JoinType::Left | JoinType::Full); + let mut fields: Vec = l + .fields() + .iter() + .map(|f| Field::new(&f.name, f.data_type, f.nullable || widen_left)) + .collect(); + fields.extend( + r.fields() + .iter() + .map(|f| Field::new(&f.name, f.data_type, f.nullable || widen_right)), + ); + Arc::new(Schema::new(fields)) + } + LogicalPlan::Values { schema, .. } => Arc::clone(schema), + }) + } + + pub fn children(&self) -> Vec<&LogicalPlan> { + match self { + LogicalPlan::Scan { .. } | LogicalPlan::Values { .. } => vec![], + LogicalPlan::Filter { input, .. } + | LogicalPlan::Projection { input, .. } + | LogicalPlan::Aggregate { input, .. } + | LogicalPlan::Sort { input, .. } + | LogicalPlan::Limit { input, .. } + | LogicalPlan::Distinct { input } => vec![input], + LogicalPlan::Join { left, right, .. } => vec![left, right], + } + } + + /// Replaces this node's children, keeping everything else. The optimiser + /// rebuilds trees bottom-up through this. + pub fn with_children(&self, mut new: Vec) -> Result { + let expected = self.children().len(); + if new.len() != expected { + return Err(Error::internal(format!( + "{} expects {expected} children, got {}", + self.name(), + new.len() + ))); + } + Ok(match self { + LogicalPlan::Scan { .. } | LogicalPlan::Values { .. } => self.clone(), + LogicalPlan::Filter { predicate, .. } => LogicalPlan::Filter { + input: Box::new(new.remove(0)), + predicate: predicate.clone(), + }, + LogicalPlan::Projection { exprs, .. } => LogicalPlan::Projection { + input: Box::new(new.remove(0)), + exprs: exprs.clone(), + }, + LogicalPlan::Aggregate { + group_exprs, + aggregates, + .. + } => LogicalPlan::Aggregate { + input: Box::new(new.remove(0)), + group_exprs: group_exprs.clone(), + aggregates: aggregates.clone(), + }, + LogicalPlan::Sort { exprs, .. } => LogicalPlan::Sort { + input: Box::new(new.remove(0)), + exprs: exprs.clone(), + }, + LogicalPlan::Limit { limit, offset, .. } => LogicalPlan::Limit { + input: Box::new(new.remove(0)), + limit: *limit, + offset: *offset, + }, + LogicalPlan::Distinct { .. } => LogicalPlan::Distinct { + input: Box::new(new.remove(0)), + }, + LogicalPlan::Join { + join_type, + on, + filter, + .. + } => { + let left = Box::new(new.remove(0)); + let right = Box::new(new.remove(0)); + LogicalPlan::Join { + left, + right, + join_type: *join_type, + on: on.clone(), + filter: filter.clone(), + } + } + }) + } + + pub fn name(&self) -> &'static str { + match self { + LogicalPlan::Scan { .. } => "Scan", + LogicalPlan::Filter { .. } => "Filter", + LogicalPlan::Projection { .. } => "Projection", + LogicalPlan::Aggregate { .. } => "Aggregate", + LogicalPlan::Join { .. } => "Join", + LogicalPlan::Sort { .. } => "Sort", + LogicalPlan::Limit { .. } => "Limit", + LogicalPlan::Distinct { .. } => "Distinct", + LogicalPlan::Values { .. } => "Values", + } + } + + /// One line describing this node, without its children. This is what + /// `EXPLAIN` prints. + pub fn describe(&self) -> String { + match self { + LogicalPlan::Scan { + table, + source_schema, + projection, + pushed_filters, + .. + } => { + let cols = match projection { + None => "*".to_string(), + Some(ix) => ix + .iter() + .filter_map(|i| source_schema.field(*i).ok()) + .map(|f| f.name.clone()) + .collect::>() + .join(", "), + }; + let mut s = format!("Scan: {table} [{cols}]"); + if !pushed_filters.is_empty() { + let fs: Vec = pushed_filters.iter().map(BoundExpr::to_string).collect(); + s.push_str(&format!(", pushed: {}", fs.join(" AND "))); + } + s + } + LogicalPlan::Filter { predicate, .. } => format!("Filter: {predicate}"), + LogicalPlan::Projection { exprs, .. } => { + let items: Vec = exprs + .iter() + .map(|(e, n)| { + if &e.output_name() == n { + e.to_string() + } else { + format!("{e} AS {n}") + } + }) + .collect(); + format!("Projection: {}", items.join(", ")) + } + LogicalPlan::Aggregate { + group_exprs, + aggregates, + .. + } => { + let g: Vec = group_exprs.iter().map(|(e, _)| e.to_string()).collect(); + let a: Vec = aggregates.iter().map(BoundAggregate::to_string).collect(); + if g.is_empty() { + format!("Aggregate: {}", a.join(", ")) + } else { + format!("Aggregate: group by [{}], {}", g.join(", "), a.join(", ")) + } + } + LogicalPlan::Join { + join_type, + on, + filter, + .. + } => { + let keys: Vec = on.iter().map(|(l, r)| format!("{l} = {r}")).collect(); + let mut s = format!("{join_type} Join"); + if !keys.is_empty() { + s.push_str(&format!(": on {}", keys.join(" AND "))); + } + if let Some(f) = filter { + s.push_str(&format!(", filter {f}")); + } + s + } + LogicalPlan::Sort { exprs, .. } => { + let items: Vec = exprs + .iter() + .map(|s| { + format!( + "{} {}{}", + s.expr, + if s.ascending { "ASC" } else { "DESC" }, + if s.nulls_first { + " NULLS FIRST" + } else { + " NULLS LAST" + } + ) + }) + .collect(); + format!("Sort: {}", items.join(", ")) + } + LogicalPlan::Limit { limit, offset, .. } => match (limit, offset) { + (Some(l), 0) => format!("Limit: {l}"), + (Some(l), o) => format!("Limit: {l}, offset {o}"), + (None, o) => format!("Offset: {o}"), + }, + LogicalPlan::Distinct { .. } => "Distinct".to_string(), + LogicalPlan::Values { rows, .. } => format!("Values: {} rows", rows.len()), + } + } + + /// The whole tree, indented. + pub fn explain(&self) -> String { + let mut out = String::new(); + self.explain_into(&mut out, 0); + out + } + + fn explain_into(&self, out: &mut String, depth: usize) { + for _ in 0..depth { + out.push_str(" "); + } + out.push_str(&self.describe()); + out.push('\n'); + for c in self.children() { + c.explain_into(out, depth + 1); + } + } + + /// Number of nodes in the tree, used by tests to assert a rule removed one. + pub fn node_count(&self) -> usize { + 1 + self + .children() + .iter() + .map(|c| c.node_count()) + .sum::() + } +} + +impl fmt::Display for LogicalPlan { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.explain().trim_end()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::expr::AggFunc; + use qf_common::{DataType, Value}; + use qf_sql::ast::BinaryOp; + + fn table_schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("region", DataType::Utf8, false), + Field::new("amount", DataType::Float64, true), + ])) + } + + fn scan() -> LogicalPlan { + LogicalPlan::Scan { + table: "sales".into(), + source_schema: table_schema(), + projection: None, + pushed_filters: vec![], + stats: TableStats::default(), + } + } + + #[test] + fn a_scan_reports_its_table_schema_and_a_projected_scan_reports_a_subset() { + assert_eq!(scan().schema().unwrap().len(), 3); + let projected = LogicalPlan::Scan { + table: "sales".into(), + source_schema: table_schema(), + projection: Some(vec![2, 0]), + pushed_filters: vec![], + stats: TableStats::default(), + }; + let s = projected.schema().unwrap(); + assert_eq!(s.len(), 2); + assert_eq!(s.field(0).unwrap().name, "amount"); + assert_eq!(s.field(1).unwrap().name, "id"); + } + + #[test] + fn filters_sorts_limits_and_distinct_pass_their_input_schema_through() { + let p = BoundExpr::binary( + BoundExpr::column(0, "id", DataType::Int64), + BinaryOp::Gt, + BoundExpr::Literal(Value::Int64(1)), + ) + .unwrap(); + for node in [ + LogicalPlan::Filter { + input: Box::new(scan()), + predicate: p.clone(), + }, + LogicalPlan::Sort { + input: Box::new(scan()), + exprs: vec![], + }, + LogicalPlan::Limit { + input: Box::new(scan()), + limit: Some(1), + offset: 0, + }, + LogicalPlan::Distinct { + input: Box::new(scan()), + }, + ] { + assert_eq!(node.schema().unwrap().len(), 3, "{}", node.name()); + } + } + + #[test] + fn a_projection_names_its_output_columns() { + let p = LogicalPlan::Projection { + input: Box::new(scan()), + exprs: vec![( + BoundExpr::binary( + BoundExpr::column(2, "amount", DataType::Float64), + BinaryOp::Multiply, + BoundExpr::Literal(Value::Int64(2)), + ) + .unwrap(), + "doubled".into(), + )], + }; + let s = p.schema().unwrap(); + assert_eq!(s.len(), 1); + assert_eq!(s.field(0).unwrap().name, "doubled"); + assert_eq!(s.field(0).unwrap().data_type, DataType::Float64); + } + + #[test] + fn an_aggregate_outputs_its_group_columns_then_its_aggregates() { + let a = LogicalPlan::Aggregate { + input: Box::new(scan()), + group_exprs: vec![( + BoundExpr::column(1, "region", DataType::Utf8), + "region".into(), + )], + aggregates: vec![BoundAggregate { + func: AggFunc::Count, + arg: None, + distinct: false, + output_name: "n".into(), + data_type: DataType::Int64, + }], + }; + let s = a.schema().unwrap(); + assert_eq!(s.len(), 2); + assert_eq!(s.field(0).unwrap().name, "region"); + assert_eq!(s.field(1).unwrap().name, "n"); + assert_eq!(s.field(1).unwrap().data_type, DataType::Int64); + } + + #[test] + fn an_outer_join_widens_the_nullability_of_the_side_that_may_not_match() { + let join = |t| LogicalPlan::Join { + left: Box::new(scan()), + right: Box::new(scan()), + join_type: t, + on: vec![], + filter: None, + }; + // `id` is declared NOT NULL on both sides. + let inner = join(JoinType::Inner).schema().unwrap(); + assert!(!inner.field(0).unwrap().nullable); + assert!(!inner.field(3).unwrap().nullable); + + let left = join(JoinType::Left).schema().unwrap(); + assert!(!left.field(0).unwrap().nullable); + assert!(left.field(3).unwrap().nullable, "right side may be padded"); + + let right = join(JoinType::Right).schema().unwrap(); + assert!(right.field(0).unwrap().nullable, "left side may be padded"); + + let full = join(JoinType::Full).schema().unwrap(); + assert!(full.field(0).unwrap().nullable); + assert!(full.field(3).unwrap().nullable); + assert_eq!(full.len(), 6); + } + + #[test] + fn children_can_be_swapped_out_without_losing_the_node_s_own_settings() { + let filter = LogicalPlan::Filter { + input: Box::new(scan()), + predicate: BoundExpr::Literal(Value::Boolean(true)), + }; + let rebuilt = filter.with_children(vec![scan()]).unwrap(); + assert_eq!(rebuilt, filter); + assert!(filter.with_children(vec![]).is_err()); + assert!(scan().with_children(vec![]).unwrap() == scan()); + } + + #[test] + fn with_children_covers_every_node_shape() { + let nodes = vec![ + LogicalPlan::Projection { + input: Box::new(scan()), + exprs: vec![(BoundExpr::column(0, "id", DataType::Int64), "id".into())], + }, + LogicalPlan::Aggregate { + input: Box::new(scan()), + group_exprs: vec![], + aggregates: vec![], + }, + LogicalPlan::Sort { + input: Box::new(scan()), + exprs: vec![SortExpr { + expr: BoundExpr::column(0, "id", DataType::Int64), + ascending: true, + nulls_first: false, + }], + }, + LogicalPlan::Limit { + input: Box::new(scan()), + limit: Some(3), + offset: 1, + }, + LogicalPlan::Distinct { + input: Box::new(scan()), + }, + ]; + for n in nodes { + let children: Vec = n.children().into_iter().cloned().collect(); + assert_eq!(n.with_children(children).unwrap(), n, "{}", n.name()); + } + let j = LogicalPlan::Join { + left: Box::new(scan()), + right: Box::new(scan()), + join_type: JoinType::Inner, + on: vec![], + filter: None, + }; + assert_eq!(j.with_children(vec![scan(), scan()]).unwrap(), j); + } + + #[test] + fn explain_indents_children_under_their_parent() { + let plan = LogicalPlan::Limit { + input: Box::new(LogicalPlan::Filter { + input: Box::new(scan()), + predicate: BoundExpr::Literal(Value::Boolean(true)), + }), + limit: Some(10), + offset: 0, + }; + let text = plan.explain(); + let lines: Vec<&str> = text.lines().collect(); + assert_eq!(lines[0], "Limit: 10"); + assert_eq!(lines[1], " Filter: true"); + assert_eq!(lines[2], " Scan: sales [*]"); + assert_eq!(plan.node_count(), 3); + assert_eq!(plan.to_string(), text.trim_end()); + } + + #[test] + fn every_node_describes_itself() { + let scan_with_pushdown = LogicalPlan::Scan { + table: "sales".into(), + source_schema: table_schema(), + projection: Some(vec![0]), + pushed_filters: vec![BoundExpr::binary( + BoundExpr::column(0, "id", DataType::Int64), + BinaryOp::Gt, + BoundExpr::Literal(Value::Int64(5)), + ) + .unwrap()], + stats: TableStats::default(), + }; + assert!(scan_with_pushdown.describe().contains("pushed:")); + assert!(scan_with_pushdown.describe().contains("[id]")); + + assert_eq!( + LogicalPlan::Values { + schema: table_schema(), + rows: vec![vec![], vec![]] + } + .describe(), + "Values: 2 rows" + ); + assert_eq!( + LogicalPlan::Limit { + input: Box::new(scan()), + limit: None, + offset: 4 + } + .describe(), + "Offset: 4" + ); + assert_eq!( + LogicalPlan::Limit { + input: Box::new(scan()), + limit: Some(2), + offset: 4 + } + .describe(), + "Limit: 2, offset 4" + ); + assert_eq!( + LogicalPlan::Distinct { + input: Box::new(scan()) + } + .describe(), + "Distinct" + ); + assert!(LogicalPlan::Sort { + input: Box::new(scan()), + exprs: vec![SortExpr { + expr: BoundExpr::column(0, "id", DataType::Int64), + ascending: false, + nulls_first: true, + }], + } + .describe() + .contains("DESC NULLS FIRST")); + } + + #[test] + fn a_join_describes_its_keys_and_any_leftover_filter() { + let j = LogicalPlan::Join { + left: Box::new(scan()), + right: Box::new(scan()), + join_type: JoinType::Left, + on: vec![( + BoundExpr::column(0, "id", DataType::Int64), + BoundExpr::column(3, "id", DataType::Int64), + )], + filter: Some(BoundExpr::Literal(Value::Boolean(true))), + }; + let d = j.describe(); + assert!(d.starts_with("LEFT Join: on")); + assert!(d.contains("filter true")); + assert_eq!( + LogicalPlan::Join { + left: Box::new(scan()), + right: Box::new(scan()), + join_type: JoinType::Cross, + on: vec![], + filter: None, + } + .describe(), + "CROSS Join" + ); + } + + #[test] + fn a_projection_shows_an_alias_only_when_it_differs_from_the_expression() { + let p = LogicalPlan::Projection { + input: Box::new(scan()), + exprs: vec![ + (BoundExpr::column(0, "id", DataType::Int64), "id".into()), + (BoundExpr::column(1, "region", DataType::Utf8), "r".into()), + ], + }; + let d = p.describe(); + assert!(d.contains("id#0,")); + assert!(d.contains("region#1 AS r")); + } + + #[test] + fn an_aggregate_without_grouping_describes_only_its_aggregates() { + let a = LogicalPlan::Aggregate { + input: Box::new(scan()), + group_exprs: vec![], + aggregates: vec![BoundAggregate { + func: AggFunc::Count, + arg: None, + distinct: false, + output_name: "n".into(), + data_type: DataType::Int64, + }], + }; + assert_eq!(a.describe(), "Aggregate: count(*)"); + } +} diff --git a/crates/qf-plan/src/optimizer.rs b/crates/qf-plan/src/optimizer.rs new file mode 100644 index 0000000..42c2bed --- /dev/null +++ b/crates/qf-plan/src/optimizer.rs @@ -0,0 +1,1758 @@ +//! Plan rewrites. +//! +//! Each rule is a plan-to-plan function that must preserve two things: the +//! rows the plan produces, and the schema it produces them in. The second is +//! what makes the rules composable — a rule may move a filter or reorder a +//! join, but the node above it must not be able to tell, or every rule would +//! have to know about every other. +//! +//! The rules, in the order they run: +//! +//! 1. **Constant folding** — evaluate what is already known, so later rules see +//! `false` rather than `1 = 2`. +//! 2. **Predicate pushdown** — move filters as close to the data as possible, +//! ending inside the scan, where a zone map can skip whole row groups. +//! 3. **Join reordering** — order an inner-join chain by estimated size. +//! 4. **Projection pushdown** — stop reading columns nothing asks for. +//! 5. **Constant folding again** — the earlier rules leave new constants behind. + +use crate::cost::estimate_rows; +use crate::expr::BoundExpr; +use crate::logical::LogicalPlan; +use qf_common::{DataType, Error, Result, Value}; +use qf_sql::ast::{BinaryOp, JoinType, UnaryOp}; +use std::collections::BTreeSet; + +/// Runs every rule, in order. +pub fn optimize(plan: LogicalPlan) -> Result { + let plan = fold_constants(plan)?; + let plan = push_down_predicates(plan)?; + let plan = reorder_joins(plan)?; + let plan = push_down_projections(plan)?; + fold_constants(plan) +} + +/// Names of the rules, for `EXPLAIN`'s header. +pub const RULES: &[&str] = &[ + "constant folding", + "predicate pushdown", + "join reordering", + "projection pushdown", +]; + +// ---------------------------------------------------------------- folding + +/// Evaluates subexpressions whose value is already known and applies the +/// boolean identities. +/// +/// Worth doing on its own (`WHERE 1 = 1` should not cost a comparison per row) +/// but the real payoff is downstream: predicate pushdown can drop a conjunct +/// entirely once it folds to `true`. +pub fn fold_constants(plan: LogicalPlan) -> Result { + let plan = map_children(plan, fold_constants)?; + Ok(match plan { + LogicalPlan::Filter { input, predicate } => { + let folded = fold_expr(&predicate)?; + match folded.as_literal() { + // `WHERE true` is not a filter at all. + Some(Value::Boolean(true)) => *input, + _ => LogicalPlan::Filter { + input, + predicate: folded, + }, + } + } + LogicalPlan::Projection { input, exprs } => LogicalPlan::Projection { + input, + exprs: exprs + .into_iter() + .map(|(e, n)| Ok((fold_expr(&e)?, n))) + .collect::>>()?, + }, + LogicalPlan::Scan { + table, + source_schema, + projection, + pushed_filters, + stats, + } => LogicalPlan::Scan { + table, + source_schema, + projection, + pushed_filters: pushed_filters + .iter() + .map(fold_expr) + .collect::>>()? + .into_iter() + .filter(|f| !matches!(f.as_literal(), Some(Value::Boolean(true)))) + .collect(), + stats, + }, + other => other, + }) +} + +fn fold_expr(expr: &BoundExpr) -> Result { + Ok(match expr { + BoundExpr::Binary { + left, + op, + right, + data_type, + } => { + let (l, r) = (fold_expr(left)?, fold_expr(right)?); + if let Some(v) = fold_binary(&l, *op, &r)? { + return Ok(BoundExpr::Literal(v)); + } + if let Some(simplified) = simplify_boolean(&l, *op, &r) { + return Ok(simplified); + } + BoundExpr::Binary { + left: Box::new(l), + op: *op, + right: Box::new(r), + data_type: *data_type, + } + } + BoundExpr::Unary { + op, + expr, + data_type, + } => { + let inner = fold_expr(expr)?; + match (op, inner.as_literal()) { + (UnaryOp::Not, Some(Value::Boolean(b))) => BoundExpr::Literal(Value::Boolean(!b)), + (UnaryOp::Neg, Some(Value::Int64(i))) => BoundExpr::Literal(Value::Int64(-i)), + (UnaryOp::Neg, Some(Value::Float64(x))) => BoundExpr::Literal(Value::Float64(-x)), + // Double negation cancels. + (UnaryOp::Not, None) => match &inner { + BoundExpr::Unary { + op: UnaryOp::Not, + expr: nested, + .. + } => nested.as_ref().clone(), + _ => BoundExpr::Unary { + op: *op, + expr: Box::new(inner), + data_type: *data_type, + }, + }, + _ => BoundExpr::Unary { + op: *op, + expr: Box::new(inner), + data_type: *data_type, + }, + } + } + BoundExpr::Cast { expr, data_type } => { + let inner = fold_expr(expr)?; + match inner.as_literal() { + Some(v) => BoundExpr::Literal(v.cast_to(*data_type)?), + None => BoundExpr::Cast { + expr: Box::new(inner), + data_type: *data_type, + }, + } + } + BoundExpr::IsNull { expr, negated } => { + let inner = fold_expr(expr)?; + match inner.as_literal() { + Some(v) => BoundExpr::Literal(Value::Boolean(v.is_null() != *negated)), + None => BoundExpr::IsNull { + expr: Box::new(inner), + negated: *negated, + }, + } + } + BoundExpr::InList { + expr, + list, + negated, + } => BoundExpr::InList { + expr: Box::new(fold_expr(expr)?), + list: list.iter().map(fold_expr).collect::>>()?, + negated: *negated, + }, + BoundExpr::Like { + expr, + pattern, + negated, + } => BoundExpr::Like { + expr: Box::new(fold_expr(expr)?), + pattern: Box::new(fold_expr(pattern)?), + negated: *negated, + }, + BoundExpr::Case { + branches, + else_result, + data_type, + } => { + let mut kept = Vec::new(); + for (when, then) in branches { + let w = fold_expr(when)?; + let t = fold_expr(then)?; + match w.as_literal() { + // A branch that can never fire is dead. + Some(Value::Boolean(false)) | Some(Value::Null) => continue, + // A branch that always fires makes everything after it dead. + Some(Value::Boolean(true)) => { + if kept.is_empty() { + return Ok(t); + } + kept.push((w, t)); + break; + } + _ => kept.push((w, t)), + } + } + let else_folded = match else_result { + Some(e) => Some(fold_expr(e)?), + None => None, + }; + if kept.is_empty() { + return Ok(else_folded.unwrap_or(BoundExpr::Literal(Value::Null))); + } + BoundExpr::Case { + branches: kept, + else_result: else_folded.map(Box::new), + data_type: *data_type, + } + } + other => other.clone(), + }) +} + +fn fold_binary(left: &BoundExpr, op: BinaryOp, right: &BoundExpr) -> Result> { + let (Some(l), Some(r)) = (left.as_literal(), right.as_literal()) else { + return Ok(None); + }; + if l.is_null() || r.is_null() { + // Any operation with NULL is NULL, except that AND/OR have their own + // three-valued rules, which `simplify_boolean` handles. + if op.is_logical() { + return Ok(None); + } + return Ok(Some(Value::Null)); + } + Ok(Some(match op { + BinaryOp::And => { + Value::Boolean(l.as_bool()?.unwrap_or(false) && r.as_bool()?.unwrap_or(false)) + } + BinaryOp::Or => { + Value::Boolean(l.as_bool()?.unwrap_or(false) || r.as_bool()?.unwrap_or(false)) + } + op if op.is_comparison() => { + let Some(ord) = l.sql_compare(r)? else { + return Ok(Some(Value::Null)); + }; + Value::Boolean(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(), + }) + } + _ => return fold_arithmetic(l, op, r), + })) +} + +fn fold_arithmetic(l: &Value, op: BinaryOp, r: &Value) -> Result> { + // Division by zero is left alone rather than folded: whatever the runtime + // does with it, folding must not change the answer at plan time. + if let (Some(_), Some(0.0)) = (l.as_f64(), r.as_f64()) { + if matches!(op, BinaryOp::Divide | BinaryOp::Modulo) { + return Ok(None); + } + } + Ok(Some(match (l, r) { + (Value::Int64(a), Value::Int64(b)) => match op { + // Overflow is left unfolded so the runtime's behaviour is the only + // behaviour. + BinaryOp::Plus => match a.checked_add(*b) { + Some(v) => Value::Int64(v), + None => return Ok(None), + }, + BinaryOp::Minus => match a.checked_sub(*b) { + Some(v) => Value::Int64(v), + None => return Ok(None), + }, + BinaryOp::Multiply => match a.checked_mul(*b) { + Some(v) => Value::Int64(v), + None => return Ok(None), + }, + BinaryOp::Divide => Value::Float64(*a as f64 / *b as f64), + BinaryOp::Modulo => Value::Int64(a % b), + _ => return Ok(None), + }, + _ => { + let (Some(a), Some(b)) = (l.as_f64(), r.as_f64()) else { + return Ok(None); + }; + match op { + BinaryOp::Plus => Value::Float64(a + b), + BinaryOp::Minus => Value::Float64(a - b), + BinaryOp::Multiply => Value::Float64(a * b), + BinaryOp::Divide => Value::Float64(a / b), + BinaryOp::Modulo => Value::Float64(a % b), + _ => return Ok(None), + } + } + })) +} + +/// `x AND true`, `x OR false` and their absorbing counterparts. +fn simplify_boolean(left: &BoundExpr, op: BinaryOp, right: &BoundExpr) -> Option { + let literal = |e: &BoundExpr| match e.as_literal() { + Some(Value::Boolean(b)) => Some(*b), + _ => None, + }; + match op { + BinaryOp::And => match (literal(left), literal(right)) { + (Some(true), _) => Some(right.clone()), + (_, Some(true)) => Some(left.clone()), + (Some(false), _) | (_, Some(false)) => Some(BoundExpr::Literal(Value::Boolean(false))), + _ => None, + }, + BinaryOp::Or => match (literal(left), literal(right)) { + (Some(false), _) => Some(right.clone()), + (_, Some(false)) => Some(left.clone()), + (Some(true), _) | (_, Some(true)) => Some(BoundExpr::Literal(Value::Boolean(true))), + _ => None, + }, + _ => None, + } +} + +// ------------------------------------------------------- predicate pushdown + +/// Moves filters down the tree, ending inside the scan. +/// +/// A filter left above a scan still reads every row. Pushed into the scan it +/// becomes a zone-map test, and whole row groups stop being read at all — the +/// difference between saving CPU and saving I/O. +pub fn push_down_predicates(plan: LogicalPlan) -> Result { + let plan = map_children(plan, push_down_predicates)?; + let LogicalPlan::Filter { input, predicate } = plan else { + return Ok(plan); + }; + push_filter(*input, split_and(predicate)) +} + +fn push_filter(input: LogicalPlan, conjuncts: Vec) -> Result { + if conjuncts.is_empty() { + return Ok(input); + } + match input { + // Two stacked filters become one. + LogicalPlan::Filter { + input: inner, + predicate, + } => { + let mut all = split_and(predicate); + all.extend(conjuncts); + push_filter(*inner, all) + } + + LogicalPlan::Scan { + table, + source_schema, + projection, + mut pushed_filters, + stats, + } => { + pushed_filters.extend(conjuncts); + Ok(LogicalPlan::Scan { + table, + source_schema, + projection, + pushed_filters, + stats, + }) + } + + // Rewrite the predicate in terms of the projection's own input and + // push it underneath. Only safe when every column it reads is a plain + // column of the input — substituting a computed expression could + // duplicate work rather than save it. + LogicalPlan::Projection { + input: inner, + exprs, + } => { + let mut pushable = Vec::new(); + let mut stuck = Vec::new(); + for c in conjuncts { + match substitute(&c, &exprs) { + Some(rewritten) => pushable.push(rewritten), + None => stuck.push(c), + } + } + let below = push_filter(*inner, pushable)?; + let projection = LogicalPlan::Projection { + input: Box::new(below), + exprs, + }; + Ok(rebuild_filter(projection, stuck)) + } + + // Sorting and then filtering is the same as filtering and then + // sorting, and the second reads fewer rows. + LogicalPlan::Sort { + input: inner, + exprs, + } => Ok(LogicalPlan::Sort { + input: Box::new(push_filter(*inner, conjuncts)?), + exprs, + }), + + LogicalPlan::Join { + left, + right, + join_type, + on, + filter, + } => push_filter_into_join(*left, *right, join_type, on, filter, conjuncts), + + // Everything else keeps the filter above it. Pushing below a Limit + // would change which rows the limit sees; pushing below an Aggregate + // would filter rows the aggregate was supposed to see. + other => Ok(rebuild_filter(other, conjuncts)), + } +} + +fn push_filter_into_join( + left: LogicalPlan, + right: LogicalPlan, + join_type: JoinType, + mut on: Vec<(BoundExpr, BoundExpr)>, + join_filter: Option, + conjuncts: Vec, +) -> Result { + let left_width = left.schema()?.len(); + // Which sides may receive a predicate. Pushing into the null-padded side + // of an outer join changes the answer: a row that would have been padded + // with NULLs disappears instead. + let (can_push_left, can_push_right) = match join_type { + JoinType::Inner | JoinType::Cross => (true, true), + JoinType::Left => (true, false), + JoinType::Right => (false, true), + JoinType::Full => (false, false), + }; + + let mut to_left = Vec::new(); + let mut to_right = Vec::new(); + let mut stuck = Vec::new(); + let mut join_type = join_type; + + for c in conjuncts { + let indices = c.column_indices(); + let touches_left = indices.iter().any(|i| *i < left_width); + let touches_right = indices.iter().any(|i| *i >= left_width); + + if !touches_right && can_push_left { + to_left.push(c); + continue; + } + if !touches_left && can_push_right { + to_right.push(c.remap_columns(&|i| i.checked_sub(left_width))?); + continue; + } + // A cross join with an equality across it in the WHERE clause is an + // inner join that was written the old way. Turning it into a join key + // is the difference between a hash join and a full cross product. + if join_type == JoinType::Cross && touches_left && touches_right { + if let BoundExpr::Binary { + left: l, + op: BinaryOp::Eq, + right: r, + .. + } = &c + { + let l_only_left = l.column_indices().iter().all(|i| *i < left_width); + let l_only_right = l.column_indices().iter().all(|i| *i >= left_width); + let r_only_left = r.column_indices().iter().all(|i| *i < left_width); + let r_only_right = r.column_indices().iter().all(|i| *i >= left_width); + if l_only_left && r_only_right { + on.push(( + l.as_ref().clone(), + r.remap_columns(&|i| i.checked_sub(left_width))?, + )); + join_type = JoinType::Inner; + continue; + } + if r_only_left && l_only_right { + on.push(( + r.as_ref().clone(), + l.remap_columns(&|i| i.checked_sub(left_width))?, + )); + join_type = JoinType::Inner; + continue; + } + } + } + stuck.push(c); + } + + let new_left = push_filter(left, to_left)?; + let new_right = push_filter(right, to_right)?; + let join = LogicalPlan::Join { + left: Box::new(new_left), + right: Box::new(new_right), + join_type, + on, + filter: join_filter, + }; + Ok(rebuild_filter(join, stuck)) +} + +/// Rewrites a predicate expressed over a projection's output into one over its +/// input, or reports that it cannot. +fn substitute(expr: &BoundExpr, exprs: &[(BoundExpr, String)]) -> Option { + let mut ok = true; + expr.walk(&mut |e| { + if let BoundExpr::Column { index, .. } = e { + match exprs.get(*index) { + Some((BoundExpr::Column { .. }, _)) => {} + _ => ok = false, + } + } + }); + if !ok { + return None; + } + rewrite_columns(expr, &|index| exprs.get(index).map(|(e, _)| e.clone())).ok() +} + +/// Replaces each column reference with whatever `lookup` returns for it. +fn rewrite_columns( + expr: &BoundExpr, + lookup: &dyn Fn(usize) -> Option, +) -> Result { + if let BoundExpr::Column { index, name, .. } = expr { + return lookup(*index) + .ok_or_else(|| Error::internal(format!("cannot rewrite column `{name}`"))); + } + // Every other shape just rebuilds with rewritten children, which + // `remap_columns` already does structurally — but it can only renumber, + // not substitute, so the walk is repeated here. + Ok(match expr { + BoundExpr::Binary { + left, + op, + right, + data_type, + } => BoundExpr::Binary { + left: Box::new(rewrite_columns(left, lookup)?), + op: *op, + right: Box::new(rewrite_columns(right, lookup)?), + data_type: *data_type, + }, + BoundExpr::Unary { + op, + expr, + data_type, + } => BoundExpr::Unary { + op: *op, + expr: Box::new(rewrite_columns(expr, lookup)?), + data_type: *data_type, + }, + BoundExpr::Cast { expr, data_type } => BoundExpr::Cast { + expr: Box::new(rewrite_columns(expr, lookup)?), + data_type: *data_type, + }, + BoundExpr::IsNull { expr, negated } => BoundExpr::IsNull { + expr: Box::new(rewrite_columns(expr, lookup)?), + negated: *negated, + }, + BoundExpr::InList { + expr, + list, + negated, + } => BoundExpr::InList { + expr: Box::new(rewrite_columns(expr, lookup)?), + list: list + .iter() + .map(|e| rewrite_columns(e, lookup)) + .collect::>>()?, + negated: *negated, + }, + BoundExpr::Like { + expr, + pattern, + negated, + } => BoundExpr::Like { + expr: Box::new(rewrite_columns(expr, lookup)?), + pattern: Box::new(rewrite_columns(pattern, lookup)?), + negated: *negated, + }, + BoundExpr::Case { + branches, + else_result, + data_type, + } => BoundExpr::Case { + branches: branches + .iter() + .map(|(w, t)| Ok((rewrite_columns(w, lookup)?, rewrite_columns(t, lookup)?))) + .collect::>>()?, + else_result: match else_result { + Some(e) => Some(Box::new(rewrite_columns(e, lookup)?)), + None => None, + }, + data_type: *data_type, + }, + other => other.clone(), + }) +} + +fn rebuild_filter(input: LogicalPlan, conjuncts: Vec) -> LogicalPlan { + match join_and(conjuncts) { + None => input, + Some(predicate) => LogicalPlan::Filter { + input: Box::new(input), + predicate, + }, + } +} + +fn split_and(expr: BoundExpr) -> Vec { + match expr { + BoundExpr::Binary { + left, + op: BinaryOp::And, + right, + .. + } => { + let mut out = split_and(*left); + out.extend(split_and(*right)); + out + } + other => vec![other], + } +} + +fn join_and(mut parts: Vec) -> Option { + if parts.is_empty() { + return None; + } + let first = parts.remove(0); + Some(parts.into_iter().fold(first, |a, b| BoundExpr::Binary { + left: Box::new(a), + op: BinaryOp::And, + right: Box::new(b), + data_type: DataType::Boolean, + })) +} + +// ------------------------------------------------------- projection pushdown + +/// Narrows every scan to the columns the query actually reads. +/// +/// In a columnar format this is the single largest win available: a query +/// touching two of forty columns should read two column chunks per row group, +/// not forty. Doing it means renumbering every expression above the scan, +/// which is the part that has to be exactly right — a mistake here reads the +/// wrong column and reports a plausible wrong answer. +pub fn push_down_projections(plan: LogicalPlan) -> Result { + let all: BTreeSet = (0..plan.schema()?.len()).collect(); + let (pruned, mapping) = prune(plan, &all)?; + // The top of the plan must still produce its columns in the original + // order; if pruning reordered them, restore it. + if mapping.iter().enumerate().all(|(i, m)| i == *m) { + return Ok(pruned); + } + let schema = pruned.schema()?; + let exprs = mapping + .iter() + .enumerate() + .map(|(new, _)| { + let f = schema.field(new)?; + Ok(( + BoundExpr::column(new, f.name.clone(), f.data_type), + f.name.clone(), + )) + }) + .collect::>>()?; + Ok(LogicalPlan::Projection { + input: Box::new(pruned), + exprs, + }) +} + +/// Rewrites `plan` to produce at least the columns in `needed`, returning the +/// new plan and the original output positions it now produces, in order. +fn prune(plan: LogicalPlan, needed: &BTreeSet) -> Result<(LogicalPlan, Vec)> { + match plan { + LogicalPlan::Scan { + table, + source_schema, + projection, + pushed_filters, + stats, + } => { + let mut wanted: BTreeSet = needed.clone(); + for f in &pushed_filters { + wanted.extend(f.column_indices()); + } + // A scan of nothing still has to produce rows to count, so keep + // one column rather than none. + if wanted.is_empty() && !source_schema.is_empty() { + wanted.insert(0); + } + let kept: Vec = wanted.into_iter().collect(); + let source_indices: Vec = match &projection { + Some(p) => kept.iter().map(|i| p[*i]).collect(), + None => kept.clone(), + }; + let position = |old: usize| kept.iter().position(|k| *k == old); + let new_filters = pushed_filters + .iter() + .map(|f| f.remap_columns(&position)) + .collect::>>()?; + Ok(( + LogicalPlan::Scan { + table, + source_schema, + projection: Some(source_indices), + pushed_filters: new_filters, + stats, + }, + kept, + )) + } + + LogicalPlan::Filter { input, predicate } => { + let mut child_needs = needed.clone(); + child_needs.extend(predicate.column_indices()); + let (new_input, mapping) = prune(*input, &child_needs)?; + let position = |old: usize| mapping.iter().position(|m| *m == old); + Ok(( + LogicalPlan::Filter { + input: Box::new(new_input), + predicate: predicate.remap_columns(&position)?, + }, + mapping, + )) + } + + LogicalPlan::Projection { input, exprs } => { + let kept: Vec = (0..exprs.len()).filter(|i| needed.contains(i)).collect(); + let kept = if kept.is_empty() { vec![0] } else { kept }; + let mut child_needs = BTreeSet::new(); + for i in &kept { + child_needs.extend(exprs[*i].0.column_indices()); + } + let (new_input, mapping) = prune(*input, &child_needs)?; + let position = |old: usize| mapping.iter().position(|m| *m == old); + let new_exprs = kept + .iter() + .map(|i| Ok((exprs[*i].0.remap_columns(&position)?, exprs[*i].1.clone()))) + .collect::>>()?; + Ok(( + LogicalPlan::Projection { + input: Box::new(new_input), + exprs: new_exprs, + }, + kept, + )) + } + + LogicalPlan::Aggregate { + input, + group_exprs, + aggregates, + } => { + // Every group column and aggregate argument is needed; the + // aggregate's own output is not narrowed, because dropping an + // aggregate would change what the groups mean to the node above. + let mut child_needs = BTreeSet::new(); + for (e, _) in &group_exprs { + child_needs.extend(e.column_indices()); + } + for a in &aggregates { + if let Some(arg) = &a.arg { + child_needs.extend(arg.column_indices()); + } + } + let (new_input, mapping) = prune(*input, &child_needs)?; + let position = |old: usize| mapping.iter().position(|m| *m == old); + let new_groups = group_exprs + .iter() + .map(|(e, n)| Ok((e.remap_columns(&position)?, n.clone()))) + .collect::>>()?; + let new_aggs = aggregates + .iter() + .map(|a| { + let mut a = a.clone(); + a.arg = match &a.arg { + Some(arg) => Some(arg.remap_columns(&position)?), + None => None, + }; + Ok(a) + }) + .collect::>>()?; + let width = new_groups.len() + new_aggs.len(); + Ok(( + LogicalPlan::Aggregate { + input: Box::new(new_input), + group_exprs: new_groups, + aggregates: new_aggs, + }, + (0..width).collect(), + )) + } + + LogicalPlan::Join { + left, + right, + join_type, + on, + filter, + } => { + let left_width = left.schema()?.len(); + let mut left_needs = BTreeSet::new(); + let mut right_needs = BTreeSet::new(); + for i in needed { + if *i < left_width { + left_needs.insert(*i); + } else { + right_needs.insert(*i - left_width); + } + } + for (l, r) in &on { + left_needs.extend(l.column_indices()); + right_needs.extend(r.column_indices()); + } + if let Some(f) = &filter { + for i in f.column_indices() { + if i < left_width { + left_needs.insert(i); + } else { + right_needs.insert(i - left_width); + } + } + } + let (new_left, left_map) = prune(*left, &left_needs)?; + let (new_right, right_map) = prune(*right, &right_needs)?; + let left_pos = |old: usize| left_map.iter().position(|m| *m == old); + let right_pos = |old: usize| right_map.iter().position(|m| *m == old); + + let new_on = on + .iter() + .map(|(l, r)| Ok((l.remap_columns(&left_pos)?, r.remap_columns(&right_pos)?))) + .collect::>>()?; + let new_width = left_map.len(); + let new_filter = match &filter { + Some(f) => Some(f.remap_columns(&|i| { + if i < left_width { + left_pos(i) + } else { + right_pos(i - left_width).map(|p| p + new_width) + } + })?), + None => None, + }; + + let mut mapping = left_map; + mapping.extend(right_map.iter().map(|i| i + left_width)); + Ok(( + LogicalPlan::Join { + left: Box::new(new_left), + right: Box::new(new_right), + join_type, + on: new_on, + filter: new_filter, + }, + mapping, + )) + } + + LogicalPlan::Sort { input, exprs } => { + let mut child_needs = needed.clone(); + for s in &exprs { + child_needs.extend(s.expr.column_indices()); + } + let (new_input, mapping) = prune(*input, &child_needs)?; + let position = |old: usize| mapping.iter().position(|m| *m == old); + let new_exprs = exprs + .iter() + .map(|s| { + Ok(crate::logical::SortExpr { + expr: s.expr.remap_columns(&position)?, + ascending: s.ascending, + nulls_first: s.nulls_first, + }) + }) + .collect::>>()?; + Ok(( + LogicalPlan::Sort { + input: Box::new(new_input), + exprs: new_exprs, + }, + mapping, + )) + } + + LogicalPlan::Limit { + input, + limit, + offset, + } => { + let (new_input, mapping) = prune(*input, needed)?; + Ok(( + LogicalPlan::Limit { + input: Box::new(new_input), + limit, + offset, + }, + mapping, + )) + } + + // Narrowing the input of a DISTINCT would change which rows are + // duplicates, so its input keeps every column. + LogicalPlan::Distinct { input } => { + let all: BTreeSet = (0..input.schema()?.len()).collect(); + let (new_input, mapping) = prune(*input, &all)?; + Ok(( + LogicalPlan::Distinct { + input: Box::new(new_input), + }, + mapping, + )) + } + + LogicalPlan::Values { schema, rows } => { + let width = schema.len(); + Ok((LogicalPlan::Values { schema, rows }, (0..width).collect())) + } + } +} + +// ---------------------------------------------------------- join reordering + +/// Orders an inner-join chain so the smallest relations join first. +/// +/// The rule is greedy: start from the smallest estimated relation, then keep +/// adding the smallest relation that has a join key to what is already joined. +/// Preferring a relation with a key matters more than preferring a small one — +/// joining an unrelated relation means a cross product, and no later choice +/// recovers from that. +/// +/// Reordering permutes the join's output columns, so the rule wraps the +/// reordered join in a projection that puts them back. The node above cannot +/// tell the difference, which is what makes the rewrite safe to apply +/// anywhere in the tree. +pub fn reorder_joins(plan: LogicalPlan) -> Result { + let plan = map_children(plan, reorder_joins)?; + let LogicalPlan::Join { + join_type: JoinType::Inner, + .. + } = &plan + else { + return Ok(plan); + }; + + let mut relations = Vec::new(); + let mut conditions = Vec::new(); + let mut residuals = Vec::new(); + flatten_join(&plan, 0, &mut relations, &mut conditions, &mut residuals)?; + if relations.len() < 3 { + // Two relations have only one order, and the build-side choice that + // does matter for them is made by the physical planner. + return Ok(plan); + } + + let order = choose_order(&relations, &conditions); + if order.iter().enumerate().all(|(i, r)| i == *r) { + return Ok(plan); + } + rebuild_join(plan, relations, conditions, residuals, order) +} + +struct Relation { + plan: LogicalPlan, + /// Where this relation's columns start in the original flattened order. + offset: usize, + width: usize, + rows: f64, +} + +/// A join key in the flattened coordinate space. +struct Condition { + left: BoundExpr, + right: BoundExpr, +} + +fn flatten_join( + plan: &LogicalPlan, + offset: usize, + relations: &mut Vec, + conditions: &mut Vec, + residuals: &mut Vec, +) -> Result { + match plan { + LogicalPlan::Join { + left, + right, + join_type: JoinType::Inner, + on, + filter, + } => { + let left_width = flatten_join(left, offset, relations, conditions, residuals)?; + let right_offset = offset + left_width; + let right_width = flatten_join(right, right_offset, relations, conditions, residuals)?; + for (l, r) in on { + conditions.push(Condition { + left: l.remap_columns(&|i| Some(i + offset))?, + right: r.remap_columns(&|i| Some(i + right_offset))?, + }); + } + if let Some(f) = filter { + residuals.push(f.remap_columns(&|i| Some(i + offset))?); + } + Ok(left_width + right_width) + } + other => { + let width = other.schema()?.len(); + relations.push(Relation { + plan: other.clone(), + offset, + width, + rows: estimate_rows(other), + }); + Ok(width) + } + } +} + +fn relation_of(relations: &[Relation], column: usize) -> Option { + relations + .iter() + .position(|r| column >= r.offset && column < r.offset + r.width) +} + +fn condition_relations(relations: &[Relation], c: &Condition) -> (Option, Option) { + let side = |e: &BoundExpr| { + let mut found = None; + for i in e.column_indices() { + match (found, relation_of(relations, i)) { + (None, r) => found = r, + (Some(a), Some(b)) if a != b => return None, + _ => {} + } + } + found + }; + (side(&c.left), side(&c.right)) +} + +fn choose_order(relations: &[Relation], conditions: &[Condition]) -> Vec { + let mut remaining: Vec = (0..relations.len()).collect(); + // Start from the smallest relation: the first join's build side is the one + // that has to fit in memory. + remaining.sort_by(|a, b| { + relations[*a] + .rows + .partial_cmp(&relations[*b].rows) + .unwrap_or(std::cmp::Ordering::Equal) + }); + let mut order = vec![remaining.remove(0)]; + + while !remaining.is_empty() { + let connected: Vec = remaining + .iter() + .copied() + .filter(|r| { + conditions.iter().any(|c| { + let (a, b) = condition_relations(relations, c); + matches!((a, b), (Some(x), Some(y)) + if (order.contains(&x) && y == *r) || (order.contains(&y) && x == *r)) + }) + }) + .collect(); + // `remaining` is already sorted by size, so the first connected + // relation is also the smallest connected one. + let pick = connected.first().copied().unwrap_or(remaining[0]); + order.push(pick); + remaining.retain(|r| *r != pick); + } + order +} + +fn rebuild_join( + original: LogicalPlan, + relations: Vec, + conditions: Vec, + residuals: Vec, + order: Vec, +) -> Result { + // Where each original column ends up in the reordered join. + let mut position = vec![usize::MAX; relations.iter().map(|r| r.width).sum()]; + let mut cursor = 0; + for r in &order { + let rel = &relations[*r]; + for k in 0..rel.width { + position[rel.offset + k] = cursor + k; + } + cursor += rel.width; + } + let map = |i: usize| position.get(i).copied().filter(|p| *p != usize::MAX); + + let mut joined: Vec = vec![order[0]]; + let mut plan = relations[order[0]].plan.clone(); + let mut used = vec![false; conditions.len()]; + let mut width = relations[order[0]].width; + + for step in &order[1..] { + let rel = &relations[*step]; + let mut keys = Vec::new(); + for (i, c) in conditions.iter().enumerate() { + if used[i] { + continue; + } + let (a, b) = condition_relations(&relations, c); + let (Some(a), Some(b)) = (a, b) else { continue }; + let (left_expr, right_expr) = if joined.contains(&a) && b == *step { + (&c.left, &c.right) + } else if joined.contains(&b) && a == *step { + (&c.right, &c.left) + } else { + continue; + }; + used[i] = true; + keys.push(( + left_expr.remap_columns(&map)?, + // The right side's positions are relative to the new relation. + right_expr.remap_columns(&|i| map(i).and_then(|p| p.checked_sub(width)))?, + )); + } + plan = LogicalPlan::Join { + left: Box::new(plan), + right: Box::new(rel.plan.clone()), + join_type: if keys.is_empty() { + JoinType::Cross + } else { + JoinType::Inner + }, + on: keys, + filter: None, + }; + joined.push(*step); + width += rel.width; + } + + // Conditions that spanned relations in a way the greedy order could not + // consume become a filter above the join rather than being dropped. + let mut leftover: Vec = Vec::new(); + for (i, c) in conditions.iter().enumerate() { + if !used[i] { + leftover.push(BoundExpr::Binary { + left: Box::new(c.left.remap_columns(&map)?), + op: BinaryOp::Eq, + right: Box::new(c.right.remap_columns(&map)?), + data_type: DataType::Boolean, + }); + } + } + for r in &residuals { + leftover.push(r.remap_columns(&map)?); + } + let plan = rebuild_filter(plan, leftover); + + // Restore the original column order so nothing above notices. + let original_schema = original.schema()?; + let exprs = (0..original_schema.len()) + .map(|i| { + let f = original_schema.field(i)?; + let p = position[i]; + Ok(( + BoundExpr::column(p, f.name.clone(), f.data_type), + f.name.clone(), + )) + }) + .collect::>>()?; + Ok(LogicalPlan::Projection { + input: Box::new(plan), + exprs, + }) +} + +// ------------------------------------------------------------------ helpers + +fn map_children( + plan: LogicalPlan, + f: impl Fn(LogicalPlan) -> Result + Copy, +) -> Result { + let children: Vec = plan.children().into_iter().cloned().collect(); + if children.is_empty() { + return Ok(plan); + } + let new = children.into_iter().map(f).collect::>>()?; + plan.with_children(new) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::binder::bind; + use qf_common::{Field, Schema}; + use qf_storage::catalog::Catalog; + use qf_storage::stats::{ColumnStats, TableStats}; + use std::sync::Arc; + + /// A catalog whose tables carry hand-set statistics, so the cost-based + /// rules have something to reason about without building real data. + fn catalog() -> Catalog { + let mut c = Catalog::new(); + add( + &mut c, + "orders", + &["id", "customer_id", "product_id", "amount"], + 100_000, + &[100_000, 5_000, 900, 50_000], + ); + add( + &mut c, + "customers", + &["id", "name", "region"], + 5_000, + &[5_000, 4_900, 8], + ); + add( + &mut c, + "products", + &["id", "title", "price"], + 900, + &[900, 900, 400], + ); + c + } + + fn add(c: &mut Catalog, name: &str, columns: &[&str], rows: usize, distinct: &[usize]) { + let schema = Arc::new(Schema::new( + columns + .iter() + .map(|n| Field::new(*n, DataType::Int64, true)) + .collect(), + )); + let table = qf_storage::catalog::Table { + name: name.to_string(), + schema, + source: qf_storage::catalog::TableSource::Memory(vec![]), + stats: TableStats { + row_count: rows, + columns: distinct + .iter() + .map(|d| ColumnStats { + min: Value::Int64(0), + max: Value::Int64(1_000_000), + null_count: 0, + row_count: rows, + distinct_count: *d, + }) + .collect(), + }, + }; + c.replace(table); + } + + fn raw(sql: &str) -> LogicalPlan { + match qf_sql::parse(sql).unwrap() { + qf_sql::ast::Statement::Query(q) => bind(&q, &catalog()).unwrap(), + other => panic!("expected a query, got {other:?}"), + } + } + + fn opt(sql: &str) -> LogicalPlan { + optimize(raw(sql)).unwrap() + } + + fn plan_text(sql: &str) -> String { + opt(sql).explain() + } + + fn find<'p>( + plan: &'p LogicalPlan, + pred: &dyn Fn(&LogicalPlan) -> bool, + ) -> Option<&'p LogicalPlan> { + if pred(plan) { + return Some(plan); + } + plan.children().into_iter().find_map(|c| find(c, pred)) + } + + fn scans(plan: &LogicalPlan) -> Vec<&LogicalPlan> { + let mut out = Vec::new(); + collect(plan, &mut out); + fn collect<'p>(p: &'p LogicalPlan, out: &mut Vec<&'p LogicalPlan>) { + if matches!(p, LogicalPlan::Scan { .. }) { + out.push(p); + } + for c in p.children() { + collect(c, out); + } + } + out + } + + /// The invariant every rule has to hold: optimising must not change the + /// shape of the answer. + fn assert_schema_preserved(sql: &str) { + let before = raw(sql).schema().unwrap(); + let after = opt(sql).schema().unwrap(); + assert_eq!(before, after, "schema changed while optimising `{sql}`"); + } + + #[test] + fn optimising_never_changes_the_output_schema() { + for sql in [ + "SELECT id FROM orders", + "SELECT * FROM orders WHERE amount > 10", + "SELECT o.id, c.name FROM orders o JOIN customers c ON o.customer_id = c.id", + "SELECT c.region, count(*) FROM orders o JOIN customers c ON o.customer_id = c.id GROUP BY c.region", + "SELECT DISTINCT customer_id FROM orders ORDER BY customer_id LIMIT 5", + "SELECT o.id FROM orders o JOIN customers c ON o.customer_id = c.id JOIN products p ON o.product_id = p.id", + "SELECT o.amount * 2 AS doubled FROM orders o WHERE o.amount BETWEEN 1 AND 5", + "SELECT 1 + 1", + ] { + assert_schema_preserved(sql); + } + } + + // ---- constant folding ---- + + #[test] + fn arithmetic_between_literals_is_evaluated_once_at_plan_time() { + let text = plan_text("SELECT 2 * 3 + 1 FROM orders"); + assert!(text.contains("Projection: 7"), "{text}"); + } + + #[test] + fn a_predicate_that_is_always_true_disappears_entirely() { + let text = plan_text("SELECT id FROM orders WHERE 1 = 1"); + assert!(!text.contains("Filter"), "{text}"); + assert!(!text.contains("pushed"), "{text}"); + } + + #[test] + fn a_predicate_that_is_always_false_survives_so_the_query_returns_nothing() { + let text = plan_text("SELECT id FROM orders WHERE 1 = 2"); + assert!(text.contains("false"), "{text}"); + } + + #[test] + fn boolean_identities_collapse() { + assert!(plan_text("SELECT id FROM orders WHERE id = 1 AND true").contains("(id#0 = 1)")); + let and_false = plan_text("SELECT id FROM orders WHERE id = 1 AND false"); + assert!(and_false.contains("false"), "{and_false}"); + assert!(!plan_text("SELECT id FROM orders WHERE id = 1 OR true").contains("id#0 = 1")); + assert!(plan_text("SELECT id FROM orders WHERE id = 1 OR false").contains("(id#0 = 1)")); + } + + #[test] + fn double_negation_cancels_and_a_literal_negation_folds() { + assert!(plan_text("SELECT id FROM orders WHERE NOT NOT (id = 1)").contains("(id#0 = 1)")); + assert!(plan_text("SELECT -(-3) FROM orders").contains("Projection: 3")); + assert!(plan_text("SELECT NOT true FROM orders").contains("Projection: false")); + } + + #[test] + fn casts_and_null_tests_over_literals_fold() { + assert!(plan_text("SELECT CAST('42' AS INT) FROM orders").contains("Projection: 42")); + assert!(plan_text("SELECT NULL IS NULL FROM orders").contains("Projection: true")); + assert!(plan_text("SELECT 1 IS NOT NULL FROM orders").contains("Projection: true")); + } + + #[test] + fn a_case_whose_conditions_are_known_collapses_to_the_branch_that_fires() { + assert!( + plan_text("SELECT CASE WHEN true THEN 1 ELSE 2 END FROM orders") + .contains("Projection: 1") + ); + assert!( + plan_text("SELECT CASE WHEN false THEN 1 ELSE 2 END FROM orders") + .contains("Projection: 2") + ); + assert!( + plan_text("SELECT CASE WHEN false THEN 1 END FROM orders").contains("Projection: NULL") + ); + } + + #[test] + fn division_by_zero_and_overflow_are_left_for_the_runtime_to_decide() { + // Folding these would mean the plan-time answer and the run-time + // answer could differ, which is worse than not folding. + let div = plan_text("SELECT 1 / 0 FROM orders"); + assert!(div.contains('/'), "{div}"); + let overflow = plan_text("SELECT 9223372036854775807 + 1 FROM orders"); + assert!(overflow.contains('+'), "{overflow}"); + } + + #[test] + fn arithmetic_involving_null_folds_to_null() { + assert!(plan_text("SELECT 1 + NULL FROM orders").contains("Projection: NULL")); + } + + // ---- predicate pushdown ---- + + #[test] + fn a_filter_ends_up_inside_the_scan_where_a_zone_map_can_use_it() { + let plan = opt("SELECT id FROM orders WHERE amount > 100"); + assert!(find(&plan, &|p| matches!(p, LogicalPlan::Filter { .. })).is_none()); + match scans(&plan)[0] { + LogicalPlan::Scan { pushed_filters, .. } => { + assert_eq!(pushed_filters.len(), 1); + assert!(pushed_filters[0].to_string().contains("> 100")); + } + other => panic!("unexpected {other:?}"), + } + } + + #[test] + fn a_conjunction_is_split_so_each_part_can_travel_separately() { + let plan = opt( + "SELECT o.id FROM orders o JOIN customers c ON o.customer_id = c.id \ + WHERE o.amount > 100 AND c.region = 3", + ); + // One predicate reached each side of the join, not both stuck above it. + for s in scans(&plan) { + match s { + LogicalPlan::Scan { + table, + pushed_filters, + .. + } => assert_eq!(pushed_filters.len(), 1, "{table} kept nothing"), + _ => unreachable!(), + } + } + assert!(find(&plan, &|p| matches!(p, LogicalPlan::Filter { .. })).is_none()); + } + + #[test] + fn a_predicate_spanning_both_sides_of_a_join_stays_above_it() { + let plan = opt( + "SELECT o.id FROM orders o JOIN customers c ON o.customer_id = c.id \ + WHERE o.amount > c.id", + ); + assert!(find(&plan, &|p| matches!(p, LogicalPlan::Filter { .. })).is_some()); + } + + #[test] + fn a_cross_join_with_an_equality_in_where_becomes_a_hash_joinable_inner_join() { + // The old-style comma join must not stay a cross product. + let plan = opt("SELECT o.id FROM orders o, customers c WHERE o.customer_id = c.id"); + match find(&plan, &|p| matches!(p, LogicalPlan::Join { .. })).unwrap() { + LogicalPlan::Join { join_type, on, .. } => { + assert_eq!(*join_type, JoinType::Inner); + assert_eq!(on.len(), 1); + } + other => panic!("unexpected {other:?}"), + } + } + + #[test] + fn a_reversed_equality_in_where_also_becomes_a_join_key() { + let plan = opt("SELECT o.id FROM orders o, customers c WHERE c.id = o.customer_id"); + match find(&plan, &|p| matches!(p, LogicalPlan::Join { .. })).unwrap() { + LogicalPlan::Join { on, join_type, .. } => { + assert_eq!(*join_type, JoinType::Inner); + // The left key must be orders.customer_id and the right one + // customers.id, even though the equality was written the other + // way round. + assert_eq!(on.len(), 1); + assert!(on[0].0.to_string().starts_with("customer_id#")); + assert!(on[0].1.to_string().starts_with("id#")); + } + other => panic!("unexpected {other:?}"), + } + } + + #[test] + fn a_predicate_is_never_pushed_into_the_padded_side_of_an_outer_join() { + // Pushing `c.region = 3` into the right side of a LEFT JOIN would + // delete rows that should have come back padded with NULLs. + let plan = opt( + "SELECT o.id FROM orders o LEFT JOIN customers c ON o.customer_id = c.id \ + WHERE c.region = 3", + ); + let customers = scans(&plan) + .into_iter() + .find(|s| matches!(s, LogicalPlan::Scan { table, .. } if table == "customers")) + .unwrap(); + match customers { + LogicalPlan::Scan { pushed_filters, .. } => assert!(pushed_filters.is_empty()), + _ => unreachable!(), + } + assert!(find(&plan, &|p| matches!(p, LogicalPlan::Filter { .. })).is_some()); + } + + #[test] + fn a_predicate_on_the_preserved_side_of_a_left_join_does_push_down() { + let plan = opt( + "SELECT o.id FROM orders o LEFT JOIN customers c ON o.customer_id = c.id \ + WHERE o.amount > 5", + ); + let orders = scans(&plan) + .into_iter() + .find(|s| matches!(s, LogicalPlan::Scan { table, .. } if table == "orders")) + .unwrap(); + match orders { + LogicalPlan::Scan { pushed_filters, .. } => assert_eq!(pushed_filters.len(), 1), + _ => unreachable!(), + } + } + + #[test] + fn a_full_outer_join_accepts_no_pushdown_at_all() { + let plan = opt( + "SELECT o.id FROM orders o FULL JOIN customers c ON o.customer_id = c.id \ + WHERE o.amount > 5", + ); + for s in scans(&plan) { + match s { + LogicalPlan::Scan { pushed_filters, .. } => assert!(pushed_filters.is_empty()), + _ => unreachable!(), + } + } + } + + #[test] + fn a_filter_travels_through_a_projection_of_plain_columns() { + let plan = opt("SELECT id FROM orders WHERE id > 5 ORDER BY id"); + match scans(&plan)[0] { + LogicalPlan::Scan { pushed_filters, .. } => assert_eq!(pushed_filters.len(), 1), + _ => unreachable!(), + } + } + + #[test] + fn a_filter_does_not_push_below_a_limit_or_an_aggregate() { + // Filtering before a LIMIT would change which rows the limit sees, + // so the filter has to stay above it. + let limited = optimize(LogicalPlan::Filter { + input: Box::new(LogicalPlan::Limit { + input: Box::new(raw("SELECT id FROM orders")), + limit: Some(10), + offset: 0, + }), + predicate: BoundExpr::binary( + BoundExpr::column(0, "id", DataType::Int64), + BinaryOp::Gt, + BoundExpr::Literal(Value::Int64(5)), + ) + .unwrap(), + }) + .unwrap(); + assert!(matches!(limited, LogicalPlan::Filter { .. })); + } + + #[test] + fn having_stays_above_the_aggregate_rather_than_becoming_a_scan_filter() { + let plan = opt( + "SELECT customer_id, count(*) FROM orders GROUP BY customer_id HAVING count(*) > 2", + ); + assert!(find(&plan, &|p| matches!(p, LogicalPlan::Filter { .. })).is_some()); + for s in scans(&plan) { + match s { + LogicalPlan::Scan { pushed_filters, .. } => assert!(pushed_filters.is_empty()), + _ => unreachable!(), + } + } + } + + // ---- projection pushdown ---- + + #[test] + fn a_scan_reads_only_the_columns_the_query_mentions() { + let plan = opt("SELECT id FROM orders WHERE amount > 5"); + match scans(&plan)[0] { + LogicalPlan::Scan { + projection: Some(p), + source_schema, + .. + } => { + // `id` and `amount` only — not `customer_id` or `product_id`. + assert_eq!(p.len(), 2, "read {p:?} of {source_schema}"); + assert!(p.contains(&0) && p.contains(&3)); + } + other => panic!("unexpected {other:?}"), + } + } + + #[test] + fn each_side_of_a_join_reads_only_its_own_needed_columns() { + let plan = opt( + "SELECT c.name FROM orders o JOIN customers c ON o.customer_id = c.id \ + WHERE o.amount > 5", + ); + for s in scans(&plan) { + match s { + LogicalPlan::Scan { + table, + projection: Some(p), + .. + } => { + // orders reads id + amount; customers reads id + name. + assert_eq!(p.len(), 2, "{table} read {p:?}"); + } + other => panic!("unexpected {other:?}"), + } + } + } + + #[test] + fn a_count_only_query_still_reads_one_column_so_it_has_rows_to_count() { + let plan = opt("SELECT count(*) FROM orders"); + match scans(&plan)[0] { + LogicalPlan::Scan { + projection: Some(p), + .. + } => assert_eq!(p.len(), 1), + other => panic!("unexpected {other:?}"), + } + } + + #[test] + fn pruning_renumbers_the_expressions_above_the_scan() { + // `amount` is source column 3 but becomes position 1 after pruning; + // the predicate has to follow it. + let plan = opt("SELECT id FROM orders WHERE amount > 5"); + match scans(&plan)[0] { + LogicalPlan::Scan { pushed_filters, .. } => { + assert_eq!(pushed_filters[0].column_indices(), vec![1]); + } + other => panic!("unexpected {other:?}"), + } + } + + #[test] + fn distinct_keeps_every_column_because_pruning_would_change_which_rows_repeat() { + let plan = opt("SELECT DISTINCT id, customer_id FROM orders"); + match scans(&plan)[0] { + LogicalPlan::Scan { + projection: Some(p), + .. + } => assert_eq!(p.len(), 2), + other => panic!("unexpected {other:?}"), + } + } + + #[test] + fn a_column_used_only_for_sorting_is_still_read() { + let plan = opt("SELECT id FROM orders ORDER BY amount"); + match scans(&plan)[0] { + LogicalPlan::Scan { + projection: Some(p), + .. + } => assert_eq!(p.len(), 2), + other => panic!("unexpected {other:?}"), + } + } + + // ---- join reordering ---- + + #[test] + fn a_three_way_join_is_reordered_smallest_first() { + // Written orders(100k) ⋈ customers(5k) ⋈ products(900); the smallest + // relation should end up at the bottom of the tree. + let plan = opt("SELECT o.id FROM orders o \ + JOIN customers c ON o.customer_id = c.id \ + JOIN products p ON o.product_id = p.id"); + let deepest = deepest_scan(&plan); + assert_eq!(deepest, "products", "plan was:\n{}", plan.explain()); + } + + #[test] + fn reordering_preserves_the_output_column_order() { + let sql = "SELECT o.id, c.name, p.title FROM orders o \ + JOIN customers c ON o.customer_id = c.id \ + JOIN products p ON o.product_id = p.id"; + let before = raw(sql).schema().unwrap(); + let after = opt(sql).schema().unwrap(); + assert_eq!(before, after); + assert_eq!(after.field(0).unwrap().name, "id"); + assert_eq!(after.field(1).unwrap().name, "name"); + assert_eq!(after.field(2).unwrap().name, "title"); + } + + #[test] + fn reordering_keeps_every_join_key() { + let plan = opt("SELECT o.id FROM orders o \ + JOIN customers c ON o.customer_id = c.id \ + JOIN products p ON o.product_id = p.id"); + let mut keys = 0; + count_keys(&plan, &mut keys); + fn count_keys(p: &LogicalPlan, n: &mut usize) { + if let LogicalPlan::Join { on, .. } = p { + *n += on.len(); + } + for c in p.children() { + count_keys(c, n); + } + } + assert_eq!(keys, 2, "a join key was lost:\n{}", plan.explain()); + } + + #[test] + fn reordering_never_introduces_a_cross_product_between_related_tables() { + let plan = opt("SELECT o.id FROM orders o \ + JOIN customers c ON o.customer_id = c.id \ + JOIN products p ON o.product_id = p.id"); + let mut crosses = 0; + count_crosses(&plan, &mut crosses); + fn count_crosses(p: &LogicalPlan, n: &mut usize) { + if let LogicalPlan::Join { + join_type: JoinType::Cross, + .. + } = p + { + *n += 1; + } + for c in p.children() { + count_crosses(c, n); + } + } + assert_eq!(crosses, 0, "plan was:\n{}", plan.explain()); + } + + #[test] + fn a_two_table_join_is_left_alone() { + let plan = opt("SELECT o.id FROM orders o JOIN customers c ON o.customer_id = c.id"); + // No restoring projection was needed, so the join sits directly under + // the query's own projection. + let text = plan.explain(); + assert_eq!(text.matches("Projection").count(), 1, "{text}"); + } + + #[test] + fn an_outer_join_chain_is_not_reordered() { + // Reordering outer joins changes which side gets padded. + let sql = "SELECT o.id FROM orders o \ + LEFT JOIN customers c ON o.customer_id = c.id \ + LEFT JOIN products p ON o.product_id = p.id"; + assert_eq!(deepest_scan(&opt(sql)), "orders"); + } + + fn deepest_scan(plan: &LogicalPlan) -> String { + fn walk(p: &LogicalPlan, depth: usize, best: &mut (usize, String)) { + if let LogicalPlan::Scan { table, .. } = p { + if depth > best.0 { + *best = (depth, table.clone()); + } + } + for c in p.children() { + walk(c, depth + 1, best); + } + } + let mut best = (0, String::new()); + walk(plan, 0, &mut best); + best.1 + } + + #[test] + fn the_rule_list_is_reported_for_explain() { + assert_eq!(RULES.len(), 4); + assert!(RULES.contains(&"predicate pushdown")); + } +}