diff --git a/src/binder/aggregate.rs b/src/binder/aggregate.rs index 1103d6e5..40ad6f21 100644 --- a/src/binder/aggregate.rs +++ b/src/binder/aggregate.rs @@ -17,8 +17,8 @@ use std::collections::HashSet; use super::{Binder, QueryBindStep}; use crate::errors::DatabaseError; use crate::expression::function::scala::ScalarFunction; -use crate::expression::visitor::{walk_expr, Visitor}; -use crate::expression::visitor_mut::{walk_mut_expr, VisitorMut}; +use crate::expression::visitor::{walk_expr, ExprVisitor}; +use crate::expression::visitor_mut::{walk_mut_expr, ExprVisitorMut}; use crate::planner::LogicalPlan; use crate::storage::Transaction; use crate::types::value::DataValue; @@ -269,7 +269,7 @@ impl> Binder<'_, '_, T, A> HashSet::from_iter(group_raw_exprs.iter().copied()); for expr in select_items { - if expr.has_agg_call() { + if expr.has_agg_call()? { continue; } group_raw_set.remove(expr); @@ -347,7 +347,7 @@ impl<'a> HavingOrderByValidator<'a> { } } -impl<'expr> Visitor<'expr> for HavingOrderByValidator<'_> { +impl<'expr> ExprVisitor<'expr> for HavingOrderByValidator<'_> { fn visit(&mut self, expr: &'expr ScalarExpression) -> Result<(), DatabaseError> { match expr { ScalarExpression::AggCall { .. } => { @@ -427,7 +427,7 @@ impl<'a, 'p> AggregateOutputBinder<'a, 'p> { } } -impl<'a> VisitorMut<'a> for AggregateOutputBinder<'_, '_> { +impl<'a> ExprVisitorMut<'a> for AggregateOutputBinder<'_, '_> { fn visit(&mut self, expr: &'a mut ScalarExpression) -> Result<(), DatabaseError> { if let ScalarExpression::Alias { expr: inner_expr, @@ -453,7 +453,7 @@ mod tests { use crate::catalog::{ColumnCatalog, ColumnDesc, ColumnRef}; use crate::errors::DatabaseError; use crate::expression::agg::AggKind; - use crate::expression::visitor_mut::VisitorMut; + use crate::expression::visitor_mut::ExprVisitorMut; use crate::expression::{AliasType, BinaryOperator, ScalarExpression}; use crate::planner::PlanArena; use crate::storage::Storage; diff --git a/src/binder/distinct.rs b/src/binder/distinct.rs index cc54bbbd..02aa2441 100644 --- a/src/binder/distinct.rs +++ b/src/binder/distinct.rs @@ -14,7 +14,7 @@ use crate::binder::{Binder, QueryBindStep}; use crate::errors::DatabaseError; -use crate::expression::visitor_mut::{walk_mut_expr, VisitorMut}; +use crate::expression::visitor_mut::{walk_mut_expr, ExprVisitorMut}; use crate::expression::ScalarExpression; use crate::planner::operator::aggregate::AggregateOperator; use crate::planner::operator::sort::SortField; @@ -101,7 +101,7 @@ impl<'a, 'p> DistinctOutputBinder<'a, 'p> { } } -impl<'a> VisitorMut<'a> for DistinctOutputBinder<'_, '_> { +impl<'a> ExprVisitorMut<'a> for DistinctOutputBinder<'_, '_> { fn visit(&mut self, expr: &'a mut ScalarExpression) -> Result<(), DatabaseError> { if let ScalarExpression::Alias { expr: inner_expr, @@ -124,7 +124,7 @@ mod tests { use super::DistinctOutputBinder; use crate::catalog::{ColumnCatalog, ColumnDesc, ColumnRef}; use crate::errors::DatabaseError; - use crate::expression::visitor_mut::VisitorMut; + use crate::expression::visitor_mut::ExprVisitorMut; use crate::expression::{AliasType, ScalarExpression}; use crate::planner::PlanArena; use crate::types::LogicalType; diff --git a/src/binder/parser.rs b/src/binder/parser.rs index 67f77910..62b8499a 100644 --- a/src/binder/parser.rs +++ b/src/binder/parser.rs @@ -25,7 +25,7 @@ use crate::db::{BindSource, DBTransaction, Database, DatabaseIter, TransactionIt use crate::errors::{DatabaseError, SqlErrorSpan}; use crate::expression; use crate::expression::simplify::ConstantCalculator; -use crate::expression::visitor_mut::VisitorMut; +use crate::expression::visitor_mut::ExprVisitorMut; use crate::expression::{AliasType, ScalarExpression}; use crate::iter_ext::Itertools; use crate::parser::parse_sql; @@ -338,7 +338,7 @@ struct UpdateExprTargetRemapper<'a, 'p> { arena: &'a PlanArena<'p>, } -impl VisitorMut<'_> for UpdateExprTargetRemapper<'_, '_> { +impl ExprVisitorMut<'_> for UpdateExprTargetRemapper<'_, '_> { fn visit_column_ref( &mut self, column: &mut ColumnRef, @@ -669,7 +669,7 @@ where ) -> Result { let mut expr = self.binder.bind_expr(&expr, self.arena)?; - if expr.any_referenced_column(self.arena, |_, _| true) { + if expr.any_referenced_column(self.arena, |_, _| true)? { return Err(DatabaseError::UnsupportedStmt( "column is not allowed to exist in default".to_string(), )); @@ -783,7 +783,7 @@ where ColumnOption::Default(expr) => { let mut expr = self.binder.bind_expr(&expr, self.arena)?; - if expr.any_referenced_column(self.arena, |_, _| true) { + if expr.any_referenced_column(self.arena, |_, _| true)? { return Err(DatabaseError::UnsupportedStmt( "column is not allowed to exist in `default`".to_string(), )); diff --git a/src/binder/select.rs b/src/binder/select.rs index ed43b3a5..ce8d69a7 100644 --- a/src/binder/select.rs +++ b/src/binder/select.rs @@ -31,7 +31,7 @@ use super::{Binder, BinderContext, QueryBindStep, SetOperatorKind, Source, SubQu use crate::catalog::{ColumnRef, ColumnRelation, TableName}; use crate::errors::DatabaseError; use crate::execution::dql::join::joins_nullable; -use crate::expression::visitor_mut::{walk_mut_expr, PositionShift, VisitorMut}; +use crate::expression::visitor_mut::{walk_mut_expr, ExprVisitorMut, PositionShift}; use crate::expression::{AliasType, BinaryOperator}; use crate::iter_ext::Itertools; use crate::planner::operator::function_scan::FunctionScanOperator; @@ -51,7 +51,7 @@ struct RightSidePositionGlobalizer<'a, 'p> { arena: &'a crate::planner::PlanArena<'p>, } -impl<'a> VisitorMut<'a> for RightSidePositionGlobalizer<'_, '_> { +impl<'a> ExprVisitorMut<'a> for RightSidePositionGlobalizer<'_, '_> { fn visit_column_ref( &mut self, column: &'a mut ColumnRef, @@ -80,7 +80,7 @@ struct SplitScopePositionRebinder<'a, 'p> { arena: &'a crate::planner::PlanArena<'p>, } -impl VisitorMut<'_> for SplitScopePositionRebinder<'_, '_> { +impl ExprVisitorMut<'_> for SplitScopePositionRebinder<'_, '_> { fn visit_column_ref( &mut self, column: &mut ColumnRef, @@ -109,7 +109,7 @@ struct MarkerPositionGlobalizer<'a, 'p> { arena: &'a crate::planner::PlanArena<'p>, } -impl VisitorMut<'_> for MarkerPositionGlobalizer<'_, '_> { +impl ExprVisitorMut<'_> for MarkerPositionGlobalizer<'_, '_> { fn visit_column_ref( &mut self, column: &mut ColumnRef, @@ -154,7 +154,7 @@ impl<'a, 'p> ProjectionOutputBinder<'a, 'p> { } } -impl<'a> VisitorMut<'a> for ProjectionOutputBinder<'_, '_> { +impl<'a> ExprVisitorMut<'a> for ProjectionOutputBinder<'_, '_> { fn visit(&mut self, expr: &'a mut ScalarExpression) -> Result<(), DatabaseError> { if let Some(output_ref) = self.output_ref(expr) { *expr = output_ref; @@ -432,8 +432,10 @@ where #[cfg(feature = "orm")] pub fn finish(self) -> Result { - if self.select_list.iter().any(ScalarExpression::has_agg_call) { - return self.aggregate_without_group()?.finish(); + for expr in &self.select_list { + if expr.has_agg_call()? { + return self.aggregate_without_group()?.finish(); + } } self.binder .bind_project(self.plan, self.select_list, self.arena) @@ -801,7 +803,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<' arena: &'a crate::planner::PlanArena<'p>, } - impl VisitorMut<'_> for AppendedRightOutputBinder<'_, '_> { + impl ExprVisitorMut<'_> for AppendedRightOutputBinder<'_, '_> { fn visit_column_ref( &mut self, column: &mut ColumnRef, @@ -1495,7 +1497,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<' plan: &mut LogicalPlan, predicates: &[ScalarExpression], arena: &mut crate::planner::PlanArena, - ) -> Vec { + ) -> Result, DatabaseError> { let output_schema = plan.output_schema(arena).clone(); let output_len = output_schema.len(); if let LogicalPlan { @@ -1505,36 +1507,38 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<' } = plan { let Childrens::Only(child) = childrens.as_mut() else { - return Vec::new(); + return Ok(Vec::new()); }; let child_schema = child.output_schema(arena); let mut appended_outputs = Vec::new(); - op.exprs.extend( - child_schema - .iter() - .enumerate() - .filter(|(_, column)| { - !output_schema.contains(column) - && predicates.iter().any(|expr| { - expr.any_referenced_column(arena, |arena, candidate| { - arena.same_column(*candidate, **column) - }) - }) - }) - .map(|(position, column)| { - appended_outputs.push(AppendedRightOutput { - column: *column, - child_position: position, - output_position: output_len + appended_outputs.len(), - }); - ScalarExpression::column_expr(*column, position) - }), - ); + for (position, column) in child_schema.iter().enumerate() { + if output_schema.contains(column) { + continue; + } + let mut referenced = false; + for expr in predicates { + if expr.any_referenced_column(arena, |arena, candidate| { + arena.same_column(*candidate, *column) + })? { + referenced = true; + break; + } + } + if referenced { + op.exprs + .push(ScalarExpression::column_expr(*column, position)); + appended_outputs.push(AppendedRightOutput { + column: *column, + child_position: position, + output_position: output_len + appended_outputs.len(), + }); + } + } plan.reset_output_schema_cache(); - return appended_outputs; + return Ok(appended_outputs); } - Vec::new() + Ok(Vec::new()) } #[allow(clippy::too_many_arguments)] @@ -1565,7 +1569,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<' if correlated { let appended_right_outputs = - Self::ensure_mark_apply_right_outputs(&mut plan, &apply_predicates, arena); + Self::ensure_mark_apply_right_outputs(&mut plan, &apply_predicates, arena)?; if !appended_right_outputs.is_empty() { Self::localize_appended_right_outputs( apply_predicates.iter_mut(), @@ -1618,25 +1622,28 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<' plan: &LogicalPlan, left_schema: &Schema, arena: &mut crate::planner::PlanArena, - ) -> bool { + ) -> Result { if !plan .operator .visit_referenced_columns(arena, &mut |arena, column| { !left_schema .iter() .any(|left| arena.same_column(*left, *column)) - }) + })? { - return true; + return Ok(true); } match plan.childrens.as_ref() { Childrens::Only(child) => Self::plan_has_correlated_refs(child, left_schema, arena), Childrens::Twins { left, right } => { - Self::plan_has_correlated_refs(left, left_schema, arena) - || Self::plan_has_correlated_refs(right, left_schema, arena) + if Self::plan_has_correlated_refs(left, left_schema, arena)? { + Ok(true) + } else { + Self::plan_has_correlated_refs(right, left_schema, arena) + } } - Childrens::None => false, + Childrens::None => Ok(false), } } @@ -1644,7 +1651,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<' expr: &ScalarExpression, left_schema: &Schema, arena: &mut crate::planner::PlanArena, - ) -> bool { + ) -> Result { expr.any_referenced_column(arena, |arena, column| { left_schema .iter() @@ -1688,7 +1695,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<' match plan.childrens.as_ref() { Childrens::Only(_) => {} Childrens::Twins { .. } => { - if Self::plan_has_correlated_refs(&plan, left_schema, arena) { + if Self::plan_has_correlated_refs(&plan, left_schema, arena)? { return Err(DatabaseError::UnsupportedStmt( "correlated EXISTS/NOT EXISTS does not support set or join subqueries" .to_string(), @@ -1715,7 +1722,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<' let mut predicates = Vec::new(); Self::split_conjuncts(op.predicate, &mut predicates); for predicate in predicates { - if Self::expr_has_correlated_refs(&predicate, left_schema, arena) { + if Self::expr_has_correlated_refs(&predicate, left_schema, arena)? { correlated_filters.push(predicate); } else { local_filters.push(predicate); @@ -1775,7 +1782,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<' arena, ), plan => { - if Self::plan_has_correlated_refs(&plan, left_schema, arena) { + if Self::plan_has_correlated_refs(&plan, left_schema, arena)? { Err(DatabaseError::UnsupportedStmt( "correlated EXISTS/NOT EXISTS only supports filter-based subqueries" .to_string(), @@ -2142,9 +2149,11 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<' // example: baz > 1 if left_expr.all_referenced_columns(arena, |_, column| { fn_or_contains(*column) - }) && right_expr.all_referenced_columns(arena, |_, column| { - fn_or_contains(*column) - }) { + })? && right_expr + .all_referenced_columns(arena, |_, column| { + fn_or_contains(*column) + })? + { accum_filter.push(ScalarExpression::Binary { left_expr, right_expr, @@ -2186,9 +2195,10 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<' } _ => { if left_expr - .all_referenced_columns(arena, |_, column| fn_or_contains(*column)) - && right_expr - .all_referenced_columns(arena, |_, column| fn_or_contains(*column)) + .all_referenced_columns(arena, |_, column| fn_or_contains(*column))? + && right_expr.all_referenced_columns(arena, |_, column| { + fn_or_contains(*column) + })? { accum_filter.push(ScalarExpression::Binary { left_expr, @@ -2202,7 +2212,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<' } } expr => { - if expr.all_referenced_columns(arena, |_, column| fn_or_contains(*column)) { + if expr.all_referenced_columns(arena, |_, column| fn_or_contains(*column))? { // example: baz > 1 accum_filter.push(expr); } @@ -2219,7 +2229,7 @@ mod tests { use crate::binder::test::build_t1_table; use crate::catalog::{ColumnCatalog, ColumnDesc}; use crate::errors::DatabaseError; - use crate::expression::visitor_mut::VisitorMut; + use crate::expression::visitor_mut::ExprVisitorMut; use crate::expression::{AliasType, ScalarExpression}; use crate::planner::operator::join::{JoinCondition, JoinType}; use crate::planner::operator::mark_apply::{ diff --git a/src/catalog/column.rs b/src/catalog/column.rs index 139c6115..03fc60f1 100644 --- a/src/catalog/column.rs +++ b/src/catalog/column.rs @@ -244,7 +244,7 @@ impl ColumnDesc { if let Some(expr) = &default { let table_arena = crate::planner::TableArenaCell::default(); let plan_arena = crate::planner::PlanArena::new(&table_arena); - if expr.has_table_ref_column(&plan_arena) { + if expr.has_table_ref_column(&plan_arena)? { return Err(DatabaseError::DefaultNotColumnRef); } } diff --git a/src/catalog/view.rs b/src/catalog/view.rs index 84c23251..503c9b7b 100644 --- a/src/catalog/view.rs +++ b/src/catalog/view.rs @@ -27,7 +27,11 @@ pub struct View { } impl View { - pub(crate) fn visit_column_refs(&self, arena: &mut A, f: &mut F) + pub(crate) fn visit_column_refs( + &self, + arena: &mut A, + f: &mut F, + ) -> Result<(), crate::errors::DatabaseError> where A: MetaArena, F: FnMut(&ColumnRef) + ?Sized, @@ -35,7 +39,7 @@ impl View { for column in &self.schema { f(column); } - self.plan.visit_column_refs(arena, f); + self.plan.visit_column_refs(arena, f) } } diff --git a/src/db.rs b/src/db.rs index e43334fd..65ad0f8d 100644 --- a/src/db.rs +++ b/src/db.rs @@ -464,7 +464,7 @@ impl State { Ok(()) } - fn recycle_table_arena(&mut self) { + fn recycle_table_arena(&mut self) -> Result<(), DatabaseError> { let mut live_columns = HashSet::new(); { let mut live = |column: &crate::catalog::ColumnRef| { @@ -479,7 +479,7 @@ impl State { let table_arena = self.table_arena.borrow_mut(); for view in self.view_cache.values() { - view.visit_column_refs(table_arena, &mut live); + view.visit_column_refs(table_arena, &mut live)?; } for function in self.table_functions.values() { @@ -491,6 +491,7 @@ impl State { self.table_arena .borrow_mut() .recycle_unreferenced_positions(live_columns); + Ok(()) } pub(crate) fn build_plan<'a, 'txn, A: AsRef<[(&'static str, DataValue)]>, F>( @@ -762,7 +763,7 @@ impl Database { .map_err(|err| err.with_sql_context(context))?; } if catalog_changed { - unsafe { (&mut *state).recycle_table_arena() }; + unsafe { (&mut *state).recycle_table_arena() }?; } Ok(()) } @@ -783,7 +784,7 @@ impl Database { } CatalogKind::TableFunction(function) => { self.state.load_table_function(function)?; - self.state.recycle_table_arena(); + self.state.recycle_table_arena()?; Ok(()) } CatalogKind::Table(name) => { @@ -808,7 +809,7 @@ impl Database { } } self.state.table_cache.insert(name, table); - self.state.recycle_table_arena(); + self.state.recycle_table_arena()?; Ok(()) } CatalogKind::View(name) => { @@ -825,7 +826,7 @@ impl Database { )? .ok_or(DatabaseError::ViewNotFound)?; self.state.view_cache.insert(name, view); - self.state.recycle_table_arena(); + self.state.recycle_table_arena()?; Ok(()) } } @@ -1581,7 +1582,7 @@ pub(crate) mod test { filter.visit_referenced_columns(&mut source_plan_arena, &mut |_, column| { referenced_columns.push(*column); true - }); + })?; assert_eq!(referenced_columns.len(), 1); assert_eq!(source_plan_arena.column(referenced_columns[0]).name(), "y"); diff --git a/src/expression/mod.rs b/src/expression/mod.rs index 17ac3f6f..48340e2a 100644 --- a/src/expression/mod.rs +++ b/src/expression/mod.rs @@ -17,8 +17,8 @@ use crate::catalog::{ColumnCatalog, ColumnDesc, ColumnRef}; use crate::errors::DatabaseError; use crate::expression::function::scala::ScalarFunction; use crate::expression::function::table::TableFunction; -use crate::expression::visitor::{walk_expr, Visitor}; -use crate::expression::visitor_mut::VisitorMut; +use crate::expression::visitor::{walk_expr, ExprVisitor}; +use crate::expression::visitor_mut::ExprVisitorMut; use crate::iter_ext::Itertools; use crate::planner::operator::sort::SortField; use crate::planner::{MetaArena, PlanArena}; @@ -277,7 +277,7 @@ pub struct BindEvaluator<'a, 'p> { pub(crate) arena: &'a PlanArena<'p>, } -impl VisitorMut<'_> for BindEvaluator<'_, '_> { +impl ExprVisitorMut<'_> for BindEvaluator<'_, '_> { fn visit_type_cast( &mut self, expr: &'_ mut ScalarExpression, @@ -361,7 +361,7 @@ pub struct HasCountStar { pub value: bool, } -impl Visitor<'_> for HasCountStar { +impl ExprVisitor<'_> for HasCountStar { fn visit_agg( &mut self, _distinct: bool, @@ -529,14 +529,14 @@ impl ScalarExpression { &self, arena: &mut A, f: &mut impl FnMut(&mut A, &ColumnRef) -> bool, - ) -> bool { + ) -> Result { struct ColumnRefVisitor<'a, A, F> { f: &'a mut F, keep_going: bool, arena: &'a mut A, } - impl Visitor<'_> for ColumnRefVisitor<'_, A, F> + impl ExprVisitor<'_> for ColumnRefVisitor<'_, A, F> where A: MetaArena, F: FnMut(&mut A, &ColumnRef) -> bool, @@ -559,22 +559,22 @@ impl ScalarExpression { keep_going: true, arena, }; - visitor.visit(self).unwrap(); - visitor.keep_going + visitor.visit(self)?; + Ok(visitor.keep_going) } pub fn any_referenced_column( &self, arena: &PlanArena, mut predicate: impl FnMut(&PlanArena, &ColumnRef) -> bool, - ) -> bool { + ) -> Result { struct ColumnRefVisitor<'a, 'p, F> { f: &'a mut F, any: bool, arena: &'a PlanArena<'p>, } - impl bool> Visitor<'_> for ColumnRefVisitor<'_, '_, F> { + impl bool> ExprVisitor<'_> for ColumnRefVisitor<'_, '_, F> { fn visit(&mut self, expr: &ScalarExpression) -> Result<(), DatabaseError> { if !self.any { walk_expr(self, expr)?; @@ -593,22 +593,22 @@ impl ScalarExpression { any: false, arena, }; - visitor.visit(self).unwrap(); - visitor.any + visitor.visit(self)?; + Ok(visitor.any) } pub fn all_referenced_columns( &self, arena: &PlanArena, mut predicate: impl FnMut(&PlanArena, &ColumnRef) -> bool, - ) -> bool { + ) -> Result { struct ColumnRefVisitor<'a, 'p, F> { f: &'a mut F, all: bool, arena: &'a PlanArena<'p>, } - impl bool> Visitor<'_> for ColumnRefVisitor<'_, '_, F> { + impl bool> ExprVisitor<'_> for ColumnRefVisitor<'_, '_, F> { fn visit(&mut self, expr: &ScalarExpression) -> Result<(), DatabaseError> { if self.all { walk_expr(self, expr)?; @@ -627,16 +627,16 @@ impl ScalarExpression { all: true, arena, }; - visitor.visit(self).unwrap(); - visitor.all + visitor.visit(self)?; + Ok(visitor.all) } - pub fn has_table_ref_column(&self, arena: &PlanArena) -> bool { + pub fn has_table_ref_column(&self, arena: &PlanArena) -> Result { struct TableRefChecker<'arena, 'table> { found: bool, arena: &'arena PlanArena<'table>, } - impl Visitor<'_> for TableRefChecker<'_, '_> { + impl ExprVisitor<'_> for TableRefChecker<'_, '_> { fn visit_column_ref(&mut self, col: &ColumnRef) -> Result<(), DatabaseError> { let col = self.arena.column(*col); if col.table_name().is_some() && col.id().is_some() { @@ -649,15 +649,15 @@ impl ScalarExpression { found: false, arena, }; - checker.visit(self).unwrap(); - checker.found + checker.visit(self)?; + Ok(checker.found) } - pub fn has_agg_call(&self) -> bool { + pub fn has_agg_call(&self) -> Result { struct AggCallChecker { has_agg: bool, } - impl<'a> Visitor<'a> for AggCallChecker { + impl<'a> ExprVisitor<'a> for AggCallChecker { fn visit(&mut self, expr: &'a ScalarExpression) -> Result<(), DatabaseError> { if self.has_agg { return Ok(()); @@ -679,8 +679,8 @@ impl ScalarExpression { } } let mut checker = AggCallChecker { has_agg: false }; - checker.visit(self).unwrap(); - checker.has_agg + checker.visit(self)?; + Ok(checker.has_agg) } fn output_name_by(&self, fn_display: &impl Fn(ColumnRef) -> N) -> String { diff --git a/src/expression/range_detacher.rs b/src/expression/range_detacher.rs index 99e70efc..a5b67cac 100644 --- a/src/expression/range_detacher.rs +++ b/src/expression/range_detacher.rs @@ -2354,4 +2354,83 @@ mod test { ); Ok(()) } + + #[test] + fn test_range_merge_preserves_open_closed_boundaries() { + let gt_one = Range::Scope { + min: Bound::Excluded(DataValue::Int32(1)), + max: Bound::Unbounded, + }; + assert_eq!( + RangeDetacher::merge_binary( + BinaryOperator::Or, + gt_one.clone(), + Range::Eq(DataValue::Int32(1)), + ), + Some(Range::Scope { + min: Bound::Included(DataValue::Int32(1)), + max: Bound::Unbounded, + }) + ); + assert_eq!( + RangeDetacher::merge_binary( + BinaryOperator::And, + gt_one, + Range::Eq(DataValue::Int32(1)), + ), + Some(Range::Dummy) + ); + + let disjoint = RangeDetacher::merge_binary( + BinaryOperator::Or, + Range::Scope { + min: Bound::Included(DataValue::Int32(1)), + max: Bound::Included(DataValue::Int32(2)), + }, + Range::Scope { + min: Bound::Included(DataValue::Int32(4)), + max: Bound::Included(DataValue::Int32(5)), + }, + ); + assert_eq!( + disjoint, + Some(Range::SortedRanges(vec![ + Range::Scope { + min: Bound::Included(DataValue::Int32(1)), + max: Bound::Included(DataValue::Int32(2)), + }, + Range::Scope { + min: Bound::Included(DataValue::Int32(4)), + max: Bound::Included(DataValue::Int32(5)), + }, + ])) + ); + } + + #[test] + fn test_combining_eqs_builds_tuple_ranges_and_rejects_non_eq_prefixes() { + let suffix = Range::Scope { + min: Bound::Included(DataValue::Int32(10)), + max: Bound::Excluded(DataValue::Int32(20)), + }; + let prefixes = [ + Range::Eq(DataValue::Int32(1)), + Range::SortedRanges(vec![ + Range::Eq(DataValue::Int32(2)), + Range::Eq(DataValue::Int32(3)), + ]), + ]; + let combined = suffix.combining_eqs(&prefixes).unwrap(); + assert!(matches!(combined, Range::SortedRanges(ref ranges) if ranges.len() == 2)); + assert!(prefixes[0].only_eq()); + assert!(prefixes[1].only_eq()); + + assert!(suffix + .combining_eqs(&[Range::Scope { + min: Bound::Unbounded, + max: Bound::Unbounded, + }]) + .is_none()); + assert!(!suffix.only_eq()); + } } diff --git a/src/expression/simplify.rs b/src/expression/simplify.rs index 537b1244..e6ffbaba 100644 --- a/src/expression/simplify.rs +++ b/src/expression/simplify.rs @@ -14,7 +14,7 @@ use crate::catalog::ColumnRef; use crate::errors::DatabaseError; -use crate::expression::visitor_mut::{walk_mut_expr, VisitorMut}; +use crate::expression::visitor_mut::{walk_mut_expr, ExprVisitorMut}; use crate::expression::{BinaryOperator, ScalarExpression, UnaryOperator}; use crate::planner::PlanArena; use crate::types::evaluator::{binary_create, unary_create}; @@ -55,7 +55,7 @@ impl<'a, 'p> ConstantCalculator<'a, 'p> { } } -impl VisitorMut<'_> for ConstantCalculator<'_, '_> { +impl ExprVisitorMut<'_> for ConstantCalculator<'_, '_> { fn visit(&mut self, expr: &'_ mut ScalarExpression) -> Result<(), DatabaseError> { match expr { ScalarExpression::Unary { @@ -122,7 +122,7 @@ pub struct Simplify { replaces: Vec, } -impl VisitorMut<'_> for Simplify { +impl ExprVisitorMut<'_> for Simplify { fn visit(&mut self, expr: &'_ mut ScalarExpression) -> Result<(), DatabaseError> { match expr { ScalarExpression::Unary { diff --git a/src/expression/visitor.rs b/src/expression/visitor.rs index c897eb1b..c8ca08d4 100644 --- a/src/expression/visitor.rs +++ b/src/expression/visitor.rs @@ -23,7 +23,7 @@ use crate::types::evaluator::{BinaryEvaluatorRef, CastEvaluatorRef, UnaryEvaluat use crate::types::value::DataValue; use crate::types::LogicalType; -pub trait Visitor<'a>: Sized { +pub trait ExprVisitor<'a>: Sized { fn visit(&mut self, expr: &'a ScalarExpression) -> Result<(), DatabaseError> { walk_expr(self, expr) } @@ -265,7 +265,7 @@ pub trait Visitor<'a>: Sized { } } -pub fn walk_expr<'a, V: Visitor<'a>>( +pub fn walk_expr<'a, V: ExprVisitor<'a>>( visitor: &mut V, expr: &'a ScalarExpression, ) -> Result<(), DatabaseError> { diff --git a/src/expression/visitor_mut.rs b/src/expression/visitor_mut.rs index 96d5f07c..dbddafd4 100644 --- a/src/expression/visitor_mut.rs +++ b/src/expression/visitor_mut.rs @@ -27,7 +27,7 @@ pub(crate) struct PositionShift { pub(crate) delta: isize, } -impl VisitorMut<'_> for PositionShift { +impl ExprVisitorMut<'_> for PositionShift { fn visit_column_ref( &mut self, _column: &mut ColumnRef, @@ -42,7 +42,7 @@ impl VisitorMut<'_> for PositionShift { } } -pub trait VisitorMut<'a>: Sized { +pub trait ExprVisitorMut<'a>: Sized { fn visit(&mut self, expr: &'a mut ScalarExpression) -> Result<(), DatabaseError> { walk_mut_expr(self, expr) } @@ -288,7 +288,7 @@ pub trait VisitorMut<'a>: Sized { } } -pub fn walk_mut_expr<'a, V: VisitorMut<'a>>( +pub fn walk_mut_expr<'a, V: ExprVisitorMut<'a>>( visitor: &mut V, expr: &'a mut ScalarExpression, ) -> Result<(), DatabaseError> { diff --git a/src/optimizer/core/cm_sketch.rs b/src/optimizer/core/cm_sketch.rs index 0bccba05..afe2d8d8 100644 --- a/src/optimizer/core/cm_sketch.rs +++ b/src/optimizer/core/cm_sketch.rs @@ -401,9 +401,27 @@ impl ReferenceSerialization for CountMinSketch { #[cfg(all(test, not(target_arch = "wasm32")))] mod tests { use crate::expression::range_detacher::Range; - use crate::optimizer::core::cm_sketch::CountMinSketch; + use crate::optimizer::core::cm_sketch::{ + CountMinSketch, CountMinSketchMeta, CountMinSketchPage, + }; + use crate::planner::TableArena; + use crate::serdes::{ReferenceSerialization, ReferenceTables}; + use crate::storage::rocksdb::RocksTransaction; use crate::types::value::DataValue; use std::collections::Bound; + use std::io::Cursor; + + fn assert_invalid_storage_parts( + meta: CountMinSketchMeta, + pages: Vec, + expected: &str, + ) { + let err = CountMinSketch::::from_storage_parts(meta, pages).unwrap_err(); + assert!( + err.to_string().contains(expected), + "unexpected error: {err}" + ); + } #[test] fn test_increment() { @@ -472,4 +490,85 @@ mod tests { rebuilt.estimate(&DataValue::Int32(9)) ); } + + #[test] + fn test_storage_parts_validate_metadata_and_page_layout() { + let cms = CountMinSketch::::with_relative_error(0.95, 0.5).unwrap(); + let (meta, pages) = cms.into_storage_parts(2); + let pages = pages.collect::>(); + + for invalid_meta in [ + CountMinSketchMeta { + width: 0, + ..meta.clone() + }, + CountMinSketchMeta { + k_num: 0, + ..meta.clone() + }, + CountMinSketchMeta { + page_len: 0, + ..meta.clone() + }, + ] { + assert_invalid_storage_parts(invalid_meta, pages.clone(), "storage meta is invalid"); + } + assert_invalid_storage_parts( + CountMinSketchMeta { + width: 3, + ..meta.clone() + }, + pages.clone(), + "width must be a power of two", + ); + + let mut invalid_pages = pages.clone(); + invalid_pages[0].row_idx = meta.k_num; + assert_invalid_storage_parts(meta.clone(), invalid_pages, "row index out of bounds"); + + let mut invalid_pages = pages.clone(); + invalid_pages[0].page_idx = 1; + assert_invalid_storage_parts(meta.clone(), invalid_pages, "page sequence is invalid"); + + let mut invalid_pages = pages.clone(); + invalid_pages[0].counters.push(0); + assert_invalid_storage_parts(meta.clone(), invalid_pages, "page is too large"); + + assert_invalid_storage_parts( + meta, + pages[..pages.len() - 1].to_vec(), + "row width mismatch", + ); + } + + #[test] + fn test_configuration_clear_and_reference_serialization() { + for probability in [-0.1, 1.0, f64::NAN] { + assert!(CountMinSketch::::with_relative_error(probability, 0.1).is_err()); + } + for relative_error in [0.0, -0.1, f64::NAN, f64::MIN_POSITIVE] { + assert!(CountMinSketch::::with_relative_error(0.95, relative_error).is_err()); + } + + let mut cms = CountMinSketch::::with_relative_error(0.95, 0.1).unwrap(); + cms.add(&7, usize::MAX); + cms.add(&7, 1); + assert_eq!(cms.estimate(&7), usize::MAX); + cms.clear(); + assert_eq!(cms.estimate(&7), 0); + + cms.add(&7, 23); + let mut bytes = Vec::new(); + let mut tables = ReferenceTables::new(); + let mut arena = TableArena::default(); + cms.encode(&mut bytes, false, &mut tables, &arena).unwrap(); + let decoded = CountMinSketch::::decode::( + &mut Cursor::new(bytes), + None, + &tables, + &mut arena, + ) + .unwrap(); + assert_eq!(decoded.estimate(&7), 23); + } } diff --git a/src/optimizer/heuristic/optimizer.rs b/src/optimizer/heuristic/optimizer.rs index 6c7b1d68..84b2e3d6 100644 --- a/src/optimizer/heuristic/optimizer.rs +++ b/src/optimizer/heuristic/optimizer.rs @@ -32,7 +32,8 @@ use crate::planner::{Childrens, LogicalPlan, PlanArena}; use std::array; use std::ops::Not; -type ScanHintApplier<'a> = dyn Fn(&mut TableScanOperator, &PlanArena) + 'a; +type ScanHintApplier<'a> = + dyn Fn(&mut TableScanOperator, &PlanArena) -> Result<(), DatabaseError> + 'a; pub struct HepOptimizer<'a> { before_batches: &'a [HepBatch], @@ -74,9 +75,10 @@ impl<'a> HepOptimizer<'a> { if let Some(loader) = loader { if self.implementation_index.is_empty().not() { - let apply_no_sort_hints = |_scan_op: &mut TableScanOperator, _arena: &PlanArena| {}; + let apply_no_sort_hints = + |_scan_op: &mut TableScanOperator, _arena: &PlanArena| Ok(()); let apply_no_stream_distinct_hints = - |_scan_op: &mut TableScanOperator, _arena: &PlanArena| {}; + |_scan_op: &mut TableScanOperator, _arena: &PlanArena| Ok(()); Self::annotate_hints_and_physical_options( &mut self.plan, loader, @@ -249,8 +251,8 @@ impl<'a> HepOptimizer<'a> { arena: &mut PlanArena, ) -> Result<(), DatabaseError> { if let Operator::TableScan(scan_op) = &mut plan.operator { - inherited_sort_hints(scan_op, arena); - inherited_stream_distinct_hints(scan_op, arena); + inherited_sort_hints(scan_op, arena)?; + inherited_stream_distinct_hints(scan_op, arena)?; } { @@ -337,19 +339,19 @@ impl<'a> HepOptimizer<'a> { match operator { Operator::Sort(op) => { let child_sort_hints = |scan_op: &mut TableScanOperator, arena: &PlanArena| { - inherited_sort_hints(scan_op, arena); + inherited_sort_hints(scan_op, arena)?; apply_scan_order_hint( scan_op, ScanOrderHint::sort_fields(&op.sort_fields), OrderHintKind::SortElimination, arena, - ); + ) }; f(&child_sort_hints) } _ if propagate_hints => f(inherited_sort_hints), _ => { - let no_sort_hints = |_scan_op: &mut TableScanOperator, _arena: &PlanArena| {}; + let no_sort_hints = |_scan_op: &mut TableScanOperator, _arena: &PlanArena| Ok(()); f(&no_sort_hints) } } @@ -380,14 +382,14 @@ impl<'a> HepOptimizer<'a> { ScanOrderHint::distinct_groupby(&op.groupby_exprs), OrderHintKind::StreamDistinct, arena, - ); + ) }; f(&child_stream_distinct_hints) } _ if propagate_hints => f(inherited_stream_distinct_hints), _ => { let no_stream_distinct_hints = - |_scan_op: &mut TableScanOperator, _arena: &PlanArena| {}; + |_scan_op: &mut TableScanOperator, _arena: &PlanArena| Ok(()); f(&no_stream_distinct_hints) } } diff --git a/src/optimizer/rule/normalization/column_pruning.rs b/src/optimizer/rule/normalization/column_pruning.rs index 1a749c9c..bdfacf63 100644 --- a/src/optimizer/rule/normalization/column_pruning.rs +++ b/src/optimizer/rule/normalization/column_pruning.rs @@ -15,11 +15,15 @@ use crate::catalog::ColumnRef; use crate::errors::DatabaseError; use crate::expression::agg::AggKind; -use crate::expression::visitor::Visitor; +use crate::expression::visitor::ExprVisitor; use crate::expression::{AliasType, HasCountStar, ScalarExpression}; use crate::optimizer::core::rule::NormalizationRule; -use crate::optimizer::rule::normalization::{remap_expr_positions, remap_exprs_positions}; +use crate::optimizer::rule::normalization::{ + remap_expr_positions, remap_exprs_positions, PositionRemapper, +}; use crate::planner::operator::join::JoinCondition; +use crate::planner::operator::visitor::{OperatorExprVisitor, OperatorVisitor}; +use crate::planner::operator::visitor_mut::{OperatorExprVisitorMut, OperatorVisitorMut}; use crate::planner::operator::Operator; use crate::planner::{Childrens, LogicalPlan}; use crate::types::value::{DataValue, Utf8Type}; @@ -79,6 +83,29 @@ impl ReferencedColumns { } } +struct ReferencedColumnCollector<'a, 'p> { + referenced_columns: &'a mut ReferencedColumns, + arena: &'a crate::planner::PlanArena<'p>, +} + +impl ExprVisitor<'_> for ReferencedColumnCollector<'_, '_> { + fn visit_column_ref( + &mut self, + column: &crate::catalog::ColumnRef, + ) -> Result<(), DatabaseError> { + self.referenced_columns.insert(*column, self.arena); + Ok(()) + } + + fn visit_alias( + &mut self, + expr: &ScalarExpression, + _ty: &AliasType, + ) -> Result<(), DatabaseError> { + self.visit(expr) + } +} + impl ApplyOutcome { fn with_arena_capacity(arena: &crate::planner::PlanArena) -> Self { Self { @@ -94,119 +121,88 @@ impl ColumnPruning { referenced_columns: &mut ReferencedColumns, arena: &mut crate::planner::PlanArena, ) -> Result<(), DatabaseError> { - match operator { - Operator::Aggregate(op) => { - Self::extend_expr_referenced_columns( - op.agg_calls.iter().chain(op.groupby_exprs.iter()), - referenced_columns, - arena, - )?; - } - Operator::Filter(op) => { - Self::extend_expr_referenced_columns([&op.predicate], referenced_columns, arena)?; - } - Operator::Join(op) => { - if let JoinCondition::On { on, filter } = &op.on { - for (left_expr, right_expr) in on { - Self::extend_expr_referenced_columns( - [left_expr, right_expr], - referenced_columns, - arena, - )?; - } - if let Some(filter_expr) = filter { - Self::extend_expr_referenced_columns( - [filter_expr], - referenced_columns, - arena, - )?; - } - } - } - Operator::Project(op) => { - Self::extend_expr_referenced_columns(op.exprs.iter(), referenced_columns, arena)?; - } - Operator::MarkApply(op) => { - Self::extend_expr_referenced_columns( - op.predicates().iter(), - referenced_columns, - arena, - )?; - referenced_columns.insert(*op.output_column(), arena); - } - Operator::TableScan(op) => { - referenced_columns.extend(op.columns.iter().copied(), arena); - } - Operator::FunctionScan(op) => { - Self::extend_expr_referenced_columns( - op.table_function.args.iter(), - referenced_columns, - arena, - )?; - } - Operator::Sort(op) => { - Self::extend_expr_referenced_columns( - op.sort_fields.iter().map(|field| &field.expr), - referenced_columns, - arena, - )?; + let mut collector = ReferencedColumnCollector { + referenced_columns, + arena, + }; + OperatorExprVisitor::new(&mut collector).visit_operator(operator)?; + + struct ReferencedOperatorColumnCollector<'a, 'p> { + referenced_columns: &'a mut ReferencedColumns, + arena: &'a crate::planner::PlanArena<'p>, + } + + impl<'a> OperatorVisitor<'a> for ReferencedOperatorColumnCollector<'_, '_> { + fn visit_mark_apply( + &mut self, + op: &'a crate::planner::operator::mark_apply::MarkApplyOperator, + ) -> Result<(), DatabaseError> { + self.referenced_columns + .insert(*op.output_column(), self.arena); + Ok(()) } - Operator::TopK(op) => { - Self::extend_expr_referenced_columns( - op.sort_fields.iter().map(|field| &field.expr), - referenced_columns, - arena, - )?; + + fn visit_table_scan( + &mut self, + op: &'a crate::planner::operator::table_scan::TableScanOperator, + ) -> Result<(), DatabaseError> { + self.referenced_columns + .extend(op.columns.iter().copied(), self.arena); + Ok(()) } - Operator::Values(op) => { - referenced_columns.extend(op.schema_ref.iter().copied(), arena); + + fn visit_values( + &mut self, + op: &'a crate::planner::operator::values::ValuesOperator, + ) -> Result<(), DatabaseError> { + self.referenced_columns + .extend(op.schema_ref.iter().copied(), self.arena); + Ok(()) } - Operator::Union(op) => { - referenced_columns.extend( + + fn visit_union( + &mut self, + op: &'a crate::planner::operator::union::UnionOperator, + ) -> Result<(), DatabaseError> { + self.referenced_columns.extend( op.left_schema_ref .iter() - .chain(op._right_schema_ref.iter()) + .chain(&op._right_schema_ref) .copied(), - arena, + self.arena, ); + Ok(()) } - Operator::SetMembership(op) => { - referenced_columns.extend( + + fn visit_set_membership( + &mut self, + op: &'a crate::planner::operator::set_membership::SetMembershipOperator, + ) -> Result<(), DatabaseError> { + self.referenced_columns.extend( op.left_schema_ref .iter() - .chain(op._right_schema_ref.iter()) + .chain(&op._right_schema_ref) .copied(), - arena, + self.arena, ); + Ok(()) } - Operator::Delete(op) => { - referenced_columns.extend(op.primary_keys.iter().copied(), arena); + + fn visit_delete( + &mut self, + op: &'a crate::planner::operator::delete::DeleteOperator, + ) -> Result<(), DatabaseError> { + self.referenced_columns + .extend(op.primary_keys.iter().copied(), self.arena); + Ok(()) } - Operator::Dummy - | Operator::Limit(_) - | Operator::ScalarApply(_) - | Operator::ScalarSubquery(_) - | Operator::Analyze(_) - | Operator::ShowTable - | Operator::ShowView - | Operator::Explain - | Operator::Describe(_) - | Operator::Insert(_) - | Operator::Update(_) - | Operator::AddColumn(_) - | Operator::ChangeColumn(_) - | Operator::DropColumn(_) - | Operator::CreateTable(_) - | Operator::CreateIndex(_) - | Operator::CreateView(_) - | Operator::DropTable(_) - | Operator::DropView(_) - | Operator::DropIndex(_) - | Operator::Truncate(_) => {} - #[cfg(feature = "copy")] - Operator::CopyFromFile(_) | Operator::CopyToFile(_) => {} } - Ok(()) + + ReferencedOperatorColumnCollector { + referenced_columns, + arena, + } + .visit_operator(operator) } fn extend_expr_referenced_columns<'a>( @@ -214,29 +210,6 @@ impl ColumnPruning { referenced_columns: &mut ReferencedColumns, arena: &mut crate::planner::PlanArena, ) -> Result<(), DatabaseError> { - struct ReferencedColumnCollector<'a, 'p> { - referenced_columns: &'a mut ReferencedColumns, - arena: &'a crate::planner::PlanArena<'p>, - } - - impl Visitor<'_> for ReferencedColumnCollector<'_, '_> { - fn visit_column_ref( - &mut self, - column: &crate::catalog::ColumnRef, - ) -> Result<(), DatabaseError> { - self.referenced_columns.insert(*column, self.arena); - Ok(()) - } - - fn visit_alias( - &mut self, - expr: &ScalarExpression, - _ty: &AliasType, - ) -> Result<(), DatabaseError> { - self.visit(expr) - } - } - let mut collector = ReferencedColumnCollector { referenced_columns, arena, @@ -280,75 +253,8 @@ impl ColumnPruning { operator: &mut Operator, removed_positions: &[usize], ) -> Result<(), DatabaseError> { - match operator { - Operator::Aggregate(op) => { - Self::remap_exprs_after_child_change( - op.agg_calls.iter_mut().chain(op.groupby_exprs.iter_mut()), - removed_positions, - )?; - } - Operator::Filter(op) => { - remap_expr_positions(&mut op.predicate, removed_positions)?; - } - Operator::Project(op) => { - remap_exprs_positions(op.exprs.iter_mut(), removed_positions)?; - } - Operator::MarkApply(op) => { - Self::remap_exprs_after_child_change( - op.predicates_mut().iter_mut(), - removed_positions, - )?; - } - Operator::ScalarApply(_) => {} - Operator::ScalarSubquery(_) => {} - Operator::Sort(op) => { - Self::remap_exprs_after_child_change( - op.sort_fields.iter_mut().map(|field| &mut field.expr), - removed_positions, - )?; - } - Operator::TopK(op) => { - Self::remap_exprs_after_child_change( - op.sort_fields.iter_mut().map(|field| &mut field.expr), - removed_positions, - )?; - } - Operator::Update(op) => { - Self::remap_exprs_after_child_change( - op.value_exprs.iter_mut().map(|(_, expr)| expr), - removed_positions, - )?; - } - Operator::Limit(_) - | Operator::Explain - | Operator::Insert(_) - | Operator::Delete(_) - | Operator::Analyze(_) - | Operator::Dummy - | Operator::TableScan(_) - | Operator::Join(_) - | Operator::Values(_) - | Operator::FunctionScan(_) - | Operator::ShowTable - | Operator::ShowView - | Operator::Describe(_) - | Operator::Union(_) - | Operator::SetMembership(_) - | Operator::AddColumn(_) - | Operator::ChangeColumn(_) - | Operator::DropColumn(_) - | Operator::CreateTable(_) - | Operator::CreateIndex(_) - | Operator::CreateView(_) - | Operator::DropTable(_) - | Operator::DropView(_) - | Operator::DropIndex(_) - | Operator::Truncate(_) => {} - #[cfg(feature = "copy")] - Operator::CopyFromFile(_) | Operator::CopyToFile(_) => {} - } - - Ok(()) + OperatorExprVisitorMut::new(&mut PositionRemapper { removed_positions }) + .visit_operator(operator) } fn remap_exprs_after_child_change<'a>( @@ -1066,4 +972,60 @@ mod tests { Ok(()) } + + #[test] + fn test_column_pruning_keeps_join_filter_columns_on_both_sides() -> Result<(), DatabaseError> { + let table_state = build_t1_table()?; + let mut arena = PlanArena::new(&table_state.table_arena); + let best_plan = optimize_column_pruning( + &table_state, + &mut arena, + "select c1 from t1 join t2 on c1 = c3 and c2 > c4", + )?; + + assert_single_scan_columns(&best_plan, "t1", &arena, &["c1", "c2"]); + assert_single_scan_columns(&best_plan, "t2", &arena, &["c3", "c4"]); + + Ok(()) + } + + #[test] + fn test_column_pruning_keeps_mark_apply_predicate_columns() -> Result<(), DatabaseError> { + let table_state = build_t1_table()?; + let mut arena = PlanArena::new(&table_state.table_arena); + let best_plan = optimize_column_pruning( + &table_state, + &mut arena, + "select c1 from t1 where c2 in (select c4 from t2)", + )?; + + assert!(contains_operator(&best_plan, |op| matches!( + op, + Operator::MarkApply(_) + ))); + assert_single_scan_columns(&best_plan, "t1", &arena, &["c1", "c2"]); + assert_single_scan_columns(&best_plan, "t2", &arena, &["c4"]); + + Ok(()) + } + + #[test] + fn test_column_pruning_preserves_union_input_mapping() -> Result<(), DatabaseError> { + let table_state = build_t1_table()?; + let mut arena = PlanArena::new(&table_state.table_arena); + let best_plan = optimize_column_pruning( + &table_state, + &mut arena, + "select c1 from t1 union all select c3 from t2", + )?; + + assert!(contains_operator(&best_plan, |op| matches!( + op, + Operator::Union(_) + ))); + assert_single_scan_columns(&best_plan, "t1", &arena, &["c1"]); + assert_single_scan_columns(&best_plan, "t2", &arena, &["c3"]); + + Ok(()) + } } diff --git a/src/optimizer/rule/normalization/compilation_in_advance.rs b/src/optimizer/rule/normalization/compilation_in_advance.rs index 779d76f6..a62f2e1f 100644 --- a/src/optimizer/rule/normalization/compilation_in_advance.rs +++ b/src/optimizer/rule/normalization/compilation_in_advance.rs @@ -13,10 +13,9 @@ // limitations under the License. use crate::errors::DatabaseError; -use crate::expression::visitor_mut::VisitorMut; use crate::expression::BindEvaluator; use crate::optimizer::core::rule::NormalizationRule; -use crate::planner::operator::join::JoinCondition; +use crate::planner::operator::visitor_mut::{OperatorExprVisitorMut, OperatorVisitorMut}; use crate::planner::operator::Operator; use crate::planner::{Childrens, LogicalPlan, PlanArena}; @@ -27,94 +26,8 @@ pub(crate) fn evaluator_bind_current( plan: &mut LogicalPlan, arena: &PlanArena, ) -> Result<(), DatabaseError> { - let operator = &mut plan.operator; let mut evaluator = BindEvaluator { arena }; - - match operator { - Operator::Join(op) => { - match &mut op.on { - JoinCondition::On { on, filter } => { - for (left_expr, right_expr) in on { - evaluator.visit(left_expr)?; - evaluator.visit(right_expr)?; - } - if let Some(expr) = filter { - evaluator.visit(expr)?; - } - } - JoinCondition::None => {} - } - - return Ok(()); - } - Operator::Aggregate(op) => { - for expr in op.agg_calls.iter_mut().chain(op.groupby_exprs.iter_mut()) { - evaluator.visit(expr)?; - } - } - Operator::Filter(op) => { - evaluator.visit(&mut op.predicate)?; - } - Operator::Project(op) => { - for expr in op.exprs.iter_mut() { - evaluator.visit(expr)?; - } - } - Operator::MarkApply(op) => { - for predicate in op.predicates_mut().iter_mut() { - evaluator.visit(predicate)?; - } - } - Operator::ScalarApply(_) => {} - Operator::Sort(op) => { - for sort_field in op.sort_fields.iter_mut() { - evaluator.visit(&mut sort_field.expr)?; - } - } - Operator::TopK(op) => { - for sort_field in op.sort_fields.iter_mut() { - evaluator.visit(&mut sort_field.expr)?; - } - } - Operator::FunctionScan(op) => { - for expr in op.table_function.args.iter_mut() { - evaluator.visit(expr)?; - } - } - Operator::Update(op) => { - for (_, expr) in op.value_exprs.iter_mut() { - evaluator.visit(expr)?; - } - } - Operator::Dummy - | Operator::TableScan(_) - | Operator::Limit(_) - | Operator::ScalarSubquery(_) - | Operator::Values(_) - | Operator::ShowTable - | Operator::ShowView - | Operator::Explain - | Operator::Describe(_) - | Operator::Insert(_) - | Operator::Delete(_) - | Operator::Analyze(_) - | Operator::AddColumn(_) - | Operator::ChangeColumn(_) - | Operator::DropColumn(_) - | Operator::CreateTable(_) - | Operator::CreateIndex(_) - | Operator::CreateView(_) - | Operator::DropTable(_) - | Operator::DropView(_) - | Operator::DropIndex(_) - | Operator::Truncate(_) - | Operator::Union(_) - | Operator::SetMembership(_) => (), - #[cfg(feature = "copy")] - Operator::CopyFromFile(_) | Operator::CopyToFile(_) => (), - } - - Ok(()) + OperatorExprVisitorMut::new(&mut evaluator).visit_operator(&mut plan.operator) } impl EvaluatorBind { @@ -151,3 +64,181 @@ impl NormalizationRule for EvaluatorBind { Ok(true) } } + +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::*; + use crate::catalog::{ColumnCatalog, ColumnDesc}; + use crate::expression::function::table::{ + ArcTableFunctionImpl, TableFunction, TableFunctionCatalog, TableFunctionImpl, + }; + use crate::expression::{BinaryOperator, ScalarExpression}; + use crate::function::numbers::Numbers; + use crate::optimizer::core::rule::NormalizationRule; + use crate::planner::operator::filter::FilterOperator; + use crate::planner::operator::function_scan::FunctionScanOperator; + use crate::planner::operator::join::{JoinCondition, JoinOperator, JoinType}; + use crate::planner::operator::mark_apply::MarkApplyOperator; + use crate::planner::operator::sort::{SortField, SortOperator}; + use crate::planner::operator::top_k::TopKOperator; + use crate::planner::operator::union::UnionOperator; + use crate::planner::operator::update::UpdateOperator; + use crate::planner::operator::Operator; + use crate::planner::TableArenaCell; + use crate::types::value::DataValue; + use crate::types::LogicalType; + + fn unbound_binary(left: ScalarExpression, right: ScalarExpression) -> ScalarExpression { + ScalarExpression::Binary { + op: BinaryOperator::Plus, + left_expr: Box::new(left), + right_expr: Box::new(right), + evaluator: None, + ty: LogicalType::Integer, + } + } + + fn is_bound(expr: &ScalarExpression) -> bool { + matches!( + expr, + ScalarExpression::Binary { + evaluator: Some(_), + .. + } + ) + } + + #[test] + fn binds_each_expression_container() -> Result<(), DatabaseError> { + let table_arena = TableArenaCell::default(); + let numbers = Numbers::new(); + let mut schema = Vec::new(); + numbers.output_schema_into(table_arena.borrow_mut(), &mut schema); + let mut arena = PlanArena::new(&table_arena); + let column = arena.alloc_column(ColumnCatalog::new( + "value".to_string(), + false, + ColumnDesc::new(LogicalType::Integer, None, false, None)?, + )); + let expr = || { + unbound_binary( + ScalarExpression::column_expr(column, 0), + DataValue::Int32(1).into(), + ) + }; + + let mut operators = vec![ + Operator::Filter(FilterOperator { + predicate: expr(), + is_optimized: false, + having: false, + }), + Operator::Sort(SortOperator { + sort_fields: vec![SortField::from(expr())], + limit: None, + }), + Operator::TopK(TopKOperator { + sort_fields: vec![SortField::from(expr())], + limit: 1, + offset: None, + }), + Operator::MarkApply(MarkApplyOperator::new_exists(column, vec![expr()])), + Operator::Update(UpdateOperator { + table_name: "t1".into(), + value_exprs: vec![(column, expr())], + }), + ]; + + operators.push(Operator::FunctionScan(FunctionScanOperator { + table_function: TableFunction { + args: vec![expr()], + catalog: TableFunctionCatalog { + schema, + inner: ArcTableFunctionImpl(numbers), + }, + }, + })); + + for operator in &mut operators { + let mut plan = LogicalPlan::new(operator.clone(), Childrens::None); + evaluator_bind_current(&mut plan, &arena)?; + match &plan.operator { + Operator::Filter(op) => assert!(is_bound(&op.predicate)), + Operator::Sort(op) => assert!(is_bound(&op.sort_fields[0].expr)), + Operator::TopK(op) => assert!(is_bound(&op.sort_fields[0].expr)), + Operator::MarkApply(op) => assert!(is_bound(&op.predicates()[0])), + Operator::Update(op) => assert!(is_bound(&op.value_exprs[0].1)), + Operator::FunctionScan(op) => assert!(is_bound(&op.table_function.args[0])), + _ => unreachable!(), + } + } + + Ok(()) + } + + #[test] + fn binds_join_conditions_and_selected_twin_subtrees() -> Result<(), DatabaseError> { + let table_arena = TableArenaCell::default(); + let mut arena = PlanArena::new(&table_arena); + let column = arena.alloc_column(ColumnCatalog::new( + "value".to_string(), + false, + ColumnDesc::new(LogicalType::Integer, None, false, None)?, + )); + let expr = || { + unbound_binary( + ScalarExpression::column_expr(column, 0), + DataValue::Int32(1).into(), + ) + }; + let filter = || { + LogicalPlan::new( + Operator::Filter(FilterOperator { + predicate: expr(), + is_optimized: false, + having: false, + }), + Childrens::None, + ) + }; + + let mut join = LogicalPlan::new( + Operator::Join(JoinOperator { + join_type: JoinType::Inner, + on: JoinCondition::On { + on: vec![(expr(), expr())], + filter: Some(expr()), + }, + }), + Childrens::Twins { + left: Box::new(filter()), + right: Box::new(filter()), + }, + ); + assert!(EvaluatorBind.apply(&mut join, &mut arena)?); + let Operator::Join(op) = &join.operator else { + unreachable!() + }; + let JoinCondition::On { + on, + filter: join_filter, + } = &op.on + else { + unreachable!() + }; + assert!(is_bound(&on[0].0)); + assert!(is_bound(&on[0].1)); + assert!(is_bound(join_filter.as_ref().unwrap())); + assert!(join.childrens.iter().all(|child| { + matches!(&child.operator, Operator::Filter(op) if is_bound(&op.predicate)) + })); + + let mut union = UnionOperator::build(vec![column], vec![column], filter(), filter()); + EvaluatorBind.apply(&mut union, &mut arena)?; + assert!(union.childrens.iter().all(|child| { + matches!(&child.operator, Operator::Filter(op) if is_bound(&op.predicate)) + })); + + Ok(()) + } +} diff --git a/src/optimizer/rule/normalization/elimination.rs b/src/optimizer/rule/normalization/elimination.rs index 336f6cea..527f5053 100644 --- a/src/optimizer/rule/normalization/elimination.rs +++ b/src/optimizer/rule/normalization/elimination.rs @@ -44,7 +44,7 @@ impl NormalizationRule for EliminateRedundantSort { Some(child) => child, None => return Ok(false), }; - mark_sort_preserving_indexes(child, &sort_fields, arena); + mark_sort_preserving_indexes(child, &sort_fields, arena)?; let can_remove = ensure_index_order(child, &sort_fields, arena); if !can_remove { @@ -109,8 +109,8 @@ fn mark_sort_preserving_indexes( plan: &mut LogicalPlan, required: &[SortField], arena: &crate::planner::PlanArena, -) { - mark_order_hint(plan, required, OrderHintKind::SortElimination, arena); +) -> Result<(), DatabaseError> { + mark_order_hint(plan, required, OrderHintKind::SortElimination, arena) } #[derive(Copy, Clone)] @@ -140,9 +140,9 @@ fn mark_order_hint( required: &[SortField], hint: OrderHintKind, arena: &crate::planner::PlanArena, -) { +) -> Result<(), DatabaseError> { if required.is_empty() { - return; + return Ok(()); } match &mut plan.operator { @@ -152,14 +152,15 @@ fn mark_order_hint( | Operator::TopK(_) | Operator::Sort(_) => { if let Childrens::Only(child) = plan.childrens.as_mut() { - mark_order_hint(child, required, hint, arena); + mark_order_hint(child, required, hint, arena)?; } } Operator::TableScan(scan_op) => { - apply_scan_order_hint(scan_op, ScanOrderHint::sort_fields(required), hint, arena); + apply_scan_order_hint(scan_op, ScanOrderHint::sort_fields(required), hint, arena)?; } _ => {} } + Ok(()) } pub(crate) fn apply_scan_order_hint( @@ -167,27 +168,25 @@ pub(crate) fn apply_scan_order_hint( required: ScanOrderHint<'_>, hint: OrderHintKind, arena: &crate::planner::PlanArena, -) { - let required_from_table = match required { - ScanOrderHint::SortFields(fields) => fields.iter().all(|field| { - field.expr.all_referenced_columns(arena, |arena, column| { - scan_op - .columns - .iter() - .any(|scan_column| arena.same_column(*scan_column, *column)) - }) - }), - ScanOrderHint::DistinctGroupBy(groupby_exprs) => groupby_exprs.iter().all(|expr| { - expr.all_referenced_columns(arena, |arena, column| { - scan_op - .columns - .iter() - .any(|scan_column| arena.same_column(*scan_column, *column)) - }) - }), - }; +) -> Result<(), DatabaseError> { + let mut required_from_table = true; + for index in 0..hint_len(required) { + let expr = match required { + ScanOrderHint::SortFields(fields) => &fields[index].expr, + ScanOrderHint::DistinctGroupBy(groupby_exprs) => &groupby_exprs[index], + }; + if !expr.all_referenced_columns(arena, |arena, column| { + scan_op + .columns + .iter() + .any(|scan_column| arena.same_column(*scan_column, *column)) + })? { + required_from_table = false; + break; + } + } if !required_from_table { - return; + return Ok(()); } for index_info in scan_op.index_infos.iter_mut() { if hint_covers(required, &index_info.sort_option, arena) { @@ -210,6 +209,7 @@ pub(crate) fn apply_scan_order_hint( } } } + Ok(()) } fn hint_len(required: ScanOrderHint<'_>) -> usize { @@ -733,7 +733,7 @@ mod tests { let c1 = make_sort_field(&mut arena, "c1"); let c2 = make_sort_field(&mut arena, "c2"); let mut plan = build_plan(&mut arena, vec![c2.clone()], vec![c1, c2.clone()], 1); - super::mark_sort_preserving_indexes(&mut plan, &[c2], &arena); + super::mark_sort_preserving_indexes(&mut plan, &[c2], &arena)?; let rule = EliminateRedundantSort; assert!(rule.apply(&mut plan, &mut arena)?); @@ -801,7 +801,7 @@ mod tests { Operator::Sort(sort_op) => sort_op.sort_fields.clone(), _ => unreachable!("expected sort operator"), }; - super::mark_sort_preserving_indexes(&mut plan, &sort_fields, &arena); + super::mark_sort_preserving_indexes(&mut plan, &sort_fields, &arena)?; let table_plan = plan.childrens.pop_only(); match table_plan.operator { @@ -832,7 +832,7 @@ mod tests { &required, super::OrderHintKind::StreamDistinct, &arena, - ); + )?; } let child = plan.childrens.pop_only(); @@ -893,7 +893,7 @@ mod tests { vec![c1.clone(), c2.clone()], 0, ); - super::mark_sort_preserving_indexes(&mut plan, &[c2], &arena); + super::mark_sort_preserving_indexes(&mut plan, &[c2], &arena)?; let rule = EliminateRedundantSort; assert!(!rule.apply(&mut plan, &mut arena)?); @@ -955,7 +955,7 @@ mod tests { Operator::Sort(sort_op) => sort_op.sort_fields.clone(), _ => unreachable!("expected sort operator"), }; - super::mark_sort_preserving_indexes(&mut plan, &sort_fields, &arena); + super::mark_sort_preserving_indexes(&mut plan, &sort_fields, &arena)?; let rule = EliminateRedundantSort; assert!(rule.apply(&mut plan, &mut arena)?); assert!(matches!(plan.operator, Operator::Filter(_))); diff --git a/src/optimizer/rule/normalization/mod.rs b/src/optimizer/rule/normalization/mod.rs index 1f9df2bc..0df08ea5 100644 --- a/src/optimizer/rule/normalization/mod.rs +++ b/src/optimizer/rule/normalization/mod.rs @@ -13,7 +13,7 @@ // limitations under the License. use crate::errors::DatabaseError; -use crate::expression::visitor_mut::{walk_mut_expr, VisitorMut}; +use crate::expression::visitor_mut::{walk_mut_expr, ExprVisitorMut}; use crate::expression::{AliasType, ScalarExpression}; use crate::optimizer::core::rule::NormalizationRule; use crate::optimizer::rule::normalization::column_pruning::ColumnPruning; @@ -250,7 +250,7 @@ struct PositionRemapper<'a> { removed_positions: &'a [usize], } -impl<'a> VisitorMut<'a> for PositionRemapper<'_> { +impl<'a> ExprVisitorMut<'a> for PositionRemapper<'_> { fn visit(&mut self, expr: &'a mut ScalarExpression) -> Result<(), DatabaseError> { match expr { ScalarExpression::ColumnRef { position, .. } => { diff --git a/src/optimizer/rule/normalization/parameterized_index.rs b/src/optimizer/rule/normalization/parameterized_index.rs index be0042f5..cffb6390 100644 --- a/src/optimizer/rule/normalization/parameterized_index.rs +++ b/src/optimizer/rule/normalization/parameterized_index.rs @@ -39,7 +39,7 @@ impl NormalizationRule for ParameterizeMarkApply { left.output_schema(arena), right.output_schema(arena), arena, - ); + )?; let new_probe = probe.and_then(|(right_column, left_expr)| { parameterize_right_subtree(right, &right_column, arena).then_some(left_expr) }); @@ -60,17 +60,26 @@ fn find_parameterized_probe( left_schema: &Schema, right_schema: &Schema, arena: &crate::planner::PlanArena, -) -> Option<(ColumnRef, ScalarExpression)> { +) -> Result, DatabaseError> { match kind { - MarkApplyKind::Exists => predicates.iter().find_map(|predicate| { - extract_parameterized_probe(predicate, left_schema, right_schema, arena) - }), + MarkApplyKind::Exists => { + for predicate in predicates { + if let Some(probe) = + extract_parameterized_probe(predicate, left_schema, right_schema, arena)? + { + return Ok(Some(probe)); + } + } + Ok(None) + } MarkApplyKind::Quantified(MarkApplyQuantifier::Any) => { - predicates.first().and_then(|predicate| { + if let Some(predicate) = predicates.first() { extract_parameterized_probe(predicate, left_schema, right_schema, arena) - }) + } else { + Ok(None) + } } - MarkApplyKind::Quantified(MarkApplyQuantifier::All) => None, + MarkApplyKind::Quantified(MarkApplyQuantifier::All) => Ok(None), } } @@ -79,21 +88,23 @@ fn extract_parameterized_probe( left_schema: &Schema, right_schema: &Schema, arena: &crate::planner::PlanArena, -) -> Option<(ColumnRef, ScalarExpression)> { +) -> Result, DatabaseError> { match predicate.unpack_alias_ref() { ScalarExpression::Binary { op: BinaryOperator::Eq, left_expr, right_expr, .. - } => extract_parameterized_probe_side( - left_expr, - right_expr, - left_schema, - right_schema, - arena, - ) - .or_else(|| { + } => { + if let Some(probe) = extract_parameterized_probe_side( + left_expr, + right_expr, + left_schema, + right_schema, + arena, + )? { + return Ok(Some(probe)); + } extract_parameterized_probe_side( right_expr, left_expr, @@ -101,8 +112,8 @@ fn extract_parameterized_probe( right_schema, arena, ) - }), - _ => None, + } + _ => Ok(None), } } @@ -112,24 +123,26 @@ fn extract_parameterized_probe_side( left_schema: &Schema, right_schema: &Schema, arena: &crate::planner::PlanArena, -) -> Option<(ColumnRef, ScalarExpression)> { - let (right_column, _) = right_expr.unpack_alias_ref().unpack_bound_col(false)?; +) -> Result, DatabaseError> { + let Some((right_column, _)) = right_expr.unpack_alias_ref().unpack_bound_col(false) else { + return Ok(None); + }; if !schema_contains_column(right_schema, &right_column, arena) { - return None; + return Ok(None); } if !left_expr.all_referenced_columns(arena, |arena, candidate| { schema_contains_column(left_schema, candidate, arena) - }) { - return None; + })? { + return Ok(None); } if left_expr.any_referenced_column(arena, |arena, candidate| { schema_contains_column(right_schema, candidate, arena) - }) { - return None; + })? { + return Ok(None); } - Some((right_column, left_expr.clone())) + Ok(Some((right_column, left_expr.clone()))) } fn parameterize_right_subtree( @@ -221,3 +234,144 @@ fn schema_contains_column( .iter() .any(|candidate| arena.same_column(*candidate, *column)) } + +// GRCOV_EXCL_START +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::*; + use crate::catalog::{ColumnCatalog, ColumnDesc}; + use crate::optimizer::core::rule::NormalizationRule; + use crate::planner::operator::filter::FilterOperator; + use crate::planner::{PlanArena, TableArenaCell}; + use crate::types::LogicalType; + + fn column(arena: &mut PlanArena, name: &str) -> ColumnRef { + arena.alloc_column(ColumnCatalog::new( + name.to_string(), + true, + ColumnDesc::new(LogicalType::Integer, None, false, None).unwrap(), + )) + } + + fn eq(left: ScalarExpression, right: ScalarExpression) -> ScalarExpression { + ScalarExpression::Binary { + op: BinaryOperator::Eq, + left_expr: Box::new(left), + right_expr: Box::new(right), + evaluator: None, + ty: LogicalType::Boolean, + } + } + + #[test] + fn probe_detection_covers_quantifiers_and_rejected_sides() -> Result<(), DatabaseError> { + let table_arena = TableArenaCell::default(); + let mut arena = PlanArena::new(&table_arena); + let left = column(&mut arena, "left"); + let right = column(&mut arena, "right"); + let outside = column(&mut arena, "outside"); + let left_schema = vec![left]; + let right_schema = vec![right]; + let overlapping_schema = vec![left, right]; + + assert!(find_parameterized_probe( + MarkApplyKind::Quantified(MarkApplyQuantifier::Any), + &[], + &left_schema, + &right_schema, + &arena, + )? + .is_none()); + assert!(find_parameterized_probe( + MarkApplyKind::Quantified(MarkApplyQuantifier::All), + &[eq( + ScalarExpression::column_expr(right, 0), + ScalarExpression::column_expr(left, 0), + )], + &left_schema, + &right_schema, + &arena, + )? + .is_none()); + + let predicates = vec![ + ScalarExpression::from(true), + eq( + ScalarExpression::column_expr(right, 0), + ScalarExpression::column_expr(left, 0), + ), + ]; + let probe = find_parameterized_probe( + MarkApplyKind::Exists, + &predicates, + &left_schema, + &right_schema, + &arena, + )? + .expect("right = left should be parameterizable"); + assert_eq!(probe.0, right); + + assert!(extract_parameterized_probe( + &eq( + ScalarExpression::column_expr(right, 0), + ScalarExpression::column_expr(outside, 0), + ), + &left_schema, + &right_schema, + &arena, + )? + .is_none()); + assert!(extract_parameterized_probe( + &eq( + ScalarExpression::column_expr(right, 0), + ScalarExpression::column_expr(right, 0), + ), + &overlapping_schema, + &right_schema, + &arena, + )? + .is_none()); + assert!(extract_parameterized_probe( + &ScalarExpression::from(false), + &left_schema, + &right_schema, + &arena, + )? + .is_none()); + + Ok(()) + } + + #[test] + fn parameterization_rejects_unsupported_operator_and_child_shapes() -> Result<(), DatabaseError> + { + let table_arena = TableArenaCell::default(); + let mut arena = PlanArena::new(&table_arena); + let column = column(&mut arena, "value"); + let mut plan = LogicalPlan::new(Operator::Dummy, Childrens::None); + + assert!(!ParameterizeMarkApply.apply(&mut plan, &mut arena)?); + assert!(!parameterize_right_subtree(&mut plan, &column, &arena)); + + let mut filter = LogicalPlan::new( + Operator::Filter(FilterOperator { + predicate: ScalarExpression::from(true), + is_optimized: false, + having: false, + }), + Childrens::None, + ); + assert!(!parameterize_right_subtree(&mut filter, &column, &arena)); + + assert_eq!( + index_priority(IndexType::PrimaryKey { is_multiple: false }), + 0 + ); + assert_eq!(index_priority(IndexType::Unique), 1); + assert_eq!(index_priority(IndexType::Composite), 2); + assert_eq!(index_priority(IndexType::Normal), 3); + + Ok(()) + } +} +// GRCOV_EXCL_STOP diff --git a/src/optimizer/rule/normalization/pushdown_predicates.rs b/src/optimizer/rule/normalization/pushdown_predicates.rs index d27eef69..e9e93ed6 100644 --- a/src/optimizer/rule/normalization/pushdown_predicates.rs +++ b/src/optimizer/rule/normalization/pushdown_predicates.rs @@ -14,9 +14,8 @@ use crate::errors::DatabaseError; use crate::expression::range_detacher::{DetachedPredicate, Range, RangeDetacher}; -use crate::expression::visitor_mut::{PositionShift, VisitorMut}; +use crate::expression::visitor_mut::{ExprVisitorMut, PositionShift}; use crate::expression::{BinaryOperator, ScalarExpression}; -use crate::iter_ext::Itertools; use crate::optimizer::core::rule::NormalizationRule; use crate::optimizer::plan_utils::{replace_with_only_child, wrap_child_with}; use crate::planner::operator::filter::FilterOperator; @@ -30,25 +29,79 @@ use std::ops::Bound; use std::{borrow::Cow, mem, slice}; const EMPTY_SCHEMA: [crate::catalog::ColumnRef; 0] = []; - -fn split_conjunctive_predicates(expr: &ScalarExpression) -> Vec { +type ClassifiedJoinFilters = ( + Option, + Option, + Option, + usize, +); + +fn split_conjunctive_predicates( + expr: &ScalarExpression, + f: &mut impl FnMut(ScalarExpression) -> Result<(), DatabaseError>, +) -> Result<(), DatabaseError> { match expr { ScalarExpression::Binary { op: BinaryOperator::And, left_expr, right_expr, .. - } => split_conjunctive_predicates(left_expr) - .into_iter() - .chain(split_conjunctive_predicates(right_expr)) - .collect_vec(), - _ => vec![expr.clone()], + } => { + split_conjunctive_predicates(left_expr, f)?; + split_conjunctive_predicates(right_expr, f) + } + _ => f(expr.clone()), } } +fn classify_join_filters( + predicate: &ScalarExpression, + childrens: &mut Childrens, + arena: &mut crate::planner::PlanArena, +) -> Result { + let (left_columns, right_columns): ( + &[crate::catalog::ColumnRef], + &[crate::catalog::ColumnRef], + ) = match childrens { + Childrens::Only(left) => (left.output_schema(arena), &EMPTY_SCHEMA), + Childrens::Twins { left, right } => (left.output_schema(arena), right.output_schema(arena)), + Childrens::None => (&EMPTY_SCHEMA, &EMPTY_SCHEMA), + }; + let left_len = left_columns.len(); + let append_filter = |slot: &mut Option, expr| { + *slot = Some(match slot.take() { + Some(current) => ScalarExpression::Binary { + op: BinaryOperator::And, + left_expr: Box::new(current), + right_expr: Box::new(expr), + evaluator: None, + ty: LogicalType::Boolean, + }, + None => expr, + }); + }; + let mut left_filter = None; + let mut right_filter = None; + let mut common_filter = None; + split_conjunctive_predicates(predicate, &mut |expr| { + if expr.all_referenced_columns(arena, |_, column| left_columns.contains(column))? { + append_filter(&mut left_filter, expr); + } else if expr.all_referenced_columns(arena, |_, column| right_columns.contains(column))? { + append_filter(&mut right_filter, expr); + } else { + append_filter(&mut common_filter, expr); + } + Ok(()) + })?; + Ok((left_filter, right_filter, common_filter, left_len)) +} + /// reduce filters into a filter, and then build a new LogicalFilter node with input child. /// if filters is empty, return the input child. -fn reduce_filters(filters: Vec, having: bool) -> Option { +fn reduce_filters( + filters: impl IntoIterator, + having: bool, +) -> Option { filters .into_iter() .reduce(|a, b| ScalarExpression::Binary { @@ -65,23 +118,6 @@ fn reduce_filters(filters: Vec, having: bool) -> Option Result<(), DatabaseError> { - if filters.is_empty() { - return Ok(()); - } - - let mut localizer = PositionShift { - delta: -(left_len as isize), - }; - for expr in filters { - localizer.visit(expr)?; - } - Ok(()) -} - /// Comments copied from Spark Catalyst PushPredicateThroughJoin /// /// Pushes down `Filter` operators where the `condition` can be @@ -128,65 +164,55 @@ impl NormalizationRule for PushPredicateThroughJoin { return Ok(false); } - let filter_exprs = split_conjunctive_predicates(&filter_op.predicate); - let left_columns: &[crate::catalog::ColumnRef] = match join_plan.childrens.as_mut() { - Childrens::Only(left) => left.output_schema(arena), - Childrens::Twins { left, .. } => left.output_schema(arena), - Childrens::None => &EMPTY_SCHEMA, - }; - let (left_filters, rest): (Vec<_>, Vec<_>) = filter_exprs.into_iter().partition(|f| { - f.all_referenced_columns(arena, |_, column| left_columns.contains(column)) - }); - let left_len = left_columns.len(); - - let right_columns: &[crate::catalog::ColumnRef] = match join_plan.childrens.as_mut() { - Childrens::Twins { right, .. } => right.output_schema(arena), - _ => &EMPTY_SCHEMA, - }; - let (right_filters, common_filters): (Vec<_>, Vec<_>) = - rest.into_iter().partition(|f| { - f.all_referenced_columns(arena, |_, column| right_columns.contains(column)) - }); + let (left_filter, mut right_filter, common_filter, left_len) = + classify_join_filters(&filter_op.predicate, join_plan.childrens.as_mut(), arena)?; let mut new_ops = (None, None, None); - let replace_filters = match join_type { + match join_type { JoinType::Inner => { - if let Some(left_filter_op) = reduce_filters(left_filters, filter_op.having) { + if let Some(left_filter_op) = reduce_filters(left_filter, filter_op.having) { new_ops.0 = Some(Operator::Filter(left_filter_op)); } - let mut right_filters = right_filters; - localize_right_filters(&mut right_filters, left_len)?; - if let Some(right_filter_op) = reduce_filters(right_filters, filter_op.having) { + if let Some(expr) = &mut right_filter { + PositionShift { + delta: -(left_len as isize), + } + .visit(expr)?; + } + if let Some(right_filter_op) = reduce_filters(right_filter, filter_op.having) { new_ops.1 = Some(Operator::Filter(right_filter_op)); } - - common_filters + new_ops.2 = + reduce_filters(common_filter, filter_op.having).map(Operator::Filter); } JoinType::LeftOuter => { - if let Some(left_filter_op) = reduce_filters(left_filters, filter_op.having) { + if let Some(left_filter_op) = reduce_filters(left_filter, filter_op.having) { new_ops.0 = Some(Operator::Filter(left_filter_op)); } - - common_filters - .into_iter() - .chain(right_filters) - .collect_vec() + new_ops.2 = reduce_filters( + common_filter.into_iter().chain(right_filter), + filter_op.having, + ) + .map(Operator::Filter); } JoinType::RightOuter => { - let mut right_filters = right_filters; - localize_right_filters(&mut right_filters, left_len)?; - if let Some(right_filter_op) = reduce_filters(right_filters, filter_op.having) { + if let Some(expr) = &mut right_filter { + PositionShift { + delta: -(left_len as isize), + } + .visit(expr)?; + } + if let Some(right_filter_op) = reduce_filters(right_filter, filter_op.having) { new_ops.1 = Some(Operator::Filter(right_filter_op)); } - - common_filters.into_iter().chain(left_filters).collect_vec() + new_ops.2 = reduce_filters( + common_filter.into_iter().chain(left_filter), + filter_op.having, + ) + .map(Operator::Filter); } - _ => vec![], - }; - - if let Some(replace_filter_op) = reduce_filters(replace_filters, filter_op.having) { - new_ops.2 = Some(Operator::Filter(replace_filter_op)); + _ => {} } if let Some(left_op) = new_ops.0 { @@ -427,25 +453,8 @@ impl NormalizationRule for PushJoinPredicateIntoScan { (join_op.join_type, filter_expr) }; - let filter_exprs = split_conjunctive_predicates(&filter_expr); - let left_columns: &[crate::catalog::ColumnRef] = match plan.childrens.as_mut() { - Childrens::Only(left) => left.output_schema(arena), - Childrens::Twins { left, .. } => left.output_schema(arena), - Childrens::None => &EMPTY_SCHEMA, - }; - let (left_filters, rest): (Vec<_>, Vec<_>) = filter_exprs.into_iter().partition(|expr| { - expr.all_referenced_columns(arena, |_, column| left_columns.contains(column)) - }); - let left_len = left_columns.len(); - - let right_columns: &[crate::catalog::ColumnRef] = match plan.childrens.as_mut() { - Childrens::Twins { right, .. } => right.output_schema(arena), - _ => &EMPTY_SCHEMA, - }; - let (right_filters, common_filters): (Vec<_>, Vec<_>) = - rest.into_iter().partition(|expr| { - expr.all_referenced_columns(arena, |_, column| right_columns.contains(column)) - }); + let (left_filter, mut right_filter, common_filter, left_len) = + classify_join_filters(&filter_expr, plan.childrens.as_mut(), arena)?; let (push_left, push_right) = match join_type { JoinType::Inner => (true, true), @@ -455,30 +464,29 @@ impl NormalizationRule for PushJoinPredicateIntoScan { }; let mut new_ops = (None, None); - let mut remaining_filters = common_filters; - - let (left_push, left_remain) = if push_left { - (left_filters, Vec::new()) + let left_remain = if push_left { + if let Some(filter_op) = reduce_filters(left_filter, false) { + new_ops.0 = Some(Operator::Filter(filter_op)); + } + None } else { - (Vec::new(), left_filters) + left_filter }; - if let Some(filter_op) = reduce_filters(left_push, false) { - new_ops.0 = Some(Operator::Filter(filter_op)); - } else { - remaining_filters.extend(left_remain); - } - let (mut right_push, right_remain) = if push_right { - (right_filters, Vec::new()) + let right_remain = if push_right { + if let Some(expr) = &mut right_filter { + PositionShift { + delta: -(left_len as isize), + } + .visit(expr)?; + } + if let Some(filter_op) = reduce_filters(right_filter, false) { + new_ops.1 = Some(Operator::Filter(filter_op)); + } + None } else { - (Vec::new(), right_filters) + right_filter }; - localize_right_filters(&mut right_push, left_len)?; - if let Some(filter_op) = reduce_filters(right_push, false) { - new_ops.1 = Some(Operator::Filter(filter_op)); - } else { - remaining_filters.extend(right_remain); - } let mut applied = false; if let Some(left_op) = new_ops.0 { @@ -488,7 +496,14 @@ impl NormalizationRule for PushJoinPredicateIntoScan { applied |= wrap_child_with(plan, 1, right_op); } - let mut join_filter = reduce_filters(remaining_filters, false).map(|op| op.predicate); + let mut join_filter = reduce_filters( + common_filter + .into_iter() + .chain(left_remain) + .chain(right_remain), + false, + ) + .map(|op| op.predicate); let filter_changed = match &join_filter { Some(expr) => expr != &filter_expr, None => true, diff --git a/src/optimizer/rule/normalization/simplification.rs b/src/optimizer/rule/normalization/simplification.rs index 50b81f8e..308b1409 100644 --- a/src/optimizer/rule/normalization/simplification.rs +++ b/src/optimizer/rule/normalization/simplification.rs @@ -14,9 +14,9 @@ use crate::errors::DatabaseError; use crate::expression::simplify::{ConstantCalculator, Simplify}; -use crate::expression::visitor_mut::VisitorMut; +use crate::expression::visitor_mut::ExprVisitorMut; use crate::optimizer::core::rule::NormalizationRule; -use crate::planner::operator::join::JoinCondition; +use crate::planner::operator::visitor_mut::{OperatorExprVisitorMut, OperatorVisitorMut}; use crate::planner::operator::Operator; use crate::planner::{Childrens, LogicalPlan}; @@ -27,43 +27,8 @@ pub(crate) fn constant_calculation_current( plan: &mut LogicalPlan, arena: &crate::planner::PlanArena, ) -> Result<(), DatabaseError> { - let operator = &mut plan.operator; let mut calculator = ConstantCalculator::new(arena); - - match operator { - Operator::Aggregate(op) => { - for expr in op.agg_calls.iter_mut().chain(op.groupby_exprs.iter_mut()) { - calculator.visit(expr)?; - } - } - Operator::Filter(op) => { - calculator.visit(&mut op.predicate)?; - } - Operator::Join(op) => { - if let JoinCondition::On { on, filter } = &mut op.on { - for (left_expr, right_expr) in on { - calculator.visit(left_expr)?; - calculator.visit(right_expr)?; - } - if let Some(expr) = filter { - calculator.visit(expr)?; - } - } - } - Operator::Project(op) => { - for expr in &mut op.exprs { - calculator.visit(expr)?; - } - } - Operator::Sort(op) => { - for field in &mut op.sort_fields { - calculator.visit(&mut field.expr)?; - } - } - _ => (), - } - - Ok(()) + OperatorExprVisitorMut::new(&mut calculator).visit_operator(&mut plan.operator) } impl ConstantCalculation { diff --git a/src/planner/mod.rs b/src/planner/mod.rs index ca6eda78..559831fe 100644 --- a/src/planner/mod.rs +++ b/src/planner/mod.rs @@ -16,6 +16,7 @@ mod arena; pub mod operator; use crate::catalog::TableName; +use crate::errors::DatabaseError; use crate::planner::operator::set_membership::SetMembershipOperator; use crate::planner::operator::union::UnionOperator; use crate::planner::operator::values::ValuesOperator; @@ -128,7 +129,11 @@ impl LogicalPlan { tables } - pub(crate) fn visit_column_refs(&self, arena: &mut A, f: &mut F) + pub(crate) fn visit_column_refs( + &self, + arena: &mut A, + f: &mut F, + ) -> Result<(), DatabaseError> where A: MetaArena, F: FnMut(&crate::catalog::ColumnRef) + ?Sized, @@ -137,10 +142,11 @@ impl LogicalPlan { .visit_referenced_columns(arena, &mut |_, column| { f(column); true - }); + })?; for child in self.childrens.iter() { - child.visit_column_refs(arena, f); + child.visit_column_refs(arena, f)?; } + Ok(()) } pub fn output_schema<'plan>( diff --git a/src/planner/operator/mark_apply.rs b/src/planner/operator/mark_apply.rs index bbe5c32d..68758d70 100644 --- a/src/planner/operator/mark_apply.rs +++ b/src/planner/operator/mark_apply.rs @@ -35,9 +35,9 @@ pub enum MarkApplyKind { #[derive(Debug, PartialEq, Eq, Clone, Hash, ReferenceSerialization)] pub struct MarkApplyOperator { pub kind: MarkApplyKind, - predicates: Vec, + pub predicates: Vec, output_column: ColumnRef, - parameterized_probe: Option, + pub parameterized_probe: Option, } impl MarkApplyOperator { diff --git a/src/planner/operator/mod.rs b/src/planner/operator/mod.rs index 076b00a1..e3b03ac4 100644 --- a/src/planner/operator/mod.rs +++ b/src/planner/operator/mod.rs @@ -44,6 +44,8 @@ pub mod truncate; pub mod union; pub mod update; pub mod values; +pub mod visitor; +pub mod visitor_mut; use self::{ aggregate::AggregateOperator, alter_table::add_column::AddColumnOperator, @@ -53,6 +55,10 @@ use self::{ table_scan::TableScanOperator, }; use crate::catalog::ColumnRef; +use crate::errors::DatabaseError; +use crate::expression::visitor::{walk_expr, ExprVisitor}; +use crate::expression::ScalarExpression; +use crate::planner::operator::alter_table::change_column::DefaultChange as ColumnDefaultChange; use crate::planner::operator::alter_table::drop_column::DropColumnOperator; use crate::planner::operator::analyze::AnalyzeOperator; #[cfg(feature = "copy")] @@ -77,6 +83,7 @@ use crate::planner::operator::truncate::TruncateOperator; use crate::planner::operator::union::UnionOperator; use crate::planner::operator::update::UpdateOperator; use crate::planner::operator::values::ValuesOperator; +use crate::planner::operator::visitor::OperatorVisitor; use crate::planner::{MetaArena, PlanArena}; use crate::types::index::IndexInfo; use kite_sql_serde_macros::ReferenceSerialization; @@ -200,120 +207,224 @@ impl Operator { &self, arena: &mut A, f: &mut impl FnMut(&mut A, &ColumnRef) -> bool, - ) -> bool { - match self { - Operator::Aggregate(op) => op - .agg_calls - .iter() - .chain(op.groupby_exprs.iter()) - .all(|expr| expr.visit_referenced_columns(arena, f)), - Operator::ScalarApply(_) => true, - Operator::MarkApply(op) => op - .predicates() - .iter() - .all(|expr| expr.visit_referenced_columns(arena, f)), - Operator::Filter(op) => op.predicate.visit_referenced_columns(arena, f), - Operator::Join(op) => { + ) -> Result { + struct ReferencedColumnVisitor<'a, A, F> { + arena: &'a mut A, + f: &'a mut F, + keep_going: bool, + } + + impl<'expr, A, F> ExprVisitor<'expr> for ReferencedColumnVisitor<'_, A, F> + where + F: FnMut(&mut A, &ColumnRef) -> bool, + { + fn visit(&mut self, expr: &'expr ScalarExpression) -> Result<(), DatabaseError> { + if self.keep_going { + walk_expr(self, expr)?; + } + Ok(()) + } + + fn visit_column_ref(&mut self, column: &'expr ColumnRef) -> Result<(), DatabaseError> { + if self.keep_going { + self.keep_going = (self.f)(self.arena, column); + } + Ok(()) + } + } + + impl<'operator, A, F> OperatorVisitor<'operator> for ReferencedColumnVisitor<'_, A, F> + where + F: FnMut(&mut A, &ColumnRef) -> bool, + { + fn visit_aggregate( + &mut self, + op: &'operator AggregateOperator, + ) -> Result<(), DatabaseError> { + for expr in op.agg_calls.iter().chain(&op.groupby_exprs) { + ExprVisitor::visit(self, expr)?; + } + Ok(()) + } + + fn visit_mark_apply( + &mut self, + op: &'operator MarkApplyOperator, + ) -> Result<(), DatabaseError> { + for expr in &op.predicates { + ExprVisitor::visit(self, expr)?; + } + if let Some(expr) = &op.parameterized_probe { + ExprVisitor::visit(self, expr)?; + } + Ok(()) + } + + fn visit_filter(&mut self, op: &'operator FilterOperator) -> Result<(), DatabaseError> { + ExprVisitor::visit(self, &op.predicate) + } + + fn visit_join(&mut self, op: &'operator JoinOperator) -> Result<(), DatabaseError> { if let JoinCondition::On { on, filter } = &op.on { for (left_expr, right_expr) in on { - if !left_expr.visit_referenced_columns(arena, f) - || !right_expr.visit_referenced_columns(arena, f) - { - return false; - } + ExprVisitor::visit(self, left_expr)?; + ExprVisitor::visit(self, right_expr)?; } - - if let Some(filter_expr) = filter { - return filter_expr.visit_referenced_columns(arena, f); + if let Some(expr) = filter { + ExprVisitor::visit(self, expr)?; } } - true + Ok(()) + } + + fn visit_project( + &mut self, + op: &'operator ProjectOperator, + ) -> Result<(), DatabaseError> { + for expr in &op.exprs { + ExprVisitor::visit(self, expr)?; + } + Ok(()) } - Operator::Project(op) => op - .exprs - .iter() - .all(|expr| expr.visit_referenced_columns(arena, f)), - Operator::ScalarSubquery(_) => true, - Operator::TableScan(op) => op.columns.iter().all(|column| f(arena, column)), - Operator::FunctionScan(op) => op - .table_function - .args - .iter() - .all(|expr| expr.visit_referenced_columns(arena, f)), - Operator::Sort(op) => op - .sort_fields - .iter() - .map(|field| &field.expr) - .all(|expr| expr.visit_referenced_columns(arena, f)), - Operator::TopK(op) => op - .sort_fields - .iter() - .map(|field| &field.expr) - .all(|expr| expr.visit_referenced_columns(arena, f)), - Operator::Values(ValuesOperator { schema_ref, .. }) => { - schema_ref.iter().all(|column| f(arena, column)) + + fn visit_table_scan( + &mut self, + op: &'operator TableScanOperator, + ) -> Result<(), DatabaseError> { + for column in &op.columns { + self.visit_column_ref(column)?; + } + Ok(()) + } + + fn visit_function_scan( + &mut self, + op: &'operator FunctionScanOperator, + ) -> Result<(), DatabaseError> { + for expr in &op.table_function.args { + ExprVisitor::visit(self, expr)?; + } + Ok(()) + } + + fn visit_sort(&mut self, op: &'operator SortOperator) -> Result<(), DatabaseError> { + for field in &op.sort_fields { + ExprVisitor::visit(self, &field.expr)?; + } + Ok(()) + } + + fn visit_top_k(&mut self, op: &'operator TopKOperator) -> Result<(), DatabaseError> { + for field in &op.sort_fields { + ExprVisitor::visit(self, &field.expr)?; + } + Ok(()) + } + + fn visit_values(&mut self, op: &'operator ValuesOperator) -> Result<(), DatabaseError> { + for column in &op.schema_ref { + self.visit_column_ref(column)?; + } + Ok(()) + } + + fn visit_union(&mut self, op: &'operator UnionOperator) -> Result<(), DatabaseError> { + for column in op.left_schema_ref.iter().chain(&op._right_schema_ref) { + self.visit_column_ref(column)?; + } + Ok(()) + } + + fn visit_set_membership( + &mut self, + op: &'operator SetMembershipOperator, + ) -> Result<(), DatabaseError> { + for column in op.left_schema_ref.iter().chain(&op._right_schema_ref) { + self.visit_column_ref(column)?; + } + Ok(()) + } + + fn visit_delete(&mut self, op: &'operator DeleteOperator) -> Result<(), DatabaseError> { + for column in &op.primary_keys { + self.visit_column_ref(column)?; + } + Ok(()) + } + + fn visit_update(&mut self, op: &'operator UpdateOperator) -> Result<(), DatabaseError> { + for (_, expr) in &op.value_exprs { + ExprVisitor::visit(self, expr)?; + } + Ok(()) + } + + fn visit_add_column( + &mut self, + op: &'operator AddColumnOperator, + ) -> Result<(), DatabaseError> { + if let Some(expr) = &op.column.desc().default { + ExprVisitor::visit(self, expr)?; + } + Ok(()) + } + + fn visit_change_column( + &mut self, + op: &'operator ChangeColumnOperator, + ) -> Result<(), DatabaseError> { + if let ColumnDefaultChange::Set(expr) = &op.default_change { + ExprVisitor::visit(self, expr)?; + } + Ok(()) + } + + fn visit_create_table( + &mut self, + op: &'operator CreateTableOperator, + ) -> Result<(), DatabaseError> { + for column in &op.columns { + if let Some(expr) = &column.desc().default { + ExprVisitor::visit(self, expr)?; + } + } + Ok(()) } - Operator::Union(UnionOperator { - left_schema_ref, - _right_schema_ref, - }) - | Operator::SetMembership(SetMembershipOperator { - left_schema_ref, - _right_schema_ref, - .. - }) => left_schema_ref - .iter() - .chain(_right_schema_ref.iter()) - .all(|column| f(arena, column)), - Operator::Analyze(_) => true, - Operator::Delete(op) => op.primary_keys.iter().all(|column| f(arena, column)), - Operator::Dummy - | Operator::Limit(_) - | Operator::ShowTable - | Operator::ShowView - | Operator::Explain - | Operator::Describe(_) - | Operator::Insert(_) - | Operator::Update(_) - | Operator::AddColumn(_) - | Operator::ChangeColumn(_) - | Operator::DropColumn(_) - | Operator::CreateTable(_) - | Operator::CreateIndex(_) - | Operator::CreateView(_) - | Operator::DropTable(_) - | Operator::DropView(_) - | Operator::DropIndex(_) - | Operator::Truncate(_) => true, - #[cfg(feature = "copy")] - Operator::CopyFromFile(_) | Operator::CopyToFile(_) => true, } + + let mut visitor = ReferencedColumnVisitor { + arena, + f, + keep_going: true, + }; + visitor.visit_operator(self)?; + Ok(visitor.keep_going) } pub fn any_referenced_column( &self, arena: &mut PlanArena, mut predicate: impl FnMut(&ColumnRef) -> bool, - ) -> bool { + ) -> Result { let mut found = false; self.visit_referenced_columns(arena, &mut |_, column| { found = predicate(column); !found - }); - found + })?; + Ok(found) } pub fn all_referenced_columns( &self, arena: &mut PlanArena, mut predicate: impl FnMut(&ColumnRef) -> bool, - ) -> bool { + ) -> Result { let mut all = true; self.visit_referenced_columns(arena, &mut |_, column| { all = predicate(column); all - }); - all + })?; + Ok(all) } } @@ -482,13 +593,16 @@ mod tests { ScalarExpression::column_expr(column, position) } - fn referenced_columns(operator: &Operator, arena: &mut PlanArena) -> Vec { + fn referenced_columns( + operator: &Operator, + arena: &mut PlanArena, + ) -> Result, DatabaseError> { let mut columns = Vec::new(); operator.visit_referenced_columns(arena, &mut |_, column| { columns.push(*column); true - }); - columns + })?; + Ok(columns) } #[test] @@ -553,10 +667,15 @@ mod tests { PlanImpl::IndexScan(Box::new(index_info())).to_string(), "IndexScan By #4 => EMPTY" ); + #[cfg(feature = "copy")] + { + assert_eq!(PlanImpl::CopyFromFile.to_string(), "CopyFromFile"); + assert_eq!(PlanImpl::CopyToFile.to_string(), "CopyToFile"); + } } #[test] - fn referenced_column_helpers_stop_on_predicate_result() { + fn referenced_column_helpers_stop_on_predicate_result() -> Result<(), DatabaseError> { let table_arena = TableArenaCell::default(); let mut arena = PlanArena::new(&table_arena); let left = column("left", &mut arena); @@ -566,24 +685,26 @@ mod tests { schema_ref: vec![left, right], }); - assert!(values.any_referenced_column(&mut arena, |column| *column == right)); + assert!(values.any_referenced_column(&mut arena, |column| *column == right)?); assert!(!values - .any_referenced_column(&mut arena, |column| { *column != left && *column != right })); - assert!(values - .all_referenced_columns(&mut arena, |column| { *column == left || *column == right })); - assert!(!values.all_referenced_columns(&mut arena, |column| *column == left)); + .any_referenced_column(&mut arena, |column| { *column != left && *column != right })?); + assert!(values.all_referenced_columns(&mut arena, |column| { + *column == left || *column == right + })?); + assert!(!values.all_referenced_columns(&mut arena, |column| *column == left)?); let delete = Operator::Delete(DeleteOperator { table_name: "users".into(), primary_keys: vec![left], }); - assert!(delete.any_referenced_column(&mut arena, |column| *column == left)); - assert!(Operator::Dummy.all_referenced_columns(&mut arena, |_| false)); - assert!(!Operator::Dummy.any_referenced_column(&mut arena, |_| true)); + assert!(delete.any_referenced_column(&mut arena, |column| *column == left)?); + assert!(Operator::Dummy.all_referenced_columns(&mut arena, |_| false)?); + assert!(!Operator::Dummy.any_referenced_column(&mut arena, |_| true)?); + Ok(()) } #[test] - fn referenced_column_visitor_covers_expression_driven_variants() { + fn referenced_column_visitor_covers_expression_driven_variants() -> Result<(), DatabaseError> { let table_arena = TableArenaCell::default(); let mut arena = PlanArena::new(&table_arena); let a = column("a", &mut arena); @@ -596,18 +717,21 @@ mod tests { groupby_exprs: vec![column_expr(b, 1)], is_distinct: false, }); - assert_eq!(referenced_columns(&aggregate, &mut arena), vec![a, b]); + assert_eq!(referenced_columns(&aggregate, &mut arena)?, vec![a, b]); - let mark_apply = - Operator::MarkApply(MarkApplyOperator::new_exists(d, vec![column_expr(c, 2)])); - assert_eq!(referenced_columns(&mark_apply, &mut arena), vec![c]); + let mut mark_apply = MarkApplyOperator::new_exists(d, vec![column_expr(c, 2)]); + mark_apply.set_parameterized_probe(Some(column_expr(d, 3))); + assert_eq!( + referenced_columns(&Operator::MarkApply(mark_apply), &mut arena)?, + vec![c, d] + ); let filter = Operator::Filter(FilterOperator { predicate: column_expr(a, 0), is_optimized: false, having: false, }); - assert_eq!(referenced_columns(&filter, &mut arena), vec![a]); + assert_eq!(referenced_columns(&filter, &mut arena)?, vec![a]); let join = Operator::Join(JoinOperator { join_type: join::JoinType::Inner, @@ -616,13 +740,61 @@ mod tests { filter: Some(column_expr(c, 2)), }, }); - assert_eq!(referenced_columns(&join, &mut arena), vec![a, b, c]); - assert!(!join.all_referenced_columns(&mut arena, |column| *column == a)); + assert_eq!(referenced_columns(&join, &mut arena)?, vec![a, b, c]); + assert!(!join.all_referenced_columns(&mut arena, |column| *column == a)?); let project = Operator::Project(ProjectOperator { exprs: vec![column_expr(b, 1), column_expr(c, 2)], }); - assert_eq!(referenced_columns(&project, &mut arena), vec![b, c]); + assert_eq!(referenced_columns(&project, &mut arena)?, vec![b, c]); + + let update = Operator::Update(UpdateOperator { + table_name: "users".into(), + value_exprs: vec![(b, column_expr(a, 0))], + }); + assert_eq!(referenced_columns(&update, &mut arena)?, vec![a]); + + let add_column = Operator::AddColumn(AddColumnOperator { + table_name: "users".into(), + if_not_exists: false, + column: ColumnCatalog::new( + "added".to_string(), + true, + ColumnDesc::new( + LogicalType::Integer, + None, + false, + Some(ScalarExpression::from(1_i32)), + )?, + ), + }); + assert!(referenced_columns(&add_column, &mut arena)?.is_empty()); + + let change_column = Operator::ChangeColumn(ChangeColumnOperator { + table_name: "users".into(), + old_column_name: "old".to_string(), + new_column_name: "new".to_string(), + data_type: LogicalType::Integer, + default_change: DefaultChange::Set(column_expr(b, 1)), + not_null_change: NotNullChange::NoChange, + }); + assert_eq!(referenced_columns(&change_column, &mut arena)?, vec![b]); + + let create_table = Operator::CreateTable(CreateTableOperator { + table_name: "created".into(), + columns: vec![ColumnCatalog::new( + "value".to_string(), + true, + ColumnDesc::new( + LogicalType::Integer, + None, + false, + Some(ScalarExpression::from(2_i32)), + )?, + )], + if_not_exists: false, + }); + assert!(referenced_columns(&create_table, &mut arena)?.is_empty()); let table_scan = Operator::TableScan(TableScanOperator { table_name: "users".into(), @@ -631,7 +803,7 @@ mod tests { index_infos: Vec::new(), with_pk: false, }); - assert_eq!(referenced_columns(&table_scan, &mut arena), vec![a, d]); + assert_eq!(referenced_columns(&table_scan, &mut arena)?, vec![a, d]); let function_scan = Operator::FunctionScan(FunctionScanOperator { table_function: TableFunction { @@ -642,39 +814,39 @@ mod tests { }, }, }); - assert_eq!(referenced_columns(&function_scan, &mut arena), vec![c]); + assert_eq!(referenced_columns(&function_scan, &mut arena)?, vec![c]); let sort = Operator::Sort(SortOperator { sort_fields: vec![SortField::from(column_expr(a, 0))], limit: None, }); - assert_eq!(referenced_columns(&sort, &mut arena), vec![a]); + assert_eq!(referenced_columns(&sort, &mut arena)?, vec![a]); let top_k = Operator::TopK(TopKOperator { sort_fields: vec![SortField::from(column_expr(b, 1))], limit: 3, offset: None, }); - assert_eq!(referenced_columns(&top_k, &mut arena), vec![b]); + assert_eq!(referenced_columns(&top_k, &mut arena)?, vec![b]); let union = Operator::Union(UnionOperator { left_schema_ref: vec![a], _right_schema_ref: vec![b], }); - assert_eq!(referenced_columns(&union, &mut arena), vec![a, b]); + assert_eq!(referenced_columns(&union, &mut arena)?, vec![a, b]); let set_membership = Operator::SetMembership(SetMembershipOperator { kind: SetMembershipKind::Intersect, left_schema_ref: vec![c], _right_schema_ref: vec![d], }); - assert_eq!(referenced_columns(&set_membership, &mut arena), vec![c, d]); + assert_eq!(referenced_columns(&set_membership, &mut arena)?, vec![c, d]); let delete = Operator::Delete(DeleteOperator { table_name: "users".into(), primary_keys: vec![a], }); - assert_eq!(referenced_columns(&delete, &mut arena), vec![a]); + assert_eq!(referenced_columns(&delete, &mut arena)?, vec![a]); let no_reference_operators = [ Operator::ScalarApply(ScalarApplyOperator), @@ -686,8 +858,9 @@ mod tests { }), ]; for operator in no_reference_operators { - assert!(referenced_columns(&operator, &mut arena).is_empty()); + assert!(referenced_columns(&operator, &mut arena)?.is_empty()); } + Ok(()) } #[test] @@ -1008,7 +1181,7 @@ mod tests { #[cfg(feature = "copy")] #[test] - fn copy_from_file_display_formats_source_table_and_schema() { + fn copy_display_formats_source_target_table_and_schema() { use crate::binder::copy::{ExtSource, FileFormat}; use std::path::PathBuf; @@ -1035,6 +1208,21 @@ mod tests { operator.to_string(), "Copy /tmp/users.csv -> users [#0, #1]" ); + assert_eq!( + Operator::CopyToFile(CopyToFileOperator { + target: ExtSource { + path: PathBuf::from("/tmp/output.csv"), + format: FileFormat::Csv { + delimiter: ',', + quote: '"', + escape: None, + header: false, + }, + }, + }) + .to_string(), + "Copy To /tmp/output.csv" + ); } } // GRCOV_EXCL_STOP diff --git a/src/planner/operator/visitor.rs b/src/planner/operator/visitor.rs new file mode 100644 index 00000000..06911e8c --- /dev/null +++ b/src/planner/operator/visitor.rs @@ -0,0 +1,586 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::alter_table::change_column::DefaultChange; +use super::*; +use crate::errors::DatabaseError; +use crate::expression::visitor::ExprVisitor; + +pub trait OperatorVisitor<'a>: Sized { + fn visit_operator(&mut self, operator: &'a Operator) -> Result<(), DatabaseError> { + walk_operator(self, operator) + } + + fn visit_dummy(&mut self) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_aggregate(&mut self, _op: &'a AggregateOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_scalar_apply(&mut self, _op: &'a ScalarApplyOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_mark_apply(&mut self, _op: &'a MarkApplyOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_filter(&mut self, _op: &'a FilterOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_join(&mut self, _op: &'a JoinOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_project(&mut self, _op: &'a ProjectOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_scalar_subquery( + &mut self, + _op: &'a ScalarSubqueryOperator, + ) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_table_scan(&mut self, _op: &'a TableScanOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_function_scan(&mut self, _op: &'a FunctionScanOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_sort(&mut self, _op: &'a SortOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_limit(&mut self, _op: &'a LimitOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_top_k(&mut self, _op: &'a TopKOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_values(&mut self, _op: &'a ValuesOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_show_table(&mut self) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_show_view(&mut self) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_explain(&mut self) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_describe(&mut self, _op: &'a DescribeOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_set_membership( + &mut self, + _op: &'a SetMembershipOperator, + ) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_union(&mut self, _op: &'a UnionOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_insert(&mut self, _op: &'a InsertOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_update(&mut self, _op: &'a UpdateOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_delete(&mut self, _op: &'a DeleteOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_analyze(&mut self, _op: &'a AnalyzeOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_add_column(&mut self, _op: &'a AddColumnOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_change_column(&mut self, _op: &'a ChangeColumnOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_drop_column(&mut self, _op: &'a DropColumnOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_create_table(&mut self, _op: &'a CreateTableOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_create_index(&mut self, _op: &'a CreateIndexOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_create_view(&mut self, _op: &'a CreateViewOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_drop_table(&mut self, _op: &'a DropTableOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_drop_view(&mut self, _op: &'a DropViewOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_drop_index(&mut self, _op: &'a DropIndexOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_truncate(&mut self, _op: &'a TruncateOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + #[cfg(feature = "copy")] + fn visit_copy_from_file(&mut self, _op: &'a CopyFromFileOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + #[cfg(feature = "copy")] + fn visit_copy_to_file(&mut self, _op: &'a CopyToFileOperator) -> Result<(), DatabaseError> { + Ok(()) + } +} + +pub struct OperatorExprVisitor<'a, V> { + visitor: &'a mut V, +} + +impl<'a, V> OperatorExprVisitor<'a, V> { + pub fn new(visitor: &'a mut V) -> Self { + Self { visitor } + } +} + +impl<'a, V: ExprVisitor<'a>> OperatorVisitor<'a> for OperatorExprVisitor<'_, V> { + fn visit_aggregate(&mut self, op: &'a AggregateOperator) -> Result<(), DatabaseError> { + for expr in op.agg_calls.iter().chain(&op.groupby_exprs) { + ExprVisitor::visit(self.visitor, expr)?; + } + Ok(()) + } + + fn visit_mark_apply(&mut self, op: &'a MarkApplyOperator) -> Result<(), DatabaseError> { + for expr in &op.predicates { + ExprVisitor::visit(self.visitor, expr)?; + } + if let Some(expr) = &op.parameterized_probe { + ExprVisitor::visit(self.visitor, expr)?; + } + Ok(()) + } + + fn visit_filter(&mut self, op: &'a FilterOperator) -> Result<(), DatabaseError> { + ExprVisitor::visit(self.visitor, &op.predicate) + } + + fn visit_join(&mut self, op: &'a JoinOperator) -> Result<(), DatabaseError> { + if let JoinCondition::On { on, filter } = &op.on { + for (left_expr, right_expr) in on { + ExprVisitor::visit(self.visitor, left_expr)?; + ExprVisitor::visit(self.visitor, right_expr)?; + } + if let Some(expr) = filter { + ExprVisitor::visit(self.visitor, expr)?; + } + } + Ok(()) + } + + fn visit_project(&mut self, op: &'a ProjectOperator) -> Result<(), DatabaseError> { + for expr in &op.exprs { + ExprVisitor::visit(self.visitor, expr)?; + } + Ok(()) + } + + fn visit_table_scan(&mut self, op: &'a TableScanOperator) -> Result<(), DatabaseError> { + for index_info in &op.index_infos { + if let SortOption::OrderBy { fields, .. } = &index_info.sort_option { + for field in fields { + ExprVisitor::visit(self.visitor, &field.expr)?; + } + } + if let Some(expr) = &index_info.residual_predicate { + ExprVisitor::visit(self.visitor, expr)?; + } + } + Ok(()) + } + + fn visit_function_scan(&mut self, op: &'a FunctionScanOperator) -> Result<(), DatabaseError> { + for expr in &op.table_function.args { + ExprVisitor::visit(self.visitor, expr)?; + } + Ok(()) + } + + fn visit_sort(&mut self, op: &'a SortOperator) -> Result<(), DatabaseError> { + for field in &op.sort_fields { + ExprVisitor::visit(self.visitor, &field.expr)?; + } + Ok(()) + } + + fn visit_top_k(&mut self, op: &'a TopKOperator) -> Result<(), DatabaseError> { + for field in &op.sort_fields { + ExprVisitor::visit(self.visitor, &field.expr)?; + } + Ok(()) + } + + fn visit_update(&mut self, op: &'a UpdateOperator) -> Result<(), DatabaseError> { + for (_, expr) in &op.value_exprs { + ExprVisitor::visit(self.visitor, expr)?; + } + Ok(()) + } + + fn visit_add_column(&mut self, op: &'a AddColumnOperator) -> Result<(), DatabaseError> { + if let Some(expr) = &op.column.desc().default { + ExprVisitor::visit(self.visitor, expr)?; + } + Ok(()) + } + + fn visit_change_column(&mut self, op: &'a ChangeColumnOperator) -> Result<(), DatabaseError> { + if let DefaultChange::Set(expr) = &op.default_change { + ExprVisitor::visit(self.visitor, expr)?; + } + Ok(()) + } + + fn visit_create_table(&mut self, op: &'a CreateTableOperator) -> Result<(), DatabaseError> { + for column in &op.columns { + if let Some(expr) = &column.desc().default { + ExprVisitor::visit(self.visitor, expr)?; + } + } + Ok(()) + } +} + +pub fn walk_operator<'a, V: OperatorVisitor<'a>>( + visitor: &mut V, + operator: &'a Operator, +) -> Result<(), DatabaseError> { + match operator { + Operator::Dummy => visitor.visit_dummy(), + Operator::Aggregate(op) => visitor.visit_aggregate(op), + Operator::ScalarApply(op) => visitor.visit_scalar_apply(op), + Operator::MarkApply(op) => visitor.visit_mark_apply(op), + Operator::Filter(op) => visitor.visit_filter(op), + Operator::Join(op) => visitor.visit_join(op), + Operator::Project(op) => visitor.visit_project(op), + Operator::ScalarSubquery(op) => visitor.visit_scalar_subquery(op), + Operator::TableScan(op) => visitor.visit_table_scan(op), + Operator::FunctionScan(op) => visitor.visit_function_scan(op), + Operator::Sort(op) => visitor.visit_sort(op), + Operator::Limit(op) => visitor.visit_limit(op), + Operator::TopK(op) => visitor.visit_top_k(op), + Operator::Values(op) => visitor.visit_values(op), + Operator::ShowTable => visitor.visit_show_table(), + Operator::ShowView => visitor.visit_show_view(), + Operator::Explain => visitor.visit_explain(), + Operator::Describe(op) => visitor.visit_describe(op), + Operator::SetMembership(op) => visitor.visit_set_membership(op), + Operator::Union(op) => visitor.visit_union(op), + Operator::Insert(op) => visitor.visit_insert(op), + Operator::Update(op) => visitor.visit_update(op), + Operator::Delete(op) => visitor.visit_delete(op), + Operator::Analyze(op) => visitor.visit_analyze(op), + Operator::AddColumn(op) => visitor.visit_add_column(op), + Operator::ChangeColumn(op) => visitor.visit_change_column(op), + Operator::DropColumn(op) => visitor.visit_drop_column(op), + Operator::CreateTable(op) => visitor.visit_create_table(op), + Operator::CreateIndex(op) => visitor.visit_create_index(op), + Operator::CreateView(op) => visitor.visit_create_view(op), + Operator::DropTable(op) => visitor.visit_drop_table(op), + Operator::DropView(op) => visitor.visit_drop_view(op), + Operator::DropIndex(op) => visitor.visit_drop_index(op), + Operator::Truncate(op) => visitor.visit_truncate(op), + #[cfg(feature = "copy")] + Operator::CopyFromFile(op) => visitor.visit_copy_from_file(op), + #[cfg(feature = "copy")] + Operator::CopyToFile(op) => visitor.visit_copy_to_file(op), + } +} + +#[cfg(all(test, not(target_arch = "wasm32")))] +pub(crate) mod tests { + use super::*; + #[cfg(feature = "copy")] + use crate::binder::copy::{ExtSource, FileFormat}; + use crate::catalog::view::View; + use crate::catalog::{ColumnCatalog, ColumnDesc, ColumnRef}; + use crate::expression::function::table::{ + ArcTableFunctionImpl, TableFunction, TableFunctionCatalog, + }; + use crate::expression::visitor::{walk_expr, ExprVisitor}; + use crate::expression::ScalarExpression; + use crate::function::numbers::Numbers; + use crate::planner::operator::alter_table::change_column::NotNullChange; + use crate::planner::operator::join::{JoinOperator, JoinType}; + use crate::planner::operator::mark_apply::MarkApplyOperator; + use crate::planner::operator::set_membership::SetMembershipKind; + use crate::planner::{Childrens, LogicalPlan}; + use crate::types::index::{IndexInfo, IndexMetaRef, IndexType}; + use crate::types::value::DataValue; + use crate::types::LogicalType; + + fn index_info() -> IndexInfo { + IndexInfo { + meta: IndexMetaRef::new(0), + sort_option: SortOption::OrderBy { + fields: vec![SortField::from(ScalarExpression::from(10_i32))], + ignore_prefix_len: 0, + }, + lookup: None, + residual_predicate: Some(11_i32.into()), + covered_deserializers: None, + cover_mapping: None, + sort_elimination_hint: None, + stream_distinct_hint: None, + } + } + + pub(crate) fn all_operators() -> Result, DatabaseError> { + let column_ref = ColumnRef::new(0); + let column = ColumnCatalog::new( + "value".to_string(), + false, + ColumnDesc::new(LogicalType::Integer, None, false, Some(12_i32.into()))?, + ); + let mut mark_apply = MarkApplyOperator::new_exists(column_ref, vec![3_i32.into()]); + mark_apply.set_parameterized_probe(Some(4_i32.into())); + let table_function = TableFunction { + args: vec![8_i32.into()], + catalog: TableFunctionCatalog { + schema: Vec::new(), + inner: ArcTableFunctionImpl(Numbers::new()), + }, + }; + let operators = vec![ + Operator::Dummy, + Operator::Aggregate(AggregateOperator { + groupby_exprs: vec![1_i32.into()], + agg_calls: vec![2_i32.into()], + is_distinct: false, + }), + Operator::ScalarApply(ScalarApplyOperator), + Operator::MarkApply(mark_apply), + Operator::Filter(FilterOperator { + predicate: 5_i32.into(), + is_optimized: false, + having: false, + }), + Operator::Join(JoinOperator { + on: JoinCondition::On { + on: vec![(6_i32.into(), 7_i32.into())], + filter: Some(8_i32.into()), + }, + join_type: JoinType::Inner, + }), + Operator::Project(ProjectOperator { + exprs: vec![9_i32.into()], + }), + Operator::ScalarSubquery(ScalarSubqueryOperator), + Operator::TableScan(TableScanOperator { + table_name: "t1".into(), + columns: vec![column_ref], + limit: (None, None), + index_infos: vec![index_info()], + with_pk: false, + }), + Operator::FunctionScan(FunctionScanOperator { table_function }), + Operator::Sort(SortOperator { + sort_fields: vec![SortField::from(ScalarExpression::from(13_i32))], + limit: None, + }), + Operator::Limit(LimitOperator { + offset: None, + limit: Some(1), + }), + Operator::TopK(TopKOperator { + sort_fields: vec![SortField::from(ScalarExpression::from(14_i32))], + limit: 1, + offset: None, + }), + Operator::Values(ValuesOperator { + rows: vec![vec![DataValue::Int32(1)]], + schema_ref: vec![column_ref], + }), + Operator::ShowTable, + Operator::ShowView, + Operator::Explain, + Operator::Describe(DescribeOperator { + table_name: "t1".into(), + }), + Operator::SetMembership(SetMembershipOperator { + kind: SetMembershipKind::Intersect, + left_schema_ref: vec![column_ref], + _right_schema_ref: vec![column_ref], + }), + Operator::Union(UnionOperator { + left_schema_ref: vec![column_ref], + _right_schema_ref: vec![column_ref], + }), + Operator::Insert(InsertOperator { + table_name: "t1".into(), + is_overwrite: false, + is_mapping_by_name: false, + }), + Operator::Update(UpdateOperator { + table_name: "t1".into(), + value_exprs: vec![(column_ref, 15_i32.into())], + }), + Operator::Delete(DeleteOperator { + table_name: "t1".into(), + primary_keys: vec![column_ref], + }), + Operator::Analyze(AnalyzeOperator { + table_name: "t1".into(), + index_metas: vec![IndexMetaRef::new(0)], + histogram_buckets: Some(1), + }), + Operator::AddColumn(AddColumnOperator { + table_name: "t1".into(), + if_not_exists: false, + column: column.clone(), + }), + Operator::ChangeColumn(ChangeColumnOperator { + table_name: "t1".into(), + old_column_name: "value".to_string(), + new_column_name: "value".to_string(), + data_type: LogicalType::Integer, + default_change: DefaultChange::Set(16_i32.into()), + not_null_change: NotNullChange::NoChange, + }), + Operator::DropColumn(DropColumnOperator { + table_name: "t1".into(), + column_name: "value".to_string(), + if_exists: false, + }), + Operator::CreateTable(CreateTableOperator { + table_name: "t1".into(), + columns: vec![column], + if_not_exists: false, + }), + Operator::CreateIndex(CreateIndexOperator { + table_name: "t1".into(), + columns: vec![column_ref], + index_name: "idx".to_string(), + if_not_exists: false, + ty: IndexType::Normal, + }), + Operator::CreateView(CreateViewOperator { + view: View { + name: "v1".into(), + plan: Box::new(LogicalPlan::new(Operator::Dummy, Childrens::None)), + schema: vec![column_ref], + }, + or_replace: false, + }), + Operator::DropTable(DropTableOperator { + table_name: "t1".into(), + if_exists: false, + }), + Operator::DropView(DropViewOperator { + view_name: "v1".into(), + if_exists: false, + }), + Operator::DropIndex(DropIndexOperator { + table_name: "t1".into(), + index_name: "idx".to_string(), + if_exists: false, + }), + Operator::Truncate(TruncateOperator { + table_name: "t1".into(), + }), + ]; + #[cfg(feature = "copy")] + let operators = { + let mut with_copy_operators = operators; + let source = ExtSource { + path: "data.csv".into(), + format: FileFormat::Csv { + delimiter: ',', + quote: '"', + escape: None, + header: false, + }, + }; + with_copy_operators.push(Operator::CopyFromFile(CopyFromFileOperator { + table: "t1".into(), + source: source.clone(), + schema_ref: vec![column_ref], + })); + with_copy_operators.push(Operator::CopyToFile(CopyToFileOperator { target: source })); + with_copy_operators + }; + Ok(operators) + } + + #[derive(Default)] + struct ExpressionCounter(usize); + + impl<'a> ExprVisitor<'a> for ExpressionCounter { + fn visit(&mut self, expr: &'a ScalarExpression) -> Result<(), DatabaseError> { + self.0 += 1; + walk_expr(self, expr) + } + } + + #[test] + fn dispatches_all_variants_and_visits_expressions() -> Result<(), DatabaseError> { + struct NoopVisitor; + impl OperatorVisitor<'_> for NoopVisitor {} + + let operators = all_operators()?; + for operator in &operators { + NoopVisitor.visit_operator(operator)?; + } + + let mut counter = ExpressionCounter::default(); + let mut visitor = OperatorExprVisitor::new(&mut counter); + for operator in &operators { + visitor.visit_operator(operator)?; + } + assert_eq!(counter.0, 18); + + Ok(()) + } +} diff --git a/src/planner/operator/visitor_mut.rs b/src/planner/operator/visitor_mut.rs new file mode 100644 index 00000000..c5b740cb --- /dev/null +++ b/src/planner/operator/visitor_mut.rs @@ -0,0 +1,405 @@ +// Copyright 2024 KipData/KiteSQL +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::alter_table::change_column::DefaultChange; +use super::*; +use crate::errors::DatabaseError; +use crate::expression::visitor_mut::ExprVisitorMut; + +pub trait OperatorVisitorMut<'a>: Sized { + fn visit_operator(&mut self, operator: &'a mut Operator) -> Result<(), DatabaseError> { + walk_mut_operator(self, operator) + } + + fn visit_dummy(&mut self) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_aggregate(&mut self, _op: &'a mut AggregateOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_scalar_apply( + &mut self, + _op: &'a mut ScalarApplyOperator, + ) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_mark_apply(&mut self, _op: &'a mut MarkApplyOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_filter(&mut self, _op: &'a mut FilterOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_join(&mut self, _op: &'a mut JoinOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_project(&mut self, _op: &'a mut ProjectOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_scalar_subquery( + &mut self, + _op: &'a mut ScalarSubqueryOperator, + ) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_table_scan(&mut self, _op: &'a mut TableScanOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_function_scan( + &mut self, + _op: &'a mut FunctionScanOperator, + ) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_sort(&mut self, _op: &'a mut SortOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_limit(&mut self, _op: &'a mut LimitOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_top_k(&mut self, _op: &'a mut TopKOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_values(&mut self, _op: &'a mut ValuesOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_show_table(&mut self) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_show_view(&mut self) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_explain(&mut self) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_describe(&mut self, _op: &'a mut DescribeOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_set_membership( + &mut self, + _op: &'a mut SetMembershipOperator, + ) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_union(&mut self, _op: &'a mut UnionOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_insert(&mut self, _op: &'a mut InsertOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_update(&mut self, _op: &'a mut UpdateOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_delete(&mut self, _op: &'a mut DeleteOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_analyze(&mut self, _op: &'a mut AnalyzeOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_add_column(&mut self, _op: &'a mut AddColumnOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_change_column( + &mut self, + _op: &'a mut ChangeColumnOperator, + ) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_drop_column(&mut self, _op: &'a mut DropColumnOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_create_table( + &mut self, + _op: &'a mut CreateTableOperator, + ) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_create_index( + &mut self, + _op: &'a mut CreateIndexOperator, + ) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_create_view(&mut self, _op: &'a mut CreateViewOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_drop_table(&mut self, _op: &'a mut DropTableOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_drop_view(&mut self, _op: &'a mut DropViewOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_drop_index(&mut self, _op: &'a mut DropIndexOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + fn visit_truncate(&mut self, _op: &'a mut TruncateOperator) -> Result<(), DatabaseError> { + Ok(()) + } + + #[cfg(feature = "copy")] + fn visit_copy_from_file( + &mut self, + _op: &'a mut CopyFromFileOperator, + ) -> Result<(), DatabaseError> { + Ok(()) + } + + #[cfg(feature = "copy")] + fn visit_copy_to_file(&mut self, _op: &'a mut CopyToFileOperator) -> Result<(), DatabaseError> { + Ok(()) + } +} + +pub struct OperatorExprVisitorMut<'a, V> { + visitor: &'a mut V, +} + +impl<'a, V> OperatorExprVisitorMut<'a, V> { + pub fn new(visitor: &'a mut V) -> Self { + Self { visitor } + } +} + +impl<'a, V: ExprVisitorMut<'a>> OperatorVisitorMut<'a> for OperatorExprVisitorMut<'_, V> { + fn visit_aggregate(&mut self, op: &'a mut AggregateOperator) -> Result<(), DatabaseError> { + for expr in op.agg_calls.iter_mut().chain(&mut op.groupby_exprs) { + ExprVisitorMut::visit(self.visitor, expr)?; + } + Ok(()) + } + + fn visit_mark_apply(&mut self, op: &'a mut MarkApplyOperator) -> Result<(), DatabaseError> { + for expr in &mut op.predicates { + ExprVisitorMut::visit(self.visitor, expr)?; + } + if let Some(expr) = &mut op.parameterized_probe { + ExprVisitorMut::visit(self.visitor, expr)?; + } + Ok(()) + } + + fn visit_filter(&mut self, op: &'a mut FilterOperator) -> Result<(), DatabaseError> { + ExprVisitorMut::visit(self.visitor, &mut op.predicate) + } + + fn visit_join(&mut self, op: &'a mut JoinOperator) -> Result<(), DatabaseError> { + if let JoinCondition::On { on, filter } = &mut op.on { + for (left_expr, right_expr) in on { + ExprVisitorMut::visit(self.visitor, left_expr)?; + ExprVisitorMut::visit(self.visitor, right_expr)?; + } + if let Some(expr) = filter { + ExprVisitorMut::visit(self.visitor, expr)?; + } + } + Ok(()) + } + + fn visit_project(&mut self, op: &'a mut ProjectOperator) -> Result<(), DatabaseError> { + for expr in &mut op.exprs { + ExprVisitorMut::visit(self.visitor, expr)?; + } + Ok(()) + } + + fn visit_table_scan(&mut self, op: &'a mut TableScanOperator) -> Result<(), DatabaseError> { + for index_info in &mut op.index_infos { + if let SortOption::OrderBy { fields, .. } = &mut index_info.sort_option { + for field in fields { + ExprVisitorMut::visit(self.visitor, &mut field.expr)?; + } + } + if let Some(expr) = &mut index_info.residual_predicate { + ExprVisitorMut::visit(self.visitor, expr)?; + } + } + Ok(()) + } + + fn visit_function_scan( + &mut self, + op: &'a mut FunctionScanOperator, + ) -> Result<(), DatabaseError> { + for expr in &mut op.table_function.args { + ExprVisitorMut::visit(self.visitor, expr)?; + } + Ok(()) + } + + fn visit_sort(&mut self, op: &'a mut SortOperator) -> Result<(), DatabaseError> { + for field in &mut op.sort_fields { + ExprVisitorMut::visit(self.visitor, &mut field.expr)?; + } + Ok(()) + } + + fn visit_top_k(&mut self, op: &'a mut TopKOperator) -> Result<(), DatabaseError> { + for field in &mut op.sort_fields { + ExprVisitorMut::visit(self.visitor, &mut field.expr)?; + } + Ok(()) + } + + fn visit_update(&mut self, op: &'a mut UpdateOperator) -> Result<(), DatabaseError> { + for (_, expr) in &mut op.value_exprs { + ExprVisitorMut::visit(self.visitor, expr)?; + } + Ok(()) + } + + fn visit_add_column(&mut self, op: &'a mut AddColumnOperator) -> Result<(), DatabaseError> { + if let Some(expr) = &mut op.column.desc_mut().default { + ExprVisitorMut::visit(self.visitor, expr)?; + } + Ok(()) + } + + fn visit_change_column( + &mut self, + op: &'a mut ChangeColumnOperator, + ) -> Result<(), DatabaseError> { + if let DefaultChange::Set(expr) = &mut op.default_change { + ExprVisitorMut::visit(self.visitor, expr)?; + } + Ok(()) + } + + fn visit_create_table(&mut self, op: &'a mut CreateTableOperator) -> Result<(), DatabaseError> { + for column in &mut op.columns { + if let Some(expr) = &mut column.desc_mut().default { + ExprVisitorMut::visit(self.visitor, expr)?; + } + } + Ok(()) + } +} + +pub fn walk_mut_operator<'a, V: OperatorVisitorMut<'a>>( + visitor: &mut V, + operator: &'a mut Operator, +) -> Result<(), DatabaseError> { + match operator { + Operator::Dummy => visitor.visit_dummy(), + Operator::Aggregate(op) => visitor.visit_aggregate(op), + Operator::ScalarApply(op) => visitor.visit_scalar_apply(op), + Operator::MarkApply(op) => visitor.visit_mark_apply(op), + Operator::Filter(op) => visitor.visit_filter(op), + Operator::Join(op) => visitor.visit_join(op), + Operator::Project(op) => visitor.visit_project(op), + Operator::ScalarSubquery(op) => visitor.visit_scalar_subquery(op), + Operator::TableScan(op) => visitor.visit_table_scan(op), + Operator::FunctionScan(op) => visitor.visit_function_scan(op), + Operator::Sort(op) => visitor.visit_sort(op), + Operator::Limit(op) => visitor.visit_limit(op), + Operator::TopK(op) => visitor.visit_top_k(op), + Operator::Values(op) => visitor.visit_values(op), + Operator::ShowTable => visitor.visit_show_table(), + Operator::ShowView => visitor.visit_show_view(), + Operator::Explain => visitor.visit_explain(), + Operator::Describe(op) => visitor.visit_describe(op), + Operator::SetMembership(op) => visitor.visit_set_membership(op), + Operator::Union(op) => visitor.visit_union(op), + Operator::Insert(op) => visitor.visit_insert(op), + Operator::Update(op) => visitor.visit_update(op), + Operator::Delete(op) => visitor.visit_delete(op), + Operator::Analyze(op) => visitor.visit_analyze(op), + Operator::AddColumn(op) => visitor.visit_add_column(op), + Operator::ChangeColumn(op) => visitor.visit_change_column(op), + Operator::DropColumn(op) => visitor.visit_drop_column(op), + Operator::CreateTable(op) => visitor.visit_create_table(op), + Operator::CreateIndex(op) => visitor.visit_create_index(op), + Operator::CreateView(op) => visitor.visit_create_view(op), + Operator::DropTable(op) => visitor.visit_drop_table(op), + Operator::DropView(op) => visitor.visit_drop_view(op), + Operator::DropIndex(op) => visitor.visit_drop_index(op), + Operator::Truncate(op) => visitor.visit_truncate(op), + #[cfg(feature = "copy")] + Operator::CopyFromFile(op) => visitor.visit_copy_from_file(op), + #[cfg(feature = "copy")] + Operator::CopyToFile(op) => visitor.visit_copy_to_file(op), + } +} + +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::*; + use crate::planner::operator::visitor::tests::all_operators; + use crate::types::value::DataValue; + + struct IncrementConstants(usize); + + impl ExprVisitorMut<'_> for IncrementConstants { + fn visit_constant(&mut self, value: &mut DataValue) -> Result<(), DatabaseError> { + if let DataValue::Int32(value) = value { + *value += 1; + } + self.0 += 1; + Ok(()) + } + } + + #[test] + fn dispatches_all_variants_and_mutates_expressions() -> Result<(), DatabaseError> { + struct NoopVisitor; + impl OperatorVisitorMut<'_> for NoopVisitor {} + + let mut operators = all_operators()?; + for operator in &mut operators { + NoopVisitor.visit_operator(operator)?; + } + + let mut counter = IncrementConstants(0); + { + let mut visitor = OperatorExprVisitorMut::new(&mut counter); + for operator in &mut operators { + visitor.visit_operator(operator)?; + } + } + assert_eq!(counter.0, 18); + + Ok(()) + } +} diff --git a/src/serdes/data_value.rs b/src/serdes/data_value.rs index e54641d8..99fce987 100644 --- a/src/serdes/data_value.rs +++ b/src/serdes/data_value.rs @@ -398,6 +398,7 @@ fn read_char_length_units(reader: &mut R) -> Result Result<(), DatabaseError> { let sources = vec![ @@ -465,4 +474,78 @@ pub(crate) mod test { Ok(()) } + + #[test] + fn invalid_tags_and_truncated_values_are_rejected() { + assert_invalid_value(&[u8::MAX], "invalid data value tag"); + assert_invalid_value(&[TAG_BOOLEAN, 2], "invalid bool value"); + assert_invalid_value(&[TAG_INT64, 1], "failed to fill whole buffer"); + + let mut utf8 = vec![TAG_UTF8]; + utf8.extend(0u32.to_le_bytes()); + utf8.extend([0, 2]); + assert_invalid_value(&utf8, "invalid option tag"); + + let mut utf8 = vec![TAG_UTF8]; + utf8.extend(0u32.to_le_bytes()); + utf8.push(2); + assert_invalid_value(&utf8, "invalid utf8 type tag"); + + let mut utf8 = vec![TAG_UTF8]; + utf8.extend(0u32.to_le_bytes()); + utf8.extend([0, 0, 2]); + assert_invalid_value(&utf8, "invalid char length units tag"); + + let mut time64 = vec![TAG_TIME64]; + time64.extend(0i64.to_le_bytes()); + time64.extend(0u64.to_le_bytes()); + time64.push(2); + assert_invalid_value(&time64, "invalid bool value"); + + let mut tuple = vec![TAG_TUPLE]; + tuple.extend(0u32.to_le_bytes()); + tuple.push(2); + assert_invalid_value(&tuple, "invalid bool value"); + } + + #[test] + fn utf8_type_reference_serialization_validates_tags() -> Result<(), DatabaseError> { + let values = [ + Utf8Type::Variable(None), + Utf8Type::Variable(Some(32)), + Utf8Type::Fixed(8), + ]; + let mut tables = ReferenceTables::new(); + let mut arena = crate::planner::TableArena::default(); + + for (index, value) in values.into_iter().enumerate() { + let mut bytes = Vec::new(); + value.encode(&mut bytes, false, &mut tables, &arena)?; + let decoded = Utf8Type::decode::( + &mut Cursor::new(bytes), + None, + &tables, + &mut arena, + )?; + assert!(matches!( + (index, decoded), + (0, Utf8Type::Variable(None)) + | (1, Utf8Type::Variable(Some(32))) + | (2, Utf8Type::Fixed(8)) + )); + } + + let result = Utf8Type::decode::( + &mut Cursor::new([2]), + None, + &tables, + &mut arena, + ); + let Err(err) = result else { + panic!("invalid tag should fail") + }; + assert!(err.to_string().contains("invalid utf8 type tag")); + + Ok(()) + } } diff --git a/src/serdes/function.rs b/src/serdes/function.rs index fa177d12..1627af02 100644 --- a/src/serdes/function.rs +++ b/src/serdes/function.rs @@ -91,3 +91,91 @@ impl ReferenceSerialization for ArcTableFunctionImpl { Ok(function.inner.clone()) } } + +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use super::*; + use crate::db::{ScalaFunctions, TableFunctions}; + use crate::planner::TableArena; + use crate::serdes::ReferenceDecodeContext; + use crate::storage::rocksdb::RocksTransaction; + use crate::types::LogicalType; + use std::io::Cursor; + + fn encoded_summary(name: &str) -> Vec { + let summary = FunctionSummary { + name: name.into(), + arg_types: vec![LogicalType::Integer], + }; + let mut bytes = Vec::new(); + summary + .encode( + &mut bytes, + false, + &mut ReferenceTables::new(), + &TableArena::default(), + ) + .unwrap(); + bytes + } + + #[test] + fn scalar_function_decode_requires_context_and_registration() { + let bytes = encoded_summary("missing_scalar"); + let tables = ReferenceTables::new(); + let mut arena = TableArena::default(); + + let err = ArcScalarFunctionImpl::decode::( + &mut Cursor::new(&bytes), + None, + &tables, + &mut arena, + ) + .unwrap_err(); + assert!(err + .to_string() + .contains("scalar function decode context missing")); + + let scalars = ScalaFunctions::default(); + let table_functions = TableFunctions::default(); + let context = ReferenceDecodeContext::with_functions(None, &scalars, &table_functions); + let err = ArcScalarFunctionImpl::decode::( + &mut Cursor::new(&bytes), + Some(&context), + &tables, + &mut arena, + ) + .unwrap_err(); + assert!(err.to_string().contains("scalar function not found")); + } + + #[test] + fn table_function_decode_requires_context_and_registration() { + let bytes = encoded_summary("missing_table"); + let tables = ReferenceTables::new(); + let mut arena = TableArena::default(); + + let err = ArcTableFunctionImpl::decode::( + &mut Cursor::new(&bytes), + None, + &tables, + &mut arena, + ) + .unwrap_err(); + assert!(err + .to_string() + .contains("table function decode context missing")); + + let scalars = ScalaFunctions::default(); + let table_functions = TableFunctions::default(); + let context = ReferenceDecodeContext::with_functions(None, &scalars, &table_functions); + let err = ArcTableFunctionImpl::decode::( + &mut Cursor::new(&bytes), + Some(&context), + &tables, + &mut arena, + ) + .unwrap_err(); + assert!(err.to_string().contains("table function not found")); + } +} diff --git a/tests/slt/join.slt b/tests/slt/join.slt index 7dbfed87..8a765660 100644 --- a/tests/slt/join.slt +++ b/tests/slt/join.slt @@ -193,3 +193,10 @@ select a.* from a where v1 >= (select 1) 1 2 2 2 3 3 +# A table alias column list must match the source width. +statement error +select * from a as aliased(only_one) + +# Every USING column must exist on both sides. +statement error +select * from a join b using (missing) diff --git a/tests/slt/set_operation.slt b/tests/slt/set_operation.slt index eeba3617..38a1ffcb 100644 --- a/tests/slt/set_operation.slt +++ b/tests/slt/set_operation.slt @@ -117,6 +117,10 @@ null statement error select v from set_left union select v from set_right order by set_left.v +# Set-operation inputs must expose the same number of columns. +statement error +select v from set_left union all select id, v from set_right + statement ok drop table set_right diff --git a/tests/slt/subquery.slt b/tests/slt/subquery.slt index 51d66ccf..e6e55741 100644 --- a/tests/slt/subquery.slt +++ b/tests/slt/subquery.slt @@ -487,6 +487,21 @@ where exists ( select id from orders where amount = 300 ); +# EXISTS/IN subqueries are not supported in the SELECT list. +statement error +select exists(select id from orders) from users; + +# Correlated scalar subqueries are not supported in WHERE. +statement error +select * from users +where id = (select user_id from orders where amount = age); + +# Mark-apply and scalar-join subqueries cannot share one WHERE scope yet. +statement error +select * from users +where exists(select 1 from orders) + and id = (select max(user_id) from orders); + statement ok drop table users;