chryso_optimizer/
physical_rules.rs

1use chryso_planner::{LogicalPlan, PhysicalPlan};
2
3pub trait PhysicalRule {
4    fn name(&self) -> &str;
5    fn apply(&self, logical: &LogicalPlan, inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan>;
6}
7
8pub struct PhysicalRuleSet {
9    rules: Vec<Box<dyn PhysicalRule + Send + Sync>>,
10}
11
12impl PhysicalRuleSet {
13    pub fn new() -> Self {
14        Self { rules: Vec::new() }
15    }
16
17    pub fn with_rule(mut self, rule: impl PhysicalRule + Send + Sync + 'static) -> Self {
18        self.rules.push(Box::new(rule));
19        self
20    }
21
22    pub fn apply_all(&self, logical: &LogicalPlan, inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
23        let mut plans = Vec::new();
24        for rule in &self.rules {
25            plans.extend(rule.apply(logical, inputs));
26        }
27        plans
28    }
29}
30
31impl Default for PhysicalRuleSet {
32    fn default() -> Self {
33        PhysicalRuleSet::new()
34            .with_rule(ScanRule)
35            .with_rule(IndexScanRule)
36            .with_rule(DmlRule)
37            .with_rule(DerivedRule)
38            .with_rule(FilterRule)
39            .with_rule(ProjectionRule)
40            .with_rule(JoinRule)
41            .with_rule(AggregateRule)
42            .with_rule(DistinctRule)
43            .with_rule(TopNRule)
44            .with_rule(SortRule)
45            .with_rule(LimitRule)
46    }
47}
48
49pub struct ScanRule;
50
51impl PhysicalRule for ScanRule {
52    fn name(&self) -> &str {
53        "scan_rule"
54    }
55
56    fn apply(&self, logical: &LogicalPlan, _inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
57        let LogicalPlan::Scan { table } = logical else {
58            return Vec::new();
59        };
60        vec![PhysicalPlan::TableScan {
61            table: table.clone(),
62        }]
63    }
64}
65
66pub struct IndexScanRule;
67
68impl PhysicalRule for IndexScanRule {
69    fn name(&self) -> &str {
70        "index_scan_rule"
71    }
72
73    fn apply(&self, logical: &LogicalPlan, _inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
74        let LogicalPlan::IndexScan {
75            table,
76            index,
77            predicate,
78        } = logical
79        else {
80            return Vec::new();
81        };
82        vec![PhysicalPlan::IndexScan {
83            table: table.clone(),
84            index: index.clone(),
85            predicate: predicate.clone(),
86        }]
87    }
88}
89
90pub struct DmlRule;
91
92impl PhysicalRule for DmlRule {
93    fn name(&self) -> &str {
94        "dml_rule"
95    }
96
97    fn apply(&self, logical: &LogicalPlan, _inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
98        let LogicalPlan::Dml { sql } = logical else {
99            return Vec::new();
100        };
101        vec![PhysicalPlan::Dml { sql: sql.clone() }]
102    }
103}
104
105pub struct DerivedRule;
106
107impl PhysicalRule for DerivedRule {
108    fn name(&self) -> &str {
109        "derived_rule"
110    }
111
112    fn apply(&self, logical: &LogicalPlan, inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
113        let LogicalPlan::Derived {
114            alias,
115            column_aliases,
116            ..
117        } = logical
118        else {
119            return Vec::new();
120        };
121        let Some(input) = inputs.first() else {
122            return Vec::new();
123        };
124        vec![PhysicalPlan::Derived {
125            input: Box::new(input.clone()),
126            alias: alias.clone(),
127            column_aliases: column_aliases.clone(),
128        }]
129    }
130}
131
132pub struct FilterRule;
133
134impl PhysicalRule for FilterRule {
135    fn name(&self) -> &str {
136        "filter_rule"
137    }
138
139    fn apply(&self, logical: &LogicalPlan, inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
140        let LogicalPlan::Filter { predicate, .. } = logical else {
141            return Vec::new();
142        };
143        let Some(input) = inputs.first() else {
144            return Vec::new();
145        };
146        vec![PhysicalPlan::Filter {
147            predicate: predicate.clone(),
148            input: Box::new(input.clone()),
149        }]
150    }
151}
152
153pub struct ProjectionRule;
154
155impl PhysicalRule for ProjectionRule {
156    fn name(&self) -> &str {
157        "projection_rule"
158    }
159
160    fn apply(&self, logical: &LogicalPlan, inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
161        let LogicalPlan::Projection { exprs, .. } = logical else {
162            return Vec::new();
163        };
164        let Some(input) = inputs.first() else {
165            return Vec::new();
166        };
167        vec![PhysicalPlan::Projection {
168            exprs: exprs.clone(),
169            input: Box::new(input.clone()),
170        }]
171    }
172}
173
174pub struct JoinRule;
175
176impl PhysicalRule for JoinRule {
177    fn name(&self) -> &str {
178        "join_rule"
179    }
180
181    fn apply(&self, logical: &LogicalPlan, inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
182        let LogicalPlan::Join { join_type, on, .. } = logical else {
183            return Vec::new();
184        };
185        if inputs.len() < 2 {
186            return Vec::new();
187        }
188        vec![
189            PhysicalPlan::Join {
190                join_type: *join_type,
191                algorithm: chryso_planner::JoinAlgorithm::Hash,
192                left: Box::new(inputs[0].clone()),
193                right: Box::new(inputs[1].clone()),
194                on: on.clone(),
195            },
196            PhysicalPlan::Join {
197                join_type: *join_type,
198                algorithm: chryso_planner::JoinAlgorithm::NestedLoop,
199                left: Box::new(inputs[0].clone()),
200                right: Box::new(inputs[1].clone()),
201                on: on.clone(),
202            },
203        ]
204    }
205}
206
207pub struct AggregateRule;
208
209impl PhysicalRule for AggregateRule {
210    fn name(&self) -> &str {
211        "aggregate_rule"
212    }
213
214    fn apply(&self, logical: &LogicalPlan, inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
215        let LogicalPlan::Aggregate {
216            group_exprs,
217            aggr_exprs,
218            ..
219        } = logical
220        else {
221            return Vec::new();
222        };
223        let Some(input) = inputs.first() else {
224            return Vec::new();
225        };
226        vec![PhysicalPlan::Aggregate {
227            group_exprs: group_exprs.clone(),
228            aggr_exprs: aggr_exprs.clone(),
229            input: Box::new(input.clone()),
230        }]
231    }
232}
233
234pub struct DistinctRule;
235
236impl PhysicalRule for DistinctRule {
237    fn name(&self) -> &str {
238        "distinct_rule"
239    }
240
241    fn apply(&self, logical: &LogicalPlan, inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
242        let LogicalPlan::Distinct { .. } = logical else {
243            return Vec::new();
244        };
245        let Some(input) = inputs.first() else {
246            return Vec::new();
247        };
248        vec![PhysicalPlan::Distinct {
249            input: Box::new(input.clone()),
250        }]
251    }
252}
253
254pub struct SortRule;
255
256impl PhysicalRule for SortRule {
257    fn name(&self) -> &str {
258        "sort_rule"
259    }
260
261    fn apply(&self, logical: &LogicalPlan, inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
262        let LogicalPlan::Sort { order_by, .. } = logical else {
263            return Vec::new();
264        };
265        let Some(input) = inputs.first() else {
266            return Vec::new();
267        };
268        vec![PhysicalPlan::Sort {
269            order_by: order_by.clone(),
270            input: Box::new(input.clone()),
271        }]
272    }
273}
274
275pub struct TopNRule;
276
277impl PhysicalRule for TopNRule {
278    fn name(&self) -> &str {
279        "topn_rule"
280    }
281
282    fn apply(&self, logical: &LogicalPlan, inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
283        let LogicalPlan::TopN {
284            order_by, limit, ..
285        } = logical
286        else {
287            return Vec::new();
288        };
289        let Some(input) = inputs.first() else {
290            return Vec::new();
291        };
292        vec![PhysicalPlan::TopN {
293            order_by: order_by.clone(),
294            limit: *limit,
295            input: Box::new(input.clone()),
296        }]
297    }
298}
299
300pub struct LimitRule;
301
302impl PhysicalRule for LimitRule {
303    fn name(&self) -> &str {
304        "limit_rule"
305    }
306
307    fn apply(&self, logical: &LogicalPlan, inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
308        let LogicalPlan::Limit { limit, offset, .. } = logical else {
309            return Vec::new();
310        };
311        let Some(input) = inputs.first() else {
312            return Vec::new();
313        };
314        vec![PhysicalPlan::Limit {
315            limit: *limit,
316            offset: *offset,
317            input: Box::new(input.clone()),
318        }]
319    }
320}
321
322#[cfg(test)]
323mod tests {
324    use super::{JoinRule, PhysicalRule};
325    use chryso_planner::{LogicalPlan, PhysicalPlan};
326
327    #[test]
328    fn join_rule_produces_two_algorithms() {
329        let logical = LogicalPlan::Join {
330            join_type: chryso_core::ast::JoinType::Inner,
331            left: Box::new(LogicalPlan::Scan {
332                table: "t1".to_string(),
333            }),
334            right: Box::new(LogicalPlan::Scan {
335                table: "t2".to_string(),
336            }),
337            on: chryso_core::ast::Expr::Identifier("t1.id = t2.id".to_string()),
338        };
339        let inputs = vec![
340            PhysicalPlan::TableScan {
341                table: "t1".to_string(),
342            },
343            PhysicalPlan::TableScan {
344                table: "t2".to_string(),
345            },
346        ];
347        let rule = JoinRule;
348        let plans = rule.apply(&logical, &inputs);
349        assert_eq!(plans.len(), 2);
350    }
351}
chryso_optimizer/physical_rules.rs

chryso_optimizer/
physical_rules.rs
