quill_sql/optimizer/
logical_optimizer.rs

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

quill_sql/optimizer/
logical_optimizer.rs
