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 { alias, column_aliases, .. } = logical else {
114            return Vec::new();
115        };
116        let Some(input) = inputs.first() else {
117            return Vec::new();
118        };
119        vec![PhysicalPlan::Derived {
120            input: Box::new(input.clone()),
121            alias: alias.clone(),
122            column_aliases: column_aliases.clone(),
123        }]
124    }
125}
126
127pub struct FilterRule;
128
129impl PhysicalRule for FilterRule {
130    fn name(&self) -> &str {
131        "filter_rule"
132    }
133
134    fn apply(&self, logical: &LogicalPlan, inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
135        let LogicalPlan::Filter { predicate, .. } = logical else {
136            return Vec::new();
137        };
138        let Some(input) = inputs.first() else {
139            return Vec::new();
140        };
141        vec![PhysicalPlan::Filter {
142            predicate: predicate.clone(),
143            input: Box::new(input.clone()),
144        }]
145    }
146}
147
148pub struct ProjectionRule;
149
150impl PhysicalRule for ProjectionRule {
151    fn name(&self) -> &str {
152        "projection_rule"
153    }
154
155    fn apply(&self, logical: &LogicalPlan, inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
156        let LogicalPlan::Projection { exprs, .. } = logical else {
157            return Vec::new();
158        };
159        let Some(input) = inputs.first() else {
160            return Vec::new();
161        };
162        vec![PhysicalPlan::Projection {
163            exprs: exprs.clone(),
164            input: Box::new(input.clone()),
165        }]
166    }
167}
168
169pub struct JoinRule;
170
171impl PhysicalRule for JoinRule {
172    fn name(&self) -> &str {
173        "join_rule"
174    }
175
176    fn apply(&self, logical: &LogicalPlan, inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
177        let LogicalPlan::Join {
178            join_type, on, ..
179        } = logical
180        else {
181            return Vec::new();
182        };
183        if inputs.len() < 2 {
184            return Vec::new();
185        }
186        vec![
187            PhysicalPlan::Join {
188                join_type: *join_type,
189                algorithm: chryso_planner::JoinAlgorithm::Hash,
190                left: Box::new(inputs[0].clone()),
191                right: Box::new(inputs[1].clone()),
192                on: on.clone(),
193            },
194            PhysicalPlan::Join {
195            join_type: *join_type,
196            algorithm: chryso_planner::JoinAlgorithm::NestedLoop,
197            left: Box::new(inputs[0].clone()),
198            right: Box::new(inputs[1].clone()),
199            on: on.clone(),
200        },
201        ]
202    }
203}
204
205pub struct AggregateRule;
206
207impl PhysicalRule for AggregateRule {
208    fn name(&self) -> &str {
209        "aggregate_rule"
210    }
211
212    fn apply(&self, logical: &LogicalPlan, inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
213        let LogicalPlan::Aggregate {
214            group_exprs,
215            aggr_exprs,
216            ..
217        } = logical
218        else {
219            return Vec::new();
220        };
221        let Some(input) = inputs.first() else {
222            return Vec::new();
223        };
224        vec![PhysicalPlan::Aggregate {
225            group_exprs: group_exprs.clone(),
226            aggr_exprs: aggr_exprs.clone(),
227            input: Box::new(input.clone()),
228        }]
229    }
230}
231
232pub struct DistinctRule;
233
234impl PhysicalRule for DistinctRule {
235    fn name(&self) -> &str {
236        "distinct_rule"
237    }
238
239    fn apply(&self, logical: &LogicalPlan, inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
240        let LogicalPlan::Distinct { .. } = logical else {
241            return Vec::new();
242        };
243        let Some(input) = inputs.first() else {
244            return Vec::new();
245        };
246        vec![PhysicalPlan::Distinct {
247            input: Box::new(input.clone()),
248        }]
249    }
250}
251
252pub struct SortRule;
253
254impl PhysicalRule for SortRule {
255    fn name(&self) -> &str {
256        "sort_rule"
257    }
258
259    fn apply(&self, logical: &LogicalPlan, inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
260        let LogicalPlan::Sort { order_by, .. } = logical else {
261            return Vec::new();
262        };
263        let Some(input) = inputs.first() else {
264            return Vec::new();
265        };
266        vec![PhysicalPlan::Sort {
267            order_by: order_by.clone(),
268            input: Box::new(input.clone()),
269        }]
270    }
271}
272
273pub struct TopNRule;
274
275impl PhysicalRule for TopNRule {
276    fn name(&self) -> &str {
277        "topn_rule"
278    }
279
280    fn apply(&self, logical: &LogicalPlan, inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
281        let LogicalPlan::TopN { order_by, limit, .. } = logical else {
282            return Vec::new();
283        };
284        let Some(input) = inputs.first() else {
285            return Vec::new();
286        };
287        vec![PhysicalPlan::TopN {
288            order_by: order_by.clone(),
289            limit: *limit,
290            input: Box::new(input.clone()),
291        }]
292    }
293}
294
295pub struct LimitRule;
296
297impl PhysicalRule for LimitRule {
298    fn name(&self) -> &str {
299        "limit_rule"
300    }
301
302    fn apply(&self, logical: &LogicalPlan, inputs: &[PhysicalPlan]) -> Vec<PhysicalPlan> {
303        let LogicalPlan::Limit { limit, offset, .. } = logical else {
304            return Vec::new();
305        };
306        let Some(input) = inputs.first() else {
307            return Vec::new();
308        };
309        vec![PhysicalPlan::Limit {
310            limit: *limit,
311            offset: *offset,
312            input: Box::new(input.clone()),
313        }]
314    }
315}
316
317#[cfg(test)]
318mod tests {
319    use super::{JoinRule, PhysicalRule};
320    use chryso_planner::{LogicalPlan, PhysicalPlan};
321
322    #[test]
323    fn join_rule_produces_two_algorithms() {
324        let logical = LogicalPlan::Join {
325            join_type: chryso_core::ast::JoinType::Inner,
326            left: Box::new(LogicalPlan::Scan {
327                table: "t1".to_string(),
328            }),
329            right: Box::new(LogicalPlan::Scan {
330                table: "t2".to_string(),
331            }),
332            on: chryso_core::ast::Expr::Identifier("t1.id = t2.id".to_string()),
333        };
334        let inputs = vec![
335            PhysicalPlan::TableScan {
336                table: "t1".to_string(),
337            },
338            PhysicalPlan::TableScan {
339                table: "t2".to_string(),
340            },
341        ];
342        let rule = JoinRule;
343        let plans = rule.apply(&logical, &inputs);
344        assert_eq!(plans.len(), 2);
345    }
346}
chryso_optimizer/physical_rules.rs

chryso_optimizer/
physical_rules.rs
