Improve xBestIndex call site and allow for proper handling of join and where constraints

This commit is contained in:
PThorpe92 2025-04-17 13:15:20 -04:00
parent 245e7f94f6
commit a25a02efe1
No known key found for this signature in database
GPG key ID: 66DB3FBACBDD05CC
3 changed files with 106 additions and 84 deletions

View file

@ -21,8 +21,8 @@ use super::{
optimizer::Optimizable,
order_by::{order_by_sorter_insert, sorter_insert},
plan::{
try_convert_to_constraint_info, IterationDirection, Operation, Search, SeekDef, SelectPlan,
SelectQueryType, TableReference, WhereTerm,
convert_where_to_vtab_constraint, IterationDirection, Operation, Search, SeekDef,
SelectPlan, SelectQueryType, TableReference, WhereTerm,
},
};
@ -291,31 +291,33 @@ pub fn open_loop(
});
}
} else if let Some(vtab) = table.virtual_table() {
// Virtual tables may be used either as VTab or TVF
let (start_reg, count, maybe_idx_str, maybe_idx_int) = if vtab
.kind
.eq(&VTabKind::VirtualTable)
{
// Build converted constraints from the predicates.
let mut converted_constraints = Vec::with_capacity(predicates.len());
for (i, pred) in predicates.iter().enumerate() {
if let Some(cinfo) =
try_convert_to_constraint_info(pred, table_index, i)
{
converted_constraints.push((cinfo, pred));
}
}
let constraints: Vec<_> =
converted_constraints.iter().map(|(c, _)| *c).collect();
let order_by = [OrderByInfo {
column_index: *t_ctx
.result_column_indexes_in_orderby_sorter
.first()
.unwrap_or(&0) as u32,
desc: matches!(iter_dir, IterationDirection::Backwards),
}];
// Call xBestIndex method on the underlying vtable.
let index_info = vtab.best_index(&constraints, &order_by);
// Virtualtable (nonTVF) modules can receive constraints via xBestIndex.
// They return information with which to pass to VFilter operation.
// We forward every predicate that touches vtab columns.
//
// vtab.col = literal (always usable)
// vtab.col = outer_table.col (usable, because outer_table is already positioned)
// vtab.col = later_table.col (forwarded with usable = false)
//
// xBestIndex decides which ones it wants by setting argvIndex and whether the
// core layer may omit them (omit = true).
// We then materialise the RHS/LHS into registers before issuing VFilter.
let converted_constraints = predicates
.iter()
.filter(|p| p.should_eval_at_loop(table_index))
.enumerate()
.filter_map(|(i, p)| {
// Build ConstraintInfo from the predicates
convert_where_to_vtab_constraint(p, table_index, i)
})
.collect::<Vec<_>>();
// TODO: get proper order_by information to pass to the vtab.
// maybe encode more info on t_ctx? we need: [col_idx, is_descending]
let index_info = vtab.best_index(&converted_constraints, &[]);
// Determine the number of VFilter arguments (constraints with an argv_index).
let args_needed = index_info
@ -328,13 +330,12 @@ pub fn open_loop(
// For each constraint used by best_index, translate the opposite side.
for (i, usage) in index_info.constraint_usages.iter().enumerate() {
if let Some(argv_index) = usage.argv_index {
if let Some((_, pred)) = converted_constraints.get(i) {
if let ast::Expr::Binary(lhs, _, rhs) = &pred.expr {
let expr = match (&**lhs, &**rhs) {
(ast::Expr::Column { .. }, lit) => lit,
(lit, ast::Expr::Column { .. }) => lit,
_ => continue,
};
if let Some(cinfo) = converted_constraints.get(i) {
let (pred_idx, is_rhs) = cinfo.unpack_plan_info();
if let ast::Expr::Binary(lhs, _, rhs) =
&predicates[pred_idx].expr
{
let expr = if is_rhs { rhs } else { lhs };
// argv_index is 1-based; adjust to get the proper register offset.
let target_reg = start_reg + (argv_index - 1) as usize;
translate_expr(
@ -344,6 +345,9 @@ pub fn open_loop(
target_reg,
&t_ctx.resolver,
)?;
if cinfo.usable && usage.omit {
t_ctx.omit_predicates.push(pred_idx)
}
}
}
}
@ -359,15 +363,6 @@ pub fn open_loop(
} else {
None
};
// Record (in t_ctx) the indices of predicates that best_index tells us to omit.
// Here we insert directly into t_ctx.omit_predicates
for (j, usage) in index_info.constraint_usages.iter().enumerate() {
if usage.argv_index.is_some() && usage.omit {
if let Some(constraint) = constraints.get(j) {
t_ctx.omit_predicates.push(constraint.pred_idx);
}
}
}
(
start_reg,
args_needed,

View file

@ -97,6 +97,18 @@ fn reverse_operator(op: &Operator) -> Option<Operator> {
}
}
fn to_ext_constraint_op(op: &Operator) -> Option<ConstraintOp> {
match op {
Operator::Equals => Some(ConstraintOp::Eq),
Operator::Less => Some(ConstraintOp::Lt),
Operator::LessEquals => Some(ConstraintOp::Le),
Operator::Greater => Some(ConstraintOp::Gt),
Operator::GreaterEquals => Some(ConstraintOp::Ge),
Operator::NotEquals => Some(ConstraintOp::Ne),
_ => None,
}
}
/// This function takes a WhereTerm for a select involving a VTab at index 'table_index'.
/// It determines whether or not it involves the given table and whether or not it can
/// be converted into a ConstraintInfo which can be passed to the vtab module's xBestIndex
@ -106,7 +118,7 @@ fn reverse_operator(op: &Operator) -> Option<Operator> {
/// 'SELECT key, value FROM vtab WHERE key = 'some_key';
/// we need to send the OwnedValue('some_key') as an argument to VFilter, and possibly omit it from
/// the filtration in the vdbe layer.
pub fn try_convert_to_constraint_info(
pub fn convert_where_to_vtab_constraint(
term: &WhereTerm,
table_index: usize,
pred_idx: usize,
@ -114,53 +126,63 @@ pub fn try_convert_to_constraint_info(
if term.from_outer_join {
return None;
}
let Expr::Binary(lhs, op, rhs) = &term.expr else {
return None;
};
let (col_expr, _, op) = match (&**lhs, &**rhs) {
(Expr::Column { table, .. }, rhs) if can_pushdown_predicate(rhs) => {
if table != &table_index {
return None;
let expr_is_ready = |e: &Expr| -> bool { can_pushdown_predicate(e, table_index) };
let (vcol_idx, op_for_vtab, usable, is_rhs) = match (&**lhs, &**rhs) {
(
Expr::Column {
table: tbl_l,
column: col_l,
..
},
Expr::Column {
table: tbl_r,
column: col_r,
..
},
) => {
// one side must be the virtual table
let vtab_on_l = *tbl_l == table_index;
let vtab_on_r = *tbl_r == table_index;
if vtab_on_l == vtab_on_r {
return None; // either both or none -> not convertible
}
(lhs, rhs, op)
}
(lhs, Expr::Column { table, .. }) if can_pushdown_predicate(lhs) => {
if table != &table_index {
return None;
if vtab_on_l {
// vtab on left side: operator unchanged
let usable = *tbl_r < table_index; // usable if the other table is already positioned
(col_l, op, usable, false)
} else {
// vtab on right side of the expr: reverse operator
let usable = *tbl_l < table_index;
(col_r, &reverse_operator(op).unwrap_or(*op), usable, true)
}
// if the column is on the rhs, swap the operands and possibly
// the operator if it's a logical comparison.
(rhs, lhs, &reverse_operator(op).unwrap_or(*op))
}
_ => {
return None;
(Expr::Column { table, column, .. }, other) if *table == table_index => {
(
column,
op,
expr_is_ready(other), // literal / earliertable / deterministic func ?
false,
)
}
};
(other, Expr::Column { table, column, .. }) if *table == table_index => (
column,
&reverse_operator(op).unwrap_or(*op),
expr_is_ready(other),
true,
),
let Expr::Column { column, .. } = **col_expr else {
return None;
};
let column_index = column as u32;
let constraint_op = match op {
Operator::Equals => ConstraintOp::Eq,
Operator::Less => ConstraintOp::Lt,
Operator::LessEquals => ConstraintOp::Le,
Operator::Greater => ConstraintOp::Gt,
Operator::GreaterEquals => ConstraintOp::Ge,
Operator::NotEquals => ConstraintOp::Ne,
Operator::Is => ConstraintOp::Is,
Operator::IsNot => ConstraintOp::IsNot,
_ => return None,
_ => return None, // does not involve the virtual table at all
};
Some(ConstraintInfo {
column_index,
op: constraint_op,
usable: true,
pred_idx,
column_index: *vcol_idx as u32,
op: to_ext_constraint_op(op_for_vtab)?,
usable,
plan_info: ConstraintInfo::pack_plan_info(pred_idx as u32, is_rhs),
})
}
/// The loop index where to evaluate the condition.

View file

@ -568,12 +568,15 @@ pub fn columns_from_create_table_body(body: &ast::CreateTableBody) -> crate::Res
/// This function checks if a given expression is a constant value that can be pushed down to the database engine.
/// It is expected to be called with the other half of a binary expression with an Expr::Column
pub fn can_pushdown_predicate(expr: &Expr) -> bool {
pub fn can_pushdown_predicate(expr: &Expr, table_idx: usize) -> bool {
match expr {
Expr::Literal(_) => true,
Expr::Binary(lhs, _, rhs) => can_pushdown_predicate(lhs) && can_pushdown_predicate(rhs),
Expr::Parenthesized(exprs) => can_pushdown_predicate(exprs.first().unwrap()),
Expr::Unary(_, expr) => can_pushdown_predicate(expr),
Expr::Column { table, .. } => *table <= table_idx,
Expr::Binary(lhs, _, rhs) => {
can_pushdown_predicate(lhs, table_idx) && can_pushdown_predicate(rhs, table_idx)
}
Expr::Parenthesized(exprs) => can_pushdown_predicate(exprs.first().unwrap(), table_idx),
Expr::Unary(_, expr) => can_pushdown_predicate(expr, table_idx),
Expr::FunctionCall { args, name, .. } => {
let function = crate::function::Func::resolve_function(
&name.0,
@ -582,13 +585,15 @@ pub fn can_pushdown_predicate(expr: &Expr) -> bool {
// is deterministic
matches!(function, Ok(Func::Scalar(_)))
}
Expr::Like { lhs, rhs, .. } => can_pushdown_predicate(lhs) && can_pushdown_predicate(rhs),
Expr::Like { lhs, rhs, .. } => {
can_pushdown_predicate(lhs, table_idx) && can_pushdown_predicate(rhs, table_idx)
}
Expr::Between {
lhs, start, end, ..
} => {
can_pushdown_predicate(lhs)
&& can_pushdown_predicate(start)
&& can_pushdown_predicate(end)
can_pushdown_predicate(lhs, table_idx)
&& can_pushdown_predicate(start, table_idx)
&& can_pushdown_predicate(end, table_idx)
}
_ => false,
}