uqa_planner/optimizer/
api.rs1use super::{
10 optimize_unified_plan, reorder_unified_plan_joins, AggregateClassifier, JoinGraphError,
11 RelationStats, ScalarExpr, SourcePlan, UnifiedPlan,
12};
13use crate::LocalAccessEstimate;
14
15#[derive(Debug, thiserror::Error)]
17pub enum OptimizerError {
18 #[error(transparent)]
19 Expression(#[from] uqa_sql::SQLError),
20 #[error(transparent)]
21 JoinGraph(#[from] JoinGraphError),
22}
23
24pub type OptimizerResult<T> = Result<T, OptimizerError>;
25
26pub type ConstantEvaluator = fn(&ScalarExpr) -> Result<uqa_core::Value, uqa_sql::SQLError>;
28
29#[derive(Debug, Clone)]
30pub struct OptimizerConfig {
31 pub enable_filter_pushdown: bool,
32 pub enable_boolean_simplify: bool,
33 pub enable_vector_threshold_merge: bool,
34 pub enable_join_reordering: bool,
35 pub constant_evaluator: ConstantEvaluator,
37}
38
39impl OptimizerConfig {
40 #[must_use]
41 pub const fn new(constant_evaluator: ConstantEvaluator) -> Self {
42 Self {
43 enable_filter_pushdown: true,
44 enable_boolean_simplify: true,
45 enable_vector_threshold_merge: true,
46 enable_join_reordering: true,
47 constant_evaluator,
48 }
49 }
50}
51
52pub trait SourceStatistics {
57 fn relation_statistics(&self, table: &str) -> Option<RelationStats>;
58
59 fn source_access_estimate(&self, _source: &SourcePlan) -> Option<LocalAccessEstimate> {
62 None
63 }
64
65 fn local_access_estimate(
68 &self,
69 _table: &str,
70 _predicate: &ScalarExpr,
71 ) -> Option<LocalAccessEstimate> {
72 None
73 }
74}
75
76impl<F> SourceStatistics for F
77where
78 F: Fn(&str) -> Option<RelationStats>,
79{
80 fn relation_statistics(&self, table: &str) -> Option<RelationStats> {
81 self(table)
82 }
83}
84
85struct NoSourceStatistics;
86
87impl SourceStatistics for NoSourceStatistics {
88 fn relation_statistics(&self, _table: &str) -> Option<RelationStats> {
89 None
90 }
91}
92
93struct NoRegisteredAggregates;
94
95impl AggregateClassifier for NoRegisteredAggregates {
96 fn is_registered_aggregate(&self, _name: &str) -> bool {
97 false
98 }
99}
100
101pub fn optimize(plan: UnifiedPlan, config: &OptimizerConfig) -> OptimizerResult<UnifiedPlan> {
103 optimize_with_aggregates_and_statistics(
104 plan,
105 config,
106 &NoRegisteredAggregates,
107 &NoSourceStatistics,
108 )
109}
110
111pub fn optimize_with_aggregates(
113 plan: UnifiedPlan,
114 config: &OptimizerConfig,
115 aggregates: &dyn AggregateClassifier,
116) -> OptimizerResult<UnifiedPlan> {
117 optimize_with_aggregates_and_statistics(plan, config, aggregates, &NoSourceStatistics)
118}
119
120pub fn optimize_with_statistics(
122 plan: UnifiedPlan,
123 config: &OptimizerConfig,
124 statistics: &dyn SourceStatistics,
125) -> OptimizerResult<UnifiedPlan> {
126 optimize_with_aggregates_and_statistics(plan, config, &NoRegisteredAggregates, statistics)
127}
128
129pub fn optimize_with_aggregates_and_statistics(
132 mut plan: UnifiedPlan,
133 config: &OptimizerConfig,
134 aggregates: &dyn AggregateClassifier,
135 statistics: &dyn SourceStatistics,
136) -> OptimizerResult<UnifiedPlan> {
137 optimize_unified_plan(&mut plan, config, aggregates)?;
138 if config.enable_join_reordering {
139 reorder_unified_plan_joins(&mut plan, statistics)?;
140 }
141 Ok(plan)
142}