Skip to main content

oxisql_parse/optimizer/
mod.rs

1//! Rule-based SQL query optimizer for [`LogicalPlan`].
2//!
3//! The optimizer applies a configurable sequence of rewrite passes to a
4//! [`LogicalPlan`] tree.  Each pass implements [`OptPass`] and is applied in
5//! order.  The default pipeline (built by [`Optimizer::new`]) includes:
6//!
7//! 1. [`PredicatePushdown`]  — push filters toward leaf scans
8//! 2. [`ProjectionPruning`]  — trim projections to the required column set
9//! 3. [`ConstantFolding`]    — evaluate literal expressions at compile time
10//! 4. [`LimitPushThrough`]   — push `LIMIT N` into single-table scans
11//! 5. [`JoinAlgorithmPass`]  — annotate `Join` nodes with an algorithm hint
12
13pub mod cse;
14pub mod expr_util;
15mod folding;
16pub mod join_algo;
17pub mod join_reorder;
18mod limit;
19mod pruning;
20mod pushdown;
21pub mod simplify;
22
23pub use cse::CommonSubexprElimination;
24pub use folding::ConstantFolding;
25pub use join_algo::{JoinAlgoHint, JoinAlgorithmPass};
26pub use limit::LimitPushThrough;
27pub use pruning::ProjectionPruning;
28pub use pushdown::PredicatePushdown;
29pub use simplify::PredicateSimplification;
30
31use crate::LogicalPlan;
32
33// ── Trait ────────────────────────────────────────────────────────────────────
34
35/// A single optimizer rewrite rule.
36pub trait OptPass: Send + Sync {
37    /// Short name identifying this pass (for tracing / debugging).
38    fn name(&self) -> &'static str;
39
40    /// Apply the rewrite to `plan`, returning the (possibly modified) plan.
41    fn apply(&self, plan: LogicalPlan) -> LogicalPlan;
42}
43
44// ── Optimizer ────────────────────────────────────────────────────────────────
45
46/// Multi-pass rule-based optimizer.
47///
48/// Passes are applied left-to-right in the order they were registered.  The
49/// default pipeline is constructed by [`Optimizer::new`].
50pub struct Optimizer {
51    passes: Vec<Box<dyn OptPass>>,
52}
53
54impl Optimizer {
55    /// Build the default five-pass optimizer pipeline.
56    ///
57    /// The `ProjectionPruning` pass is created with an empty `required` set,
58    /// which means it operates as a no-op unless a caller explicitly builds an
59    /// optimizer with a non-empty requirement list via [`Optimizer::with_passes`].
60    pub fn new() -> Self {
61        let passes: Vec<Box<dyn OptPass>> = vec![
62            Box::new(PredicatePushdown),
63            Box::new(PredicateSimplification),
64            Box::new(CommonSubexprElimination),
65            Box::new(ProjectionPruning { required: vec![] }),
66            Box::new(ConstantFolding),
67            Box::new(LimitPushThrough),
68            Box::new(JoinAlgorithmPass),
69        ];
70        Self { passes }
71    }
72
73    /// Build an optimizer with a custom ordered list of passes.
74    pub fn with_passes(passes: Vec<Box<dyn OptPass>>) -> Self {
75        Self { passes }
76    }
77
78    /// Run all registered passes over `plan` in registration order.
79    pub fn optimize(&self, plan: LogicalPlan) -> LogicalPlan {
80        self.passes.iter().fold(plan, |p, pass| pass.apply(p))
81    }
82}
83
84impl Default for Optimizer {
85    fn default() -> Self {
86        Self::new()
87    }
88}
89
90impl Optimizer {
91    /// Create an [`Optimizer`] with the standard pipeline plus cost-aware passes.
92    ///
93    /// Cost-aware passes (such as `JoinReorder`) own an `Arc<CostModel>` and
94    /// require this constructor to be registered.  Currently this returns the
95    /// same five-pass pipeline as [`Optimizer::new`]; cost-aware passes will
96    /// attach themselves here as they are added in subsequent phases.
97    pub fn with_cost_model(model: crate::cost::CostModel) -> Self {
98        use std::sync::Arc;
99        let model = Arc::new(model);
100        let mut opt = Self::new();
101        opt.passes
102            .push(Box::new(crate::optimizer::join_reorder::JoinReorder::new(
103                model,
104            )));
105        opt
106    }
107}
108
109// ── Public convenience wrapper ────────────────────────────────────────────────
110
111/// Run all five default optimizer passes over `plan`.
112///
113/// This is the top-level convenience wrapper that creates a default
114/// [`Optimizer`] and applies it to the given plan.
115pub fn optimize(plan: LogicalPlan) -> LogicalPlan {
116    Optimizer::new().optimize(plan)
117}
118
119// ── Tests ─────────────────────────────────────────────────────────────────────
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124    use crate::{JoinType, LogicalPlan};
125
126    // Helper: build an unbounded Scan node.
127    fn scan(table: &str) -> LogicalPlan {
128        LogicalPlan::Scan {
129            table: table.to_string(),
130            alias: None,
131            limit: None,
132        }
133    }
134
135    // Helper: build a named-alias Scan.
136    fn scan_alias(table: &str, alias: &str) -> LogicalPlan {
137        LogicalPlan::Scan {
138            table: table.to_string(),
139            alias: Some(alias.to_string()),
140            limit: None,
141        }
142    }
143
144    // Helper: build a Filter.
145    fn filter(input: LogicalPlan, pred: &str) -> LogicalPlan {
146        LogicalPlan::Filter {
147            input: Box::new(input),
148            predicate: pred.to_string(),
149        }
150    }
151
152    // Helper: build an inner Join.
153    fn join(left: LogicalPlan, right: LogicalPlan, on: &str) -> LogicalPlan {
154        LogicalPlan::Join {
155            left: Box::new(left),
156            right: Box::new(right),
157            on: on.to_string(),
158            join_type: JoinType::Inner,
159            algo_hint: None,
160        }
161    }
162
163    /// 1. A predicate referencing only the left side of an INNER join is pushed
164    ///    below the join onto the left scan.
165    #[test]
166    fn test_predicate_pushdown_moves_filter_below_join() {
167        // Filter(Join(Scan(a AS a), Scan(b AS b)), "a.x > 1")
168        let plan = filter(
169            join(scan_alias("a", "a"), scan_alias("b", "b"), "a.id = b.id"),
170            "a.x > 1",
171        );
172        let result = PredicatePushdown.apply(plan);
173
174        match result {
175            LogicalPlan::Join { left, .. } => match *left {
176                LogicalPlan::Filter { predicate, .. } => {
177                    assert_eq!(predicate, "a.x > 1");
178                }
179                other => panic!("expected Filter on left side, got {other:?}"),
180            },
181            other => panic!("expected Join at root after pushdown, got {other:?}"),
182        }
183    }
184
185    /// 2. Projection pruning reduces the column list to only required columns.
186    #[test]
187    fn test_projection_pruning_reduces_columns() {
188        let plan = LogicalPlan::Project {
189            input: Box::new(scan("t")),
190            columns: vec!["col1".to_string(), "col2".to_string(), "col3".to_string()],
191        };
192        let pass = ProjectionPruning {
193            required: vec!["col1".to_string()],
194        };
195        let result = pass.apply(plan);
196        match result {
197            LogicalPlan::Project { columns, .. } => {
198                assert_eq!(columns, vec!["col1".to_string()]);
199            }
200            other => panic!("expected Project, got {other:?}"),
201        }
202    }
203
204    /// 3. Constant-folding a `1 = 1` predicate removes the Filter entirely.
205    #[test]
206    fn test_constant_folding_true_removes_filter() {
207        let plan = filter(scan("t"), "1 = 1");
208        let result = ConstantFolding.apply(plan);
209        match result {
210            LogicalPlan::Scan { table, .. } => assert_eq!(table, "t"),
211            other => panic!("expected Scan after folding 1=1, got {other:?}"),
212        }
213    }
214
215    /// 4. Constant-folding a `1 = 2` predicate replaces the subtree with Empty.
216    #[test]
217    fn test_constant_folding_false_yields_empty() {
218        let plan = filter(scan("t"), "1 = 2");
219        let result = ConstantFolding.apply(plan);
220        assert!(
221            matches!(result, LogicalPlan::Empty),
222            "expected Empty after folding 1=2, got {result:?}"
223        );
224    }
225
226    /// 5. LimitPushThrough pushes `Limit(10)` into the Scan's `limit` field,
227    ///    removing the Limit wrapper node.
228    #[test]
229    fn test_limit_push_through_scan() {
230        let plan = LogicalPlan::Limit {
231            input: Box::new(scan("t")),
232            count: Some(10),
233            offset: None,
234        };
235        let result = LimitPushThrough.apply(plan);
236        match result {
237            LogicalPlan::Scan { limit, .. } => assert_eq!(limit, Some(10)),
238            other => panic!("expected Scan with limit after push-through, got {other:?}"),
239        }
240    }
241
242    /// 6. The full four-pass pipeline composes without errors.
243    #[test]
244    fn test_optimizer_pipeline_composes() {
245        // Build: Limit(Filter(Join(Scan(a AS a), Scan(b AS b)), "a.id > 0"), 5)
246        let base = filter(
247            join(scan_alias("a", "a"), scan_alias("b", "b"), "a.id = b.id"),
248            "a.id > 0",
249        );
250        let limited = LogicalPlan::Limit {
251            input: Box::new(base),
252            count: Some(5),
253            offset: None,
254        };
255        let result = Optimizer::new().optimize(limited);
256        let text = crate::explain(&result);
257        assert!(!text.is_empty(), "explain should produce non-empty output");
258        assert!(
259            text.contains("Join") || text.contains("Scan"),
260            "expected join or scan in plan: {text}"
261        );
262    }
263}