use uqa_sql::SQLError;
use super::{
choose_access_path, optimize_command, optimize_scalar_slot, prioritize_access_predicates,
rewrite_implicit_hybrid_fusion, source_allows_unqualified_signals, AggregateClassifier,
ComputePlan, OptimizerConfig, QueryBlockPlan, QueryPlan, RelationalPlan, SourcePlan,
UnifiedPlan,
};
pub(super) fn optimize_unified_plan(
plan: &mut UnifiedPlan,
config: &OptimizerConfig,
aggregates: &dyn AggregateClassifier,
) -> Result<(), SQLError> {
match plan {
UnifiedPlan::Query(query) => optimize_query(query, config, aggregates)?,
UnifiedPlan::Command(command) => optimize_command(command, config, aggregates)?,
}
Ok(())
}
pub(super) fn optimize_cte(
body: &mut crate::CtePlanBody,
config: &OptimizerConfig,
aggregates: &dyn AggregateClassifier,
) -> Result<(), SQLError> {
match body {
crate::CtePlanBody::Query(query) => optimize_query(query, config, aggregates)?,
crate::CtePlanBody::Command(command) => optimize_command(command, config, aggregates)?,
}
Ok(())
}
pub(super) fn optimize_query(
query: &mut QueryPlan,
config: &OptimizerConfig,
aggregates: &dyn AggregateClassifier,
) -> Result<(), SQLError> {
super::subqueries::prune_query(query);
let reachable = uqa_sql::semantics::reachable_plan_cte_names(query);
for cte in query
.ctes
.iter_mut()
.filter(|cte| reachable.contains(&cte.name))
{
optimize_cte(&mut cte.body, config, aggregates)?;
}
match &mut query.root {
RelationalPlan::QueryBlock(block) => optimize_query_block(block, config, aggregates)?,
RelationalPlan::SetOp {
left,
right,
order_by,
limit,
offset,
subqueries,
..
} => {
optimize_query(left, config, aggregates)?;
optimize_query(right, config, aggregates)?;
for order in &mut *order_by {
optimize_scalar_slot(&mut order.expr, config)?;
}
if let Some(limit) = limit {
optimize_scalar_slot(limit, config)?;
}
if let Some(offset) = offset {
optimize_scalar_slot(offset, config)?;
}
let live = super::subqueries::live_slots(
order_by
.iter()
.map(|order| &order.expr)
.chain(limit.as_deref())
.chain(offset.as_deref()),
)?;
super::subqueries::optimize_live(subqueries, &live, config, aggregates)?;
}
RelationalPlan::Values { rows, subqueries } => {
for row in &mut *rows {
for expression in row {
optimize_scalar_slot(expression, config)?;
}
}
let live = super::subqueries::live_slots(rows.iter().flatten())?;
super::subqueries::optimize_live(subqueries, &live, config, aggregates)?;
}
}
Ok(())
}
pub(super) fn optimize_query_block(
block: &mut QueryBlockPlan,
config: &OptimizerConfig,
aggregates: &dyn AggregateClassifier,
) -> Result<(), SQLError> {
super::derived_projection::prune(block, aggregates);
if let Some(source) = &mut block.from {
optimize_source(source, config, aggregates)?;
}
super::source_constants::propagate_source_constants(
block,
config
.routine_inlining
.as_ref()
.map(|context| context.types),
config
.routine_inlining
.as_ref()
.map(|context| context.volatility),
)?;
super::optimize_projections(&mut block.projections, config)?;
if let Some(predicate) = &mut block.r#where {
optimize_scalar_slot(predicate, config)?;
let allow_unqualified_signals = source_allows_unqualified_signals(block.from.as_ref());
rewrite_implicit_hybrid_fusion(predicate, allow_unqualified_signals);
if config.enable_filter_pushdown {
prioritize_access_predicates(predicate);
}
}
for expression in &mut block.group_by {
optimize_scalar_slot(expression, config)?;
}
for set in &mut block.grouping_sets {
for expression in set {
optimize_scalar_slot(expression, config)?;
}
}
if let Some(having) = &mut block.having {
optimize_scalar_slot(having, config)?;
}
for order in &mut block.order_by {
optimize_scalar_slot(&mut order.expr, config)?;
}
if let Some(limit) = &mut block.limit {
optimize_scalar_slot(limit, config)?;
}
if let Some(offset) = &mut block.offset {
optimize_scalar_slot(offset, config)?;
}
for expression in &mut block.distinct_on {
optimize_scalar_slot(expression, config)?;
}
for window in &mut block.windows {
for expression in window.spec.expressions_mut() {
optimize_scalar_slot(expression, config)?;
}
}
super::source_constants::simplify_constant_group(
block,
config
.routine_inlining
.as_ref()
.map(|context| context.types),
);
let live_subqueries = super::subqueries::live_slots(block.expressions())?;
super::subqueries::optimize_live(&mut block.subqueries, &live_subqueries, config, aggregates)?;
let is_aggregate = |name: &str| {
crate::unified_plan::is_builtin_aggregate(name) || aggregates.is_registered_aggregate(name)
};
let has_aggregate = !block.group_by.is_empty()
|| !block.grouping_sets.is_empty()
|| block.having.is_some()
|| block
.projections
.iter()
.any(|projection| projection.expr.contains_aggregate(&is_aggregate))
|| block
.windows
.iter()
.flat_map(|window| window.spec.expressions())
.any(|expression| expression.contains_aggregate(&is_aggregate));
let has_window = block
.projections
.iter()
.any(|projection| projection.expr.contains_window());
block.compute = if has_aggregate {
ComputePlan::Aggregate
} else if has_window {
ComputePlan::Window
} else {
ComputePlan::Project
};
block.access = choose_access_path(block);
Ok(())
}
pub(super) fn optimize_source(
source: &mut SourcePlan,
config: &OptimizerConfig,
aggregates: &dyn AggregateClassifier,
) -> Result<(), SQLError> {
match source {
SourcePlan::Join {
left, right, on, ..
} => {
optimize_source(left, config, aggregates)?;
optimize_source(right, config, aggregates)?;
if let Some(on) = on {
optimize_scalar_slot(on, config)?;
if config.enable_filter_pushdown {
prioritize_access_predicates(on);
}
}
}
SourcePlan::Values { rows, .. } => {
for row in rows {
for expression in row {
optimize_scalar_slot(expression, config)?;
}
}
}
SourcePlan::Function { args, .. } => {
for expression in args {
optimize_scalar_slot(expression, config)?;
}
}
SourcePlan::FunctionGroup { functions, .. } => {
for function in functions {
for expression in &mut function.args {
optimize_scalar_slot(expression, config)?;
}
}
}
SourcePlan::Subquery { body, .. } => optimize_query(body, config, aggregates)?,
SourcePlan::Table { .. } => {}
}
Ok(())
}