diff --git a/core/translate/optimizer/mod.rs b/core/translate/optimizer/mod.rs index 066b23567..0b6042687 100644 --- a/core/translate/optimizer/mod.rs +++ b/core/translate/optimizer/mod.rs @@ -301,6 +301,7 @@ fn add_ephemeral_table_to_update_plan( distinctness: super::plan::Distinctness::NonDistinct, values: vec![], window: None, + non_from_clause_subqueries: vec![], }; plan.ephemeral_plan = Some(ephemeral_plan); diff --git a/core/translate/plan.rs b/core/translate/plan.rs index 1a5e70529..f7f34ff29 100644 --- a/core/translate/plan.rs +++ b/core/translate/plan.rs @@ -1,5 +1,7 @@ use std::{cmp::Ordering, sync::Arc}; -use turso_parser::ast::{self, FrameBound, FrameClause, FrameExclude, FrameMode, SortOrder}; +use turso_parser::ast::{ + self, FrameBound, FrameClause, FrameExclude, FrameMode, SortOrder, SubqueryType, +}; use crate::{ function::AggFunc, @@ -359,6 +361,8 @@ pub struct SelectPlan { /// The window definition and all window functions associated with it. There is at most one /// window per SELECT. If the original query contains more, they are pushed down into subqueries. pub window: Option, + /// Subqueries that appear in any part of the query apart from the FROM clause + pub non_from_clause_subqueries: Vec, } impl SelectPlan { @@ -1324,6 +1328,84 @@ pub struct WindowFunction { pub original_expr: Expr, } +#[derive(Debug, Clone)] +pub enum SubqueryState { + /// The subquery has not been evaluated yet. + /// The 'plan' field is only optional because it is .take()'d when the the subquery + /// is translated into bytecode. + Unevaluated { plan: Option> }, + /// The subquery has been evaluated. + /// The [evaluated_at] field contains the loop index where the subquery was evaluated. + /// The query plan struct no longer exists because translating the plan currently + /// requires an ownership transfer. + Evaluated { evaluated_at: EvalAt }, +} + +#[derive(Debug, Clone)] +/// A subquery that is not part of the `FROM` clause. +/// This is used for subqueries in the WHERE clause, HAVING clause, ORDER BY clause, LIMIT clause, OFFSET clause, etc. +/// Currently only subqueries in the WHERE clause are supported. +pub struct NonFromClauseSubquery { + pub internal_id: TableInternalId, + pub query_type: SubqueryType, + pub state: SubqueryState, + pub correlated: bool, +} + +impl NonFromClauseSubquery { + /// Returns true if the subquery has been evaluated (translated into bytecode). + pub fn has_been_evaluated(&self) -> bool { + matches!(self.state, SubqueryState::Evaluated { .. }) + } + + /// Returns the loop index where the subquery should be evaluated in this particular join order. + /// If the subquery references tables from the parent query, it will be evaluated at the right-most + /// nested loop whose table it references. + pub fn get_eval_at(&self, join_order: &[JoinOrderMember]) -> Result { + let mut eval_at = EvalAt::BeforeLoop; + let SubqueryState::Unevaluated { plan } = &self.state else { + crate::bail_parse_error!("subquery has already been evaluated"); + }; + let used_outer_refs = plan + .as_ref() + .unwrap() + .table_references + .outer_query_refs() + .iter() + .filter(|t| t.is_used()); + + for outer_ref in used_outer_refs { + let Some(loop_idx) = join_order + .iter() + .position(|t| t.table_id == outer_ref.internal_id) + else { + continue; + }; + eval_at = eval_at.max(EvalAt::Loop(loop_idx)); + } + for subquery in plan.as_ref().unwrap().non_from_clause_subqueries.iter() { + let eval_at_inner = subquery.get_eval_at(join_order)?; + eval_at = eval_at.max(eval_at_inner); + } + Ok(eval_at) + } + + /// Consumes the plan and returns it, and sets the subquery to the evaluated state. + /// This is used when the subquery is translated into bytecode. + pub fn consume_plan(&mut self, evaluated_at: EvalAt) -> Box { + match &mut self.state { + SubqueryState::Unevaluated { plan } => { + let plan = plan.take().unwrap(); + self.state = SubqueryState::Evaluated { evaluated_at }; + plan + } + SubqueryState::Evaluated { .. } => { + panic!("subquery has already been evaluated"); + } + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/core/translate/select.rs b/core/translate/select.rs index 59b6ff6cb..8c6c57e84 100644 --- a/core/translate/select.rs +++ b/core/translate/select.rs @@ -273,6 +273,7 @@ fn prepare_one_select_plan( distinctness: Distinctness::from_ast(distinctness.as_ref()), values: vec![], window: None, + non_from_clause_subqueries: vec![], }; let mut windows = Vec::with_capacity(window_clause.len()); @@ -555,6 +556,7 @@ fn prepare_one_select_plan( .map(|values| values.iter().map(|value| *value.clone()).collect()) .collect(), window: None, + non_from_clause_subqueries: vec![], }; Ok(plan) diff --git a/core/translate/window.rs b/core/translate/window.rs index 7ab80207d..ad62e31fc 100644 --- a/core/translate/window.rs +++ b/core/translate/window.rs @@ -208,6 +208,7 @@ fn prepare_window_subquery( distinctness: Distinctness::NonDistinct, values: vec![], window: None, + non_from_clause_subqueries: vec![], }; prepare_window_subquery(