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 LogicalOptimizer {
41    pub fn new() -> Self {
42        let rules: Vec<Arc<dyn LogicalOptimizerRule + Sync + Send>> = vec![
43            Arc::new(PushDownFilterToScan),
44            Arc::new(PushLimitIntoScan),
45            Arc::new(EliminateLimit {}),
46            Arc::new(MergeLimit {}),
47            Arc::new(PushDownLimit {}),
48        ];
49
50        Self {
51            rules,
52            max_passes: 3,
53        }
54    }
55
56    pub fn with_rules(rules: Vec<Arc<dyn LogicalOptimizerRule + Send + Sync>>) -> Self {
57        Self {
58            rules,
59            max_passes: 3,
60        }
61    }
62
63    pub fn optimize(&self, plan: &LogicalPlan) -> QuillSQLResult<LogicalPlan> {
64        let mut new_plan = plan.clone();
65        let mut i = 0;
66        while i < self.max_passes {
67            for rule in &self.rules {
68                if let Some(optimized_plan) = self.optimize_recursively(rule, &new_plan)? {
69                    new_plan = optimized_plan;
70                }
71            }
72
73            i += 1;
74        }
75        Ok(new_plan)
76    }
77
78    pub fn optimize_recursively(
79        &self,
80        rule: &Arc<dyn LogicalOptimizerRule + Send + Sync>,
81        plan: &LogicalPlan,
82    ) -> QuillSQLResult<Option<LogicalPlan>> {
83        match rule.apply_order() {
84            Some(order) => match order {
85                ApplyOrder::TopDown => {
86                    let optimize_self_opt = rule.try_optimize(plan)?;
87                    let optimize_inputs_opt = match &optimize_self_opt {
88                        Some(optimized_plan) => self.optimize_inputs(rule, optimized_plan)?,
89                        _ => self.optimize_inputs(rule, plan)?,
90                    };
91                    Ok(optimize_inputs_opt.or(optimize_self_opt))
92                }
93                ApplyOrder::BottomUp => {
94                    let optimize_inputs_opt = self.optimize_inputs(rule, plan)?;
95                    let optimize_self_opt = match &optimize_inputs_opt {
96                        Some(optimized_plan) => rule.try_optimize(optimized_plan)?,
97                        _ => rule.try_optimize(plan)?,
98                    };
99                    Ok(optimize_self_opt.or(optimize_inputs_opt))
100                }
101            },
102            _ => rule.try_optimize(plan),
103        }
104    }
105
106    fn optimize_inputs(
107        &self,
108        rule: &Arc<dyn LogicalOptimizerRule + Send + Sync>,
109        plan: &LogicalPlan,
110    ) -> QuillSQLResult<Option<LogicalPlan>> {
111        let inputs = plan.inputs();
112        let result = inputs
113            .iter()
114            .map(|sub_plan| self.optimize_recursively(rule, sub_plan))
115            .collect::<QuillSQLResult<Vec<_>>>()?;
116        if result.is_empty() || result.iter().all(|o| o.is_none()) {
117            return Ok(None);
118        }
119
120        let new_inputs = result
121            .into_iter()
122            .zip(inputs)
123            .map(|(new_plan, old_plan)| match new_plan {
124                Some(plan) => plan,
125                None => old_plan.clone(),
126            })
127            .collect::<Vec<_>>();
128        plan.with_new_inputs(&new_inputs).map(Some)
129    }
130}
quill_sql/optimizer/logical_optimizer.rs

quill_sql/optimizer/
logical_optimizer.rs
