diff --git a/core/translate/expr.rs b/core/translate/expr.rs index 0adaee185..df7e82724 100644 --- a/core/translate/expr.rs +++ b/core/translate/expr.rs @@ -2579,3 +2579,38 @@ pub fn maybe_apply_affinity(col_type: Type, target_register: usize, program: &mu pub fn sanitize_string(input: &str) -> String { input[1..input.len() - 1].replace("''", "'").to_string() } + +/// Returns the components of a binary expression +/// e.g. t.x = 5 -> Some((t.x, =, 5)) +pub fn as_binary_components( + expr: &ast::Expr, +) -> Result> { + match unwrap_parens(expr)? { + ast::Expr::Binary(lhs, operator, rhs) + if matches!( + operator, + ast::Operator::Equals + | ast::Operator::Greater + | ast::Operator::Less + | ast::Operator::GreaterEquals + | ast::Operator::LessEquals + ) => + { + Ok(Some((lhs.as_ref(), *operator, rhs.as_ref()))) + } + _ => Ok(None), + } +} + +/// Recursively unwrap parentheses from an expression +/// e.g. (((t.x > 5))) -> t.x > 5 +fn unwrap_parens(expr: &ast::Expr) -> Result<&ast::Expr> { + match expr { + ast::Expr::Column { .. } => Ok(expr), + ast::Expr::Parenthesized(exprs) => match exprs.len() { + 1 => unwrap_parens(exprs.first().unwrap()), + _ => crate::bail_parse_error!("expected single expression in parentheses"), + }, + _ => Ok(expr), + } +} diff --git a/core/translate/main_loop.rs b/core/translate/main_loop.rs index 4b25474e5..bf30b8508 100644 --- a/core/translate/main_loop.rs +++ b/core/translate/main_loop.rs @@ -131,9 +131,7 @@ pub fn init_loop( program.emit_insn(Insn::VOpen { cursor_id }); } } - _ => { - unimplemented!() - } + _ => {} }, Operation::Search(search) => { match mode { @@ -435,13 +433,7 @@ pub fn open_loop( // Rowid equality point lookups are handled with a SeekRowid instruction which does not loop, since it is a single row lookup. if let Search::RowidEq { cmp_expr } = search { let src_reg = program.alloc_register(); - translate_expr( - program, - Some(tables), - &cmp_expr.expr, - src_reg, - &t_ctx.resolver, - )?; + translate_expr(program, Some(tables), cmp_expr, src_reg, &t_ctx.resolver)?; program.emit_insn(Insn::SeekRowid { cursor_id: table_cursor_id .expect("Search::RowidEq requires a table cursor"), diff --git a/core/translate/optimizer/OPTIMIZER.md b/core/translate/optimizer/OPTIMIZER.md new file mode 100644 index 000000000..0939fd9e5 --- /dev/null +++ b/core/translate/optimizer/OPTIMIZER.md @@ -0,0 +1,162 @@ +# Overview of the current state of the query optimizer in Limbo + +Query optimization is obviously an important part of any SQL-based database engine. This document is an overview of what we currently do. + +## Structure of the optimizer directory + +1. `mod.rs` + - Provides the high-level optimization interface through `optimize_plan()` + +2. `access_method.rs` + - Determines what is the best index to use when joining a table to a set of preceding tables + +3. `constraints.rs` - Manages query constraints: + - Extracts constraints from the WHERE clause + - Determines which constraints are usable with indexes + +4. `cost.rs` + - Calculates the cost of doing a seek vs a scan, for example + +5. `join.rs` + - Implements the System R style dynamic programming join ordering algorithm + +6. `order.rs` + - Determines if sort operations can be eliminated based on the chosen access methods and join order + +## Join reordering and optimal index selection + +**The goals of query optimization are at least the following:** + +1. Do as little page I/O as possible +2. Do as little CPU work as possible +3. Retain query correctness. + +**The most important ways to achieve no. 1 and no. 2 are:** + +1. Choose the optimal access method for each table (e.g. an index or a rowid-based seek, or a full table scan if all else fails). +2. Choose the best or near-best way to reorder the tables in the query so that those optimal access methods can be used. +3. Also factor in whether the chosen join order and indexes allow removal of any sort operations that are necessary for query correctness. + +## Limbo's optimizer + +Limbo's optimizer is an implementation of an extremely traditional [IBM System R](https://www.cs.cmu.edu/~15721-f24/slides/02-Selinger-SystemR-opt.pdf) style optimizer, +i.e. straight from the 70s! The DP algorithm is explained below. + +### Current high level flow of the optimizer + +1. **SQL rewriting** + - Rewrite certain SQL expressions to another form (not a lot currently; e.g. rewrite BETWEEN as two comparisons) + - Eliminate constant conditions: e.g. `WHERE 1` is removed, `WHERE 0` short-circuits the whole query because it is trivially false. +2. **Check whether there is an "interesting order"** that we should consider when evaluating indexes and join orders + - Is there a GROUP BY? an ORDER BY? Both? +3. **Convert WHERE clause conjucts to Constraints** + - E.g. in `WHERE t.x = 5`, the expression `5` _constrains_ table `t` to values of `x` that are exactly `5`. + - E.g. in `Where t.x = u.x`, the expression `u.x` constrains `t`, AND `t.x` constrains `u`. + - Per table, each constraint has an estimated _selectivity_ (how much it filters the result set); this affects join order calculations, see the paragraph on _Estimation_ below. + - Per table, constraints are also analyzed for whether one or multiple of them can be used as an index seek key to avoid a full scan. +4. **Compute the best join order using a dynamic programming algorithm:** + - `n` = number of tables considered + - `n=1`: find the lowest _cost_ way to access each single table, given the constraints of the query. Memoize the result. + - `n=2`: for each table found in the `n=1` step, find the best way to join that table with each other table. Memoize the result. + - `n=3`: for each 2-table subset found, find the best way to join that result to each other table. Memoize the result. + - `n=m`: for each `m-1` table subset found, find the best way to join that result to the `m'th` table + - **Use pruning to reduce search space:** + - Compute the literal query order first, and store its _cost_ as an upper threshold + - If at any point a considered join order exceeds the upper threshold, discard that search path since it cannot be better than the current best. + - For example, we have `SELECT * FROM a JOIN b JOIN c JOIN d`. Compute `JOIN(a,b,c,d)` first. If `JOIN (b,a)` is already worse than `JOIN(a,b,c,d)`, we don't have to even try `JOIN(b,a,c)`. + - Also keep track of the best plan per _subset_: + - If we find that `JOIN(b,a,c)` is better than any other permutation of the same tables, e.g. `JOIN(a,b,c)`, then we can discard _ALL_ of the other permutations for that subset. For example, we don't need to consider `JOIN(a,b,c,d)` because we know it's worse than `JOIN(b,a,c,d)`. + - This is possible due to the associativity and commutativity of INNER JOINs. + - Also keep track of the best _ordered plan_ , i.e. one that provides the "interesting order" mentioned above. + - At the end, apply a cost penalty to the best overall plan + - If it is now worse than the best sorted plan, then choose the sorted plan as the best plan for the query. + - This allows us to eliminate a sorting operation. + - If the best overall plan is still best even with the sorting penalty, then keep it. A sorting operation is later applied to sort the rows according to the desired order. +5. **Mutate the plan's `join_order` and `Operation`s to match the computed best plan.** + +### Estimation of cost and cardinalities + a note on table statistics + +Currently, in the absence of `ANALYZE`, `sqlite_stat1` etc. we assume the following: + +1. Each table has `1,000,000` rows. +2. Each equality (`=`) filter will filter out some percentage of the result set. +3. Each nonequality (e.g. `>`) will filter out some smaller percentage of the result set. +4. Each `4096` byte database page holds `50` rows, i.e. roughly `80` bytes per row +5. Sort operations have some CPU cost dependent on the number of input rows to the sort operation. + +From the above, we derive the following formula for estimating the cost of joining `t1` with `t2` + +``` +JOIN_COST = PAGE_IO(t1.rows) + t1.rows * PAGE_IO(t2.rows) +``` + +For example, let's take the query `SELECT * FROM t1 JOIN t2 USING(foo) WHERE t2.foo > 10`. Let's assume the following: + +- `t1` has `6400` rows and `t2` has `8000` rows +- there are no indexes at all +- let's ignore the CPU cost from the equation for simplicity. + +The best access method for both is a full table scan. The output cardinality of `t1` is the full table, because nothing is filtering it. Hence, the cost of `t1 JOIN t2` becomes: + +``` +JOIN_COST = PAGE_IO(t1.input_rows) + t1.output_rows * PAGE_IO(t2.input_rows) + +// plugging in the values: + +JOIN_COST = PAGE_IO(6400) + 6400 * PAGE_IO(8000) +JOIN_COST = 80 + 6400 * 100 = 640080 +``` + +Now let's consider `t2 JOIN t1`. The best access method for both is still a full scan, but since we can filter on `t2.foo > 10`, its output cardinality decreases. Let's assume only 1/4 of the rows of `t2` match the condition `t2.foo > 10`. Hence, the cost of `t2 join t1` becomes: + +``` +JOIN_COST = PAGE_IO(t2.input_rows) + t2.output_rows * PAGE_IO(t1.input_rows) + +// plugging in the values: + +JOIN_COST = PAGE_IO(8000) + 1/4 * 8000 * PAGE_IO(6400) +JOIN_COST = 100 + 2000 * 80 = 160100 +``` + +Even though `t2` is a larger table, because we were able to reduce the input set to the join operation, it's dramatically cheaper. + +#### Statistics + +Since we don't support `ANALYZE`, nor can we assume that users will call `ANALYZE` anyway, we use simple magic constants to estimate the selectivity of join predicates, row count of tables, and so on. When we have support for `ANALYZE`, we should plug the statistics from `sqlite_stat1` and friends into the optimizer to make more informed decisions. + +### Estimating the output cardinality of a join + +The output cardinality (output row count) of an operation is as follows: + +``` +OUTPUT_CARDINALITY_JOIN = INPUT_CARDINALITY_RHS * OUTPUT_CARDINALITY_RHS + +where + +INPUT_CARDINALITY_RHS = OUTPUT_CARDINALITY_LHS +``` + +example: + +``` +SELECT * FROM products p JOIN order_lines o ON p.id = o.product_id +``` +Assuming there are 100 products, i.e. just selecting all products would yield 100 rows: + +``` +OUTPUT_CARDINALITY_LHS = 100 +INPUT_CARDINALITY_RHS = 100 +``` + +Assuming p.id = o.product_id will return three orders per each product: + +``` +OUTPUT_CARDINALITY_RHS = 3 + +OUTPUT_CARDINALITY_JOIN = 100 * 3 = 300 +``` +i.e. the join is estimated to return 300 rows, 3 for each product. + +Again, in the absence of statistics, we use magic constants to estimate these cardinalities. + +Estimating them is important because in multi-way joins the output cardinality of the previous join becomes the input cardinality of the next one. \ No newline at end of file diff --git a/core/translate/optimizer/access_method.rs b/core/translate/optimizer/access_method.rs new file mode 100644 index 000000000..8888e7755 --- /dev/null +++ b/core/translate/optimizer/access_method.rs @@ -0,0 +1,148 @@ +use std::sync::Arc; + +use limbo_sqlite3_parser::ast::SortOrder; + +use crate::{ + schema::Index, + translate::plan::{IterationDirection, JoinOrderMember, TableReference}, + Result, +}; + +use super::{ + constraints::{usable_constraints_for_join_order, ConstraintRef, TableConstraints}, + cost::{estimate_cost_for_scan_or_seek, Cost, IndexInfo}, + order::OrderTarget, +}; + +#[derive(Debug, Clone)] +/// Represents a way to access a table. +pub struct AccessMethod<'a> { + /// The estimated number of page fetches. + /// We are ignoring CPU cost for now. + pub cost: Cost, + /// The direction of iteration for the access method. + /// Typically this is backwards only if it helps satisfy an [OrderTarget]. + pub iter_dir: IterationDirection, + /// The index that is being used, if any. For rowid based searches (and full table scans), this is None. + pub index: Option>, + /// The constraint references that are being used, if any. + /// An empty list of constraint refs means a scan (full table or index); + /// a non-empty list means a search. + pub constraint_refs: &'a [ConstraintRef], +} + +impl<'a> AccessMethod<'a> { + pub fn is_scan(&self) -> bool { + self.constraint_refs.is_empty() + } + + pub fn new_table_scan(input_cardinality: f64, iter_dir: IterationDirection) -> Self { + Self { + cost: estimate_cost_for_scan_or_seek(None, &[], &[], input_cardinality), + iter_dir, + index: None, + constraint_refs: &[], + } + } +} + +/// Return the best [AccessMethod] for a given join order. +pub fn find_best_access_method_for_join_order<'a>( + rhs_table: &TableReference, + rhs_constraints: &'a TableConstraints, + join_order: &[JoinOrderMember], + maybe_order_target: Option<&OrderTarget>, + input_cardinality: f64, +) -> Result> { + let table_no = join_order.last().unwrap().table_no; + let mut best_access_method = + AccessMethod::new_table_scan(input_cardinality, IterationDirection::Forwards); + let rowid_column_idx = rhs_table.columns().iter().position(|c| c.is_rowid_alias); + + // Estimate cost for each candidate index (including the rowid index) and replace best_access_method if the cost is lower. + for candidate in rhs_constraints.candidates.iter() { + let index_info = match candidate.index.as_ref() { + Some(index) => IndexInfo { + unique: index.unique, + covering: rhs_table.index_is_covering(index), + column_count: index.columns.len(), + }, + None => IndexInfo { + unique: true, // rowids are always unique + covering: false, + column_count: 1, + }, + }; + let usable_constraint_refs = usable_constraints_for_join_order( + &rhs_constraints.constraints, + &candidate.refs, + join_order, + ); + let cost = estimate_cost_for_scan_or_seek( + Some(index_info), + &rhs_constraints.constraints, + &usable_constraint_refs, + input_cardinality, + ); + + // All other things being equal, prefer an access method that satisfies the order target. + let (iter_dir, order_satisfiability_bonus) = if let Some(order_target) = maybe_order_target + { + // If the index delivers rows in the same direction (or the exact reverse direction) as the order target, then it + // satisfies the order target. + let mut all_same_direction = true; + let mut all_opposite_direction = true; + for i in 0..order_target.0.len().min(index_info.column_count) { + let correct_table = order_target.0[i].table_no == table_no; + let correct_column = { + match &candidate.index { + Some(index) => index.columns[i].pos_in_table == order_target.0[i].column_no, + None => { + rowid_column_idx.map_or(false, |idx| idx == order_target.0[i].column_no) + } + } + }; + if !correct_table || !correct_column { + all_same_direction = false; + all_opposite_direction = false; + break; + } + let correct_order = { + match &candidate.index { + Some(index) => order_target.0[i].order == index.columns[i].order, + None => order_target.0[i].order == SortOrder::Asc, + } + }; + if correct_order { + all_opposite_direction = false; + } else { + all_same_direction = false; + } + } + if all_same_direction || all_opposite_direction { + ( + if all_same_direction { + IterationDirection::Forwards + } else { + IterationDirection::Backwards + }, + Cost(1.0), + ) + } else { + (IterationDirection::Forwards, Cost(0.0)) + } + } else { + (IterationDirection::Forwards, Cost(0.0)) + }; + if cost < best_access_method.cost + order_satisfiability_bonus { + best_access_method = AccessMethod { + cost, + index: candidate.index.clone(), + iter_dir, + constraint_refs: &usable_constraint_refs, + }; + } + } + + Ok(best_access_method) +} diff --git a/core/translate/optimizer/constraints.rs b/core/translate/optimizer/constraints.rs new file mode 100644 index 000000000..a7272ce86 --- /dev/null +++ b/core/translate/optimizer/constraints.rs @@ -0,0 +1,418 @@ +use std::{cmp::Ordering, collections::HashMap, sync::Arc}; + +use crate::{ + schema::{Column, Index}, + translate::{ + expr::as_binary_components, + plan::{JoinOrderMember, TableReference, WhereTerm}, + planner::{table_mask_from_expr, TableMask}, + }, + Result, +}; +use limbo_sqlite3_parser::ast::{self, SortOrder}; + +use super::cost::ESTIMATED_HARDCODED_ROWS_PER_TABLE; + +/// Represents a single condition derived from a `WHERE` clause term +/// that constrains a specific column of a table. +/// +/// Constraints are precomputed for each table involved in a query. They are used +/// during query optimization to estimate the cost of different access paths (e.g., using an index) +/// and to determine the optimal join order. A constraint can only be applied if all tables +/// referenced in its expression (other than the constrained table itself) are already +/// available in the current join context, i.e. on the left side in the join order +/// relative to the table. +#[derive(Debug, Clone)] +/// +pub struct Constraint { + /// The position of the original `WHERE` clause term this constraint derives from, + /// and which side of the [ast::Expr::Binary] comparison contains the expression + /// that constrains the column. + /// E.g. in SELECT * FROM t WHERE t.x = 10, the constraint is (0, BinaryExprSide::Rhs) + /// because the RHS '10' is the constraining expression. + /// + /// This is tracked so we can: + /// + /// 1. Extract the constraining expression for use in an index seek key, and + /// 2. Remove the relevant binary expression from the WHERE clause, if used as an index seek key. + pub where_clause_pos: (usize, BinaryExprSide), + /// The comparison operator (e.g., `=`, `>`, `<`) used in the constraint. + pub operator: ast::Operator, + /// The zero-based index of the constrained column within the table's schema. + pub table_col_pos: usize, + /// A bitmask representing the set of tables that appear on the *constraining* side + /// of the comparison expression. For example, in SELECT * FROM t1,t2,t3 WHERE t1.x = t2.x + t3.x, + /// the lhs_mask contains t2 and t3. Thus, this constraint can only be used if t2 and t3 + /// have already been joined (i.e. are on the left side of the join order relative to t1). + pub lhs_mask: TableMask, + /// An estimated selectivity factor (0.0 to 1.0) indicating the fraction of rows + /// expected to satisfy this constraint. Used for cost and cardinality estimation. + pub selectivity: f64, +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum BinaryExprSide { + Lhs, + Rhs, +} + +impl Constraint { + /// Get the constraining expression, e.g. '2+3' from 't.x = 2+3' + pub fn get_constraining_expr(&self, where_clause: &[WhereTerm]) -> ast::Expr { + let (idx, side) = self.where_clause_pos; + let where_term = &where_clause[idx]; + let Ok(Some((lhs, _, rhs))) = as_binary_components(&where_term.expr) else { + panic!("Expected a valid binary expression"); + }; + if side == BinaryExprSide::Lhs { + lhs.clone() + } else { + rhs.clone() + } + } +} + +#[derive(Debug, Clone)] +/// A reference to a [Constraint] in a [TableConstraints]. +/// +/// This is used to track which constraints may be used as an index seek key. +pub struct ConstraintRef { + /// The position of the constraint in the [TableConstraints::constraints] vector. + pub constraint_vec_pos: usize, + /// The position of the constrained column in the index. Always 0 for rowid indices. + pub index_col_pos: usize, + /// The sort order of the constrained column in the index. Always ascending for rowid indices. + pub sort_order: SortOrder, +} + +impl ConstraintRef { + /// Convert the constraint to a column usable in a [crate::translate::plan::SeekDef::key]. + pub fn as_seek_key_column( + &self, + constraints: &[Constraint], + where_clause: &[WhereTerm], + ) -> (ast::Expr, SortOrder) { + let constraint = &constraints[self.constraint_vec_pos]; + let constraining_expr = constraint.get_constraining_expr(where_clause); + (constraining_expr, self.sort_order) + } +} + +/// A collection of [ConstraintRef]s for a given index, or if index is None, for the table's rowid index. +/// For example, given a table `T (x,y,z)` with an index `T_I (y desc,z)`, take the following query: +/// ```sql +/// SELECT * FROM T WHERE y = 10 AND z = 20; +/// ``` +/// +/// This will produce the following [ConstraintUseCandidate]: +/// +/// ConstraintUseCandidate { +/// index: Some(T_I) +/// refs: [ +/// ConstraintRef { +/// constraint_vec_pos: 0, // y = 10 +/// index_col_pos: 0, // y +/// sort_order: SortOrder::Desc, +/// }, +/// ConstraintRef { +/// constraint_vec_pos: 1, // z = 20 +/// index_col_pos: 1, // z +/// sort_order: SortOrder::Asc, +/// }, +/// ], +/// } +/// +#[derive(Debug)] +pub struct ConstraintUseCandidate { + /// The index that may be used to satisfy the constraints. If none, the table's rowid index is used. + pub index: Option>, + /// References to the constraints that may be used as an access path for the index. + pub refs: Vec, +} + +#[derive(Debug)] +/// A collection of [Constraint]s and their potential [ConstraintUseCandidate]s for a given table. +pub struct TableConstraints { + pub table_no: usize, + /// The constraints for the table, i.e. any [WhereTerm]s that reference columns from this table. + pub constraints: Vec, + /// Candidates for indexes that may use the constraints to perform a lookup. + pub candidates: Vec, +} + +/// In lieu of statistics, we estimate that an equality filter will reduce the output set to 1% of its size. +const SELECTIVITY_EQ: f64 = 0.01; +/// In lieu of statistics, we estimate that a range filter will reduce the output set to 40% of its size. +const SELECTIVITY_RANGE: f64 = 0.4; +/// In lieu of statistics, we estimate that other filters will reduce the output set to 90% of its size. +const SELECTIVITY_OTHER: f64 = 0.9; + +const SELECTIVITY_UNIQUE_EQUALITY: f64 = 1.0 / ESTIMATED_HARDCODED_ROWS_PER_TABLE as f64; + +/// Estimate the selectivity of a constraint based on the operator and the column type. +fn estimate_selectivity(column: &Column, op: ast::Operator) -> f64 { + match op { + ast::Operator::Equals => { + if column.is_rowid_alias || column.primary_key { + SELECTIVITY_UNIQUE_EQUALITY + } else { + SELECTIVITY_EQ + } + } + ast::Operator::Greater => SELECTIVITY_RANGE, + ast::Operator::GreaterEquals => SELECTIVITY_RANGE, + ast::Operator::Less => SELECTIVITY_RANGE, + ast::Operator::LessEquals => SELECTIVITY_RANGE, + _ => SELECTIVITY_OTHER, + } +} + +/// Precompute all potentially usable [Constraints] from a WHERE clause. +/// The resulting list of [TableConstraints] is then used to evaluate the best access methods for various join orders. +pub fn constraints_from_where_clause( + where_clause: &[WhereTerm], + table_references: &[TableReference], + available_indexes: &HashMap>>, +) -> Result> { + let mut constraints = Vec::new(); + + // For each table, collect all the Constraints and all potential index candidates that may use them. + for (table_no, table_reference) in table_references.iter().enumerate() { + let rowid_alias_column = table_reference + .columns() + .iter() + .position(|c| c.is_rowid_alias); + + let mut cs = TableConstraints { + table_no, + constraints: Vec::new(), + candidates: available_indexes + .get(table_reference.table.get_name()) + .map_or(Vec::new(), |indexes| { + indexes + .iter() + .map(|index| ConstraintUseCandidate { + index: Some(index.clone()), + refs: Vec::new(), + }) + .collect() + }), + }; + // Add a candidate for the rowid index, which is always available when the table has a rowid alias. + cs.candidates.push(ConstraintUseCandidate { + index: None, + refs: Vec::new(), + }); + + for (i, term) in where_clause.iter().enumerate() { + let Some((lhs, operator, rhs)) = as_binary_components(&term.expr)? else { + continue; + }; + + // Constraints originating from a LEFT JOIN must always be evaluated in that join's RHS table's loop, + // regardless of which tables the constraint references. + if let Some(outer_join_tbl) = term.from_outer_join { + if outer_join_tbl != table_no { + continue; + } + } + + // If either the LHS or RHS of the constraint is a column from the table, add the constraint. + match lhs { + ast::Expr::Column { table, column, .. } => { + if *table == table_no { + let table_column = &table_reference.table.columns()[*column]; + cs.constraints.push(Constraint { + where_clause_pos: (i, BinaryExprSide::Rhs), + operator, + table_col_pos: *column, + lhs_mask: table_mask_from_expr(rhs)?, + selectivity: estimate_selectivity(table_column, operator), + }); + } + } + ast::Expr::RowId { table, .. } => { + // A rowid alias column must exist for the 'rowid' keyword to be considered a valid reference. + // This should be a parse error at an earlier stage of the query compilation, but nevertheless, + // we check it here. + if *table == table_no && rowid_alias_column.is_some() { + let table_column = + &table_reference.table.columns()[rowid_alias_column.unwrap()]; + cs.constraints.push(Constraint { + where_clause_pos: (i, BinaryExprSide::Rhs), + operator, + table_col_pos: rowid_alias_column.unwrap(), + lhs_mask: table_mask_from_expr(rhs)?, + selectivity: estimate_selectivity(table_column, operator), + }); + } + } + _ => {} + }; + match rhs { + ast::Expr::Column { table, column, .. } => { + if *table == table_no { + let table_column = &table_reference.table.columns()[*column]; + cs.constraints.push(Constraint { + where_clause_pos: (i, BinaryExprSide::Lhs), + operator: opposite_cmp_op(operator), + table_col_pos: *column, + lhs_mask: table_mask_from_expr(lhs)?, + selectivity: estimate_selectivity(table_column, operator), + }); + } + } + ast::Expr::RowId { table, .. } => { + if *table == table_no && rowid_alias_column.is_some() { + let table_column = + &table_reference.table.columns()[rowid_alias_column.unwrap()]; + cs.constraints.push(Constraint { + where_clause_pos: (i, BinaryExprSide::Lhs), + operator: opposite_cmp_op(operator), + table_col_pos: rowid_alias_column.unwrap(), + lhs_mask: table_mask_from_expr(lhs)?, + selectivity: estimate_selectivity(table_column, operator), + }); + } + } + _ => {} + }; + } + // sort equalities first so that index keys will be properly constructed. + // see e.g.: https://www.solarwinds.com/blog/the-left-prefix-index-rule + cs.constraints.sort_by(|a, b| { + if a.operator == ast::Operator::Equals { + Ordering::Less + } else if b.operator == ast::Operator::Equals { + Ordering::Greater + } else { + Ordering::Equal + } + }); + + // For each constraint we found, add a reference to it for each index that may be able to use it. + for (i, constraint) in cs.constraints.iter().enumerate() { + if rowid_alias_column.map_or(false, |idx| constraint.table_col_pos == idx) { + let rowid_candidate = cs + .candidates + .iter_mut() + .find_map(|candidate| { + if candidate.index.is_none() { + Some(candidate) + } else { + None + } + }) + .unwrap(); + rowid_candidate.refs.push(ConstraintRef { + constraint_vec_pos: i, + index_col_pos: 0, + sort_order: SortOrder::Asc, + }); + } + for index in available_indexes + .get(table_reference.table.get_name()) + .unwrap_or(&Vec::new()) + { + if let Some(position_in_index) = + index.column_table_pos_to_index_pos(constraint.table_col_pos) + { + let index_candidate = cs + .candidates + .iter_mut() + .find_map(|candidate| { + if candidate + .index + .as_ref() + .map_or(false, |i| Arc::ptr_eq(index, i)) + { + Some(candidate) + } else { + None + } + }) + .unwrap(); + index_candidate.refs.push(ConstraintRef { + constraint_vec_pos: i, + index_col_pos: position_in_index, + sort_order: index.columns[position_in_index].order, + }); + } + } + } + + for candidate in cs.candidates.iter_mut() { + // Sort by index_col_pos, ascending -- index columns must be consumed in contiguous order. + candidate.refs.sort_by_key(|cref| cref.index_col_pos); + // Deduplicate by position, keeping first occurrence (which will be equality if one exists, since the constraints vec is sorted that way) + candidate.refs.dedup_by_key(|cref| cref.index_col_pos); + // Truncate at first gap in positions -- again, index columns must be consumed in contiguous order. + let mut last_pos = 0; + let mut i = 0; + for cref in candidate.refs.iter() { + if cref.index_col_pos != last_pos { + if cref.index_col_pos != last_pos + 1 { + break; + } + last_pos = cref.index_col_pos; + } + i += 1; + } + candidate.refs.truncate(i); + + // Truncate after the first inequality, since the left-prefix rule of indexes requires that all constraints but the last one must be equalities; + // again see: https://www.solarwinds.com/blog/the-left-prefix-index-rule + if let Some(first_inequality) = candidate.refs.iter().position(|cref| { + cs.constraints[cref.constraint_vec_pos].operator != ast::Operator::Equals + }) { + candidate.refs.truncate(first_inequality + 1); + } + } + constraints.push(cs); + } + + Ok(constraints) +} + +/// Find which [Constraint]s are usable for a given join order. +/// Returns a slice of the references to the constraints that are usable. +/// A constraint is considered usable for a given table if all of the other tables referenced by the constraint +/// are on the left side in the join order relative to the table. +pub fn usable_constraints_for_join_order<'a>( + constraints: &'a [Constraint], + refs: &'a [ConstraintRef], + join_order: &[JoinOrderMember], +) -> &'a [ConstraintRef] { + let table_no = join_order.last().unwrap().table_no; + let mut usable_until = 0; + for cref in refs.iter() { + let constraint = &constraints[cref.constraint_vec_pos]; + let other_side_refers_to_self = constraint.lhs_mask.contains_table(table_no); + if other_side_refers_to_self { + break; + } + let lhs_mask = TableMask::from_table_number_iter( + join_order + .iter() + .take(join_order.len() - 1) + .map(|j| j.table_no), + ); + let all_required_tables_are_on_left_side = lhs_mask.contains_all(&constraint.lhs_mask); + if !all_required_tables_are_on_left_side { + break; + } + usable_until += 1; + } + &refs[..usable_until] +} + +fn opposite_cmp_op(op: ast::Operator) -> ast::Operator { + match op { + ast::Operator::Equals => ast::Operator::Equals, + ast::Operator::Greater => ast::Operator::Less, + ast::Operator::GreaterEquals => ast::Operator::LessEquals, + ast::Operator::Less => ast::Operator::Greater, + ast::Operator::LessEquals => ast::Operator::GreaterEquals, + _ => panic!("unexpected operator: {:?}", op), + } +} diff --git a/core/translate/optimizer/cost.rs b/core/translate/optimizer/cost.rs new file mode 100644 index 000000000..862592b6a --- /dev/null +++ b/core/translate/optimizer/cost.rs @@ -0,0 +1,72 @@ +use super::constraints::{Constraint, ConstraintRef}; + +/// A simple newtype wrapper over a f64 that represents the cost of an operation. +/// +/// This is used to estimate the cost of scans, seeks, and joins. +#[derive(Debug, Clone, Copy, PartialEq, PartialOrd)] +pub struct Cost(pub f64); + +impl std::ops::Add for Cost { + type Output = Cost; + + fn add(self, other: Cost) -> Cost { + Cost(self.0 + other.0) + } +} + +impl std::ops::Deref for Cost { + type Target = f64; + + fn deref(&self) -> &f64 { + &self.0 + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct IndexInfo { + pub unique: bool, + pub column_count: usize, + pub covering: bool, +} + +pub const ESTIMATED_HARDCODED_ROWS_PER_TABLE: usize = 1000000; +pub const ESTIMATED_HARDCODED_ROWS_PER_PAGE: usize = 50; // roughly 80 bytes per 4096 byte page + +pub fn estimate_page_io_cost(rowcount: f64) -> Cost { + Cost((rowcount as f64 / ESTIMATED_HARDCODED_ROWS_PER_PAGE as f64).ceil()) +} + +/// Estimate the cost of a scan or seek operation. +/// +/// This is a very simple model that estimates the number of pages read +/// based on the number of rows read, ignoring any CPU costs. +pub fn estimate_cost_for_scan_or_seek( + index_info: Option, + constraints: &[Constraint], + usable_constraint_refs: &[ConstraintRef], + input_cardinality: f64, +) -> Cost { + let Some(index_info) = index_info else { + return estimate_page_io_cost( + input_cardinality * ESTIMATED_HARDCODED_ROWS_PER_TABLE as f64, + ); + }; + + let selectivity_multiplier: f64 = usable_constraint_refs + .iter() + .map(|cref| { + let constraint = &constraints[cref.constraint_vec_pos]; + constraint.selectivity + }) + .product(); + + // little cheeky bonus for covering indexes + let covering_multiplier = if index_info.covering { 0.9 } else { 1.0 }; + + estimate_page_io_cost( + selectivity_multiplier + * ESTIMATED_HARDCODED_ROWS_PER_TABLE as f64 + * input_cardinality + * covering_multiplier, + ) +} diff --git a/core/translate/optimizer/join.rs b/core/translate/optimizer/join.rs new file mode 100644 index 000000000..aceb2b889 --- /dev/null +++ b/core/translate/optimizer/join.rs @@ -0,0 +1,1284 @@ +use std::{cell::RefCell, collections::HashMap}; + +use crate::{ + translate::{ + optimizer::{cost::Cost, order::plan_satisfies_order_target}, + plan::{JoinOrderMember, TableReference}, + planner::TableMask, + }, + Result, +}; + +use super::{ + access_method::{find_best_access_method_for_join_order, AccessMethod}, + constraints::TableConstraints, + cost::ESTIMATED_HARDCODED_ROWS_PER_TABLE, + order::OrderTarget, +}; + +/// Represents an n-ary join, anywhere from 1 table to N tables. +#[derive(Debug, Clone)] +pub struct JoinN { + /// Tuple: (table_number, access_method_index) + pub data: Vec<(usize, usize)>, + /// The estimated number of rows returned by joining these n tables together. + pub output_cardinality: usize, + /// Estimated execution cost of this N-ary join. + pub cost: Cost, +} + +impl JoinN { + pub fn table_numbers(&self) -> impl Iterator + use<'_> { + self.data.iter().map(|(table_number, _)| *table_number) + } + + pub fn best_access_methods(&self) -> impl Iterator + use<'_> { + self.data + .iter() + .map(|(_, access_method_index)| *access_method_index) + } +} + +/// Join n-1 tables with the n'th table. +/// Returns None if the plan is worse than the provided cost upper bound. +pub fn join_lhs_and_rhs<'a>( + lhs: Option<&JoinN>, + rhs_table_reference: &TableReference, + rhs_constraints: &'a TableConstraints, + join_order: &[JoinOrderMember], + maybe_order_target: Option<&OrderTarget>, + access_methods_arena: &'a RefCell>>, + cost_upper_bound: Cost, +) -> Result> { + // The input cardinality for this join is the output cardinality of the previous join. + // For example, in a 2-way join, if the left table has 1000 rows, and the right table will return 2 rows for each of the left table's rows, + // then the output cardinality of the join will be 2000. + let input_cardinality = lhs.map_or(1, |l| l.output_cardinality); + + let best_access_method = find_best_access_method_for_join_order( + rhs_table_reference, + rhs_constraints, + &join_order, + maybe_order_target, + input_cardinality as f64, + )?; + + let lhs_cost = lhs.map_or(Cost(0.0), |l| l.cost); + let cost = lhs_cost + best_access_method.cost; + + if cost > cost_upper_bound { + return Ok(None); + } + + access_methods_arena.borrow_mut().push(best_access_method); + + let mut best_access_methods = Vec::with_capacity(join_order.len()); + best_access_methods.extend(lhs.map_or(vec![], |l| l.data.clone())); + + let rhs_table_number = join_order.last().unwrap().table_no; + best_access_methods.push((rhs_table_number, access_methods_arena.borrow().len() - 1)); + + let lhs_mask = lhs.map_or(TableMask::new(), |l| { + TableMask::from_table_number_iter(l.table_numbers()) + }); + // Output cardinality is reduced by the product of the selectivities of the constraints that can be used with this join order. + let output_cardinality_multiplier = rhs_constraints + .constraints + .iter() + .filter(|c| lhs_mask.contains_all(&c.lhs_mask)) + .map(|c| c.selectivity) + .product::(); + + // Produce a number of rows estimated to be returned when this table is filtered by the WHERE clause. + // If this table is the rightmost table in the join order, we multiply by the input cardinality, + // which is the output cardinality of the previous tables. + let output_cardinality = (input_cardinality as f64 + * ESTIMATED_HARDCODED_ROWS_PER_TABLE as f64 + * output_cardinality_multiplier) + .ceil() as usize; + + Ok(Some(JoinN { + data: best_access_methods, + output_cardinality, + cost, + })) +} + +/// The result of [compute_best_join_order]. +#[derive(Debug)] +pub struct BestJoinOrderResult { + /// The best plan overall. + pub best_plan: JoinN, + /// The best plan for the given order target, if it isn't the overall best. + pub best_ordered_plan: Option, +} + +/// Compute the best way to join a given set of tables. +/// Returns the best [JoinN] if one exists, otherwise returns None. +pub fn compute_best_join_order<'a>( + table_references: &[TableReference], + maybe_order_target: Option<&OrderTarget>, + constraints: &'a [TableConstraints], + access_methods_arena: &'a RefCell>>, +) -> Result> { + // Skip work if we have no tables to consider. + if table_references.is_empty() { + return Ok(None); + } + + let num_tables = table_references.len(); + + // Compute naive left-to-right plan to use as pruning threshold + let naive_plan = compute_naive_left_deep_plan( + table_references, + maybe_order_target, + access_methods_arena, + &constraints, + )?; + + // Keep track of both 1. the best plan overall (not considering sorting), and 2. the best ordered plan (which might not be the same). + // We assign Some Cost (tm) to any required sort operation, so the best ordered plan may end up being + // the one we choose, if the cost reduction from avoiding sorting brings it below the cost of the overall best one. + let mut best_ordered_plan: Option = None; + let mut best_plan_is_also_ordered = if let Some(ref order_target) = maybe_order_target { + plan_satisfies_order_target( + &naive_plan, + &access_methods_arena, + table_references, + order_target, + ) + } else { + false + }; + + // If we have one table, then the "naive left-to-right plan" is always the best. + if table_references.len() == 1 { + return Ok(Some(BestJoinOrderResult { + best_plan: naive_plan, + best_ordered_plan: None, + })); + } + let mut best_plan = naive_plan; + + // Reuse a single mutable join order to avoid allocating join orders per permutation. + let mut join_order = Vec::with_capacity(num_tables); + join_order.push(JoinOrderMember { + table_no: 0, + is_outer: false, + }); + + // Keep track of the current best cost so we can short-circuit planning for subplans + // that already exceed the cost of the current best plan. + let cost_upper_bound = best_plan.cost; + let cost_upper_bound_ordered = best_plan.cost; + + // Keep track of the best plan for a given subset of tables. + // Consider this example: we have tables a,b,c,d to join. + // if we find that 'b JOIN a' is better than 'a JOIN b', then we don't need to even try + // to do 'a JOIN b JOIN c', because we know 'b JOIN a JOIN c' is going to be better. + // This is due to the commutativity and associativity of inner joins. + let mut best_plan_memo: HashMap = + HashMap::with_capacity(2usize.pow(num_tables as u32 - 1)); + + // Dynamic programming base case: calculate the best way to access each single table, as if + // there were no other tables. + for i in 0..num_tables { + let mut mask = TableMask::new(); + mask.add_table(i); + let table_ref = &table_references[i]; + join_order[0] = JoinOrderMember { + table_no: i, + is_outer: false, + }; + assert!(join_order.len() == 1); + let rel = join_lhs_and_rhs( + None, + table_ref, + &constraints[i], + &join_order, + maybe_order_target, + access_methods_arena, + cost_upper_bound_ordered, + )?; + if let Some(rel) = rel { + best_plan_memo.insert(mask, rel); + } + } + join_order.clear(); + + // As mentioned, inner joins are commutative. Outer joins are NOT. + // Example: + // "a LEFT JOIN b" can NOT be reordered as "b LEFT JOIN a". + // If there are outer joins in the plan, ensure correct ordering. + let left_join_illegal_map = { + let left_join_count = table_references + .iter() + .filter(|t| t.join_info.as_ref().map_or(false, |j| j.outer)) + .count(); + if left_join_count == 0 { + None + } else { + // map from rhs table index to lhs table index + let mut left_join_illegal_map: HashMap = + HashMap::with_capacity(left_join_count); + for (i, _) in table_references.iter().enumerate() { + for j in i + 1..table_references.len() { + if table_references[j] + .join_info + .as_ref() + .map_or(false, |j| j.outer) + { + // bitwise OR the masks + if let Some(illegal_lhs) = left_join_illegal_map.get_mut(&i) { + illegal_lhs.add_table(j); + } else { + let mut mask = TableMask::new(); + mask.add_table(j); + left_join_illegal_map.insert(i, mask); + } + } + } + } + Some(left_join_illegal_map) + } + }; + + // Now that we have our single-table base cases, we can start considering join subsets of 2 tables and more. + // Try to join each single table to each other table. + for subset_size in 2..=num_tables { + for mask in generate_join_bitmasks(num_tables, subset_size) { + // Keep track of the best way to join this subset of tables. + // Take the (a,b,c,d) example from above: + // E.g. for "a JOIN b JOIN c", the possibilities are (a,b,c), (a,c,b), (b,a,c) and so on. + // If we find out (b,a,c) is the best way to join these three, then we ONLY need to compute + // the cost of (b,a,c,d) in the final step, because (a,b,c,d) (and all others) are guaranteed to be worse. + let mut best_for_mask: Option = None; + // also keep track of the best plan for this subset that orders the rows in an Interesting Way (tm), + // i.e. allows us to eliminate sort operations downstream. + let (mut best_ordered_for_mask, mut best_for_mask_is_also_ordered) = (None, false); + + // Try to join all subsets (masks) with all other tables. + // In this block, LHS is always (n-1) tables, and RHS is a single table. + for rhs_idx in 0..num_tables { + // If the RHS table isn't a member of this join subset, skip. + if !mask.contains_table(rhs_idx) { + continue; + } + + // If there are no other tables except RHS, skip. + let lhs_mask = mask.without_table(rhs_idx); + if lhs_mask.is_empty() { + continue; + } + + // If this join ordering would violate LEFT JOIN ordering restrictions, skip. + if let Some(illegal_lhs) = left_join_illegal_map + .as_ref() + .and_then(|deps| deps.get(&rhs_idx)) + { + let legal = !lhs_mask.intersects(illegal_lhs); + if !legal { + continue; // Don't allow RHS before its LEFT in LEFT JOIN + } + } + + // If the already cached plan for this subset was too crappy to consider, + // then joining it with RHS won't help. Skip. + let Some(lhs) = best_plan_memo.get(&lhs_mask) else { + continue; + }; + + // Build a JoinOrder out of the table bitmask we are now considering. + for table_no in lhs.table_numbers() { + join_order.push(JoinOrderMember { + table_no, + is_outer: table_references[table_no] + .join_info + .as_ref() + .map_or(false, |j| j.outer), + }); + } + join_order.push(JoinOrderMember { + table_no: rhs_idx, + is_outer: table_references[rhs_idx] + .join_info + .as_ref() + .map_or(false, |j| j.outer), + }); + assert!(join_order.len() == subset_size); + + // Calculate the best way to join LHS with RHS. + let rel = join_lhs_and_rhs( + Some(lhs), + &table_references[rhs_idx], + &constraints[rhs_idx], + &join_order, + maybe_order_target, + access_methods_arena, + cost_upper_bound_ordered, + )?; + join_order.clear(); + + let Some(rel) = rel else { + continue; + }; + + let satisfies_order_target = if let Some(ref order_target) = maybe_order_target { + plan_satisfies_order_target( + &rel, + &access_methods_arena, + table_references, + order_target, + ) + } else { + false + }; + + // If this plan is worse than our overall best, it might still be the best ordered plan. + if rel.cost >= cost_upper_bound { + // But if it isn't, skip. + if !satisfies_order_target { + continue; + } + let existing_ordered_cost: Cost = best_ordered_for_mask + .as_ref() + .map_or(Cost(f64::MAX), |p: &JoinN| p.cost); + if rel.cost < existing_ordered_cost { + best_ordered_for_mask = Some(rel); + } + } else if best_for_mask.is_none() || rel.cost < best_for_mask.as_ref().unwrap().cost + { + best_for_mask = Some(rel); + best_for_mask_is_also_ordered = satisfies_order_target; + } + } + + if let Some(rel) = best_ordered_for_mask.take() { + let cost = rel.cost; + let has_all_tables = mask.table_count() == num_tables; + if has_all_tables && cost_upper_bound_ordered > cost { + best_ordered_plan = Some(rel); + } + } + + if let Some(rel) = best_for_mask.take() { + let cost = rel.cost; + let has_all_tables = mask.table_count() == num_tables; + if has_all_tables { + if cost_upper_bound > cost { + best_plan = rel; + best_plan_is_also_ordered = best_for_mask_is_also_ordered; + } + } else { + best_plan_memo.insert(mask, rel); + } + } + } + } + + Ok(Some(BestJoinOrderResult { + best_plan, + best_ordered_plan: if best_plan_is_also_ordered { + None + } else { + best_ordered_plan + }, + })) +} + +/// Specialized version of [compute_best_join_order] that just joins tables in the order they are given +/// in the SQL query. This is used as an upper bound for any other plans -- we can give up enumerating +/// permutations if they exceed this cost during enumeration. +pub fn compute_naive_left_deep_plan<'a>( + table_references: &[TableReference], + maybe_order_target: Option<&OrderTarget>, + access_methods_arena: &'a RefCell>>, + constraints: &'a [TableConstraints], +) -> Result { + let n = table_references.len(); + assert!(n > 0); + + let join_order = table_references + .iter() + .enumerate() + .map(|(i, t)| JoinOrderMember { + table_no: i, + is_outer: t.join_info.as_ref().map_or(false, |j| j.outer), + }) + .collect::>(); + + // Start with first table + let mut best_plan = join_lhs_and_rhs( + None, + &table_references[0], + &constraints[0], + &join_order[..1], + maybe_order_target, + access_methods_arena, + Cost(f64::MAX), + )? + .expect("call to join_lhs_and_rhs in compute_naive_left_deep_plan always returns Some(JoinN)"); + + // Add remaining tables one at a time from left to right + for i in 1..n { + best_plan = join_lhs_and_rhs( + Some(&best_plan), + &table_references[i], + &constraints[i], + &join_order[..=i], + maybe_order_target, + access_methods_arena, + Cost(f64::MAX), + )? + .expect( + "call to join_lhs_and_rhs in compute_naive_left_deep_plan always returns Some(JoinN)", + ); + } + + Ok(best_plan) +} + +/// Iterator that generates all possible size k bitmasks for a given number of tables. +/// For example, given: 3 tables and k=2, the bitmasks are: +/// - 0b011 (tables 0, 1) +/// - 0b101 (tables 0, 2) +/// - 0b110 (tables 1, 2) +/// +/// This is used in the dynamic programming approach to finding the best way to join a subset of N tables. +struct JoinBitmaskIter { + current: u128, + max_exclusive: u128, +} + +impl JoinBitmaskIter { + fn new(table_number_max_exclusive: usize, how_many: usize) -> Self { + Self { + current: (1 << how_many) - 1, // Start with smallest k-bit number (e.g., 000111 for k=3) + max_exclusive: 1 << table_number_max_exclusive, + } + } +} + +impl Iterator for JoinBitmaskIter { + type Item = TableMask; + + fn next(&mut self) -> Option { + if self.current >= self.max_exclusive { + return None; + } + + let result = TableMask::from_bits(self.current); + + // Gosper's hack: compute next k-bit combination in lexicographic order + let c = self.current & (!self.current + 1); // rightmost set bit + let r = self.current + c; // add it to get a carry + let ones = self.current ^ r; // changed bits + let ones = (ones >> 2) / c; // right-adjust shifted bits + self.current = r | ones; // form the next combination + + Some(result) + } +} + +/// Generate all possible bitmasks of size `how_many` for a given number of tables. +fn generate_join_bitmasks(table_number_max_exclusive: usize, how_many: usize) -> JoinBitmaskIter { + JoinBitmaskIter::new(table_number_max_exclusive, how_many) +} + +#[cfg(test)] +mod tests { + use std::{rc::Rc, sync::Arc}; + + use limbo_sqlite3_parser::ast::{self, Expr, Operator, SortOrder}; + + use super::*; + use crate::{ + schema::{BTreeTable, Column, Index, IndexColumn, Table, Type}, + translate::{ + optimizer::constraints::{constraints_from_where_clause, BinaryExprSide}, + plan::{ColumnUsedMask, IterationDirection, JoinInfo, Operation, WhereTerm}, + planner::TableMask, + }, + }; + + #[test] + fn test_generate_bitmasks() { + let bitmasks = generate_join_bitmasks(4, 2).collect::>(); + assert!(bitmasks.contains(&TableMask(0b110))); // {0,1} -- first bit is always set to 0 so that a Mask with value 0 means "no tables are referenced". + assert!(bitmasks.contains(&TableMask(0b1010))); // {0,2} + assert!(bitmasks.contains(&TableMask(0b1100))); // {1,2} + assert!(bitmasks.contains(&TableMask(0b10010))); // {0,3} + assert!(bitmasks.contains(&TableMask(0b10100))); // {1,3} + assert!(bitmasks.contains(&TableMask(0b11000))); // {2,3} + } + + #[test] + /// Test that [compute_best_join_order] returns None when there are no table references. + fn test_compute_best_join_order_empty() { + let table_references = vec![]; + let available_indexes = HashMap::new(); + let where_clause = vec![]; + + let access_methods_arena = RefCell::new(Vec::new()); + let table_constraints = + constraints_from_where_clause(&where_clause, &table_references, &available_indexes) + .unwrap(); + + let result = compute_best_join_order( + &table_references, + None, + &table_constraints, + &access_methods_arena, + ) + .unwrap(); + assert!(result.is_none()); + } + + #[test] + /// Test that [compute_best_join_order] returns a table scan access method when the where clause is empty. + fn test_compute_best_join_order_single_table_no_indexes() { + let t1 = _create_btree_table("test_table", _create_column_list(&["id"], Type::Integer)); + let table_references = vec![_create_table_reference(t1.clone(), None)]; + let available_indexes = HashMap::new(); + let where_clause = vec![]; + + let access_methods_arena = RefCell::new(Vec::new()); + let table_constraints = + constraints_from_where_clause(&where_clause, &table_references, &available_indexes) + .unwrap(); + + // SELECT * from test_table + // expecting best_best_plan() not to do any work due to empty where clause. + let BestJoinOrderResult { best_plan, .. } = compute_best_join_order( + &table_references, + None, + &table_constraints, + &access_methods_arena, + ) + .unwrap() + .unwrap(); + // Should just be a table scan access method + let access_method = &access_methods_arena.borrow()[best_plan.data[0].1]; + assert!(access_method.is_scan()); + assert!(access_method.iter_dir == IterationDirection::Forwards); + } + + #[test] + /// Test that [compute_best_join_order] returns a RowidEq access method when the where clause has an EQ constraint on the rowid alias. + fn test_compute_best_join_order_single_table_rowid_eq() { + let t1 = _create_btree_table("test_table", vec![_create_column_rowid_alias("id")]); + let table_references = vec![_create_table_reference(t1.clone(), None)]; + + let where_clause = vec![_create_binary_expr( + _create_column_expr(0, 0, true), // table 0, column 0 (rowid) + ast::Operator::Equals, + _create_numeric_literal("42"), + )]; + + let access_methods_arena = RefCell::new(Vec::new()); + let available_indexes = HashMap::new(); + let table_constraints = + constraints_from_where_clause(&where_clause, &table_references, &available_indexes) + .unwrap(); + + // SELECT * FROM test_table WHERE id = 42 + // expecting a RowidEq access method because id is a rowid alias. + let result = compute_best_join_order( + &table_references, + None, + &table_constraints, + &access_methods_arena, + ) + .unwrap(); + assert!(result.is_some()); + let BestJoinOrderResult { best_plan, .. } = result.unwrap(); + assert_eq!(best_plan.table_numbers().collect::>(), vec![0]); + let access_method = &access_methods_arena.borrow()[best_plan.data[0].1]; + assert!(!access_method.is_scan()); + assert!(access_method.iter_dir == IterationDirection::Forwards); + assert!(access_method.constraint_refs.len() == 1); + assert!( + table_constraints[0].constraints[access_method.constraint_refs[0].constraint_vec_pos] + .where_clause_pos + == (0, BinaryExprSide::Rhs) + ); + } + + #[test] + /// Test that [compute_best_join_order] returns an IndexScan access method when the where clause has an EQ constraint on a primary key. + fn test_compute_best_join_order_single_table_pk_eq() { + let t1 = _create_btree_table( + "test_table", + vec![_create_column_of_type("id", Type::Integer)], + ); + let table_references = vec![_create_table_reference(t1.clone(), None)]; + + let where_clause = vec![_create_binary_expr( + _create_column_expr(0, 0, false), // table 0, column 0 (id) + ast::Operator::Equals, + _create_numeric_literal("42"), + )]; + + let access_methods_arena = RefCell::new(Vec::new()); + let mut available_indexes = HashMap::new(); + let index = Arc::new(Index { + name: "sqlite_autoindex_test_table_1".to_string(), + table_name: "test_table".to_string(), + columns: vec![IndexColumn { + name: "id".to_string(), + order: SortOrder::Asc, + pos_in_table: 0, + }], + unique: true, + ephemeral: false, + root_page: 1, + }); + available_indexes.insert("test_table".to_string(), vec![index]); + + let table_constraints = + constraints_from_where_clause(&where_clause, &table_references, &available_indexes) + .unwrap(); + // SELECT * FROM test_table WHERE id = 42 + // expecting an IndexScan access method because id is a primary key with an index + let result = compute_best_join_order( + &table_references, + None, + &table_constraints, + &access_methods_arena, + ) + .unwrap(); + assert!(result.is_some()); + let BestJoinOrderResult { best_plan, .. } = result.unwrap(); + assert_eq!(best_plan.table_numbers().collect::>(), vec![0]); + let access_method = &access_methods_arena.borrow()[best_plan.data[0].1]; + assert!(!access_method.is_scan()); + assert!(access_method.iter_dir == IterationDirection::Forwards); + assert!(access_method.index.as_ref().unwrap().name == "sqlite_autoindex_test_table_1"); + assert!(access_method.constraint_refs.len() == 1); + assert!( + table_constraints[0].constraints[access_method.constraint_refs[0].constraint_vec_pos] + .where_clause_pos + == (0, BinaryExprSide::Rhs) + ); + } + + #[test] + /// Test that [compute_best_join_order] moves the outer table to the inner position when an index can be used on it, but not the original inner table. + fn test_compute_best_join_order_two_tables() { + let t1 = _create_btree_table("table1", _create_column_list(&["id"], Type::Integer)); + let t2 = _create_btree_table("table2", _create_column_list(&["id"], Type::Integer)); + + let mut table_references = vec![ + _create_table_reference(t1.clone(), None), + _create_table_reference( + t2.clone(), + Some(JoinInfo { + outer: false, + using: None, + }), + ), + ]; + + const TABLE1: usize = 0; + const TABLE2: usize = 1; + + let mut available_indexes = HashMap::new(); + // Index on the outer table (table1) + let index1 = Arc::new(Index { + name: "index1".to_string(), + table_name: "table1".to_string(), + columns: vec![IndexColumn { + name: "id".to_string(), + order: SortOrder::Asc, + pos_in_table: 0, + }], + unique: true, + ephemeral: false, + root_page: 1, + }); + available_indexes.insert("table1".to_string(), vec![index1]); + + // SELECT * FROM table1 JOIN table2 WHERE table1.id = table2.id + // expecting table2 to be chosen first due to the index on table1.id + let where_clause = vec![_create_binary_expr( + _create_column_expr(TABLE1, 0, false), // table1.id + ast::Operator::Equals, + _create_column_expr(TABLE2, 0, false), // table2.id + )]; + + let access_methods_arena = RefCell::new(Vec::new()); + let table_constraints = + constraints_from_where_clause(&where_clause, &table_references, &available_indexes) + .unwrap(); + + let result = compute_best_join_order( + &mut table_references, + None, + &table_constraints, + &access_methods_arena, + ) + .unwrap(); + assert!(result.is_some()); + let BestJoinOrderResult { best_plan, .. } = result.unwrap(); + assert_eq!(best_plan.table_numbers().collect::>(), vec![1, 0]); + let access_method = &access_methods_arena.borrow()[best_plan.data[0].1]; + assert!(access_method.is_scan()); + assert!(access_method.iter_dir == IterationDirection::Forwards); + let access_method = &access_methods_arena.borrow()[best_plan.data[1].1]; + assert!(!access_method.is_scan()); + assert!(access_method.iter_dir == IterationDirection::Forwards); + assert!(access_method.index.as_ref().unwrap().name == "index1"); + assert!(access_method.constraint_refs.len() == 1); + assert!( + table_constraints[TABLE1].constraints + [access_method.constraint_refs[0].constraint_vec_pos] + .where_clause_pos + == (0, BinaryExprSide::Rhs) + ); + } + + #[test] + /// Test that [compute_best_join_order] returns a sensible order and plan for three tables, each with indexes. + fn test_compute_best_join_order_three_tables_indexed() { + let table_orders = _create_btree_table( + "orders", + vec![ + _create_column_of_type("id", Type::Integer), + _create_column_of_type("customer_id", Type::Integer), + _create_column_of_type("total", Type::Integer), + ], + ); + let table_customers = _create_btree_table( + "customers", + vec![ + _create_column_of_type("id", Type::Integer), + _create_column_of_type("name", Type::Integer), + ], + ); + let table_order_items = _create_btree_table( + "order_items", + vec![ + _create_column_of_type("id", Type::Integer), + _create_column_of_type("order_id", Type::Integer), + _create_column_of_type("product_id", Type::Integer), + _create_column_of_type("quantity", Type::Integer), + ], + ); + + let table_references = vec![ + _create_table_reference(table_orders.clone(), None), + _create_table_reference( + table_customers.clone(), + Some(JoinInfo { + outer: false, + using: None, + }), + ), + _create_table_reference( + table_order_items.clone(), + Some(JoinInfo { + outer: false, + using: None, + }), + ), + ]; + + const TABLE_NO_ORDERS: usize = 0; + const TABLE_NO_CUSTOMERS: usize = 1; + const TABLE_NO_ORDER_ITEMS: usize = 2; + + let mut available_indexes = HashMap::new(); + ["orders", "customers", "order_items"] + .iter() + .for_each(|table_name| { + // add primary key index called sqlite_autoindex__1 + let index_name = format!("sqlite_autoindex_{}_1", table_name); + let index = Arc::new(Index { + name: index_name, + table_name: table_name.to_string(), + columns: vec![IndexColumn { + name: "id".to_string(), + order: SortOrder::Asc, + pos_in_table: 0, + }], + unique: true, + ephemeral: false, + root_page: 1, + }); + available_indexes.insert(table_name.to_string(), vec![index]); + }); + let customer_id_idx = Arc::new(Index { + name: "orders_customer_id_idx".to_string(), + table_name: "orders".to_string(), + columns: vec![IndexColumn { + name: "customer_id".to_string(), + order: SortOrder::Asc, + pos_in_table: 1, + }], + unique: false, + ephemeral: false, + root_page: 1, + }); + let order_id_idx = Arc::new(Index { + name: "order_items_order_id_idx".to_string(), + table_name: "order_items".to_string(), + columns: vec![IndexColumn { + name: "order_id".to_string(), + order: SortOrder::Asc, + pos_in_table: 1, + }], + unique: false, + ephemeral: false, + root_page: 1, + }); + + available_indexes + .entry("orders".to_string()) + .and_modify(|v| v.push(customer_id_idx)); + available_indexes + .entry("order_items".to_string()) + .and_modify(|v| v.push(order_id_idx)); + + // SELECT * FROM orders JOIN customers JOIN order_items + // WHERE orders.customer_id = customers.id AND orders.id = order_items.order_id AND customers.id = 42 + // expecting customers to be chosen first due to the index on customers.id and it having a selective filter (=42) + // then orders to be chosen next due to the index on orders.customer_id + // then order_items to be chosen last due to the index on order_items.order_id + let where_clause = vec![ + // orders.customer_id = customers.id + _create_binary_expr( + _create_column_expr(TABLE_NO_ORDERS, 1, false), // orders.customer_id + ast::Operator::Equals, + _create_column_expr(TABLE_NO_CUSTOMERS, 0, false), // customers.id + ), + // orders.id = order_items.order_id + _create_binary_expr( + _create_column_expr(TABLE_NO_ORDERS, 0, false), // orders.id + ast::Operator::Equals, + _create_column_expr(TABLE_NO_ORDER_ITEMS, 1, false), // order_items.order_id + ), + // customers.id = 42 + _create_binary_expr( + _create_column_expr(TABLE_NO_CUSTOMERS, 0, false), // customers.id + ast::Operator::Equals, + _create_numeric_literal("42"), + ), + ]; + + let access_methods_arena = RefCell::new(Vec::new()); + let table_constraints = + constraints_from_where_clause(&where_clause, &table_references, &available_indexes) + .unwrap(); + + let result = compute_best_join_order( + &table_references, + None, + &table_constraints, + &access_methods_arena, + ) + .unwrap(); + assert!(result.is_some()); + let BestJoinOrderResult { best_plan, .. } = result.unwrap(); + + // Customers (due to =42 filter) -> Orders (due to index on customer_id) -> Order_items (due to index on order_id) + assert_eq!( + best_plan.table_numbers().collect::>(), + vec![TABLE_NO_CUSTOMERS, TABLE_NO_ORDERS, TABLE_NO_ORDER_ITEMS] + ); + + let access_method = &access_methods_arena.borrow()[best_plan.data[0].1]; + assert!(!access_method.is_scan()); + assert!(access_method.iter_dir == IterationDirection::Forwards); + assert!(access_method.index.as_ref().unwrap().name == "sqlite_autoindex_customers_1"); + assert!(access_method.constraint_refs.len() == 1); + let constraint = &table_constraints[TABLE_NO_CUSTOMERS].constraints + [access_method.constraint_refs[0].constraint_vec_pos]; + assert!(constraint.lhs_mask.is_empty()); + + let access_method = &access_methods_arena.borrow()[best_plan.data[1].1]; + assert!(!access_method.is_scan()); + assert!(access_method.iter_dir == IterationDirection::Forwards); + assert!(access_method.index.as_ref().unwrap().name == "orders_customer_id_idx"); + assert!(access_method.constraint_refs.len() == 1); + let constraint = &table_constraints[TABLE_NO_ORDERS].constraints + [access_method.constraint_refs[0].constraint_vec_pos]; + assert!(constraint.lhs_mask.contains_table(TABLE_NO_CUSTOMERS)); + + let access_method = &access_methods_arena.borrow()[best_plan.data[2].1]; + assert!(!access_method.is_scan()); + assert!(access_method.iter_dir == IterationDirection::Forwards); + assert!(access_method.index.as_ref().unwrap().name == "order_items_order_id_idx"); + assert!(access_method.constraint_refs.len() == 1); + let constraint = &table_constraints[TABLE_NO_ORDER_ITEMS].constraints + [access_method.constraint_refs[0].constraint_vec_pos]; + assert!(constraint.lhs_mask.contains_table(TABLE_NO_ORDERS)); + } + + struct TestColumn { + name: String, + ty: Type, + is_rowid_alias: bool, + } + + impl Default for TestColumn { + fn default() -> Self { + Self { + name: "a".to_string(), + ty: Type::Integer, + is_rowid_alias: false, + } + } + } + + #[test] + fn test_join_order_three_tables_no_indexes() { + let t1 = _create_btree_table("t1", _create_column_list(&["id", "foo"], Type::Integer)); + let t2 = _create_btree_table("t2", _create_column_list(&["id", "foo"], Type::Integer)); + let t3 = _create_btree_table("t3", _create_column_list(&["id", "foo"], Type::Integer)); + + let mut table_references = vec![ + _create_table_reference(t1.clone(), None), + _create_table_reference( + t2.clone(), + Some(JoinInfo { + outer: false, + using: None, + }), + ), + _create_table_reference( + t3.clone(), + Some(JoinInfo { + outer: false, + using: None, + }), + ), + ]; + + let where_clause = vec![ + // t2.foo = 42 (equality filter, more selective) + _create_binary_expr( + _create_column_expr(1, 1, false), // table 1, column 1 (foo) + ast::Operator::Equals, + _create_numeric_literal("42"), + ), + // t1.foo > 10 (inequality filter, less selective) + _create_binary_expr( + _create_column_expr(0, 1, false), // table 0, column 1 (foo) + ast::Operator::Greater, + _create_numeric_literal("10"), + ), + ]; + + let available_indexes = HashMap::new(); + let access_methods_arena = RefCell::new(Vec::new()); + let table_constraints = + constraints_from_where_clause(&where_clause, &table_references, &available_indexes) + .unwrap(); + + let BestJoinOrderResult { best_plan, .. } = compute_best_join_order( + &mut table_references, + None, + &table_constraints, + &access_methods_arena, + ) + .unwrap() + .unwrap(); + + // Verify that t2 is chosen first due to its equality filter + assert_eq!(best_plan.table_numbers().nth(0).unwrap(), 1); + // Verify table scan is used since there are no indexes + let access_method = &access_methods_arena.borrow()[best_plan.data[0].1]; + assert!(access_method.is_scan()); + assert!(access_method.iter_dir == IterationDirection::Forwards); + assert!(access_method.index.is_none()); + // Verify that t1 is chosen next due to its inequality filter + let access_method = &access_methods_arena.borrow()[best_plan.data[1].1]; + assert!(access_method.is_scan()); + assert!(access_method.iter_dir == IterationDirection::Forwards); + assert!(access_method.index.is_none()); + // Verify that t3 is chosen last due to no filters + let access_method = &access_methods_arena.borrow()[best_plan.data[2].1]; + assert!(access_method.is_scan()); + assert!(access_method.iter_dir == IterationDirection::Forwards); + assert!(access_method.index.is_none()); + } + + #[test] + /// Test that [compute_best_join_order] chooses a "fact table" as the outer table, + /// when it has a foreign key to all dimension tables. + fn test_compute_best_join_order_star_schema() { + const NUM_DIM_TABLES: usize = 9; + const FACT_TABLE_IDX: usize = 9; + + // Create fact table with foreign keys to all dimension tables + let mut fact_columns = vec![_create_column_rowid_alias("id")]; + for i in 0..NUM_DIM_TABLES { + fact_columns.push(_create_column_of_type( + &format!("dim{}_id", i), + Type::Integer, + )); + } + let fact_table = _create_btree_table("fact", fact_columns); + + // Create dimension tables, each with an id and value column + let dim_tables: Vec<_> = (0..NUM_DIM_TABLES) + .map(|i| { + _create_btree_table( + &format!("dim{}", i), + vec![ + _create_column_rowid_alias("id"), + _create_column_of_type("value", Type::Integer), + ], + ) + }) + .collect(); + + let mut where_clause = vec![]; + + // Add join conditions between fact and each dimension table + for i in 0..NUM_DIM_TABLES { + where_clause.push(_create_binary_expr( + _create_column_expr(FACT_TABLE_IDX, i + 1, false), // fact.dimX_id + ast::Operator::Equals, + _create_column_expr(i, 0, true), // dimX.id + )); + } + + let table_references = { + let mut refs = vec![_create_table_reference(dim_tables[0].clone(), None)]; + refs.extend(dim_tables.iter().skip(1).map(|t| { + _create_table_reference( + t.clone(), + Some(JoinInfo { + outer: false, + using: None, + }), + ) + })); + refs.push(_create_table_reference( + fact_table.clone(), + Some(JoinInfo { + outer: false, + using: None, + }), + )); + refs + }; + + let access_methods_arena = RefCell::new(Vec::new()); + let available_indexes = HashMap::new(); + let table_constraints = + constraints_from_where_clause(&where_clause, &table_references, &available_indexes) + .unwrap(); + + let result = compute_best_join_order( + &table_references, + None, + &table_constraints, + &access_methods_arena, + ) + .unwrap(); + assert!(result.is_some()); + let BestJoinOrderResult { best_plan, .. } = result.unwrap(); + + // Expected optimal order: fact table as outer, with rowid seeks in any order on each dimension table + // Verify fact table is selected as the outer table as all the other tables can use SeekRowid + assert_eq!( + best_plan.table_numbers().nth(0).unwrap(), + FACT_TABLE_IDX, + "First table should be fact (table {}) due to available index, got table {} instead", + FACT_TABLE_IDX, + best_plan.table_numbers().nth(0).unwrap() + ); + + // Verify access methods + let access_method = &access_methods_arena.borrow()[best_plan.data[0].1]; + assert!(access_method.is_scan()); + assert!(access_method.iter_dir == IterationDirection::Forwards); + assert!(access_method.index.is_none()); + assert!(access_method.constraint_refs.is_empty()); + + for (table_number, access_method_index) in best_plan.data.iter().skip(1) { + let access_method = &access_methods_arena.borrow()[*access_method_index]; + assert!(!access_method.is_scan()); + assert!(access_method.iter_dir == IterationDirection::Forwards); + assert!(access_method.index.is_none()); + assert!(access_method.constraint_refs.len() == 1); + let constraint = &table_constraints[*table_number].constraints + [access_method.constraint_refs[0].constraint_vec_pos]; + assert!(constraint.lhs_mask.contains_table(FACT_TABLE_IDX)); + assert!(constraint.operator == ast::Operator::Equals); + } + } + + #[test] + /// Test that [compute_best_join_order] figures out that the tables form a "linked list" pattern + /// where a column in each table points to an indexed column in the next table, + /// and chooses the best order based on that. + fn test_compute_best_join_order_linked_list() { + const NUM_TABLES: usize = 5; + + // Create tables t1 -> t2 -> t3 -> t4 -> t5 where there is a foreign key from each table to the next + let mut tables = Vec::with_capacity(NUM_TABLES); + for i in 0..NUM_TABLES { + let mut columns = vec![_create_column_rowid_alias("id")]; + if i < NUM_TABLES - 1 { + columns.push(_create_column_of_type(&format!("next_id"), Type::Integer)); + } + tables.push(_create_btree_table(&format!("t{}", i + 1), columns)); + } + + let available_indexes = HashMap::new(); + + // Create table references + let table_references: Vec<_> = tables + .iter() + .map(|t| _create_table_reference(t.clone(), None)) + .collect(); + + // Create where clause linking each table to the next + let mut where_clause = Vec::new(); + for i in 0..NUM_TABLES - 1 { + where_clause.push(_create_binary_expr( + _create_column_expr(i, 1, false), // ti.next_id + ast::Operator::Equals, + _create_column_expr(i + 1, 0, true), // t(i+1).id + )); + } + + let access_methods_arena = RefCell::new(Vec::new()); + let table_constraints = + constraints_from_where_clause(&where_clause, &table_references, &available_indexes) + .unwrap(); + + // Run the optimizer + let BestJoinOrderResult { best_plan, .. } = compute_best_join_order( + &table_references, + None, + &table_constraints, + &access_methods_arena, + ) + .unwrap() + .unwrap(); + + // Verify the join order is exactly t1 -> t2 -> t3 -> t4 -> t5 + for i in 0..NUM_TABLES { + assert_eq!( + best_plan.table_numbers().nth(i).unwrap(), + i, + "Expected table {} at position {}, got table {} instead", + i, + i, + best_plan.table_numbers().nth(i).unwrap() + ); + } + + // Verify access methods: + // - First table should use Table scan + let access_method = &access_methods_arena.borrow()[best_plan.data[0].1]; + assert!(access_method.is_scan()); + assert!(access_method.iter_dir == IterationDirection::Forwards); + assert!(access_method.index.is_none()); + assert!(access_method.constraint_refs.is_empty()); + + // all of the rest should use rowid equality + for i in 1..NUM_TABLES { + let access_method = &access_methods_arena.borrow()[best_plan.data[i].1]; + assert!(!access_method.is_scan()); + assert!(access_method.iter_dir == IterationDirection::Forwards); + assert!(access_method.index.is_none()); + assert!(access_method.constraint_refs.len() == 1); + let constraint = &table_constraints[i].constraints + [access_method.constraint_refs[0].constraint_vec_pos]; + assert!(constraint.lhs_mask.contains_table(i - 1)); + assert!(constraint.operator == ast::Operator::Equals); + } + } + + fn _create_column(c: &TestColumn) -> Column { + Column { + name: Some(c.name.clone()), + ty: c.ty, + ty_str: c.ty.to_string(), + is_rowid_alias: c.is_rowid_alias, + primary_key: false, + notnull: false, + default: None, + } + } + fn _create_column_of_type(name: &str, ty: Type) -> Column { + _create_column(&TestColumn { + name: name.to_string(), + ty, + is_rowid_alias: false, + }) + } + + fn _create_column_list(names: &[&str], ty: Type) -> Vec { + names + .iter() + .map(|name| _create_column_of_type(name, ty)) + .collect() + } + + fn _create_column_rowid_alias(name: &str) -> Column { + _create_column(&TestColumn { + name: name.to_string(), + ty: Type::Integer, + is_rowid_alias: true, + }) + } + + /// Creates a BTreeTable with the given name and columns + fn _create_btree_table(name: &str, columns: Vec) -> Rc { + Rc::new(BTreeTable { + root_page: 1, // Page number doesn't matter for tests + name: name.to_string(), + primary_key_columns: vec![], + columns, + has_rowid: true, + is_strict: false, + }) + } + + /// Creates a TableReference for a BTreeTable + fn _create_table_reference( + table: Rc, + join_info: Option, + ) -> TableReference { + let name = table.name.clone(); + TableReference { + table: Table::BTree(table), + op: Operation::Scan { + iter_dir: IterationDirection::Forwards, + index: None, + }, + identifier: name, + join_info, + col_used_mask: ColumnUsedMask::new(), + } + } + + /// Creates a column expression + fn _create_column_expr(table: usize, column: usize, is_rowid_alias: bool) -> Expr { + Expr::Column { + database: None, + table, + column, + is_rowid_alias, + } + } + + /// Creates a binary expression for a WHERE clause + fn _create_binary_expr(lhs: Expr, op: Operator, rhs: Expr) -> WhereTerm { + WhereTerm { + expr: Expr::Binary(Box::new(lhs), op, Box::new(rhs)), + from_outer_join: None, + } + } + + /// Creates a numeric literal expression + fn _create_numeric_literal(value: &str) -> Expr { + Expr::Literal(ast::Literal::Numeric(value.to_string())) + } +} diff --git a/core/translate/optimizer.rs b/core/translate/optimizer/mod.rs similarity index 56% rename from core/translate/optimizer.rs rename to core/translate/optimizer/mod.rs index 38cdb6343..88ec01604 100644 --- a/core/translate/optimizer.rs +++ b/core/translate/optimizer/mod.rs @@ -1,25 +1,35 @@ -use std::{cmp::Ordering, collections::HashMap, sync::Arc}; +use std::{cell::RefCell, cmp::Ordering, collections::HashMap, sync::Arc}; +use constraints::{ + constraints_from_where_clause, usable_constraints_for_join_order, Constraint, ConstraintRef, +}; +use cost::Cost; +use join::{compute_best_join_order, BestJoinOrderResult}; use limbo_sqlite3_parser::ast::{self, Expr, SortOrder}; +use order::{compute_order_target, plan_satisfies_order_target, EliminatesSort}; use crate::{ parameters::PARAM_PREFIX, schema::{Index, IndexColumn, Schema}, translate::plan::TerminationKey, types::SeekOp, - util::exprs_are_equivalent, Result, }; use super::{ emitter::Resolver, plan::{ - DeletePlan, EvalAt, GroupBy, IterationDirection, JoinOrderMember, Operation, Plan, Search, - SeekDef, SeekKey, SelectPlan, TableReference, UpdatePlan, WhereTerm, + DeletePlan, GroupBy, IterationDirection, JoinOrderMember, Operation, Plan, Search, SeekDef, + SeekKey, SelectPlan, TableReference, UpdatePlan, WhereTerm, }, - planner::determine_where_to_eval_expr, }; +pub(crate) mod access_method; +pub(crate) mod constraints; +pub(crate) mod cost; +pub(crate) mod join; +pub(crate) mod order; + pub fn optimize_plan(plan: &mut Plan, schema: &Schema) -> Result<()> { match plan { Plan::Select(plan) => optimize_select_plan(plan, schema), @@ -43,15 +53,17 @@ fn optimize_select_plan(plan: &mut SelectPlan, schema: &Schema) -> Result<()> { return Ok(()); } - use_indexes( + let best_join_order = optimize_table_access( &mut plan.table_references, &schema.indexes, &mut plan.where_clause, &mut plan.order_by, - &plan.group_by, + &mut plan.group_by, )?; - eliminate_orderby_like_groupby(plan)?; + if let Some(best_join_order) = best_join_order { + plan.join_order = best_join_order; + } Ok(()) } @@ -65,12 +77,12 @@ fn optimize_delete_plan(plan: &mut DeletePlan, schema: &Schema) -> Result<()> { return Ok(()); } - use_indexes( + let _ = optimize_table_access( &mut plan.table_references, &schema.indexes, &mut plan.where_clause, &mut plan.order_by, - &None, + &mut None, )?; Ok(()) @@ -84,12 +96,12 @@ fn optimize_update_plan(plan: &mut UpdatePlan, schema: &Schema) -> Result<()> { plan.contains_constant_false_condition = true; return Ok(()); } - use_indexes( + let _ = optimize_table_access( &mut plan.table_references, &schema.indexes, &mut plan.where_clause, &mut plan.order_by, - &None, + &mut None, )?; Ok(()) } @@ -104,281 +116,222 @@ fn optimize_subqueries(plan: &mut SelectPlan, schema: &Schema) -> Result<()> { Ok(()) } -fn eliminate_orderby_like_groupby(plan: &mut SelectPlan) -> Result<()> { - if plan.order_by.is_none() | plan.group_by.is_none() { - return Ok(()); - } - if plan.table_references.len() == 0 { - return Ok(()); - } - - let order_by_clauses = plan.order_by.as_mut().unwrap(); - // TODO: let's make the group by sorter aware of the order by directions so we dont need to skip - // descending terms. - if order_by_clauses - .iter() - .any(|(_, dir)| matches!(dir, SortOrder::Desc)) - { - return Ok(()); - } - let group_by_clauses = plan.group_by.as_mut().unwrap(); - // all order by terms must be in the group by clause for order by to be eliminated - if !order_by_clauses.iter().all(|(o_expr, _)| { - group_by_clauses - .exprs - .iter() - .any(|g_expr| exprs_are_equivalent(g_expr, o_expr)) - }) { - return Ok(()); - } - - // reorder group by terms so that they match the order by terms - // this way the group by sorter will effectively do the order by sorter's job and - // we can remove the order by clause - group_by_clauses.exprs.sort_by_key(|g_expr| { - order_by_clauses - .iter() - .position(|(o_expr, _)| exprs_are_equivalent(o_expr, g_expr)) - .unwrap_or(usize::MAX) - }); - - plan.order_by = None; - Ok(()) -} - -/// Eliminate unnecessary ORDER BY clauses. -/// Returns true if the ORDER BY clause was eliminated. -fn eliminate_unnecessary_orderby( - table_references: &mut [TableReference], - available_indexes: &HashMap>>, - order_by: &mut Option>, - group_by: &Option, -) -> Result { - let Some(order) = order_by else { - return Ok(false); - }; - let Some(first_table_reference) = table_references.first_mut() else { - return Ok(false); - }; - let Some(btree_table) = first_table_reference.btree() else { - return Ok(false); - }; - // If GROUP BY clause is present, we can't rely on already ordered columns because GROUP BY reorders the data - // This early return prevents the elimination of ORDER BY when GROUP BY exists, as sorting must be applied after grouping - // And if ORDER BY clause duplicates GROUP BY we handle it later in fn eliminate_orderby_like_groupby - if group_by.is_some() { - return Ok(false); - } - let Operation::Scan { - index, iter_dir, .. - } = &mut first_table_reference.op - else { - return Ok(false); - }; - - assert!( - index.is_none(), - "Nothing shouldve transformed the scan to use an index yet" - ); - - // Special case: if ordering by just the rowid, we can remove the ORDER BY clause - if order.len() == 1 && order[0].0.is_rowid_alias_of(0) { - *iter_dir = match order[0].1 { - SortOrder::Asc => IterationDirection::Forwards, - SortOrder::Desc => IterationDirection::Backwards, - }; - *order_by = None; - return Ok(true); - } - - // Find the best matching index for the ORDER BY columns - let table_name = &btree_table.name; - let mut best_index = (None, 0); - - for (_, indexes) in available_indexes.iter() { - for index_candidate in indexes.iter().filter(|i| &i.table_name == table_name) { - let matching_columns = index_candidate.columns.iter().enumerate().take_while(|(i, c)| { - if let Some((Expr::Column { table, column, .. }, _)) = order.get(*i) { - let col_idx_in_table = btree_table - .columns - .iter() - .position(|tc| tc.name.as_ref() == Some(&c.name)); - matches!(col_idx_in_table, Some(col_idx) if *table == 0 && *column == col_idx) - } else { - false - } - }).count(); - - if matching_columns > best_index.1 { - best_index = (Some(index_candidate), matching_columns); - } - } - } - - let Some(matching_index) = best_index.0 else { - return Ok(false); - }; - let match_count = best_index.1; - - // If we found a matching index, use it for scanning - *index = Some(matching_index.clone()); - // If the order by direction matches the index direction, we can iterate the index in forwards order. - // If they don't, we must iterate the index in backwards order. - let index_direction = &matching_index.columns.first().as_ref().unwrap().order; - *iter_dir = match (index_direction, order[0].1) { - (SortOrder::Asc, SortOrder::Asc) | (SortOrder::Desc, SortOrder::Desc) => { - IterationDirection::Forwards - } - (SortOrder::Asc, SortOrder::Desc) | (SortOrder::Desc, SortOrder::Asc) => { - IterationDirection::Backwards - } - }; - - // If the index covers all ORDER BY columns, and one of the following applies: - // - the ORDER BY directions exactly match the index orderings, - // - the ORDER by directions are the exact opposite of the index orderings, - // we can remove the ORDER BY clause. - if match_count == order.len() { - let full_match = { - let mut all_match_forward = true; - let mut all_match_reverse = true; - for (i, (_, direction)) in order.iter().enumerate() { - match (&matching_index.columns[i].order, direction) { - (SortOrder::Asc, SortOrder::Asc) | (SortOrder::Desc, SortOrder::Desc) => { - all_match_reverse = false; - } - (SortOrder::Asc, SortOrder::Desc) | (SortOrder::Desc, SortOrder::Asc) => { - all_match_forward = false; - } - } - } - all_match_forward || all_match_reverse - }; - if full_match { - *order_by = None; - } - } - - Ok(order_by.is_none()) -} - -/** - * Use indexes where possible. - * - * When this function is called, condition expressions from both the actual WHERE clause and the JOIN clauses are in the where_clause vector. - * If we find a condition that can be used to index scan, we pop it off from the where_clause vector and put it into a Search operation. - * We put it there simply because it makes it a bit easier to track during translation. - * - * In this function we also try to eliminate ORDER BY clauses if there is an index that satisfies the ORDER BY clause. - */ -fn use_indexes( +/// Optimize the join order and index selection for a query. +/// +/// This function does the following: +/// - Computes a set of [Constraint]s for each table. +/// - Using those constraints, computes the best join order for the list of [TableReference]s +/// and selects the best [crate::translate::optimizer::access_method::AccessMethod] for each table in the join order. +/// - Mutates the [Operation]s in `table_references` to use the selected access methods. +/// - Removes predicates from the `where_clause` that are now redundant due to the selected access methods. +/// - Removes sorting operations if the selected join order and access methods satisfy the [crate::translate::optimizer::order::OrderTarget]. +/// +/// Returns the join order if it was optimized, or None if the default join order was considered best. +fn optimize_table_access( table_references: &mut [TableReference], available_indexes: &HashMap>>, where_clause: &mut Vec, order_by: &mut Option>, - group_by: &Option, -) -> Result<()> { - // Try to use indexes for eliminating ORDER BY clauses - let did_eliminate_orderby = - eliminate_unnecessary_orderby(table_references, available_indexes, order_by, group_by)?; + group_by: &mut Option, +) -> Result>> { + let access_methods_arena = RefCell::new(Vec::new()); + let maybe_order_target = compute_order_target(order_by, group_by.as_mut()); + let constraints_per_table = + constraints_from_where_clause(where_clause, table_references, available_indexes)?; + let Some(best_join_order_result) = compute_best_join_order( + table_references, + maybe_order_target.as_ref(), + &constraints_per_table, + &access_methods_arena, + )? + else { + return Ok(None); + }; - let join_order = table_references - .iter() - .enumerate() - .map(|(i, t)| JoinOrderMember { - table_no: i, - is_outer: t.join_info.as_ref().map_or(false, |j| j.outer), - }) - .collect::>(); + let BestJoinOrderResult { + best_plan, + best_ordered_plan, + } = best_join_order_result; - // Try to use indexes for WHERE conditions - for (table_index, table_reference) in table_references.iter_mut().enumerate() { - if matches!(table_reference.op, Operation::Scan { .. }) { - let index = if let Operation::Scan { index, .. } = &table_reference.op { - Option::clone(index) - } else { - None - }; - match index { - // If we decided to eliminate ORDER BY using an index, let's constrain our search to only that index - Some(index) => { - let available_indexes = available_indexes - .values() - .flatten() - .filter(|i| i.name == index.name) - .cloned() - .collect::>(); - if let Some(search) = try_extract_index_search_from_where_clause( - where_clause, - table_index, - table_reference, - &available_indexes, - &join_order, - )? { - table_reference.op = Operation::Search(search); - } - } - None => { - let table_name = table_reference.table.get_name(); - - // If we can utilize the rowid alias of the table, let's preferentially always use it for now. - let mut i = 0; - while i < where_clause.len() { - if let Some(search) = try_extract_rowid_search_expression( - &mut where_clause[i], - table_index, - table_reference, - &join_order, - )? { - where_clause.remove(i); - table_reference.op = Operation::Search(search); - continue; - } else { - i += 1; - } - } - if did_eliminate_orderby && table_index == 0 { - // If we already made the decision to remove ORDER BY based on the Rowid (e.g. ORDER BY id), then skip this. - // It would be possible to analyze the index and see if the covering index would retain the ordering guarantee, - // but we just don't do that yet. - continue; - } - let placeholder = vec![]; - let mut usable_indexes_ref = &placeholder; - if let Some(indexes) = available_indexes.get(table_name) { - usable_indexes_ref = indexes; - } - if let Some(search) = try_extract_index_search_from_where_clause( - where_clause, - table_index, - table_reference, - usable_indexes_ref, - &join_order, - )? { - table_reference.op = Operation::Search(search); - } - } - } + // See if best_ordered_plan is better than the overall best_plan if we add a sorting penalty + // to the unordered plan's cost. + let best_plan = if let Some(best_ordered_plan) = best_ordered_plan { + let best_unordered_plan_cost = best_plan.cost; + let best_ordered_plan_cost = best_ordered_plan.cost; + const SORT_COST_PER_ROW_MULTIPLIER: f64 = 0.001; + let sorting_penalty = + Cost(best_plan.output_cardinality as f64 * SORT_COST_PER_ROW_MULTIPLIER); + if best_unordered_plan_cost + sorting_penalty > best_ordered_plan_cost { + best_ordered_plan + } else { + best_plan } + } else { + best_plan + }; - // Finally, if there's no other reason to use an index, if an index covers the columns used in the query, let's use it. - if let Some(indexes) = available_indexes.get(table_reference.table.get_name()) { - for index_candidate in indexes.iter() { - let is_covering = table_reference.index_is_covering(index_candidate); - if let Operation::Scan { index, .. } = &mut table_reference.op { - if index.is_some() { - continue; - } - if is_covering { - *index = Some(index_candidate.clone()); - break; - } + // Eliminate sorting if possible. + if let Some(order_target) = maybe_order_target { + let satisfies_order_target = plan_satisfies_order_target( + &best_plan, + &access_methods_arena, + table_references, + &order_target, + ); + if satisfies_order_target { + match order_target.1 { + EliminatesSort::GroupBy => { + let _ = group_by.as_mut().and_then(|g| g.sort_order.take()); + } + EliminatesSort::OrderBy => { + let _ = order_by.take(); + } + EliminatesSort::GroupByAndOrderBy => { + let _ = group_by.as_mut().and_then(|g| g.sort_order.take()); + let _ = order_by.take(); } } } } - Ok(()) + let (best_access_methods, best_table_numbers) = ( + best_plan.best_access_methods().collect::>(), + best_plan.table_numbers().collect::>(), + ); + + let best_join_order: Vec = best_table_numbers + .into_iter() + .map(|table_number| JoinOrderMember { + table_no: table_number, + is_outer: table_references[table_number] + .join_info + .as_ref() + .map_or(false, |join_info| join_info.outer), + }) + .collect(); + let mut to_remove_from_where_clause = vec![]; + + // Mutate the Operations in `table_references` to use the selected access methods. + for (i, join_order_member) in best_join_order.iter().enumerate() { + let table_number = join_order_member.table_no; + let access_method = &access_methods_arena.borrow()[best_access_methods[i]]; + if matches!( + table_references[table_number].op, + Operation::Subquery { .. } + ) { + // FIXME: Operation::Subquery shouldn't exist. It's not an operation, it's a kind of temporary table. + assert!( + access_method.is_scan(), + "nothing in the current optimizer should be able to optimize subqueries, but got {:?} for table {}", + access_method, + table_references[table_number].table.get_name() + ); + continue; + } + if access_method.is_scan() { + if access_method.index.is_some() || i == 0 { + table_references[table_number].op = Operation::Scan { + iter_dir: access_method.iter_dir, + index: access_method.index.clone(), + }; + continue; + } + // This branch means we have a full table scan for a non-outermost table. + // Try to construct an ephemeral index since it's going to be better than a scan. + let table_constraints = constraints_per_table + .iter() + .find(|c| c.table_no == table_number); + let Some(table_constraints) = table_constraints else { + table_references[table_number].op = Operation::Scan { + iter_dir: access_method.iter_dir, + index: access_method.index.clone(), + }; + continue; + }; + let temp_constraint_refs = (0..table_constraints.constraints.len()) + .map(|i| ConstraintRef { + constraint_vec_pos: i, + index_col_pos: table_constraints.constraints[i].table_col_pos, + sort_order: SortOrder::Asc, + }) + .collect::>(); + let usable_constraint_refs = usable_constraints_for_join_order( + &table_constraints.constraints, + &temp_constraint_refs, + &best_join_order[..=i], + ); + if usable_constraint_refs.is_empty() { + table_references[table_number].op = Operation::Scan { + iter_dir: access_method.iter_dir, + index: access_method.index.clone(), + }; + continue; + } + let ephemeral_index = ephemeral_index_build( + &table_references[table_number], + table_number, + &table_constraints.constraints, + &usable_constraint_refs, + ); + let ephemeral_index = Arc::new(ephemeral_index); + table_references[table_number].op = Operation::Search(Search::Seek { + index: Some(ephemeral_index), + seek_def: build_seek_def_from_constraints( + &table_constraints.constraints, + &usable_constraint_refs, + access_method.iter_dir, + where_clause, + )?, + }); + } else { + let constraint_refs = access_method.constraint_refs; + assert!(!constraint_refs.is_empty()); + for cref in constraint_refs.iter() { + let constraint = + &constraints_per_table[table_number].constraints[cref.constraint_vec_pos]; + to_remove_from_where_clause.push(constraint.where_clause_pos.0); + } + if let Some(index) = &access_method.index { + table_references[table_number].op = Operation::Search(Search::Seek { + index: Some(index.clone()), + seek_def: build_seek_def_from_constraints( + &constraints_per_table[table_number].constraints, + &constraint_refs, + access_method.iter_dir, + where_clause, + )?, + }); + continue; + } + assert!( + constraint_refs.len() == 1, + "expected exactly one constraint for rowid seek, got {:?}", + constraint_refs + ); + let constraint = &constraints_per_table[table_number].constraints + [constraint_refs[0].constraint_vec_pos]; + table_references[table_number].op = match constraint.operator { + ast::Operator::Equals => Operation::Search(Search::RowidEq { + cmp_expr: constraint.get_constraining_expr(where_clause), + }), + _ => Operation::Search(Search::Seek { + index: None, + seek_def: build_seek_def_from_constraints( + &constraints_per_table[table_number].constraints, + &constraint_refs, + access_method.iter_dir, + where_clause, + )?, + }), + }; + } + } + to_remove_from_where_clause.sort_by_key(|c| *c); + for position in to_remove_from_where_clause.iter().rev() { + where_clause.remove(*position); + } + + Ok(Some(best_join_order)) } #[derive(Debug, PartialEq, Clone)] @@ -495,21 +448,10 @@ pub trait Optimizable { .map_or(false, |c| c == AlwaysTrueOrFalse::AlwaysFalse)) } fn is_constant(&self, resolver: &Resolver<'_>) -> bool; - fn is_rowid_alias_of(&self, table_index: usize) -> bool; fn is_nonnull(&self, tables: &[TableReference]) -> bool; } impl Optimizable for ast::Expr { - fn is_rowid_alias_of(&self, table_index: usize) -> bool { - match self { - Self::Column { - table, - is_rowid_alias, - .. - } => *is_rowid_alias && *table == table_index, - _ => false, - } - } /// Returns true if the expressions is (verifiably) non-NULL. /// It might still be non-NULL even if we return false; we just /// weren't able to prove it. @@ -795,289 +737,11 @@ impl Optimizable for ast::Expr { } } -fn opposite_cmp_op(op: ast::Operator) -> ast::Operator { - match op { - ast::Operator::Equals => ast::Operator::Equals, - ast::Operator::Greater => ast::Operator::Less, - ast::Operator::GreaterEquals => ast::Operator::LessEquals, - ast::Operator::Less => ast::Operator::Greater, - ast::Operator::LessEquals => ast::Operator::GreaterEquals, - _ => panic!("unexpected operator: {:?}", op), - } -} - -/// Struct used for scoring index scans -/// Currently we just estimate cost in a really dumb way, -/// i.e. no statistics are used. -struct IndexScore { - index: Option>, - cost: f64, - constraints: Vec, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -struct IndexInfo { - unique: bool, - column_count: usize, -} - -const ESTIMATED_HARDCODED_ROWS_PER_TABLE: f64 = 1000.0; - -/// Unbelievably dumb cost estimate for rows scanned by an index scan. -fn dumb_cost_estimator( - index_info: Option, - constraints: &[IndexConstraint], - is_inner_loop: bool, - is_ephemeral: bool, -) -> f64 { - // assume that the outer table always does a full table scan :) - // this discourages building ephemeral indexes on the outer table - // (since a scan reads TABLE_ROWS rows, so an ephemeral index on the outer table would both read TABLE_ROWS rows to build the index and then seek the index) - // but encourages building it on the inner table because it's only built once but the inner loop is run as many times as the outer loop has iterations. - let loop_multiplier = if is_inner_loop { - ESTIMATED_HARDCODED_ROWS_PER_TABLE - } else { - 1.0 - }; - - // If we are building an ephemeral index, we assume we will scan the entire source table to build it. - // Non-ephemeral indexes don't need to be built. - let cost_to_build_index = is_ephemeral as usize as f64 * ESTIMATED_HARDCODED_ROWS_PER_TABLE; - - let Some(index_info) = index_info else { - return cost_to_build_index + ESTIMATED_HARDCODED_ROWS_PER_TABLE * loop_multiplier; - }; - - let final_constraint_is_range = constraints - .last() - .map_or(false, |c| c.operator != ast::Operator::Equals); - let equalities_count = constraints - .iter() - .take(if final_constraint_is_range { - constraints.len() - 1 - } else { - constraints.len() - }) - .count() as f64; - - let selectivity = match ( - index_info.unique, - index_info.column_count as f64, - equalities_count, - ) { - // no equalities: let's assume range query selectivity is 0.4. if final constraint is not range and there are no equalities, it means full table scan incoming - (_, _, 0.0) => { - if final_constraint_is_range { - 0.4 - } else { - 1.0 - } - } - // on an unique index if we have equalities across all index columns, assume very high selectivity - (true, index_cols, eq_count) if eq_count == index_cols => 0.01 * eq_count, - // some equalities: let's assume each equality has a selectivity of 0.1 and range query selectivity is 0.4 - (_, _, eq_count) => (eq_count * 0.1) * if final_constraint_is_range { 0.4 } else { 1.0 }, - }; - cost_to_build_index + selectivity * ESTIMATED_HARDCODED_ROWS_PER_TABLE * loop_multiplier -} - -/// Try to extract an index search from the WHERE clause -/// Returns an optional [Search] struct if an index search can be extracted, otherwise returns None. -pub fn try_extract_index_search_from_where_clause( - where_clause: &mut Vec, - table_index: usize, - table_reference: &TableReference, - table_indexes: &[Arc], - join_order: &[JoinOrderMember], -) -> Result> { - // If there are no WHERE terms, we can't extract a search - if where_clause.is_empty() { - return Ok(None); - } - - let iter_dir = if let Operation::Scan { iter_dir, .. } = &table_reference.op { - *iter_dir - } else { - return Ok(None); - }; - - // Find all potential index constraints - // For WHERE terms to be used to constrain an index scan, they must: - // 1. refer to columns in the table that the index is on - // 2. be a binary comparison expression - // 3. constrain the index columns in the order that they appear in the index - // - e.g. if the index is on (a,b,c) then we can use all of "a = 1 AND b = 2 AND c = 3" to constrain the index scan, - // - but if the where clause is "a = 1 and c = 3" then we can only use "a = 1". - let cost_of_full_table_scan = dumb_cost_estimator(None, &[], table_index != 0, false); - let mut constraints_cur = vec![]; - let mut best_index = IndexScore { - index: None, - cost: cost_of_full_table_scan, - constraints: vec![], - }; - - for index in table_indexes { - // Check how many terms in the where clause constrain the index in column order - find_index_constraints( - where_clause, - table_index, - index, - join_order, - &mut constraints_cur, - )?; - // naive scoring since we don't have statistics: prefer the index where we can use the most columns - // e.g. if we can use all columns of an index on (a,b), it's better than an index of (c,d,e) where we can only use c. - let cost = dumb_cost_estimator( - Some(IndexInfo { - unique: index.unique, - column_count: index.columns.len(), - }), - &constraints_cur, - table_index != 0, - false, - ); - if cost < best_index.cost { - best_index.index = Some(Arc::clone(index)); - best_index.cost = cost; - best_index.constraints.clear(); - best_index.constraints.append(&mut constraints_cur); - } - } - - // We haven't found a persistent btree index that is any better than a full table scan; - // let's see if building an ephemeral index would be better. - if best_index.index.is_none() { - let (ephemeral_cost, constraints_with_col_idx, mut constraints_without_col_idx) = - ephemeral_index_estimate_cost(where_clause, table_reference, table_index, join_order); - if ephemeral_cost < best_index.cost { - // ephemeral index makes sense, so let's build it now. - // ephemeral columns are: columns from the table_reference, constraints first, then the rest - let ephemeral_index = - ephemeral_index_build(table_reference, table_index, &constraints_with_col_idx); - best_index.index = Some(Arc::new(ephemeral_index)); - best_index.cost = ephemeral_cost; - best_index.constraints.clear(); - best_index - .constraints - .append(&mut constraints_without_col_idx); - } - } - - if best_index.index.is_none() { - return Ok(None); - } - - // Build the seek definition - let seek_def = - build_seek_def_from_index_constraints(&best_index.constraints, iter_dir, where_clause)?; - - // Remove the used terms from the where_clause since they are now part of the seek definition - // Sort terms by position in descending order to avoid shifting indices during removal - best_index.constraints.sort_by(|a, b| { - b.position_in_where_clause - .0 - .cmp(&a.position_in_where_clause.0) - }); - - for constraint in best_index.constraints.iter() { - where_clause.remove(constraint.position_in_where_clause.0); - } - - return Ok(Some(Search::Seek { - index: best_index.index, - seek_def, - })); -} - -fn ephemeral_index_estimate_cost( - where_clause: &mut Vec, - table_reference: &TableReference, - table_index: usize, - join_order: &[JoinOrderMember], -) -> (f64, Vec<(usize, IndexConstraint)>, Vec) { - let mut constraints_with_col_idx: Vec<(usize, IndexConstraint)> = where_clause - .iter() - .enumerate() - .filter(|(_, term)| is_potential_index_constraint(term, table_index, join_order)) - .filter_map(|(i, term)| { - let Ok(ast::Expr::Binary(lhs, operator, rhs)) = unwrap_parens(&term.expr) else { - panic!("expected binary expression"); - }; - if let ast::Expr::Column { table, column, .. } = lhs.as_ref() { - if *table == table_index { - return Some(( - *column, - IndexConstraint { - position_in_where_clause: (i, BinaryExprSide::Rhs), - operator: *operator, - index_column_sort_order: SortOrder::Asc, - }, - )); - } - } - if let ast::Expr::Column { table, column, .. } = rhs.as_ref() { - if *table == table_index { - return Some(( - *column, - IndexConstraint { - position_in_where_clause: (i, BinaryExprSide::Lhs), - operator: opposite_cmp_op(*operator), - index_column_sort_order: SortOrder::Asc, - }, - )); - } - } - None - }) - .collect(); - // sort equalities first - constraints_with_col_idx.sort_by(|a, _| { - if a.1.operator == ast::Operator::Equals { - Ordering::Less - } else { - Ordering::Equal - } - }); - // drop everything after the first inequality - constraints_with_col_idx.truncate( - constraints_with_col_idx - .iter() - .position(|c| c.1.operator != ast::Operator::Equals) - .unwrap_or(constraints_with_col_idx.len()), - ); - - let ephemeral_column_count = table_reference - .columns() - .iter() - .enumerate() - .filter(|(i, _)| table_reference.column_is_used(*i)) - .count(); - - let constraints_without_col_idx = constraints_with_col_idx - .iter() - .cloned() - .map(|(_, c)| c) - .collect::>(); - let ephemeral_cost = dumb_cost_estimator( - Some(IndexInfo { - unique: false, - column_count: ephemeral_column_count, - }), - &constraints_without_col_idx, - table_index != 0, - true, - ); - ( - ephemeral_cost, - constraints_with_col_idx, - constraints_without_col_idx, - ) -} - fn ephemeral_index_build( table_reference: &TableReference, table_index: usize, - index_constraints: &[(usize, IndexConstraint)], + constraints: &[Constraint], + constraint_refs: &[ConstraintRef], ) -> Index { let mut ephemeral_columns: Vec = table_reference .columns() @@ -1093,14 +757,14 @@ fn ephemeral_index_build( .collect(); // sort so that constraints first, then rest in whatever order they were in in the table ephemeral_columns.sort_by(|a, b| { - let a_constraint = index_constraints + let a_constraint = constraint_refs .iter() .enumerate() - .find(|(_, c)| c.0 == a.pos_in_table); - let b_constraint = index_constraints + .find(|(_, c)| constraints[c.constraint_vec_pos].table_col_pos == a.pos_in_table); + let b_constraint = constraint_refs .iter() .enumerate() - .find(|(_, c)| c.0 == b.pos_in_table); + .find(|(_, c)| constraints[c.constraint_vec_pos].table_col_pos == b.pos_in_table); match (a_constraint, b_constraint) { (Some(_), None) => Ordering::Less, (None, Some(_)) => Ordering::Greater, @@ -1124,230 +788,29 @@ fn ephemeral_index_build( ephemeral_index } -#[derive(Debug, Clone)] -/// A representation of an expression in a [WhereTerm] that can potentially be used as part of an index seek key. -/// For example, if there is an index on table T(x,y) and another index on table U(z), and the where clause is "WHERE x > 10 AND 20 = z", -/// the index constraints are: -/// - x > 10 ==> IndexConstraint { position_in_where_clause: (0, [BinaryExprSide::Rhs]), operator: [ast::Operator::Greater] } -/// - 20 = z ==> IndexConstraint { position_in_where_clause: (1, [BinaryExprSide::Lhs]), operator: [ast::Operator::Equals] } -pub struct IndexConstraint { - position_in_where_clause: (usize, BinaryExprSide), - operator: ast::Operator, - index_column_sort_order: SortOrder, -} - -/// Helper enum for [IndexConstraint] to indicate which side of a binary comparison expression is being compared to the index column. -/// For example, if the where clause is "WHERE x = 10" and there's an index on x, -/// the [IndexConstraint] for the where clause term "x = 10" will have a [BinaryExprSide::Rhs] -/// because the right hand side expression "10" is being compared to the index column "x". -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum BinaryExprSide { - Lhs, - Rhs, -} - -/// Recursively unwrap parentheses from an expression -/// e.g. (((t.x > 5))) -> t.x > 5 -fn unwrap_parens(expr: T) -> Result -where - T: UnwrapParens, -{ - expr.unwrap_parens() -} - -trait UnwrapParens { - fn unwrap_parens(self) -> Result - where - Self: Sized; -} - -impl UnwrapParens for &ast::Expr { - fn unwrap_parens(self) -> Result { - match self { - ast::Expr::Column { .. } => Ok(self), - ast::Expr::Parenthesized(exprs) => match exprs.len() { - 1 => unwrap_parens(exprs.first().unwrap()), - _ => crate::bail_parse_error!("expected single expression in parentheses"), - }, - _ => Ok(self), - } - } -} - -impl UnwrapParens for ast::Expr { - fn unwrap_parens(self) -> Result { - match self { - ast::Expr::Column { .. } => Ok(self), - ast::Expr::Parenthesized(mut exprs) => match exprs.len() { - 1 => unwrap_parens(exprs.pop().unwrap()), - _ => crate::bail_parse_error!("expected single expression in parentheses"), - }, - _ => Ok(self), - } - } -} - -/// Get the position of a column in an index -/// For example, if there is an index on table T(x,y) then y's position in the index is 1. -fn get_column_position_in_index( - expr: &ast::Expr, - table_index: usize, - index: &Arc, -) -> Result> { - let ast::Expr::Column { table, column, .. } = unwrap_parens(expr)? else { - return Ok(None); - }; - if *table != table_index { - return Ok(None); - } - Ok(index.column_table_pos_to_index_pos(*column)) -} - -fn is_potential_index_constraint( - term: &WhereTerm, - table_index: usize, - join_order: &[JoinOrderMember], -) -> bool { - // Skip terms that cannot be evaluated at this table's loop level - if !term.should_eval_at_loop(table_index, join_order) { - return false; - } - // Skip terms that are not binary comparisons - let Ok(ast::Expr::Binary(lhs, operator, rhs)) = unwrap_parens(&term.expr) else { - return false; - }; - // Only consider index scans for binary ops that are comparisons - if !matches!( - *operator, - ast::Operator::Equals - | ast::Operator::Greater - | ast::Operator::GreaterEquals - | ast::Operator::Less - | ast::Operator::LessEquals - ) { - return false; - } - - // If both lhs and rhs refer to columns from this table, we can't use this constraint - // because we can't use the index to satisfy the condition. - // Examples: - // - WHERE t.x > t.y - // - WHERE t.x + 1 > t.y - 5 - // - WHERE t.x = (t.x) - let Ok(eval_at_left) = determine_where_to_eval_expr(&lhs, join_order) else { - return false; - }; - let Ok(eval_at_right) = determine_where_to_eval_expr(&rhs, join_order) else { - return false; - }; - if eval_at_left == EvalAt::Loop(table_index) && eval_at_right == EvalAt::Loop(table_index) { - return false; - } - true -} - -/// Find all [IndexConstraint]s for a given WHERE clause -/// Constraints are appended as long as they constrain the index in column order. -/// E.g. for index (a,b,c) to be fully used, there must be a [WhereTerm] for each of a, b, and c. -/// If e.g. only a and c are present, then only the first column 'a' of the index will be used. -fn find_index_constraints( - where_clause: &mut Vec, - table_index: usize, - index: &Arc, - join_order: &[JoinOrderMember], - out_constraints: &mut Vec, -) -> Result<()> { - for position_in_index in 0..index.columns.len() { - let mut found = false; - for (position_in_where_clause, term) in where_clause.iter().enumerate() { - if !is_potential_index_constraint(term, table_index, join_order) { - continue; - } - - let ast::Expr::Binary(lhs, operator, rhs) = unwrap_parens(&term.expr)? else { - panic!("expected binary expression"); - }; - - // Check if lhs is a column that is in the i'th position of the index - if Some(position_in_index) == get_column_position_in_index(lhs, table_index, index)? { - out_constraints.push(IndexConstraint { - operator: *operator, - position_in_where_clause: (position_in_where_clause, BinaryExprSide::Rhs), - index_column_sort_order: index.columns[position_in_index].order, - }); - found = true; - break; - } - // Check if rhs is a column that is in the i'th position of the index - if Some(position_in_index) == get_column_position_in_index(rhs, table_index, index)? { - out_constraints.push(IndexConstraint { - operator: opposite_cmp_op(*operator), // swap the operator since e.g. if condition is 5 >= x, we want to use x <= 5 - position_in_where_clause: (position_in_where_clause, BinaryExprSide::Lhs), - index_column_sort_order: index.columns[position_in_index].order, - }); - found = true; - break; - } - } - if !found { - // Expressions must constrain index columns in index definition order. If we didn't find a constraint for the i'th index column, - // then we stop here and return the constraints we have found so far. - break; - } - } - - // In a multicolumn index, only the last term can have a nonequality expression. - // For example, imagine an index on (x,y) and the where clause is "WHERE x > 10 AND y > 20"; - // We can't use GT(x: 10,y: 20) as the seek key, because the first row greater than (x: 10,y: 20) - // might be e.g. (x: 10,y: 21), which does not satisfy the where clause, but a row after that e.g. (x: 11,y: 21) does. - // So: - // - in this case only GT(x: 10) can be used as the seek key, and we must emit a regular condition expression for y > 20 while scanning. - // On the other hand, if the where clause is "WHERE x = 10 AND y > 20", we can use GT(x=10,y=20) as the seek key, - // because any rows where (x=10,y=20) < ROW < (x=11) will match the where clause. - for i in 0..out_constraints.len() { - if out_constraints[i].operator != ast::Operator::Equals { - out_constraints.truncate(i + 1); - break; - } - } - - Ok(()) -} - -/// Build a [SeekDef] for a given list of [IndexConstraint]s -pub fn build_seek_def_from_index_constraints( - constraints: &[IndexConstraint], +/// Build a [SeekDef] for a given list of [Constraint]s +pub fn build_seek_def_from_constraints( + constraints: &[Constraint], + constraint_refs: &[ConstraintRef], iter_dir: IterationDirection, - where_clause: &mut Vec, + where_clause: &[WhereTerm], ) -> Result { assert!( - !constraints.is_empty(), - "cannot build seek def from empty list of index constraints" + !constraint_refs.is_empty(), + "cannot build seek def from empty list of constraint refs" ); // Extract the key values and operators - let mut key = Vec::with_capacity(constraints.len()); - - for constraint in constraints { - // Extract the other expression from the binary WhereTerm (i.e. the one being compared to the index column) - let (idx, side) = constraint.position_in_where_clause; - let where_term = &mut where_clause[idx]; - let ast::Expr::Binary(lhs, _, rhs) = unwrap_parens(where_term.expr.take_ownership())? - else { - crate::bail_parse_error!("expected binary expression"); - }; - let cmp_expr = if side == BinaryExprSide::Lhs { - *lhs - } else { - *rhs - }; - key.push((cmp_expr, constraint.index_column_sort_order)); - } + let key = constraint_refs + .iter() + .map(|cref| cref.as_seek_key_column(constraints, where_clause)) + .collect(); // We know all but potentially the last term is an equality, so we can use the operator of the last term // to form the SeekOp - let op = constraints.last().unwrap().operator; + let op = constraints[constraint_refs.last().unwrap().constraint_vec_pos].operator; - build_seek_def(op, iter_dir, key) + let seek_def = build_seek_def(op, iter_dir, key)?; + Ok(seek_def) } /// Build a [SeekDef] for a given comparison operator and index key. @@ -1360,7 +823,7 @@ pub fn build_seek_def_from_index_constraints( /// 2. The [TerminationKey], which specifies the key that we will use to terminate the index scan that follows the seek. /// /// There are some nuances to how, and which parts of, the index key can be used in the [SeekKey] and [TerminationKey], -/// depending on the operator and iteration direction. This function explains those nuances inline when dealing with +/// depending on the operator and iteration order. This function explains those nuances inline when dealing with /// each case. /// /// But to illustrate the general idea, consider the following examples: @@ -1501,7 +964,7 @@ fn build_seek_def( // Termination key: end at the first GE(x:10, y:20), e.g. (x=10, y=20) // // Descending index example: (x=10 AND y<20) - // Seek key: start from the first LT(x:10, y:20), e.g. (x=10, y=19), so reversed -> GT(x:10, y:20) + // Seek key: start from the first LT(x:10, y:20), e.g. (x=10, y=19) so reversed -> GT(x:10, y:20) // Termination key: end at the first LT(x:10), e.g. (x=9, y=usize::MAX), so reversed -> GE(x:10, NULL); i.e. GE the smallest possible (x=10, y) combination (NULL is always LT) (IterationDirection::Forwards, ast::Operator::Less) => { let (seek_key_len, termination_key_len, seek_op, termination_op) = @@ -1773,93 +1236,6 @@ fn build_seek_def( }) } -pub fn try_extract_rowid_search_expression( - cond: &mut WhereTerm, - table_index: usize, - table_reference: &TableReference, - join_order: &[JoinOrderMember], -) -> Result> { - let iter_dir = if let Operation::Scan { iter_dir, .. } = &table_reference.op { - *iter_dir - } else { - return Ok(None); - }; - if !cond.should_eval_at_loop(table_index, join_order) { - return Ok(None); - } - match &mut cond.expr { - ast::Expr::Binary(lhs, operator, rhs) => { - // If both lhs and rhs refer to columns from this table, we can't perform a rowid seek - // Examples: - // - WHERE t.x > t.y - // - WHERE t.x + 1 > t.y - 5 - // - WHERE t.x = (t.x) - if determine_where_to_eval_expr(lhs, join_order)? == EvalAt::Loop(table_index) - && determine_where_to_eval_expr(rhs, join_order)? == EvalAt::Loop(table_index) - { - return Ok(None); - } - if lhs.is_rowid_alias_of(table_index) { - match operator { - ast::Operator::Equals => { - let rhs_owned = rhs.take_ownership(); - return Ok(Some(Search::RowidEq { - cmp_expr: WhereTerm { - expr: rhs_owned, - from_outer_join: cond.from_outer_join, - }, - })); - } - ast::Operator::Greater - | ast::Operator::GreaterEquals - | ast::Operator::Less - | ast::Operator::LessEquals => { - let rhs_owned = rhs.take_ownership(); - let seek_def = - build_seek_def(*operator, iter_dir, vec![(rhs_owned, SortOrder::Asc)])?; - return Ok(Some(Search::Seek { - index: None, - seek_def, - })); - } - _ => {} - } - } - - if rhs.is_rowid_alias_of(table_index) { - match operator { - ast::Operator::Equals => { - let lhs_owned = lhs.take_ownership(); - return Ok(Some(Search::RowidEq { - cmp_expr: WhereTerm { - expr: lhs_owned, - from_outer_join: cond.from_outer_join, - }, - })); - } - ast::Operator::Greater - | ast::Operator::GreaterEquals - | ast::Operator::Less - | ast::Operator::LessEquals => { - let lhs_owned = lhs.take_ownership(); - let op = opposite_cmp_op(*operator); - let seek_def = - build_seek_def(op, iter_dir, vec![(lhs_owned, SortOrder::Asc)])?; - return Ok(Some(Search::Seek { - index: None, - seek_def, - })); - } - _ => {} - } - } - - Ok(None) - } - _ => Ok(None), - } -} - pub fn rewrite_expr(expr: &mut ast::Expr, param_idx: &mut usize) -> Result<()> { match expr { ast::Expr::Id(id) => { diff --git a/core/translate/optimizer/order.rs b/core/translate/optimizer/order.rs new file mode 100644 index 000000000..b6ff1ade3 --- /dev/null +++ b/core/translate/optimizer/order.rs @@ -0,0 +1,217 @@ +use std::cell::RefCell; + +use limbo_sqlite3_parser::ast::{self, SortOrder}; + +use crate::{ + translate::plan::{GroupBy, IterationDirection, TableReference}, + util::exprs_are_equivalent, +}; + +use super::{access_method::AccessMethod, join::JoinN}; + +#[derive(Debug, PartialEq, Clone)] +/// A convenience struct for representing a (table_no, column_no, [SortOrder]) tuple. +pub struct ColumnOrder { + pub table_no: usize, + pub column_no: usize, + pub order: SortOrder, +} + +#[derive(Debug, PartialEq, Clone)] +/// If an [OrderTarget] is satisfied, then [EliminatesSort] describes which part of the query no longer requires sorting. +pub enum EliminatesSort { + GroupBy, + OrderBy, + GroupByAndOrderBy, +} + +#[derive(Debug, PartialEq, Clone)] +/// An [OrderTarget] is considered in join optimization and index selection, +/// so that if a given join ordering and its access methods satisfy the [OrderTarget], +/// then the join ordering and its access methods are preferred, all other things being equal. +pub struct OrderTarget(pub Vec, pub EliminatesSort); + +impl OrderTarget { + fn maybe_from_iterator<'a>( + list: impl Iterator + Clone, + eliminates_sort: EliminatesSort, + ) -> Option { + if list.clone().count() == 0 { + return None; + } + if list + .clone() + .any(|(expr, _)| !matches!(expr, ast::Expr::Column { .. })) + { + return None; + } + Some(OrderTarget( + list.map(|(expr, order)| { + let ast::Expr::Column { table, column, .. } = expr else { + unreachable!(); + }; + ColumnOrder { + table_no: *table, + column_no: *column, + order, + } + }) + .collect(), + eliminates_sort, + )) + } +} + +/// Compute an [OrderTarget] for the join optimizer to use. +/// Ideally, a join order is both efficient in joining the tables +/// but also returns the results in an order that minimizes the amount of +/// sorting that needs to be done later (either in GROUP BY, ORDER BY, or both). +/// +/// TODO: this does not currently handle the case where we definitely cannot eliminate +/// the ORDER BY sorter, but we could still eliminate the GROUP BY sorter. +pub fn compute_order_target( + order_by: &Option>, + group_by: Option<&mut GroupBy>, +) -> Option { + match (order_by, group_by) { + // No ordering demands - we don't care what order the joined result rows are in + (None, None) => None, + // Only ORDER BY - we would like the joined result rows to be in the order specified by the ORDER BY + (Some(order_by), None) => OrderTarget::maybe_from_iterator( + order_by.iter().map(|(expr, order)| (expr, *order)), + EliminatesSort::OrderBy, + ), + // Only GROUP BY - we would like the joined result rows to be in the order specified by the GROUP BY + (None, Some(group_by)) => OrderTarget::maybe_from_iterator( + group_by.exprs.iter().map(|expr| (expr, SortOrder::Asc)), + EliminatesSort::GroupBy, + ), + // Both ORDER BY and GROUP BY: + // If the GROUP BY does not contain all the expressions in the ORDER BY, + // then we must separately sort the result rows for ORDER BY anyway. + // However, in that case we can use the GROUP BY expressions as the target order for the join, + // so that we don't have to sort twice. + // + // If the GROUP BY contains all the expressions in the ORDER BY, + // then we again can use the GROUP BY expressions as the target order for the join; + // however in this case we must take the ASC/DESC from ORDER BY into account. + (Some(order_by), Some(group_by)) => { + // Does the group by contain all expressions in the order by? + let group_by_contains_all = group_by.exprs.iter().all(|expr| { + order_by + .iter() + .any(|(order_by_expr, _)| exprs_are_equivalent(expr, order_by_expr)) + }); + // If not, let's try to target an ordering that matches the group by -- we don't care about ASC/DESC + if !group_by_contains_all { + return OrderTarget::maybe_from_iterator( + group_by.exprs.iter().map(|expr| (expr, SortOrder::Asc)), + EliminatesSort::GroupBy, + ); + } + // If yes, let's try to target an ordering that matches the GROUP BY columns, + // but the ORDER BY orderings. First, we need to reorder the GROUP BY columns to match the ORDER BY columns. + group_by.exprs.sort_by_key(|expr| { + order_by + .iter() + .position(|(order_by_expr, _)| exprs_are_equivalent(expr, order_by_expr)) + .map_or(usize::MAX, |i| i) + }); + // Iterate over GROUP BY, but take the ORDER BY orderings into account. + OrderTarget::maybe_from_iterator( + group_by + .exprs + .iter() + .zip( + order_by + .iter() + .map(|(_, dir)| dir) + .chain(std::iter::repeat(&SortOrder::Asc)), + ) + .map(|(expr, dir)| (expr, *dir)), + EliminatesSort::GroupByAndOrderBy, + ) + } + } +} + +/// Check if the plan's row iteration order matches the [OrderTarget]'s column order. +/// If yes, and this plan is selected, then a sort operation can be eliminated. +pub fn plan_satisfies_order_target( + plan: &JoinN, + access_methods_arena: &RefCell>, + table_references: &[TableReference], + order_target: &OrderTarget, +) -> bool { + let mut target_col_idx = 0; + let num_cols_in_order_target = order_target.0.len(); + for (table_no, access_method_index) in plan.data.iter() { + let target_col = &order_target.0[target_col_idx]; + let table_ref = &table_references[*table_no]; + let correct_table = target_col.table_no == *table_no; + if !correct_table { + return false; + } + + // Check if this table has an access method that provides the right ordering. + let access_method = &access_methods_arena.borrow()[*access_method_index]; + let iter_dir = access_method.iter_dir; + let index = access_method.index.as_ref(); + match index { + None => { + // No index, so the next required column must be the rowid alias column. + let rowid_alias_col = table_ref + .table + .columns() + .iter() + .position(|c| c.is_rowid_alias); + let Some(rowid_alias_col) = rowid_alias_col else { + return false; + }; + let correct_column = target_col.column_no == rowid_alias_col; + if !correct_column { + return false; + } + + // Btree table rows are always in ascending order of rowid. + let correct_order = if iter_dir == IterationDirection::Forwards { + target_col.order == SortOrder::Asc + } else { + target_col.order == SortOrder::Desc + }; + if !correct_order { + return false; + } + target_col_idx += 1; + // All order columns matched. + if target_col_idx == num_cols_in_order_target { + return true; + } + } + Some(index) => { + // All of the index columns must match the next required columns in the order target. + for index_col in index.columns.iter() { + let target_col = &order_target.0[target_col_idx]; + let correct_column = target_col.column_no == index_col.pos_in_table; + if !correct_column { + return false; + } + let correct_order = if iter_dir == IterationDirection::Forwards { + target_col.order == index_col.order + } else { + target_col.order != index_col.order + }; + if !correct_order { + return false; + } + target_col_idx += 1; + // All order columns matched. + if target_col_idx == num_cols_in_order_target { + return true; + } + } + } + } + } + false +} diff --git a/core/translate/plan.rs b/core/translate/plan.rs index a422b6fed..2a047b3a1 100644 --- a/core/translate/plan.rs +++ b/core/translate/plan.rs @@ -64,9 +64,10 @@ impl ResultSetColumn { #[derive(Debug, Clone)] pub struct GroupBy { pub exprs: Vec, + /// sort order, if a sorter is required (= the columns aren't already in the correct order) + pub sort_order: Option>, /// having clause split into a vec at 'AND' boundaries. pub having: Option>, - pub sort_order: Option>, } /// In a query plan, WHERE clause conditions and JOIN conditions are all folded into a vector of WhereTerm. @@ -781,7 +782,7 @@ pub struct TerminationKey { #[derive(Clone, Debug)] pub enum Search { /// A rowid equality point lookup. This is a special case that uses the SeekRowid bytecode instruction and does not loop. - RowidEq { cmp_expr: WhereTerm }, + RowidEq { cmp_expr: ast::Expr }, /// A search on a table btree (via `rowid`) or a secondary index search. Uses bytecode instructions like SeekGE, SeekGT etc. Seek { index: Option>, diff --git a/core/translate/planner.rs b/core/translate/planner.rs index 1d4cd568e..c6326bac2 100644 --- a/core/translate/planner.rs +++ b/core/translate/planner.rs @@ -581,13 +581,236 @@ pub fn determine_where_to_eval_term( join_order .iter() .position(|t| t.table_no == table_no) - .unwrap(), + .unwrap_or(usize::MAX), )); } return determine_where_to_eval_expr(&term.expr, join_order); } +/// A bitmask representing a set of tables in a query plan. +/// Tables are numbered by their index in [SelectPlan::table_references]. +/// In the bitmask, the first bit is unused so that a mask with all zeros +/// can represent "no tables". +/// +/// E.g. table 0 is represented by bit index 1, table 1 by bit index 2, etc. +/// +/// Usage in Join Optimization +/// +/// In join optimization, [TableMask] is used to: +/// - Generate subsets of tables for dynamic programming in join optimization +/// - Ensure tables are joined in valid orders (e.g., respecting LEFT JOIN order) +/// +/// Usage with constraints (WHERE clause) +/// +/// [TableMask] helps determine: +/// - Which tables are referenced in a constraint +/// - When a constraint can be applied as a join condition (all referenced tables must be on the left side of the table being joined) +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub struct TableMask(pub u128); + +impl std::ops::BitOrAssign for TableMask { + fn bitor_assign(&mut self, rhs: Self) { + self.0 |= rhs.0; + } +} + +impl TableMask { + /// Creates a new empty table mask. + /// + /// The initial mask represents an empty set of tables. + pub fn new() -> Self { + Self(0) + } + + /// Returns true if the mask represents an empty set of tables. + pub fn is_empty(&self) -> bool { + self.0 == 0 + } + + /// Creates a new mask that is the same as this one but without the specified table. + pub fn without_table(&self, table_no: usize) -> Self { + assert!(table_no < 127, "table_no must be less than 127"); + Self(self.0 ^ (1 << (table_no + 1))) + } + + /// Creates a table mask from raw bits. + /// + /// The bits are shifted left by 1 to maintain the convention that table 0 is at bit 1. + pub fn from_bits(bits: u128) -> Self { + Self(bits << 1) + } + + /// Creates a table mask from an iterator of table numbers. + pub fn from_table_number_iter(iter: impl Iterator) -> Self { + iter.fold(Self::new(), |mut mask, table_no| { + assert!(table_no < 127, "table_no must be less than 127"); + mask.add_table(table_no); + mask + }) + } + + /// Adds a table to the mask. + pub fn add_table(&mut self, table_no: usize) { + assert!(table_no < 127, "table_no must be less than 127"); + self.0 |= 1 << (table_no + 1); + } + + /// Returns true if the mask contains the specified table. + pub fn contains_table(&self, table_no: usize) -> bool { + assert!(table_no < 127, "table_no must be less than 127"); + self.0 & (1 << (table_no + 1)) != 0 + } + + /// Returns true if this mask contains all tables in the other mask. + pub fn contains_all(&self, other: &TableMask) -> bool { + self.0 & other.0 == other.0 + } + + /// Returns the number of tables in the mask. + pub fn table_count(&self) -> usize { + self.0.count_ones() as usize + } + + /// Returns true if this mask shares any tables with the other mask. + pub fn intersects(&self, other: &TableMask) -> bool { + self.0 & other.0 != 0 + } +} + +/// Returns a [TableMask] representing the tables referenced in the given expression. +/// Used in the optimizer for constraint analysis. +pub fn table_mask_from_expr(expr: &Expr) -> Result { + let mut mask = TableMask::new(); + match expr { + Expr::Binary(e1, _, e2) => { + mask |= table_mask_from_expr(e1)?; + mask |= table_mask_from_expr(e2)?; + } + Expr::Column { table, .. } | Expr::RowId { table, .. } => { + mask.add_table(*table); + } + Expr::Between { + lhs, + not: _, + start, + end, + } => { + mask |= table_mask_from_expr(lhs)?; + mask |= table_mask_from_expr(start)?; + mask |= table_mask_from_expr(end)?; + } + Expr::Case { + base, + when_then_pairs, + else_expr, + } => { + if let Some(base) = base { + mask |= table_mask_from_expr(base)?; + } + for (when, then) in when_then_pairs { + mask |= table_mask_from_expr(when)?; + mask |= table_mask_from_expr(then)?; + } + if let Some(else_expr) = else_expr { + mask |= table_mask_from_expr(else_expr)?; + } + } + Expr::Cast { expr, .. } => { + mask |= table_mask_from_expr(expr)?; + } + Expr::Collate(expr, _) => { + mask |= table_mask_from_expr(expr)?; + } + Expr::DoublyQualified(_, _, _) => { + crate::bail_parse_error!( + "DoublyQualified should be resolved to a Column before resolving table mask" + ); + } + Expr::Exists(_) => { + todo!(); + } + Expr::FunctionCall { + args, + order_by, + filter_over: _, + .. + } => { + if let Some(args) = args { + for arg in args.iter() { + mask |= table_mask_from_expr(arg)?; + } + } + if let Some(order_by) = order_by { + for term in order_by.iter() { + mask |= table_mask_from_expr(&term.expr)?; + } + } + } + Expr::FunctionCallStar { .. } => {} + Expr::Id(_) => panic!("Id should be resolved to a Column before resolving table mask"), + Expr::InList { lhs, not: _, rhs } => { + mask |= table_mask_from_expr(lhs)?; + if let Some(rhs) = rhs { + for rhs_expr in rhs.iter() { + mask |= table_mask_from_expr(rhs_expr)?; + } + } + } + Expr::InSelect { .. } => todo!(), + Expr::InTable { + lhs, + not: _, + rhs: _, + args, + } => { + mask |= table_mask_from_expr(lhs)?; + if let Some(args) = args { + for arg in args.iter() { + mask |= table_mask_from_expr(arg)?; + } + } + } + Expr::IsNull(expr) => { + mask |= table_mask_from_expr(expr)?; + } + Expr::Like { + lhs, + not: _, + op: _, + rhs, + escape, + } => { + mask |= table_mask_from_expr(lhs)?; + mask |= table_mask_from_expr(rhs)?; + if let Some(escape) = escape { + mask |= table_mask_from_expr(escape)?; + } + } + Expr::Literal(_) => {} + Expr::Name(_) => {} + Expr::NotNull(expr) => { + mask |= table_mask_from_expr(expr)?; + } + Expr::Parenthesized(exprs) => { + for expr in exprs.iter() { + mask |= table_mask_from_expr(expr)?; + } + } + Expr::Qualified(_, _) => { + panic!("Qualified should be resolved to a Column before resolving table mask"); + } + Expr::Raise(_, _) => todo!(), + Expr::Subquery(_) => todo!(), + Expr::Unary(_, expr) => { + mask |= table_mask_from_expr(expr)?; + } + Expr::Variable(_) => {} + } + + Ok(mask) +} + pub fn determine_where_to_eval_expr<'a>( expr: &'a Expr, join_order: &[JoinOrderMember], @@ -602,7 +825,7 @@ pub fn determine_where_to_eval_expr<'a>( let join_idx = join_order .iter() .position(|t| t.table_no == *table) - .unwrap(); + .unwrap_or(usize::MAX); eval_at = eval_at.max(EvalAt::Loop(join_idx)); } ast::Expr::Id(_) => { diff --git a/testing/groupby.test b/testing/groupby.test index 9fce2e83e..70141be0a 100644 --- a/testing/groupby.test +++ b/testing/groupby.test @@ -192,3 +192,9 @@ do_execsql_test groupby_orderby_removal_regression_test { } {1|Foster|1 2|Salazar|1 3|Perry|1} + +do_execsql_test group_by_no_sorting_required { + select age, count(1) from users group by age limit 3; +} {1|112 +2|113 +3|97} \ No newline at end of file