Skip to content

Commit 35e8a0e

Browse files
committed
refactor: flatten union plan
1 parent cab23ee commit 35e8a0e

File tree

2 files changed

+7
-146
lines changed

2 files changed

+7
-146
lines changed

src/query/sql/src/executor/physical_plans/physical_union_all.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ use crate::executor::explain::PlanStatsInfo;
2727
use crate::executor::PhysicalPlan;
2828
use crate::executor::PhysicalPlanBuilder;
2929
use crate::optimizer::SExpr;
30+
use crate::optimizer::StatInfo;
3031
use crate::plans::BoundColumnRef;
3132
use crate::plans::ScalarItem;
3233
use crate::ColumnBindingBuilder;

src/query/sql/src/planner/binder/select.rs

Lines changed: 6 additions & 146 deletions
Original file line numberDiff line numberDiff line change
@@ -523,7 +523,6 @@ impl Binder {
523523
}
524524
(SetOperator::Union, true) => self.bind_union(
525525
left.span(),
526-
right.span(),
527526
left_bind_context,
528527
right_bind_context,
529528
left_expr,
@@ -532,7 +531,6 @@ impl Binder {
532531
),
533532
(SetOperator::Union, false) => self.bind_union(
534533
left.span(),
535-
right.span(),
536534
left_bind_context,
537535
right_bind_context,
538536
left_expr,
@@ -549,48 +547,18 @@ impl Binder {
549547
pub fn bind_union(
550548
&mut self,
551549
left_span: Span,
552-
right_span: Span,
553550
left_context: BindContext,
554551
right_context: BindContext,
555552
left_expr: SExpr,
556553
right_expr: SExpr,
557554
distinct: bool,
558555
) -> Result<(SExpr, BindContext)> {
559-
let mut coercion_types = Vec::with_capacity(left_context.columns.len());
560-
for (left_col, right_col) in left_context
556+
let pairs = left_context
561557
.columns
562558
.iter()
563559
.zip(right_context.columns.iter())
564-
{
565-
if left_col.data_type != right_col.data_type {
566-
if let Some(data_type) = common_super_type(
567-
*left_col.data_type.clone(),
568-
*right_col.data_type.clone(),
569-
&BUILTIN_FUNCTIONS.default_cast_rules,
570-
) {
571-
coercion_types.push(data_type);
572-
} else {
573-
return Err(ErrorCode::SemanticError(format!(
574-
"SetOperation's types cannot be matched, left column {:?}, type: {:?}, right column {:?}, type: {:?}",
575-
left_col.column_name,
576-
left_col.data_type,
577-
right_col.column_name,
578-
right_col.data_type
579-
)));
580-
}
581-
} else {
582-
coercion_types.push(*left_col.data_type.clone());
583-
}
584-
}
585-
let (new_bind_context, pairs, left_expr, right_expr) = self.coercion_union_type(
586-
left_span,
587-
right_span,
588-
left_context,
589-
right_context,
590-
left_expr,
591-
right_expr,
592-
coercion_types,
593-
)?;
560+
.map(|(l, r)| (l.index, r.index))
561+
.collect();
594562

595563
let union_plan = UnionAll { pairs };
596564
let mut new_expr = SExpr::create_binary(
@@ -602,14 +570,14 @@ impl Binder {
602570
if distinct {
603571
new_expr = self.bind_distinct(
604572
left_span,
605-
&new_bind_context,
606-
new_bind_context.all_column_bindings(),
573+
&left_context,
574+
left_context.all_column_bindings(),
607575
&mut HashMap::new(),
608576
new_expr,
609577
)?;
610578
}
611579

612-
Ok((new_expr, new_bind_context))
580+
Ok((new_expr, left_context))
613581
}
614582

615583
pub fn bind_intersect(
@@ -703,114 +671,6 @@ impl Binder {
703671
Ok((s_expr, left_context))
704672
}
705673

706-
#[allow(clippy::type_complexity)]
707-
#[allow(clippy::too_many_arguments)]
708-
fn coercion_union_type(
709-
&self,
710-
left_span: Span,
711-
right_span: Span,
712-
left_bind_context: BindContext,
713-
right_bind_context: BindContext,
714-
mut left_expr: SExpr,
715-
mut right_expr: SExpr,
716-
coercion_types: Vec<DataType>,
717-
) -> Result<(BindContext, Vec<(IndexType, IndexType)>, SExpr, SExpr)> {
718-
let mut left_scalar_items = Vec::with_capacity(left_bind_context.columns.len());
719-
let mut right_scalar_items = Vec::with_capacity(right_bind_context.columns.len());
720-
let mut new_bind_context = BindContext::new();
721-
let mut pairs = Vec::with_capacity(left_bind_context.columns.len());
722-
for (idx, (left_col, right_col)) in left_bind_context
723-
.columns
724-
.iter()
725-
.zip(right_bind_context.columns.iter())
726-
.enumerate()
727-
{
728-
let left_index = if *left_col.data_type != coercion_types[idx] {
729-
let new_column_index = self
730-
.metadata
731-
.write()
732-
.add_derived_column(left_col.column_name.clone(), coercion_types[idx].clone());
733-
let column_binding = ColumnBindingBuilder::new(
734-
left_col.column_name.clone(),
735-
new_column_index,
736-
Box::new(coercion_types[idx].clone()),
737-
Visibility::Visible,
738-
)
739-
.build();
740-
let left_coercion_expr = CastExpr {
741-
span: left_span,
742-
is_try: false,
743-
argument: Box::new(
744-
BoundColumnRef {
745-
span: left_span,
746-
column: left_col.clone(),
747-
}
748-
.into(),
749-
),
750-
target_type: Box::new(coercion_types[idx].clone()),
751-
};
752-
left_scalar_items.push(ScalarItem {
753-
scalar: left_coercion_expr.into(),
754-
index: new_column_index,
755-
});
756-
new_bind_context.add_column_binding(column_binding);
757-
new_column_index
758-
} else {
759-
new_bind_context.add_column_binding(left_col.clone());
760-
left_col.index
761-
};
762-
let right_index = if *right_col.data_type != coercion_types[idx] {
763-
let new_column_index = self
764-
.metadata
765-
.write()
766-
.add_derived_column(right_col.column_name.clone(), coercion_types[idx].clone());
767-
let right_coercion_expr = CastExpr {
768-
span: right_span,
769-
is_try: false,
770-
argument: Box::new(
771-
BoundColumnRef {
772-
span: right_span,
773-
column: right_col.clone(),
774-
}
775-
.into(),
776-
),
777-
target_type: Box::new(coercion_types[idx].clone()),
778-
};
779-
right_scalar_items.push(ScalarItem {
780-
scalar: right_coercion_expr.into(),
781-
index: new_column_index,
782-
});
783-
new_column_index
784-
} else {
785-
right_col.index
786-
};
787-
pairs.push((left_index, right_index));
788-
}
789-
if !left_scalar_items.is_empty() {
790-
left_expr = SExpr::create_unary(
791-
Arc::new(
792-
EvalScalar {
793-
items: left_scalar_items,
794-
}
795-
.into(),
796-
),
797-
Arc::new(left_expr),
798-
);
799-
}
800-
if !right_scalar_items.is_empty() {
801-
right_expr = SExpr::create_unary(
802-
Arc::new(
803-
EvalScalar {
804-
items: right_scalar_items,
805-
}
806-
.into(),
807-
),
808-
Arc::new(right_expr),
809-
);
810-
}
811-
Ok((new_bind_context, pairs, left_expr, right_expr))
812-
}
813-
814674
#[allow(clippy::too_many_arguments)]
815675
fn analyze_lazy_materialization(
816676
&self,

0 commit comments

Comments
 (0)