Add ColumnUsedMask struct to TableReference to track columns referenced in query

This commit is contained in:
Jussi Saurio
2025-04-14 17:25:13 +03:00
parent 72dac59813
commit 5a1cfb7d15
5 changed files with 88 additions and 21 deletions

View File

@@ -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,
)?;

View File

@@ -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.

View File

@@ -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 {

View File

@@ -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);

View File

@@ -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,
)?;