From 8244f6380cd0be4a2e3c2cdfb3f84dce37c5f6cf Mon Sep 17 00:00:00 2001 From: Sahil Kalgutkar Date: Fri, 28 Aug 2026 14:30:53 -0400 Subject: [PATCH] Add the binder, cost model and optimiser Three layers between a parse tree and something runnable. The binder is the only place that resolves names, so error messages about missing or ambiguous columns all come from one function. An unqualified column matching two tables in a join is an error rather than a first-match win, and a bare column outside GROUP BY is rejected instead of returning an arbitrary row. BETWEEN and the CASE-with-operand form desugar here, so the optimiser and the evaluator only ever see comparisons. The cost model estimates cardinality from the zone maps the storage layer already writes. Its numbers only ever choose between plans, never decide what a query returns, which is why the fallbacks are allowed to be textbook constants. Four rules run over the plan. Constant folding stops short of division by zero and integer overflow so the plan-time answer cannot differ from the run-time one. Predicate pushdown splits conjunctions so each part travels separately, refuses to push into the padded side of an outer join, and turns a comma join with an equality in WHERE into a real inner join. Join reordering puts the smallest relation at the bottom of an inner-join chain, preferring a relation that has a join key over a smaller unrelated one, and restores the original column order above so nothing higher in the tree notices. Projection pushdown narrows every scan to the columns actually read and renumbers everything above it. Every rule has to preserve the output schema, and a test asserts exactly that across a spread of queries. --- crates/qf-plan/src/binder.rs | 1258 ++++++++++++++++++++++ crates/qf-plan/src/cost.rs | 589 +++++++++++ crates/qf-plan/src/expr.rs | 843 +++++++++++++++ crates/qf-plan/src/lib.rs | 6 +- crates/qf-plan/src/logical.rs | 701 ++++++++++++ crates/qf-plan/src/optimizer.rs | 1758 +++++++++++++++++++++++++++++++ 6 files changed, 5154 insertions(+), 1 deletion(-) create mode 100644 crates/qf-plan/src/binder.rs create mode 100644 crates/qf-plan/src/cost.rs create mode 100644 crates/qf-plan/src/expr.rs create mode 100644 crates/qf-plan/src/logical.rs create mode 100644 crates/qf-plan/src/optimizer.rs 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")); + } +}