From 5a1cfb7d1550d5d9f147f873a910e8f3ca390cf5 Mon Sep 17 00:00:00 2001 From: Jussi Saurio Date: Mon, 14 Apr 2025 17:25:13 +0300 Subject: [PATCH] Add ColumnUsedMask struct to TableReference to track columns referenced in query --- core/translate/delete.rs | 7 ++++--- core/translate/plan.rs | 44 +++++++++++++++++++++++++++++++++++++++ core/translate/planner.rs | 19 ++++++++++++----- core/translate/select.rs | 27 +++++++++++++++++------- core/translate/update.rs | 12 ++++++----- 5 files changed, 88 insertions(+), 21 deletions(-) diff --git a/core/translate/delete.rs b/core/translate/delete.rs index b8b92349d..fb580b8e8 100644 --- a/core/translate/delete.rs +++ b/core/translate/delete.rs @@ -7,7 +7,7 @@ use crate::vdbe::builder::{ProgramBuilder, ProgramBuilderOpts, QueryMode}; use crate::{schema::Schema, Result, SymbolTable}; use limbo_sqlite3_parser::ast::{Expr, Limit, QualifiedName}; -use super::plan::{IterationDirection, TableReference}; +use super::plan::{ColumnUsedMask, IterationDirection, TableReference}; pub fn translate_delete( query_mode: QueryMode, @@ -50,7 +50,7 @@ pub fn prepare_delete_plan( crate::bail_corrupt_error!("Table is neither a virtual table nor a btree table"); }; let name = tbl_name.name.0.as_str().to_string(); - let table_references = vec![TableReference { + let mut table_references = vec![TableReference { table, identifier: name, op: Operation::Scan { @@ -58,6 +58,7 @@ pub fn prepare_delete_plan( index: None, }, join_info: None, + col_used_mask: ColumnUsedMask::new(), }]; let mut where_predicates = vec![]; @@ -65,7 +66,7 @@ pub fn prepare_delete_plan( // Parse the WHERE clause parse_where( where_clause.map(|e| *e), - &table_references, + &mut table_references, None, &mut where_predicates, )?; diff --git a/core/translate/plan.rs b/core/translate/plan.rs index 4bf86ab8d..038dd90ee 100644 --- a/core/translate/plan.rs +++ b/core/translate/plan.rs @@ -256,6 +256,43 @@ pub struct TableReference { pub identifier: String, /// The join info for this table reference, if it is the right side of a join (which all except the first table reference have) pub join_info: Option, + /// Bitmask of columns that are referenced in the query. + /// Used to decide whether a covering index can be used. + pub col_used_mask: ColumnUsedMask, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +#[repr(transparent)] +pub struct ColumnUsedMask(u128); + +impl ColumnUsedMask { + pub fn new() -> Self { + Self(0) + } + + pub fn set(&mut self, index: usize) { + assert!( + index < 128, + "ColumnUsedMask only supports up to 128 columns" + ); + self.0 |= 1 << index; + } + + pub fn get(&self, index: usize) -> bool { + assert!( + index < 128, + "ColumnUsedMask only supports up to 128 columns" + ); + self.0 & (1 << index) != 0 + } + + pub fn contains_all_set_bits_of(&self, other: &Self) -> bool { + self.0 & other.0 == other.0 + } + + pub fn is_empty(&self) -> bool { + self.0 == 0 + } } #[derive(Clone, Debug)] @@ -331,12 +368,19 @@ impl TableReference { table, identifier: identifier.clone(), join_info, + col_used_mask: ColumnUsedMask::new(), } } pub fn columns(&self) -> &[Column] { self.table.columns() } + + /// Mark a column as used in the query. + /// This is used to determine whether a covering index can be used. + pub fn mark_column_used(&mut self, index: usize) { + self.col_used_mask.set(index); + } } /// A definition of a rowid/index search. diff --git a/core/translate/planner.rs b/core/translate/planner.rs index 2d9246666..f1d7aaeea 100644 --- a/core/translate/planner.rs +++ b/core/translate/planner.rs @@ -1,7 +1,7 @@ use super::{ plan::{ - Aggregate, EvalAt, IterationDirection, JoinInfo, Operation, Plan, ResultSetColumn, - SelectPlan, SelectQueryType, TableReference, WhereTerm, + Aggregate, ColumnUsedMask, EvalAt, IterationDirection, JoinInfo, Operation, Plan, + ResultSetColumn, SelectPlan, SelectQueryType, TableReference, WhereTerm, }, select::prepare_select_plan, SymbolTable, @@ -85,7 +85,7 @@ pub fn resolve_aggregates(expr: &Expr, aggs: &mut Vec) -> bool { pub fn bind_column_references( expr: &mut Expr, - referenced_tables: &[TableReference], + referenced_tables: &mut [TableReference], result_columns: Option<&[ResultSetColumn]>, ) -> Result<()> { match expr { @@ -128,6 +128,7 @@ pub fn bind_column_references( column: col_idx, is_rowid_alias, }; + referenced_tables[tbl_idx].mark_column_used(col_idx); return Ok(()); } @@ -178,6 +179,7 @@ pub fn bind_column_references( column: col_idx.unwrap(), is_rowid_alias: col.is_rowid_alias, }; + referenced_tables[tbl_idx].mark_column_used(col_idx.unwrap()); Ok(()) } Expr::Between { @@ -327,6 +329,7 @@ fn parse_from_clause_table<'a>( table: tbl_ref, identifier: alias.unwrap_or(normalized_qualified_name), join_info: None, + col_used_mask: ColumnUsedMask::new(), }); return Ok(()); }; @@ -409,6 +412,7 @@ fn parse_from_clause_table<'a>( join_info: None, table: Table::Virtual(vtab), identifier: alias, + col_used_mask: ColumnUsedMask::new(), }); Ok(()) @@ -539,7 +543,7 @@ pub fn parse_from<'a>( pub fn parse_where( where_clause: Option, - table_references: &[TableReference], + table_references: &mut [TableReference], result_columns: Option<&[ResultSetColumn]>, out_where_clause: &mut Vec, ) -> Result<()> { @@ -758,7 +762,7 @@ fn parse_join<'a>( let mut preds = vec![]; break_predicate_at_and_boundaries(expr, &mut preds); for predicate in preds.iter_mut() { - bind_column_references(predicate, &scope.tables, None)?; + bind_column_references(predicate, &mut scope.tables, None)?; } for pred in preds { let cur_table_idx = scope.tables.len() - 1; @@ -832,6 +836,11 @@ fn parse_join<'a>( is_rowid_alias: right_col.is_rowid_alias, }), ); + + let left_table = scope.tables.get_mut(left_table_idx).unwrap(); + left_table.mark_column_used(left_col_idx); + let right_table = scope.tables.get_mut(cur_table_idx).unwrap(); + right_table.mark_column_used(right_col_idx); let eval_at = if outer { EvalAt::Loop(cur_table_idx) } else { diff --git a/core/translate/select.rs b/core/translate/select.rs index 24a6331e5..3972bdc85 100644 --- a/core/translate/select.rs +++ b/core/translate/select.rs @@ -104,12 +104,17 @@ pub fn prepare_select_plan<'a>( match column { ResultColumn::Star => { select_star(&plan.table_references, &mut plan.result_columns); + for table in plan.table_references.iter_mut() { + for idx in 0..table.columns().len() { + table.mark_column_used(idx); + } + } } ResultColumn::TableStar(name) => { let name_normalized = normalize_ident(name.0.as_str()); let referenced_table = plan .table_references - .iter() + .iter_mut() .enumerate() .find(|(_, t)| t.identifier == name_normalized); @@ -117,23 +122,29 @@ pub fn prepare_select_plan<'a>( crate::bail_parse_error!("Table {} not found", name.0); } let (table_index, table) = referenced_table.unwrap(); - for (idx, col) in table.columns().iter().enumerate() { + let num_columns = table.columns().len(); + for idx in 0..num_columns { + let is_rowid_alias = { + let columns = table.columns(); + columns[idx].is_rowid_alias + }; plan.result_columns.push(ResultSetColumn { expr: ast::Expr::Column { database: None, // TODO: support different databases table: table_index, column: idx, - is_rowid_alias: col.is_rowid_alias, + is_rowid_alias, }, alias: None, contains_aggregates: false, }); + table.mark_column_used(idx); } } ResultColumn::Expr(ref mut expr, maybe_alias) => { bind_column_references( expr, - &plan.table_references, + &mut plan.table_references, Some(&plan.result_columns), )?; match expr { @@ -293,7 +304,7 @@ pub fn prepare_select_plan<'a>( // Parse the actual WHERE clause and add its conditions to the plan WHERE clause that already contains the join conditions. parse_where( where_clause, - &plan.table_references, + &mut plan.table_references, Some(&plan.result_columns), &mut plan.where_clause, )?; @@ -303,7 +314,7 @@ pub fn prepare_select_plan<'a>( replace_column_number_with_copy_of_column_expr(expr, &plan.result_columns)?; bind_column_references( expr, - &plan.table_references, + &mut plan.table_references, Some(&plan.result_columns), )?; } @@ -316,7 +327,7 @@ pub fn prepare_select_plan<'a>( for expr in predicates.iter_mut() { bind_column_references( expr, - &plan.table_references, + &mut plan.table_references, Some(&plan.result_columns), )?; let contains_aggregates = @@ -352,7 +363,7 @@ pub fn prepare_select_plan<'a>( bind_column_references( &mut o.expr, - &plan.table_references, + &mut plan.table_references, Some(&plan.result_columns), )?; resolve_aggregates(&o.expr, &mut plan.aggregates); diff --git a/core/translate/update.rs b/core/translate/update.rs index 62c6c6f9f..a0e32e640 100644 --- a/core/translate/update.rs +++ b/core/translate/update.rs @@ -11,7 +11,8 @@ use limbo_sqlite3_parser::ast::{self, Expr, ResultColumn, SortOrder, Update}; use super::emitter::emit_program; use super::optimizer::optimize_plan; use super::plan::{ - Direction, IterationDirection, Plan, ResultSetColumn, TableReference, UpdatePlan, + ColumnUsedMask, Direction, IterationDirection, Plan, ResultSetColumn, TableReference, + UpdatePlan, }; use super::planner::bind_column_references; use super::planner::{parse_limit, parse_where}; @@ -88,7 +89,7 @@ pub fn prepare_update_plan(schema: &Schema, body: &mut Update) -> crate::Result< }) }) .unwrap_or(IterationDirection::Forwards); - let table_references = vec![TableReference { + let mut table_references = vec![TableReference { table: match table.as_ref() { Table::Virtual(vtab) => Table::Virtual(vtab.clone()), Table::BTree(btree_table) => Table::BTree(btree_table.clone()), @@ -100,6 +101,7 @@ pub fn prepare_update_plan(schema: &Schema, body: &mut Update) -> crate::Result< index: None, }, join_info: None, + col_used_mask: ColumnUsedMask::new(), }]; let set_clauses = body .sets @@ -123,7 +125,7 @@ pub fn prepare_update_plan(schema: &Schema, body: &mut Update) -> crate::Result< )) })?; - let _ = bind_column_references(&mut set.expr, &table_references, None); + let _ = bind_column_references(&mut set.expr, &mut table_references, None); Ok((col_index, set.expr.clone())) }) .collect::, crate::LimboError>>()?; @@ -133,7 +135,7 @@ pub fn prepare_update_plan(schema: &Schema, body: &mut Update) -> crate::Result< if let Some(returning) = &mut body.returning { for rc in returning.iter_mut() { if let ResultColumn::Expr(expr, alias) = rc { - bind_column_references(expr, &table_references, None)?; + bind_column_references(expr, &mut table_references, None)?; result_columns.push(ResultSetColumn { expr: expr.clone(), alias: alias.as_ref().and_then(|a| { @@ -169,7 +171,7 @@ pub fn prepare_update_plan(schema: &Schema, body: &mut Update) -> crate::Result< // Parse the WHERE clause parse_where( body.where_clause.as_ref().map(|w| *w.clone()), - &table_references, + &mut table_references, Some(&result_columns), &mut where_clause, )?;