quill_sql/optimizer/
logical_optimizer.rs

1use crate::error::QuillSQLResult;
2use crate::optimizer::rule::{EliminateLimit, MergeLimit, PushDownLimit};
3use crate::plan::logical_plan::LogicalPlan;
4use std::sync::Arc;
5
6/// `LogicalOptimizerRule` transforms one [`LogicalPlan`] into another which
7/// computes the same results, but in a potentially more efficient
8/// way. If there are no suitable transformations for the input plan,
9/// the optimizer can simply return it as is.
10pub trait LogicalOptimizerRule {
11    /// Try and rewrite `plan` to an optimized form, returning None if the plan cannot be
12    /// optimized by this rule.
13    fn try_optimize(&self, plan: &LogicalPlan) -> QuillSQLResult<Option<LogicalPlan>>;
14
15    /// A human-readable name for this optimizer rule
16    fn name(&self) -> &str;
17
18    /// How should the rule be applied by the optimizer
19    ///
20    /// If a rule use default None, it should traverse recursively plan inside itself
21    fn apply_order(&self) -> Option<ApplyOrder> {
22        None
23    }
24}
25
26pub enum ApplyOrder {
27    TopDown,
28    BottomUp,
29}
30
31#[derive(Clone)]
32pub struct LogicalOptimizer {
33    /// All optimizer rules to apply
34    pub rules: Vec<Arc<dyn LogicalOptimizerRule + Send + Sync>>,
35    pub max_passes: usize,
36}
37
38impl LogicalOptimizer {
39    pub fn new() -> Self {
40        let rules: Vec<Arc<dyn LogicalOptimizerRule + Sync + Send>> = vec![
41            Arc::new(EliminateLimit {}),
42            Arc::new(MergeLimit {}),
43            Arc::new(PushDownLimit {}),
44        ];
45
46        Self {
47            rules,
48            max_passes: 3,
49        }
50    }
51
52    pub fn with_rules(rules: Vec<Arc<dyn LogicalOptimizerRule + Send + Sync>>) -> Self {
53        Self {
54            rules,
55            max_passes: 3,
56        }
57    }
58
59    pub fn optimize(&self, plan: &LogicalPlan) -> QuillSQLResult<LogicalPlan> {
60        let mut new_plan = plan.clone();
61        let mut i = 0;
62        while i < self.max_passes {
63            for rule in &self.rules {
64                if let Some(optimized_plan) = self.optimize_recursively(rule, &new_plan)? {
65                    new_plan = optimized_plan;
66                }
67            }
68
69            i += 1;
70        }
71        Ok(new_plan)
72    }
73
74    pub fn optimize_recursively(
75        &self,
76        rule: &Arc<dyn LogicalOptimizerRule + Send + Sync>,
77        plan: &LogicalPlan,
78    ) -> QuillSQLResult<Option<LogicalPlan>> {
79        match rule.apply_order() {
80            Some(order) => match order {
81                ApplyOrder::TopDown => {
82                    let optimize_self_opt = rule.try_optimize(plan)?;
83                    let optimize_inputs_opt = match &optimize_self_opt {
84                        Some(optimized_plan) => self.optimize_inputs(rule, optimized_plan)?,
85                        _ => self.optimize_inputs(rule, plan)?,
86                    };
87                    Ok(optimize_inputs_opt.or(optimize_self_opt))
88                }
89                ApplyOrder::BottomUp => {
90                    let optimize_inputs_opt = self.optimize_inputs(rule, plan)?;
91                    let optimize_self_opt = match &optimize_inputs_opt {
92                        Some(optimized_plan) => rule.try_optimize(optimized_plan)?,
93                        _ => rule.try_optimize(plan)?,
94                    };
95                    Ok(optimize_self_opt.or(optimize_inputs_opt))
96                }
97            },
98            _ => rule.try_optimize(plan),
99        }
100    }
101
102    fn optimize_inputs(
103        &self,
104        rule: &Arc<dyn LogicalOptimizerRule + Send + Sync>,
105        plan: &LogicalPlan,
106    ) -> QuillSQLResult<Option<LogicalPlan>> {
107        let inputs = plan.inputs();
108        let result = inputs
109            .iter()
110            .map(|sub_plan| self.optimize_recursively(rule, sub_plan))
111            .collect::<QuillSQLResult<Vec<_>>>()?;
112        if result.is_empty() || result.iter().all(|o| o.is_none()) {
113            return Ok(None);
114        }
115
116        let new_inputs = result
117            .into_iter()
118            .zip(inputs)
119            .map(|(new_plan, old_plan)| match new_plan {
120                Some(plan) => plan,
121                None => old_plan.clone(),
122            })
123            .collect::<Vec<_>>();
124        plan.with_new_inputs(&new_inputs).map(Some)
125    }
126}
quill_sql/optimizer/logical_optimizer.rs

quill_sql/optimizer/
logical_optimizer.rs
