Files
turso/core/translate/main_loop.rs
2025-10-27 16:01:39 +02:00

1581 lines
66 KiB
Rust

use turso_parser::ast::{fmt::ToTokens, SortOrder};
use std::sync::Arc;
use super::{
aggregation::{translate_aggregation_step, AggArgumentSource},
display::PlanContext,
emitter::{OperationMode, TranslateCtx},
expr::{
translate_condition_expr, translate_expr, translate_expr_no_constant_opt,
ConditionMetadata, NoConstantOptReason,
},
group_by::{group_by_agg_phase, GroupByMetadata, GroupByRowSource},
optimizer::Optimizable,
order_by::{order_by_sorter_insert, sorter_insert},
plan::{
Aggregate, GroupBy, IterationDirection, JoinOrderMember, Operation, QueryDestination,
Search, SeekDef, SelectPlan, TableReferences, WhereTerm,
},
};
use crate::translate::{
collate::get_collseq_from_expr,
emitter::UpdateRowSource,
plan::{EvalAt, NonFromClauseSubquery},
subquery::emit_non_from_clause_subquery,
window::emit_window_loop_source,
};
use crate::{
schema::{Affinity, Index, IndexColumn, Table},
translate::{
emitter::prepare_cdc_if_necessary,
plan::{DistinctCtx, Distinctness, Scan, SeekKeyComponent},
result_row::emit_select_result,
},
types::SeekOp,
vdbe::{
builder::{CursorKey, CursorType, ProgramBuilder},
insn::{CmpInsFlags, IdxInsertFlags, Insn},
BranchOffset, CursorID,
},
Result,
};
// Metadata for handling LEFT JOIN operations
#[derive(Debug)]
pub struct LeftJoinMetadata {
// integer register that holds a flag that is set to true if the current row has a match for the left join
pub reg_match_flag: usize,
// label for the instruction that sets the match flag to true
pub label_match_flag_set_true: BranchOffset,
// label for the instruction that checks if the match flag is true
pub label_match_flag_check_value: BranchOffset,
}
/// Jump labels for each loop in the query's main execution loop
#[derive(Debug, Clone, Copy)]
pub struct LoopLabels {
/// jump to the start of the loop body
pub loop_start: BranchOffset,
/// jump to the Next instruction (or equivalent)
pub next: BranchOffset,
/// jump to the end of the loop, exiting it
pub loop_end: BranchOffset,
}
impl LoopLabels {
pub fn new(program: &mut ProgramBuilder) -> Self {
Self {
loop_start: program.allocate_label(),
next: program.allocate_label(),
loop_end: program.allocate_label(),
}
}
}
pub fn init_distinct(program: &mut ProgramBuilder, plan: &SelectPlan) -> Result<DistinctCtx> {
let index_name = format!("distinct_{}", program.offset().as_offset_int()); // we don't really care about the name that much, just enough that we don't get name collisions
let mut columns = plan
.result_columns
.iter()
.enumerate()
.map(|(i, col)| IndexColumn {
name: col
.expr
.displayer(&PlanContext(&[&plan.table_references]))
.to_string(),
order: SortOrder::Asc,
pos_in_table: i,
collation: None,
default: None,
})
.collect::<Vec<_>>();
for (i, column) in columns.iter_mut().enumerate() {
column.collation =
get_collseq_from_expr(&plan.result_columns[i].expr, &plan.table_references)?;
}
let index = Arc::new(Index {
name: index_name.clone(),
table_name: String::new(),
ephemeral: true,
root_page: 0,
columns,
unique: false,
has_rowid: false,
where_clause: None,
});
let cursor_id = program.alloc_cursor_id(CursorType::BTreeIndex(index.clone()));
let ctx = DistinctCtx {
cursor_id,
ephemeral_index_name: index_name,
label_on_conflict: program.allocate_label(),
};
program.emit_insn(Insn::OpenEphemeral {
cursor_id,
is_table: false,
});
Ok(ctx)
}
/// Initialize resources needed for the source operators (tables, joins, etc)
#[allow(clippy::too_many_arguments)]
pub fn init_loop(
program: &mut ProgramBuilder,
t_ctx: &mut TranslateCtx,
tables: &TableReferences,
aggregates: &mut [Aggregate],
group_by: Option<&GroupBy>,
mode: OperationMode,
where_clause: &[WhereTerm],
join_order: &[JoinOrderMember],
subqueries: &mut [NonFromClauseSubquery],
) -> Result<()> {
assert!(
t_ctx.meta_left_joins.len() == tables.joined_tables().len(),
"meta_left_joins length does not match tables length"
);
if matches!(
&mode,
OperationMode::INSERT | OperationMode::UPDATE { .. } | OperationMode::DELETE
) {
assert!(tables.joined_tables().len() == 1);
let changed_table = &tables.joined_tables()[0].table;
let prepared =
prepare_cdc_if_necessary(program, t_ctx.resolver.schema, changed_table.get_name())?;
if let Some((cdc_cursor_id, _)) = prepared {
t_ctx.cdc_cursor_id = Some(cdc_cursor_id);
}
}
// Initialize ephemeral indexes for distinct aggregates
for (i, agg) in aggregates
.iter_mut()
.enumerate()
.filter(|(_, agg)| agg.is_distinct())
{
assert!(
agg.args.len() == 1,
"DISTINCT aggregate functions must have exactly one argument"
);
let index_name = format!(
"distinct_agg_{}_{}",
i,
agg.args[0].displayer(&PlanContext(&[tables]))
);
let index = Arc::new(Index {
name: index_name.clone(),
table_name: String::new(),
ephemeral: true,
root_page: 0,
columns: vec![IndexColumn {
name: agg.args[0].displayer(&PlanContext(&[tables])).to_string(),
order: SortOrder::Asc,
pos_in_table: 0,
collation: get_collseq_from_expr(&agg.original_expr, tables)?,
default: None, // FIXME: this should be inferred from the expression
}],
has_rowid: false,
unique: false,
where_clause: None,
});
let cursor_id = program.alloc_cursor_id(CursorType::BTreeIndex(index.clone()));
if group_by.is_none() {
// In GROUP BY, the ephemeral index is reinitialized for every group
// in the clear accumulator subroutine, so we only do it here if there is no GROUP BY.
program.emit_insn(Insn::OpenEphemeral {
cursor_id,
is_table: false,
});
}
agg.distinctness = Distinctness::Distinct {
ctx: Some(DistinctCtx {
cursor_id,
ephemeral_index_name: index_name,
label_on_conflict: program.allocate_label(),
}),
};
}
for (table_index, table) in tables.joined_tables().iter().enumerate() {
// Initialize bookkeeping for OUTER JOIN
if let Some(join_info) = table.join_info.as_ref() {
if join_info.outer {
let lj_metadata = LeftJoinMetadata {
reg_match_flag: program.alloc_register(),
label_match_flag_set_true: program.allocate_label(),
label_match_flag_check_value: program.allocate_label(),
};
t_ctx.meta_left_joins[table_index] = Some(lj_metadata);
}
}
let (table_cursor_id, index_cursor_id) =
table.open_cursors(program, mode.clone(), t_ctx.resolver.schema)?;
match &table.op {
Operation::Scan(Scan::BTreeTable { index, .. }) => match (&mode, &table.table) {
(OperationMode::SELECT, Table::BTree(btree)) => {
let root_page = btree.root_page;
if let Some(cursor_id) = table_cursor_id {
program.emit_insn(Insn::OpenRead {
cursor_id,
root_page,
db: table.database_id,
});
}
if let Some(index_cursor_id) = index_cursor_id {
program.emit_insn(Insn::OpenRead {
cursor_id: index_cursor_id,
root_page: index.as_ref().unwrap().root_page,
db: table.database_id,
});
}
}
(OperationMode::DELETE, Table::BTree(btree)) => {
let root_page = btree.root_page;
program.emit_insn(Insn::OpenWrite {
cursor_id: table_cursor_id
.expect("table cursor is always opened in OperationMode::DELETE"),
root_page: root_page.into(),
db: table.database_id,
});
if let Some(index_cursor_id) = index_cursor_id {
program.emit_insn(Insn::OpenWrite {
cursor_id: index_cursor_id,
root_page: index.as_ref().unwrap().root_page.into(),
db: table.database_id,
});
}
// For delete, we need to open all the other indexes too for writing
if let Some(indexes) = t_ctx.resolver.schema.indexes.get(&btree.name) {
for index in indexes {
if table
.op
.index()
.is_some_and(|table_index| table_index.name == index.name)
{
continue;
}
let cursor_id = program.alloc_cursor_id_keyed(
CursorKey::index(table.internal_id, index.clone()),
CursorType::BTreeIndex(index.clone()),
);
program.emit_insn(Insn::OpenWrite {
cursor_id,
root_page: index.root_page.into(),
db: table.database_id,
});
}
}
}
(OperationMode::UPDATE(update_mode), Table::BTree(btree)) => {
let root_page = btree.root_page;
match &update_mode {
UpdateRowSource::Normal => {
program.emit_insn(Insn::OpenWrite {
cursor_id: table_cursor_id.expect(
"table cursor is always opened in OperationMode::UPDATE",
),
root_page: root_page.into(),
db: table.database_id,
});
}
UpdateRowSource::PrebuiltEphemeralTable { target_table, .. } => {
let target_table_cursor_id = program
.resolve_cursor_id(&CursorKey::table(target_table.internal_id));
program.emit_insn(Insn::OpenWrite {
cursor_id: target_table_cursor_id,
root_page: target_table.btree().unwrap().root_page.into(),
db: table.database_id,
});
}
}
if let Some(index_cursor_id) = index_cursor_id {
program.emit_insn(Insn::OpenWrite {
cursor_id: index_cursor_id,
root_page: index.as_ref().unwrap().root_page.into(),
db: table.database_id,
});
}
}
_ => {}
},
Operation::Scan(Scan::VirtualTable { .. }) => {
if let Table::Virtual(tbl) = &table.table {
let is_write = matches!(
mode,
OperationMode::INSERT
| OperationMode::UPDATE { .. }
| OperationMode::DELETE
);
if is_write && tbl.readonly() {
return Err(crate::LimboError::ReadOnly);
}
if let Some(cursor_id) = table_cursor_id {
program.emit_insn(Insn::VOpen { cursor_id });
}
}
}
Operation::Scan(_) => {}
Operation::Search(search) => {
match mode {
OperationMode::SELECT => {
if let Some(table_cursor_id) = table_cursor_id {
program.emit_insn(Insn::OpenRead {
cursor_id: table_cursor_id,
root_page: table.table.get_root_page(),
db: table.database_id,
});
}
}
OperationMode::DELETE | OperationMode::UPDATE { .. } => {
let table_cursor_id = table_cursor_id.expect(
"table cursor is always opened in OperationMode::DELETE or OperationMode::UPDATE",
);
program.emit_insn(Insn::OpenWrite {
cursor_id: table_cursor_id,
root_page: table.table.get_root_page().into(),
db: table.database_id,
});
// For DELETE, we need to open all the indexes for writing
// UPDATE opens these in emit_program_for_update() separately
if matches!(mode, OperationMode::DELETE) {
if let Some(indexes) =
t_ctx.resolver.schema.indexes.get(table.table.get_name())
{
for index in indexes {
if table
.op
.index()
.is_some_and(|table_index| table_index.name == index.name)
{
continue;
}
let cursor_id = program.alloc_cursor_id_keyed(
CursorKey::index(table.internal_id, index.clone()),
CursorType::BTreeIndex(index.clone()),
);
program.emit_insn(Insn::OpenWrite {
cursor_id,
root_page: index.root_page.into(),
db: table.database_id,
});
}
}
}
}
_ => {
unimplemented!()
}
}
if let Search::Seek {
index: Some(index), ..
} = search
{
// Ephemeral index cursor are opened ad-hoc when needed.
if !index.ephemeral {
match mode {
OperationMode::SELECT => {
program.emit_insn(Insn::OpenRead {
cursor_id: index_cursor_id
.expect("index cursor is always opened in Seek with index"),
root_page: index.root_page,
db: table.database_id,
});
}
OperationMode::UPDATE { .. } | OperationMode::DELETE => {
program.emit_insn(Insn::OpenWrite {
cursor_id: index_cursor_id
.expect("index cursor is always opened in Seek with index"),
root_page: index.root_page.into(),
db: table.database_id,
});
}
_ => {
unimplemented!()
}
}
}
}
}
}
}
for subquery in subqueries.iter_mut().filter(|s| !s.has_been_evaluated()) {
let eval_at = subquery.get_eval_at(join_order)?;
if eval_at != EvalAt::BeforeLoop {
continue;
}
let plan = subquery.consume_plan(EvalAt::BeforeLoop);
emit_non_from_clause_subquery(
program,
t_ctx,
*plan,
&subquery.query_type,
subquery.correlated,
)?;
}
for cond in where_clause
.iter()
.filter(|c| c.should_eval_before_loop(join_order, subqueries))
{
let jump_target = program.allocate_label();
let meta = ConditionMetadata {
jump_if_condition_is_true: false,
jump_target_when_true: jump_target,
jump_target_when_false: t_ctx.label_main_loop_end.unwrap(),
jump_target_when_null: t_ctx.label_main_loop_end.unwrap(),
};
translate_condition_expr(program, tables, &cond.expr, meta, &t_ctx.resolver)?;
program.preassign_label_to_next_insn(jump_target);
}
Ok(())
}
/// Set up the main query execution loop
/// For example in the case of a nested table scan, this means emitting the Rewind instruction
/// for all tables involved, outermost first.
#[allow(clippy::too_many_arguments)]
pub fn open_loop(
program: &mut ProgramBuilder,
t_ctx: &mut TranslateCtx,
table_references: &TableReferences,
join_order: &[JoinOrderMember],
predicates: &[WhereTerm],
temp_cursor_id: Option<CursorID>,
mode: OperationMode,
subqueries: &mut [NonFromClauseSubquery],
) -> Result<()> {
for (join_index, join) in join_order.iter().enumerate() {
let joined_table_index = join.original_idx;
let table = &table_references.joined_tables()[joined_table_index];
let LoopLabels {
loop_start,
loop_end,
next,
} = *t_ctx
.labels_main_loop
.get(joined_table_index)
.expect("table has no loop labels");
// Each OUTER JOIN has a "match flag" that is initially set to false,
// and is set to true when a match is found for the OUTER JOIN.
// This is used to determine whether to emit actual columns or NULLs for the columns of the right table.
if let Some(join_info) = table.join_info.as_ref() {
if join_info.outer {
let lj_meta = t_ctx.meta_left_joins[joined_table_index].as_ref().unwrap();
program.emit_insn(Insn::Integer {
value: 0,
dest: lj_meta.reg_match_flag,
});
}
}
let (table_cursor_id, index_cursor_id) = table.resolve_cursors(program, mode.clone())?;
match &table.op {
Operation::Scan(scan) => {
match (scan, &table.table) {
(Scan::BTreeTable { iter_dir, .. }, Table::BTree(_)) => {
let iteration_cursor_id = temp_cursor_id.unwrap_or_else(|| {
index_cursor_id.unwrap_or_else(|| {
table_cursor_id.expect(
"Either ephemeral or index or table cursor must be opened",
)
})
});
if *iter_dir == IterationDirection::Backwards {
program.emit_insn(Insn::Last {
cursor_id: iteration_cursor_id,
pc_if_empty: loop_end,
});
} else {
program.emit_insn(Insn::Rewind {
cursor_id: iteration_cursor_id,
pc_if_empty: loop_end,
});
}
program.preassign_label_to_next_insn(loop_start);
}
(
Scan::VirtualTable {
idx_num,
idx_str,
constraints,
},
Table::Virtual(_),
) => {
let (start_reg, count, maybe_idx_str, maybe_idx_int) = {
let args_needed = constraints.len();
let start_reg = program.alloc_registers(args_needed);
for (argv_index, expr) in constraints.iter().enumerate() {
let target_reg = start_reg + argv_index;
translate_expr(
program,
Some(table_references),
expr,
target_reg,
&t_ctx.resolver,
)?;
}
// If best_index provided an idx_str, translate it.
let maybe_idx_str = if let Some(idx_str) = idx_str {
let reg = program.alloc_register();
program.emit_insn(Insn::String8 {
dest: reg,
value: idx_str.to_owned(),
});
Some(reg)
} else {
None
};
(start_reg, args_needed, maybe_idx_str, Some(*idx_num))
};
// Emit VFilter with the computed arguments.
program.emit_insn(Insn::VFilter {
cursor_id: table_cursor_id
.expect("Virtual tables do not support covering indexes"),
arg_count: count,
args_reg: start_reg,
idx_str: maybe_idx_str,
idx_num: maybe_idx_int.unwrap_or(0) as usize,
pc_if_empty: loop_end,
});
program.preassign_label_to_next_insn(loop_start);
}
(Scan::Subquery, Table::FromClauseSubquery(from_clause_subquery)) => {
let (yield_reg, coroutine_implementation_start) =
match &from_clause_subquery.plan.query_destination {
QueryDestination::CoroutineYield {
yield_reg,
coroutine_implementation_start,
} => (*yield_reg, *coroutine_implementation_start),
_ => unreachable!("Subquery table with non-subquery query type"),
};
// In case the subquery is an inner loop, it needs to be reinitialized on each iteration of the outer loop.
program.emit_insn(Insn::InitCoroutine {
yield_reg,
jump_on_definition: BranchOffset::Offset(0),
start_offset: coroutine_implementation_start,
});
program.preassign_label_to_next_insn(loop_start);
// A subquery within the main loop of a parent query has no cursor, so instead of advancing the cursor,
// it emits a Yield which jumps back to the main loop of the subquery itself to retrieve the next row.
// When the subquery coroutine completes, this instruction jumps to the label at the top of the termination_label_stack,
// which in this case is the end of the Yield-Goto loop in the parent query.
program.emit_insn(Insn::Yield {
yield_reg,
end_offset: loop_end,
});
}
_ => unreachable!(
"{:?} scan cannot be used with {:?} table",
scan, table.table
),
}
if let Some(table_cursor_id) = table_cursor_id {
if let Some(index_cursor_id) = index_cursor_id {
program.emit_insn(Insn::DeferredSeek {
index_cursor_id,
table_cursor_id,
});
}
}
}
Operation::Search(search) => {
assert!(
!matches!(table.table, Table::FromClauseSubquery(_)),
"Subqueries do not support index seeks"
);
// Open the loop for the index search.
// Rowid equality point lookups are handled with a SeekRowid instruction which does not loop, since it is a single row lookup.
match search {
Search::RowidEq { cmp_expr } => {
let src_reg = program.alloc_register();
translate_expr(
program,
Some(table_references),
cmp_expr,
src_reg,
&t_ctx.resolver,
)?;
program.emit_insn(Insn::SeekRowid {
cursor_id: table_cursor_id
.expect("Search::RowidEq requires a table cursor"),
src_reg,
target_pc: next,
});
}
Search::Seek { index, .. } => {
// Otherwise, it's an index/rowid scan, i.e. first a seek is performed and then a scan until the comparison expression is not satisfied anymore.
if let Some(index) = index {
if index.ephemeral {
let table_has_rowid = if let Table::BTree(btree) = &table.table {
btree.has_rowid
} else {
false
};
let _ = emit_autoindex(
program,
index,
table_cursor_id.expect(
"an ephemeral index must have a source table cursor",
),
index_cursor_id
.expect("an ephemeral index must have an index cursor"),
table_has_rowid,
)?;
}
}
let seek_cursor_id = temp_cursor_id.unwrap_or_else(|| {
index_cursor_id.unwrap_or_else(|| {
table_cursor_id.expect(
"Either ephemeral or index or table cursor must be opened",
)
})
});
let Search::Seek { seek_def, .. } = search else {
unreachable!(
"Rowid equality point lookup should have been handled above"
);
};
let max_registers = seek_def
.size(&seek_def.start)
.max(seek_def.size(&seek_def.end));
let start_reg = program.alloc_registers(max_registers);
emit_seek(
program,
table_references,
seek_def,
t_ctx,
seek_cursor_id,
start_reg,
loop_end,
index.as_ref(),
)?;
emit_seek_termination(
program,
table_references,
seek_def,
t_ctx,
seek_cursor_id,
start_reg,
loop_start,
loop_end,
index.as_ref(),
)?;
if let Some(index_cursor_id) = index_cursor_id {
if let Some(table_cursor_id) = table_cursor_id {
// Don't do a btree table seek until it's actually necessary to read from the table.
program.emit_insn(Insn::DeferredSeek {
index_cursor_id,
table_cursor_id,
});
}
}
}
}
}
}
for subquery in subqueries.iter_mut().filter(|s| !s.has_been_evaluated()) {
assert!(subquery.correlated, "subquery must be correlated");
let eval_at = subquery.get_eval_at(join_order)?;
if eval_at != EvalAt::Loop(join_index) {
continue;
}
let plan = subquery.consume_plan(eval_at);
emit_non_from_clause_subquery(
program,
t_ctx,
*plan,
&subquery.query_type,
subquery.correlated,
)?;
}
// First emit outer join conditions, if any.
emit_conditions(
program,
&t_ctx,
table_references,
join_order,
predicates,
join_index,
next,
true,
subqueries,
)?;
// Set the match flag to true if this is a LEFT JOIN.
// At this point of execution we are going to emit columns for the left table,
// and either emit columns or NULLs for the right table, depending on whether the null_flag is set
// for the right table's cursor.
if let Some(join_info) = table.join_info.as_ref() {
if join_info.outer {
let lj_meta = t_ctx.meta_left_joins[joined_table_index].as_ref().unwrap();
program.resolve_label(lj_meta.label_match_flag_set_true, program.offset());
program.emit_insn(Insn::Integer {
value: 1,
dest: lj_meta.reg_match_flag,
});
}
}
// Now we can emit conditions from the WHERE clause.
// If the right table produces a NULL row, control jumps to the point where the match flag is set.
// The WHERE clause conditions may reference columns from that row, so they cannot be emitted
// before the flag is set — the row may be filtered out by the WHERE clause.
emit_conditions(
program,
&t_ctx,
table_references,
join_order,
predicates,
join_index,
next,
false,
subqueries,
)?;
}
if subqueries.iter().any(|s| !s.has_been_evaluated()) {
crate::bail_parse_error!(
"all subqueries should have already been emitted, but found {} unevaluated subqueries",
subqueries
.iter()
.filter(|s| !s.has_been_evaluated())
.count()
);
}
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn emit_conditions(
program: &mut ProgramBuilder,
t_ctx: &&mut TranslateCtx,
table_references: &TableReferences,
join_order: &[JoinOrderMember],
predicates: &[WhereTerm],
join_index: usize,
next: BranchOffset,
from_outer_join: bool,
subqueries: &[NonFromClauseSubquery],
) -> Result<()> {
for cond in predicates
.iter()
.filter(|cond| cond.from_outer_join.is_some() == from_outer_join)
.filter(|cond| cond.should_eval_at_loop(join_index, join_order, subqueries))
{
let jump_target_when_true = program.allocate_label();
let condition_metadata = ConditionMetadata {
jump_if_condition_is_true: false,
jump_target_when_true,
jump_target_when_false: next,
jump_target_when_null: next,
};
translate_condition_expr(
program,
table_references,
&cond.expr,
condition_metadata,
&t_ctx.resolver,
)?;
program.preassign_label_to_next_insn(jump_target_when_true);
}
Ok(())
}
/// SQLite (and so Turso) processes joins as a nested loop.
/// The loop may emit rows to various destinations depending on the query:
/// - a GROUP BY sorter (grouping is done by sorting based on the GROUP BY keys and aggregating while the GROUP BY keys match)
/// - a GROUP BY phase with no sorting (when the rows are already in the order required by the GROUP BY keys)
/// - an AggStep (the columns are collected for aggregation, which is finished later)
/// - a Window (rows are buffered and returned according to the rules of the window definition)
/// - an ORDER BY sorter (when there is none of the above, but there is an ORDER BY)
/// - a QueryResult (there is none of the above, so the loop either emits a ResultRow, or if it's a subquery, yields to the parent query)
enum LoopEmitTarget {
GroupBy,
OrderBySorter,
AggStep,
Window,
QueryResult,
}
/// Emits the bytecode for the inner loop of a query.
/// At this point the cursors for all tables have been opened and rewound.
pub fn emit_loop(
program: &mut ProgramBuilder,
t_ctx: &mut TranslateCtx,
plan: &SelectPlan,
) -> Result<()> {
// if we have a group by, we emit a record into the group by sorter,
// or if the rows are already sorted, we do the group by aggregation phase directly.
if plan.group_by.is_some() {
return emit_loop_source(program, t_ctx, plan, LoopEmitTarget::GroupBy);
}
// if we DONT have a group by, but we have aggregates, we emit without ResultRow.
// we also do not need to sort because we are emitting a single row.
if !plan.aggregates.is_empty() {
return emit_loop_source(program, t_ctx, plan, LoopEmitTarget::AggStep);
}
// Window processing is planned so that the query plan has neither GROUP BY nor aggregates.
// If the original query contained them, they are pushed down into a subquery.
// Rows are buffered and returned according to the rules of the window definition.
if plan.window.is_some() {
return emit_loop_source(program, t_ctx, plan, LoopEmitTarget::Window);
}
// if NONE of the above applies, but we have an order by, we emit a record into the order by sorter.
if !plan.order_by.is_empty() {
return emit_loop_source(program, t_ctx, plan, LoopEmitTarget::OrderBySorter);
}
// if we have neither, we emit a ResultRow. In that case, if we have a Limit, we handle that with DecrJumpZero.
emit_loop_source(program, t_ctx, plan, LoopEmitTarget::QueryResult)
}
/// This is a helper function for inner_loop_emit,
/// which does a different thing depending on the emit target.
/// See the InnerLoopEmitTarget enum for more details.
fn emit_loop_source(
program: &mut ProgramBuilder,
t_ctx: &mut TranslateCtx,
plan: &SelectPlan,
emit_target: LoopEmitTarget,
) -> Result<()> {
match emit_target {
LoopEmitTarget::GroupBy => {
// This function either:
// - creates a sorter for GROUP BY operations by allocating registers and translating expressions for three types of columns:
// 1) GROUP BY columns (used as sorting keys)
// 2) non-aggregate, non-GROUP BY columns
// 3) aggregate function arguments
// - or if the rows produced by the loop are already sorted in the order required by the GROUP BY keys,
// the group by comparisons are done directly inside the main loop.
let aggregates = &plan.aggregates;
let GroupByMetadata {
row_source,
registers,
..
} = t_ctx.meta_group_by.as_ref().unwrap();
let start_reg = registers.reg_group_by_source_cols_start;
let mut cur_reg = start_reg;
// Collect all non-aggregate expressions in the following order:
// 1. GROUP BY expressions. These serve as sort keys.
// 2. Remaining non-aggregate expressions that are not in GROUP BY.
//
// Example:
// SELECT col1, col2, SUM(col3) FROM table GROUP BY col1
// - col1 is added first (from GROUP BY)
// - col2 is added second (non-aggregate, in SELECT, not in GROUP BY)
for (expr, _) in t_ctx.non_aggregate_expressions.iter() {
let key_reg = cur_reg;
cur_reg += 1;
translate_expr(
program,
Some(&plan.table_references),
expr,
key_reg,
&t_ctx.resolver,
)?;
}
// Step 2: Process arguments for all aggregate functions
// For each aggregate, translate all its argument expressions
for agg in aggregates.iter() {
// For a query like: SELECT group_col, SUM(val1), AVG(val2) FROM table GROUP BY group_col
// we'll process val1 and val2 here, storing them in the sorter so they're available
// when computing the aggregates after sorting by group_col
for expr in agg.args.iter() {
let agg_reg = cur_reg;
cur_reg += 1;
translate_expr(
program,
Some(&plan.table_references),
expr,
agg_reg,
&t_ctx.resolver,
)?;
}
}
match row_source {
GroupByRowSource::Sorter {
sort_cursor,
sorter_column_count,
reg_sorter_key,
..
} => {
sorter_insert(
program,
start_reg,
*sorter_column_count,
*sort_cursor,
*reg_sorter_key,
);
}
GroupByRowSource::MainLoop { .. } => group_by_agg_phase(program, t_ctx, plan)?,
}
Ok(())
}
LoopEmitTarget::OrderBySorter => {
order_by_sorter_insert(program, t_ctx, plan)?;
if let Distinctness::Distinct { ctx } = &plan.distinctness {
let distinct_ctx = ctx.as_ref().expect("distinct context must exist");
program.preassign_label_to_next_insn(distinct_ctx.label_on_conflict);
}
Ok(())
}
LoopEmitTarget::AggStep => {
let start_reg = t_ctx
.reg_agg_start
.expect("aggregate registers must be initialized");
// In planner.rs, we have collected all aggregates from the SELECT clause, including ones where the aggregate is embedded inside
// a more complex expression. Some examples: length(sum(x)), sum(x) + avg(y), sum(x) + 1, etc.
// The result of those more complex expressions depends on the final result of the aggregate, so we don't translate the complete expressions here.
// Instead, we accumulate the intermediate results of all aggreagates, and evaluate any expressions that do not contain aggregates.
for (i, agg) in plan.aggregates.iter().enumerate() {
let reg = start_reg + i;
translate_aggregation_step(
program,
&plan.table_references,
AggArgumentSource::new_from_expression(&agg.func, &agg.args, &agg.distinctness),
reg,
&t_ctx.resolver,
)?;
if let Distinctness::Distinct { ctx } = &agg.distinctness {
let ctx = ctx
.as_ref()
.expect("distinct aggregate context not populated");
program.preassign_label_to_next_insn(ctx.label_on_conflict);
}
}
let label_emit_nonagg_only_once = if let Some(flag) = t_ctx.reg_nonagg_emit_once_flag {
let if_label = program.allocate_label();
program.emit_insn(Insn::If {
reg: flag,
target_pc: if_label,
jump_if_null: false,
});
Some(if_label)
} else {
None
};
let col_start = t_ctx.reg_result_cols_start.unwrap();
// Process only non-aggregate columns
let non_agg_columns = plan
.result_columns
.iter()
.enumerate()
.filter(|(_, rc)| !rc.contains_aggregates);
for (i, rc) in non_agg_columns {
let reg = col_start + i;
translate_expr(
program,
Some(&plan.table_references),
&rc.expr,
reg,
&t_ctx.resolver,
)?;
}
if let Some(label) = label_emit_nonagg_only_once {
program.resolve_label(label, program.offset());
let flag = t_ctx.reg_nonagg_emit_once_flag.unwrap();
program.emit_int(1, flag);
}
Ok(())
}
LoopEmitTarget::QueryResult => {
assert!(
plan.aggregates.is_empty(),
"We should not get here with aggregates"
);
let offset_jump_to = t_ctx
.labels_main_loop
.first()
.map(|l| l.next)
.or(t_ctx.label_main_loop_end);
emit_select_result(
program,
&t_ctx.resolver,
plan,
t_ctx.label_main_loop_end,
offset_jump_to,
t_ctx.reg_nonagg_emit_once_flag,
t_ctx.reg_offset,
t_ctx.reg_result_cols_start.unwrap(),
t_ctx.limit_ctx,
)?;
if let Distinctness::Distinct { ctx } = &plan.distinctness {
let distinct_ctx = ctx.as_ref().expect("distinct context must exist");
program.preassign_label_to_next_insn(distinct_ctx.label_on_conflict);
}
Ok(())
}
LoopEmitTarget::Window => {
emit_window_loop_source(program, t_ctx, plan)?;
Ok(())
}
}
}
/// Closes the loop for a given source operator.
/// For example in the case of a nested table scan, this means emitting the Next instruction
/// for all tables involved, innermost first.
pub fn close_loop(
program: &mut ProgramBuilder,
t_ctx: &mut TranslateCtx,
tables: &TableReferences,
join_order: &[JoinOrderMember],
mode: OperationMode,
) -> Result<()> {
// We close the loops for all tables in reverse order, i.e. innermost first.
// OPEN t1
// OPEN t2
// OPEN t3
// <do stuff>
// CLOSE t3
// CLOSE t2
// CLOSE t1
for join in join_order.iter().rev() {
let table_index = join.original_idx;
let table = &tables.joined_tables()[table_index];
let loop_labels = *t_ctx
.labels_main_loop
.get(table_index)
.expect("source has no loop labels");
let (table_cursor_id, index_cursor_id) = table.resolve_cursors(program, mode.clone())?;
match &table.op {
Operation::Scan(scan) => {
program.resolve_label(loop_labels.next, program.offset());
match scan {
Scan::BTreeTable { iter_dir, .. } => {
let iteration_cursor_id = if let OperationMode::UPDATE(
UpdateRowSource::PrebuiltEphemeralTable {
ephemeral_table_cursor_id,
..
},
) = &mode
{
*ephemeral_table_cursor_id
} else {
index_cursor_id.unwrap_or_else(|| {
table_cursor_id.expect(
"Either ephemeral or index or table cursor must be opened",
)
})
};
if *iter_dir == IterationDirection::Backwards {
program.emit_insn(Insn::Prev {
cursor_id: iteration_cursor_id,
pc_if_prev: loop_labels.loop_start,
});
} else {
program.emit_insn(Insn::Next {
cursor_id: iteration_cursor_id,
pc_if_next: loop_labels.loop_start,
});
}
}
Scan::VirtualTable { .. } => {
program.emit_insn(Insn::VNext {
cursor_id: table_cursor_id
.expect("Virtual tables do not support covering indexes"),
pc_if_next: loop_labels.loop_start,
});
}
Scan::Subquery => {
// A subquery has no cursor to call Next on, so it just emits a Goto
// to the Yield instruction, which in turn jumps back to the main loop of the subquery,
// so that the next row from the subquery can be read.
program.emit_insn(Insn::Goto {
target_pc: loop_labels.loop_start,
});
}
}
program.preassign_label_to_next_insn(loop_labels.loop_end);
}
Operation::Search(search) => {
assert!(
!matches!(table.table, Table::FromClauseSubquery(_)),
"Subqueries do not support index seeks"
);
program.resolve_label(loop_labels.next, program.offset());
let iteration_cursor_id =
if let OperationMode::UPDATE(UpdateRowSource::PrebuiltEphemeralTable {
ephemeral_table_cursor_id,
..
}) = &mode
{
*ephemeral_table_cursor_id
} else {
index_cursor_id.unwrap_or_else(|| {
table_cursor_id
.expect("Either ephemeral or index or table cursor must be opened")
})
};
// Rowid equality point lookups are handled with a SeekRowid instruction which does not loop, so there is no need to emit a Next instruction.
if !matches!(search, Search::RowidEq { .. }) {
let iter_dir = match search {
Search::Seek { seek_def, .. } => seek_def.iter_dir,
Search::RowidEq { .. } => unreachable!(),
};
if iter_dir == IterationDirection::Backwards {
program.emit_insn(Insn::Prev {
cursor_id: iteration_cursor_id,
pc_if_prev: loop_labels.loop_start,
});
} else {
program.emit_insn(Insn::Next {
cursor_id: iteration_cursor_id,
pc_if_next: loop_labels.loop_start,
});
}
}
program.preassign_label_to_next_insn(loop_labels.loop_end);
}
}
// Handle OUTER JOIN logic. The reason this comes after the "loop end" mark is that we may need to still jump back
// and emit a row with NULLs for the right table, and then jump back to the next row of the left table.
if let Some(join_info) = table.join_info.as_ref() {
if join_info.outer {
let lj_meta = t_ctx.meta_left_joins[table_index].as_ref().unwrap();
// The left join match flag is set to 1 when there is any match on the right table
// (e.g. SELECT * FROM t1 LEFT JOIN t2 ON t1.a = t2.a).
// If the left join match flag has been set to 1, we jump to the next row on the outer table,
// i.e. continue to the next row of t1 in our example.
program.resolve_label(lj_meta.label_match_flag_check_value, program.offset());
let label_when_right_table_notnull = program.allocate_label();
program.emit_insn(Insn::IfPos {
reg: lj_meta.reg_match_flag,
target_pc: label_when_right_table_notnull,
decrement_by: 0,
});
// If the left join match flag is still 0, it means there was no match on the right table,
// but since it's a LEFT JOIN, we still need to emit a row with NULLs for the right table.
// In that case, we now enter the routine that does exactly that.
// First we set the right table cursor's "pseudo null bit" on, which means any Insn::Column will return NULL.
// This needs to be set for both the table and the index cursor, if present,
// since even if the iteration cursor is the index cursor, it might fetch values from the table cursor.
[table_cursor_id, index_cursor_id]
.iter()
.filter_map(|maybe_cursor_id| maybe_cursor_id.as_ref())
.for_each(|cursor_id| {
program.emit_insn(Insn::NullRow {
cursor_id: *cursor_id,
});
});
// Then we jump to setting the left join match flag to 1 again,
// but this time the right table cursor will set everything to null.
// This leads to emitting a row with cols from the left + nulls from the right,
// and we will end up back in the IfPos instruction above, which will then
// check the match flag again, and since it is now 1, we will jump to the
// next row in the left table.
program.emit_insn(Insn::Goto {
target_pc: lj_meta.label_match_flag_set_true,
});
program.preassign_label_to_next_insn(label_when_right_table_notnull);
}
}
}
Ok(())
}
/// Emits instructions for an index seek. See e.g. [crate::translate::plan::SeekDef]
/// for more details about the seek definition.
///
/// Index seeks always position the cursor to the first row that matches the seek key,
/// and then continue to emit rows until the termination condition is reached,
/// see [emit_seek_termination] below.
///
/// If either 1. the seek finds no rows or 2. the termination condition is reached,
/// the loop for that given table/index is fully exited.
#[allow(clippy::too_many_arguments)]
fn emit_seek(
program: &mut ProgramBuilder,
tables: &TableReferences,
seek_def: &SeekDef,
t_ctx: &mut TranslateCtx,
seek_cursor_id: usize,
start_reg: usize,
loop_end: BranchOffset,
seek_index: Option<&Arc<Index>>,
) -> Result<()> {
let is_index = seek_index.is_some();
if seek_def.prefix.is_empty() && matches!(seek_def.start.last_component, SeekKeyComponent::None)
{
// If there is no seek key, we start from the first or last row of the index,
// depending on the iteration direction.
//
// Also, if we will encounter NULLs in the index at the beginning of iteration (Forward + Asc OR Backward + Desc)
// then, we must explicitly skip them as seek always has some bound condition over indexed column (e.g. c < ?, c >= ?, ...)
//
// note that table seek has some rules to convert seek key values to the integer affinity + it has some logic to check for explicit NULL searches
// so, we emit simple Rewind/Last in case of search over table BTree
// (this is safe as table BTree keys are always non-null single integer)
match seek_def.iter_dir {
IterationDirection::Forwards => {
if seek_index.is_some_and(|index| index.columns[0].order == SortOrder::Asc) {
program.emit_null(start_reg, None);
program.emit_insn(Insn::SeekGT {
is_index,
cursor_id: seek_cursor_id,
start_reg,
num_regs: 1,
target_pc: loop_end,
});
} else {
program.emit_insn(Insn::Rewind {
cursor_id: seek_cursor_id,
pc_if_empty: loop_end,
});
}
}
IterationDirection::Backwards => {
if seek_index.is_some_and(|index| index.columns[0].order == SortOrder::Desc) {
program.emit_null(start_reg, None);
program.emit_insn(Insn::SeekLT {
is_index,
cursor_id: seek_cursor_id,
start_reg,
num_regs: 1,
target_pc: loop_end,
});
} else {
program.emit_insn(Insn::Last {
cursor_id: seek_cursor_id,
pc_if_empty: loop_end,
});
}
}
}
return Ok(());
};
// We allocated registers for the full index key, but our seek key might not use the full index key.
// See [crate::translate::optimizer::build_seek_def] for more details about in which cases we do and don't use the full index key.
for (i, key) in seek_def.iter(&seek_def.start).enumerate() {
let reg = start_reg + i;
match key {
SeekKeyComponent::Expr(expr) => {
translate_expr_no_constant_opt(
program,
Some(tables),
expr,
reg,
&t_ctx.resolver,
NoConstantOptReason::RegisterReuse,
)?;
// If the seek key column is not verifiably non-NULL, we need check whether it is NULL,
// and if so, jump to the loop end.
// This is to avoid returning rows for e.g. SELECT * FROM t WHERE t.x > NULL,
// which would erroneously return all rows from t, as NULL is lower than any non-NULL value in index key comparisons.
if !expr.is_nonnull(tables) {
program.emit_insn(Insn::IsNull {
reg,
target_pc: loop_end,
});
}
}
SeekKeyComponent::None => unreachable!("None component is not possible in iterator"),
}
}
let num_regs = seek_def.size(&seek_def.start);
match seek_def.start.op {
SeekOp::GE { eq_only } => program.emit_insn(Insn::SeekGE {
is_index,
cursor_id: seek_cursor_id,
start_reg,
num_regs,
target_pc: loop_end,
eq_only,
}),
SeekOp::GT => program.emit_insn(Insn::SeekGT {
is_index,
cursor_id: seek_cursor_id,
start_reg,
num_regs,
target_pc: loop_end,
}),
SeekOp::LE { eq_only } => program.emit_insn(Insn::SeekLE {
is_index,
cursor_id: seek_cursor_id,
start_reg,
num_regs,
target_pc: loop_end,
eq_only,
}),
SeekOp::LT => program.emit_insn(Insn::SeekLT {
is_index,
cursor_id: seek_cursor_id,
start_reg,
num_regs,
target_pc: loop_end,
}),
};
Ok(())
}
/// Emits instructions for an index seek termination. See e.g. [crate::translate::plan::SeekDef]
/// for more details about the seek definition.
///
/// Index seeks always position the cursor to the first row that matches the seek key
/// (see [emit_seek] above), and then continue to emit rows until the termination condition
/// (if any) is reached.
///
/// If the termination condition is not present, the cursor is fully scanned to the end.
#[allow(clippy::too_many_arguments)]
fn emit_seek_termination(
program: &mut ProgramBuilder,
tables: &TableReferences,
seek_def: &SeekDef,
t_ctx: &mut TranslateCtx,
seek_cursor_id: usize,
start_reg: usize,
loop_start: BranchOffset,
loop_end: BranchOffset,
seek_index: Option<&Arc<Index>>,
) -> Result<()> {
let is_index = seek_index.is_some();
if seek_def.prefix.is_empty() && matches!(seek_def.end.last_component, SeekKeyComponent::None) {
program.preassign_label_to_next_insn(loop_start);
// If we will encounter NULLs in the index at the end of iteration (Forward + Desc OR Backward + Asc)
// then, we must explicitly stop before them as seek always has some bound condition over indexed column (e.g. c < ?, c >= ?, ...)
match seek_def.iter_dir {
IterationDirection::Forwards => {
if seek_index.is_some_and(|index| index.columns[0].order == SortOrder::Desc) {
program.emit_null(start_reg, None);
program.emit_insn(Insn::IdxGE {
cursor_id: seek_cursor_id,
start_reg,
num_regs: 1,
target_pc: loop_end,
});
}
}
IterationDirection::Backwards => {
if seek_index.is_some_and(|index| index.columns[0].order == SortOrder::Asc) {
program.emit_null(start_reg, None);
program.emit_insn(Insn::IdxLE {
cursor_id: seek_cursor_id,
start_reg,
num_regs: 1,
target_pc: loop_end,
});
}
}
}
return Ok(());
};
// For all index key values apart from the last one, we are guaranteed to use the same values
// as these values were emited from common prefix, so we don't need to emit them again.
let num_regs = seek_def.size(&seek_def.end);
let last_reg = start_reg + seek_def.prefix.len();
match &seek_def.end.last_component {
SeekKeyComponent::Expr(expr) => {
translate_expr_no_constant_opt(
program,
Some(tables),
expr,
last_reg,
&t_ctx.resolver,
NoConstantOptReason::RegisterReuse,
)?;
}
SeekKeyComponent::None => {}
}
program.preassign_label_to_next_insn(loop_start);
let mut rowid_reg = None;
let mut affinity = None;
if !is_index {
rowid_reg = Some(program.alloc_register());
program.emit_insn(Insn::RowId {
cursor_id: seek_cursor_id,
dest: rowid_reg.unwrap(),
});
affinity = if let Some(table_ref) = tables
.joined_tables()
.iter()
.find(|t| t.columns().iter().any(|c| c.is_rowid_alias))
{
if let Some(rowid_col_idx) = table_ref.columns().iter().position(|c| c.is_rowid_alias) {
Some(table_ref.columns()[rowid_col_idx].affinity())
} else {
Some(Affinity::Numeric)
}
} else {
Some(Affinity::Numeric)
};
}
match (is_index, seek_def.end.op) {
(true, SeekOp::GE { .. }) => program.emit_insn(Insn::IdxGE {
cursor_id: seek_cursor_id,
start_reg,
num_regs,
target_pc: loop_end,
}),
(true, SeekOp::GT) => program.emit_insn(Insn::IdxGT {
cursor_id: seek_cursor_id,
start_reg,
num_regs,
target_pc: loop_end,
}),
(true, SeekOp::LE { .. }) => program.emit_insn(Insn::IdxLE {
cursor_id: seek_cursor_id,
start_reg,
num_regs,
target_pc: loop_end,
}),
(true, SeekOp::LT) => program.emit_insn(Insn::IdxLT {
cursor_id: seek_cursor_id,
start_reg,
num_regs,
target_pc: loop_end,
}),
(false, SeekOp::GE { .. }) => program.emit_insn(Insn::Ge {
lhs: rowid_reg.unwrap(),
rhs: start_reg,
target_pc: loop_end,
flags: CmpInsFlags::default()
.jump_if_null()
.with_affinity(affinity.unwrap()),
collation: program.curr_collation(),
}),
(false, SeekOp::GT) => program.emit_insn(Insn::Gt {
lhs: rowid_reg.unwrap(),
rhs: start_reg,
target_pc: loop_end,
flags: CmpInsFlags::default()
.jump_if_null()
.with_affinity(affinity.unwrap()),
collation: program.curr_collation(),
}),
(false, SeekOp::LE { .. }) => program.emit_insn(Insn::Le {
lhs: rowid_reg.unwrap(),
rhs: start_reg,
target_pc: loop_end,
flags: CmpInsFlags::default()
.jump_if_null()
.with_affinity(affinity.unwrap()),
collation: program.curr_collation(),
}),
(false, SeekOp::LT) => program.emit_insn(Insn::Lt {
lhs: rowid_reg.unwrap(),
rhs: start_reg,
target_pc: loop_end,
flags: CmpInsFlags::default()
.jump_if_null()
.with_affinity(affinity.unwrap()),
collation: program.curr_collation(),
}),
};
Ok(())
}
/// Open an ephemeral index cursor and build an automatic index on a table.
/// This is used as a last-resort to avoid a nested full table scan
/// Returns the cursor id of the ephemeral index cursor.
fn emit_autoindex(
program: &mut ProgramBuilder,
index: &Arc<Index>,
table_cursor_id: CursorID,
index_cursor_id: CursorID,
table_has_rowid: bool,
) -> Result<CursorID> {
assert!(index.ephemeral, "Index {} is not ephemeral", index.name);
let label_ephemeral_build_end = program.allocate_label();
// Since this typically happens in an inner loop, we only build it once.
program.emit_insn(Insn::Once {
target_pc_when_reentered: label_ephemeral_build_end,
});
program.emit_insn(Insn::OpenAutoindex {
cursor_id: index_cursor_id,
});
// Rewind source table
let label_ephemeral_build_loop_start = program.allocate_label();
program.emit_insn(Insn::Rewind {
cursor_id: table_cursor_id,
pc_if_empty: label_ephemeral_build_loop_start,
});
program.preassign_label_to_next_insn(label_ephemeral_build_loop_start);
// Emit all columns from source table that are needed in the ephemeral index.
// Also reserve a register for the rowid if the source table has rowids.
let num_regs_to_reserve = index.columns.len() + table_has_rowid as usize;
let ephemeral_cols_start_reg = program.alloc_registers(num_regs_to_reserve);
for (i, col) in index.columns.iter().enumerate() {
let reg = ephemeral_cols_start_reg + i;
program.emit_column_or_rowid(table_cursor_id, col.pos_in_table, reg);
}
if table_has_rowid {
program.emit_insn(Insn::RowId {
cursor_id: table_cursor_id,
dest: ephemeral_cols_start_reg + index.columns.len(),
});
}
let record_reg = program.alloc_register();
program.emit_insn(Insn::MakeRecord {
start_reg: ephemeral_cols_start_reg,
count: num_regs_to_reserve,
dest_reg: record_reg,
index_name: Some(index.name.clone()),
affinity_str: None,
});
program.emit_insn(Insn::IdxInsert {
cursor_id: index_cursor_id,
record_reg,
unpacked_start: Some(ephemeral_cols_start_reg),
unpacked_count: Some(num_regs_to_reserve as u16),
flags: IdxInsertFlags::new().use_seek(false),
});
program.emit_insn(Insn::Next {
cursor_id: table_cursor_id,
pc_if_next: label_ephemeral_build_loop_start,
});
program.preassign_label_to_next_insn(label_ephemeral_build_end);
Ok(index_cursor_id)
}