mirror of
https://github.com/aljazceru/turso.git
synced 2026-02-22 08:25:29 +01:00
Add ColumnUsedMask struct to TableReference to track columns referenced in query
This commit is contained in:
@@ -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,
|
||||
)?;
|
||||
|
||||
@@ -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<JoinInfo>,
|
||||
/// 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.
|
||||
|
||||
@@ -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<Aggregate>) -> 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<Expr>,
|
||||
table_references: &[TableReference],
|
||||
table_references: &mut [TableReference],
|
||||
result_columns: Option<&[ResultSetColumn]>,
|
||||
out_where_clause: &mut Vec<WhereTerm>,
|
||||
) -> 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 {
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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::<Result<Vec<(usize, Expr)>, 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,
|
||||
)?;
|
||||
|
||||
Reference in New Issue
Block a user