Skip to content

Commit dec3643

Browse files
KKouldKKould
authored
refactor: unify expression and operator traversal with visitors (#372)
* test: cover optimizer and serde edge cases * refactor(planner): unify expression and operator traversal * test: improve binder and operator coverage
1 parent a28b4ab commit dec3643

31 files changed

Lines changed: 2491 additions & 727 deletions

‎src/binder/aggregate.rs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,8 @@ use std::collections::HashSet;
1717
use super::{Binder, QueryBindStep};
1818
use crate::errors::DatabaseError;
1919
use crate::expression::function::scala::ScalarFunction;
20-
use crate::expression::visitor::{walk_expr, Visitor};
21-
use crate::expression::visitor_mut::{walk_mut_expr, VisitorMut};
20+
use crate::expression::visitor::{walk_expr, ExprVisitor};
21+
use crate::expression::visitor_mut::{walk_mut_expr, ExprVisitorMut};
2222
use crate::planner::LogicalPlan;
2323
use crate::storage::Transaction;
2424
use crate::types::value::DataValue;
@@ -269,7 +269,7 @@ impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A>
269269
HashSet::from_iter(group_raw_exprs.iter().copied());
270270

271271
for expr in select_items {
272-
if expr.has_agg_call() {
272+
if expr.has_agg_call()? {
273273
continue;
274274
}
275275
group_raw_set.remove(expr);
@@ -347,7 +347,7 @@ impl<'a> HavingOrderByValidator<'a> {
347347
}
348348
}
349349

350-
impl<'expr> Visitor<'expr> for HavingOrderByValidator<'_> {
350+
impl<'expr> ExprVisitor<'expr> for HavingOrderByValidator<'_> {
351351
fn visit(&mut self, expr: &'expr ScalarExpression) -> Result<(), DatabaseError> {
352352
match expr {
353353
ScalarExpression::AggCall { .. } => {
@@ -427,7 +427,7 @@ impl<'a, 'p> AggregateOutputBinder<'a, 'p> {
427427
}
428428
}
429429

430-
impl<'a> VisitorMut<'a> for AggregateOutputBinder<'_, '_> {
430+
impl<'a> ExprVisitorMut<'a> for AggregateOutputBinder<'_, '_> {
431431
fn visit(&mut self, expr: &'a mut ScalarExpression) -> Result<(), DatabaseError> {
432432
if let ScalarExpression::Alias {
433433
expr: inner_expr,
@@ -453,7 +453,7 @@ mod tests {
453453
use crate::catalog::{ColumnCatalog, ColumnDesc, ColumnRef};
454454
use crate::errors::DatabaseError;
455455
use crate::expression::agg::AggKind;
456-
use crate::expression::visitor_mut::VisitorMut;
456+
use crate::expression::visitor_mut::ExprVisitorMut;
457457
use crate::expression::{AliasType, BinaryOperator, ScalarExpression};
458458
use crate::planner::PlanArena;
459459
use crate::storage::Storage;

‎src/binder/distinct.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515
use crate::binder::{Binder, QueryBindStep};
1616
use crate::errors::DatabaseError;
17-
use crate::expression::visitor_mut::{walk_mut_expr, VisitorMut};
17+
use crate::expression::visitor_mut::{walk_mut_expr, ExprVisitorMut};
1818
use crate::expression::ScalarExpression;
1919
use crate::planner::operator::aggregate::AggregateOperator;
2020
use crate::planner::operator::sort::SortField;
@@ -101,7 +101,7 @@ impl<'a, 'p> DistinctOutputBinder<'a, 'p> {
101101
}
102102
}
103103

104-
impl<'a> VisitorMut<'a> for DistinctOutputBinder<'_, '_> {
104+
impl<'a> ExprVisitorMut<'a> for DistinctOutputBinder<'_, '_> {
105105
fn visit(&mut self, expr: &'a mut ScalarExpression) -> Result<(), DatabaseError> {
106106
if let ScalarExpression::Alias {
107107
expr: inner_expr,
@@ -124,7 +124,7 @@ mod tests {
124124
use super::DistinctOutputBinder;
125125
use crate::catalog::{ColumnCatalog, ColumnDesc, ColumnRef};
126126
use crate::errors::DatabaseError;
127-
use crate::expression::visitor_mut::VisitorMut;
127+
use crate::expression::visitor_mut::ExprVisitorMut;
128128
use crate::expression::{AliasType, ScalarExpression};
129129
use crate::planner::PlanArena;
130130
use crate::types::LogicalType;

‎src/binder/parser.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ use crate::db::{BindSource, DBTransaction, Database, DatabaseIter, TransactionIt
2525
use crate::errors::{DatabaseError, SqlErrorSpan};
2626
use crate::expression;
2727
use crate::expression::simplify::ConstantCalculator;
28-
use crate::expression::visitor_mut::VisitorMut;
28+
use crate::expression::visitor_mut::ExprVisitorMut;
2929
use crate::expression::{AliasType, ScalarExpression};
3030
use crate::iter_ext::Itertools;
3131
use crate::parser::parse_sql;
@@ -338,7 +338,7 @@ struct UpdateExprTargetRemapper<'a, 'p> {
338338
arena: &'a PlanArena<'p>,
339339
}
340340

341-
impl VisitorMut<'_> for UpdateExprTargetRemapper<'_, '_> {
341+
impl ExprVisitorMut<'_> for UpdateExprTargetRemapper<'_, '_> {
342342
fn visit_column_ref(
343343
&mut self,
344344
column: &mut ColumnRef,
@@ -669,7 +669,7 @@ where
669669
) -> Result<ScalarExpression, DatabaseError> {
670670
let mut expr = self.binder.bind_expr(&expr, self.arena)?;
671671

672-
if expr.any_referenced_column(self.arena, |_, _| true) {
672+
if expr.any_referenced_column(self.arena, |_, _| true)? {
673673
return Err(DatabaseError::UnsupportedStmt(
674674
"column is not allowed to exist in default".to_string(),
675675
));
@@ -783,7 +783,7 @@ where
783783
ColumnOption::Default(expr) => {
784784
let mut expr = self.binder.bind_expr(&expr, self.arena)?;
785785

786-
if expr.any_referenced_column(self.arena, |_, _| true) {
786+
if expr.any_referenced_column(self.arena, |_, _| true)? {
787787
return Err(DatabaseError::UnsupportedStmt(
788788
"column is not allowed to exist in `default`".to_string(),
789789
));

‎src/binder/select.rs‎

Lines changed: 62 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ use super::{Binder, BinderContext, QueryBindStep, SetOperatorKind, Source, SubQu
3131
use crate::catalog::{ColumnRef, ColumnRelation, TableName};
3232
use crate::errors::DatabaseError;
3333
use crate::execution::dql::join::joins_nullable;
34-
use crate::expression::visitor_mut::{walk_mut_expr, PositionShift, VisitorMut};
34+
use crate::expression::visitor_mut::{walk_mut_expr, ExprVisitorMut, PositionShift};
3535
use crate::expression::{AliasType, BinaryOperator};
3636
use crate::iter_ext::Itertools;
3737
use crate::planner::operator::function_scan::FunctionScanOperator;
@@ -51,7 +51,7 @@ struct RightSidePositionGlobalizer<'a, 'p> {
5151
arena: &'a crate::planner::PlanArena<'p>,
5252
}
5353

54-
impl<'a> VisitorMut<'a> for RightSidePositionGlobalizer<'_, '_> {
54+
impl<'a> ExprVisitorMut<'a> for RightSidePositionGlobalizer<'_, '_> {
5555
fn visit_column_ref(
5656
&mut self,
5757
column: &'a mut ColumnRef,
@@ -80,7 +80,7 @@ struct SplitScopePositionRebinder<'a, 'p> {
8080
arena: &'a crate::planner::PlanArena<'p>,
8181
}
8282

83-
impl VisitorMut<'_> for SplitScopePositionRebinder<'_, '_> {
83+
impl ExprVisitorMut<'_> for SplitScopePositionRebinder<'_, '_> {
8484
fn visit_column_ref(
8585
&mut self,
8686
column: &mut ColumnRef,
@@ -109,7 +109,7 @@ struct MarkerPositionGlobalizer<'a, 'p> {
109109
arena: &'a crate::planner::PlanArena<'p>,
110110
}
111111

112-
impl VisitorMut<'_> for MarkerPositionGlobalizer<'_, '_> {
112+
impl ExprVisitorMut<'_> for MarkerPositionGlobalizer<'_, '_> {
113113
fn visit_column_ref(
114114
&mut self,
115115
column: &mut ColumnRef,
@@ -154,7 +154,7 @@ impl<'a, 'p> ProjectionOutputBinder<'a, 'p> {
154154
}
155155
}
156156

157-
impl<'a> VisitorMut<'a> for ProjectionOutputBinder<'_, '_> {
157+
impl<'a> ExprVisitorMut<'a> for ProjectionOutputBinder<'_, '_> {
158158
fn visit(&mut self, expr: &'a mut ScalarExpression) -> Result<(), DatabaseError> {
159159
if let Some(output_ref) = self.output_ref(expr) {
160160
*expr = output_ref;
@@ -432,8 +432,10 @@ where
432432

433433
#[cfg(feature = "orm")]
434434
pub fn finish(self) -> Result<LogicalPlan, DatabaseError> {
435-
if self.select_list.iter().any(ScalarExpression::has_agg_call) {
436-
return self.aggregate_without_group()?.finish();
435+
for expr in &self.select_list {
436+
if expr.has_agg_call()? {
437+
return self.aggregate_without_group()?.finish();
438+
}
437439
}
438440
self.binder
439441
.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<'
801803
arena: &'a crate::planner::PlanArena<'p>,
802804
}
803805

804-
impl VisitorMut<'_> for AppendedRightOutputBinder<'_, '_> {
806+
impl ExprVisitorMut<'_> for AppendedRightOutputBinder<'_, '_> {
805807
fn visit_column_ref(
806808
&mut self,
807809
column: &mut ColumnRef,
@@ -1495,7 +1497,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'
14951497
plan: &mut LogicalPlan,
14961498
predicates: &[ScalarExpression],
14971499
arena: &mut crate::planner::PlanArena,
1498-
) -> Vec<AppendedRightOutput> {
1500+
) -> Result<Vec<AppendedRightOutput>, DatabaseError> {
14991501
let output_schema = plan.output_schema(arena).clone();
15001502
let output_len = output_schema.len();
15011503
if let LogicalPlan {
@@ -1505,36 +1507,38 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'
15051507
} = plan
15061508
{
15071509
let Childrens::Only(child) = childrens.as_mut() else {
1508-
return Vec::new();
1510+
return Ok(Vec::new());
15091511
};
15101512
let child_schema = child.output_schema(arena);
15111513
let mut appended_outputs = Vec::new();
1512-
op.exprs.extend(
1513-
child_schema
1514-
.iter()
1515-
.enumerate()
1516-
.filter(|(_, column)| {
1517-
!output_schema.contains(column)
1518-
&& predicates.iter().any(|expr| {
1519-
expr.any_referenced_column(arena, |arena, candidate| {
1520-
arena.same_column(*candidate, **column)
1521-
})
1522-
})
1523-
})
1524-
.map(|(position, column)| {
1525-
appended_outputs.push(AppendedRightOutput {
1526-
column: *column,
1527-
child_position: position,
1528-
output_position: output_len + appended_outputs.len(),
1529-
});
1530-
ScalarExpression::column_expr(*column, position)
1531-
}),
1532-
);
1514+
for (position, column) in child_schema.iter().enumerate() {
1515+
if output_schema.contains(column) {
1516+
continue;
1517+
}
1518+
let mut referenced = false;
1519+
for expr in predicates {
1520+
if expr.any_referenced_column(arena, |arena, candidate| {
1521+
arena.same_column(*candidate, *column)
1522+
})? {
1523+
referenced = true;
1524+
break;
1525+
}
1526+
}
1527+
if referenced {
1528+
op.exprs
1529+
.push(ScalarExpression::column_expr(*column, position));
1530+
appended_outputs.push(AppendedRightOutput {
1531+
column: *column,
1532+
child_position: position,
1533+
output_position: output_len + appended_outputs.len(),
1534+
});
1535+
}
1536+
}
15331537
plan.reset_output_schema_cache();
1534-
return appended_outputs;
1538+
return Ok(appended_outputs);
15351539
}
15361540

1537-
Vec::new()
1541+
Ok(Vec::new())
15381542
}
15391543

15401544
#[allow(clippy::too_many_arguments)]
@@ -1565,7 +1569,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'
15651569

15661570
if correlated {
15671571
let appended_right_outputs =
1568-
Self::ensure_mark_apply_right_outputs(&mut plan, &apply_predicates, arena);
1572+
Self::ensure_mark_apply_right_outputs(&mut plan, &apply_predicates, arena)?;
15691573
if !appended_right_outputs.is_empty() {
15701574
Self::localize_appended_right_outputs(
15711575
apply_predicates.iter_mut(),
@@ -1618,33 +1622,36 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'
16181622
plan: &LogicalPlan,
16191623
left_schema: &Schema,
16201624
arena: &mut crate::planner::PlanArena,
1621-
) -> bool {
1625+
) -> Result<bool, DatabaseError> {
16221626
if !plan
16231627
.operator
16241628
.visit_referenced_columns(arena, &mut |arena, column| {
16251629
!left_schema
16261630
.iter()
16271631
.any(|left| arena.same_column(*left, *column))
1628-
})
1632+
})?
16291633
{
1630-
return true;
1634+
return Ok(true);
16311635
}
16321636

16331637
match plan.childrens.as_ref() {
16341638
Childrens::Only(child) => Self::plan_has_correlated_refs(child, left_schema, arena),
16351639
Childrens::Twins { left, right } => {
1636-
Self::plan_has_correlated_refs(left, left_schema, arena)
1637-
|| Self::plan_has_correlated_refs(right, left_schema, arena)
1640+
if Self::plan_has_correlated_refs(left, left_schema, arena)? {
1641+
Ok(true)
1642+
} else {
1643+
Self::plan_has_correlated_refs(right, left_schema, arena)
1644+
}
16381645
}
1639-
Childrens::None => false,
1646+
Childrens::None => Ok(false),
16401647
}
16411648
}
16421649

16431650
fn expr_has_correlated_refs(
16441651
expr: &ScalarExpression,
16451652
left_schema: &Schema,
16461653
arena: &mut crate::planner::PlanArena,
1647-
) -> bool {
1654+
) -> Result<bool, DatabaseError> {
16481655
expr.any_referenced_column(arena, |arena, column| {
16491656
left_schema
16501657
.iter()
@@ -1688,7 +1695,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'
16881695
match plan.childrens.as_ref() {
16891696
Childrens::Only(_) => {}
16901697
Childrens::Twins { .. } => {
1691-
if Self::plan_has_correlated_refs(&plan, left_schema, arena) {
1698+
if Self::plan_has_correlated_refs(&plan, left_schema, arena)? {
16921699
return Err(DatabaseError::UnsupportedStmt(
16931700
"correlated EXISTS/NOT EXISTS does not support set or join subqueries"
16941701
.to_string(),
@@ -1715,7 +1722,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'
17151722
let mut predicates = Vec::new();
17161723
Self::split_conjuncts(op.predicate, &mut predicates);
17171724
for predicate in predicates {
1718-
if Self::expr_has_correlated_refs(&predicate, left_schema, arena) {
1725+
if Self::expr_has_correlated_refs(&predicate, left_schema, arena)? {
17191726
correlated_filters.push(predicate);
17201727
} else {
17211728
local_filters.push(predicate);
@@ -1775,7 +1782,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'
17751782
arena,
17761783
),
17771784
plan => {
1778-
if Self::plan_has_correlated_refs(&plan, left_schema, arena) {
1785+
if Self::plan_has_correlated_refs(&plan, left_schema, arena)? {
17791786
Err(DatabaseError::UnsupportedStmt(
17801787
"correlated EXISTS/NOT EXISTS only supports filter-based subqueries"
17811788
.to_string(),
@@ -2142,9 +2149,11 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'
21422149
// example: baz > 1
21432150
if left_expr.all_referenced_columns(arena, |_, column| {
21442151
fn_or_contains(*column)
2145-
}) && right_expr.all_referenced_columns(arena, |_, column| {
2146-
fn_or_contains(*column)
2147-
}) {
2152+
})? && right_expr
2153+
.all_referenced_columns(arena, |_, column| {
2154+
fn_or_contains(*column)
2155+
})?
2156+
{
21482157
accum_filter.push(ScalarExpression::Binary {
21492158
left_expr,
21502159
right_expr,
@@ -2186,9 +2195,10 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'
21862195
}
21872196
_ => {
21882197
if left_expr
2189-
.all_referenced_columns(arena, |_, column| fn_or_contains(*column))
2190-
&& right_expr
2191-
.all_referenced_columns(arena, |_, column| fn_or_contains(*column))
2198+
.all_referenced_columns(arena, |_, column| fn_or_contains(*column))?
2199+
&& right_expr.all_referenced_columns(arena, |_, column| {
2200+
fn_or_contains(*column)
2201+
})?
21922202
{
21932203
accum_filter.push(ScalarExpression::Binary {
21942204
left_expr,
@@ -2202,7 +2212,7 @@ impl<'a: 'b, 'b, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'
22022212
}
22032213
}
22042214
expr => {
2205-
if expr.all_referenced_columns(arena, |_, column| fn_or_contains(*column)) {
2215+
if expr.all_referenced_columns(arena, |_, column| fn_or_contains(*column))? {
22062216
// example: baz > 1
22072217
accum_filter.push(expr);
22082218
}
@@ -2219,7 +2229,7 @@ mod tests {
22192229
use crate::binder::test::build_t1_table;
22202230
use crate::catalog::{ColumnCatalog, ColumnDesc};
22212231
use crate::errors::DatabaseError;
2222-
use crate::expression::visitor_mut::VisitorMut;
2232+
use crate::expression::visitor_mut::ExprVisitorMut;
22232233
use crate::expression::{AliasType, ScalarExpression};
22242234
use crate::planner::operator::join::{JoinCondition, JoinType};
22252235
use crate::planner::operator::mark_apply::{

‎src/catalog/column.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -244,7 +244,7 @@ impl ColumnDesc {
244244
if let Some(expr) = &default {
245245
let table_arena = crate::planner::TableArenaCell::default();
246246
let plan_arena = crate::planner::PlanArena::new(&table_arena);
247-
if expr.has_table_ref_column(&plan_arena) {
247+
if expr.has_table_ref_column(&plan_arena)? {
248248
return Err(DatabaseError::DefaultNotColumnRef);
249249
}
250250
}

0 commit comments

Comments
 (0)