use crate::planner::{ContextProvider, PlannerContext, SqlToRel};
use datafusion_common::{DataFusionError, Result};
use datafusion_expr::{LogicalPlan, LogicalPlanBuilder};
use sqlparser::ast::{SetExpr, SetOperator, SetQuantifier};
impl<'a, S: ContextProvider> SqlToRel<'a, S> {
pub(super) fn set_expr_to_plan(
&self,
set_expr: SetExpr,
planner_context: &mut PlannerContext,
) -> Result<LogicalPlan> {
match set_expr {
SetExpr::Select(s) => self.select_to_plan(*s, planner_context),
SetExpr::Values(v) => self.sql_values_to_plan(v, planner_context),
SetExpr::SetOperation {
op,
left,
right,
set_quantifier,
} => {
let all = match set_quantifier {
SetQuantifier::All => true,
SetQuantifier::Distinct | SetQuantifier::None => false,
};
let left_plan = self.set_expr_to_plan(*left, planner_context)?;
let right_plan = self.set_expr_to_plan(*right, planner_context)?;
match (op, all) {
(SetOperator::Union, true) => LogicalPlanBuilder::from(left_plan)
.union(right_plan)?
.build(),
(SetOperator::Union, false) => LogicalPlanBuilder::from(left_plan)
.union_distinct(right_plan)?
.build(),
(SetOperator::Intersect, true) => {
LogicalPlanBuilder::intersect(left_plan, right_plan, true)
}
(SetOperator::Intersect, false) => {
LogicalPlanBuilder::intersect(left_plan, right_plan, false)
}
(SetOperator::Except, true) => {
LogicalPlanBuilder::except(left_plan, right_plan, true)
}
(SetOperator::Except, false) => {
LogicalPlanBuilder::except(left_plan, right_plan, false)
}
}
}
SetExpr::Query(q) => self.query_to_plan(*q, planner_context),
_ => Err(DataFusionError::NotImplemented(format!(
"Query {set_expr} not implemented yet"
))),
}
}
}