Skip to content

Commit

Permalink
add annotation for select_to_plan (#1714)
Browse files Browse the repository at this point in the history
  • Loading branch information
xudong963 authored Jan 31, 2022
1 parent 3494e9c commit 2512608
Showing 1 changed file with 26 additions and 8 deletions.
34 changes: 26 additions & 8 deletions datafusion/src/sql/planner.rs
Original file line number Diff line number Diff line change
Expand Up @@ -697,14 +697,13 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
}
}

/// Generate a logic plan from an SQL select
fn select_to_plan(
/// Generate a logic plan from selection clause, the function contain optimization for cross join to inner join
/// Related PR: https://github.com/apache/arrow-datafusion/pull/1566
fn plan_selection(
&self,
select: &Select,
ctes: &mut HashMap<String, LogicalPlan>,
alias: Option<String>,
plans: Vec<LogicalPlan>,
) -> Result<LogicalPlan> {
let plans = self.plan_from_tables(&select.from, ctes)?;
let plan = match &select.selection {
Some(predicate_expr) => {
// build join schema
Expand Down Expand Up @@ -822,9 +821,23 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
}
}
};
let plan = plan?;
plan
}

// The SELECT expressions, with wildcards expanded.
/// Generate a logic plan from an SQL select
fn select_to_plan(
&self,
select: &Select,
ctes: &mut HashMap<String, LogicalPlan>,
alias: Option<String>,
) -> Result<LogicalPlan> {
// process `from` clause
let plans = self.plan_from_tables(&select.from, ctes)?;

// process `where` clause
let plan = self.plan_selection(select, plans)?;

// process the SELECT expressions, with wildcards expanded.
let select_exprs = self.prepare_select_exprs(&plan, select)?;

// having and group by clause may reference aliases defined in select projection
Expand Down Expand Up @@ -873,6 +886,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
// All of the aggregate expressions (deduplicated).
let aggr_exprs = find_aggregate_exprs(&aggr_expr_haystack);

// All of the group by expressions
let group_by_exprs = select
.group_by
.iter()
Expand All @@ -891,6 +905,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
})
.collect::<Result<Vec<Expr>>>()?;

// process group by, aggregation or having
let (plan, select_exprs_post_aggr, having_expr_post_aggr_opt) = if !group_by_exprs
.is_empty()
|| !aggr_exprs.is_empty()
Expand Down Expand Up @@ -931,7 +946,7 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
plan
};

// window function
// process window function
let window_func_exprs = find_window_exprs(&select_exprs_post_aggr);

let plan = if window_func_exprs.is_empty() {
Expand All @@ -940,13 +955,16 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> {
LogicalPlanBuilder::window_plan(plan, window_func_exprs)?
};

// process distinct clause
let plan = if select.distinct {
return LogicalPlanBuilder::from(plan)
.aggregate(select_exprs_post_aggr, iter::empty::<Expr>())?
.build();
} else {
plan
};

// generate the final projection plan
project_with_alias(plan, select_exprs_post_aggr, alias)
}

Expand Down

0 comments on commit 2512608

Please sign in to comment.