quill_sql/optimizer/rule/
eliminate_limit.rs

1use crate::error::QuillSQLResult;
2use crate::optimizer::logical_optimizer::ApplyOrder;
3use crate::optimizer::LogicalOptimizerRule;
4use crate::plan::logical_plan::{EmptyRelation, LogicalPlan};
5
6pub struct EliminateLimit;
7
8impl LogicalOptimizerRule for EliminateLimit {
9    fn try_optimize(&self, plan: &LogicalPlan) -> QuillSQLResult<Option<LogicalPlan>> {
10        if let LogicalPlan::Limit(limit) = plan {
11            match limit.limit {
12                Some(fetch) => {
13                    if fetch == 0 {
14                        return Ok(Some(LogicalPlan::EmptyRelation(EmptyRelation {
15                            produce_one_row: false,
16                            schema: limit.input.schema().clone(),
17                        })));
18                    }
19                }
20                None => {
21                    if limit.offset == 0 {
22                        let input = limit.input.as_ref();
23                        // input also can be Limit, so we should apply again.
24                        return Ok(Some(
25                            self.try_optimize(input)?.unwrap_or_else(|| input.clone()),
26                        ));
27                    }
28                }
29            }
30        }
31        Ok(None)
32    }
33
34    fn name(&self) -> &str {
35        "EliminateLimit"
36    }
37
38    fn apply_order(&self) -> Option<ApplyOrder> {
39        Some(ApplyOrder::BottomUp)
40    }
41}
42
43#[cfg(test)]
44mod tests {
45    use crate::database::Database;
46    use crate::optimizer::rule::EliminateLimit;
47    use crate::optimizer::LogicalOptimizer;
48    use crate::plan::logical_plan::LogicalPlan;
49    use std::sync::Arc;
50
51    fn build_optimizer() -> LogicalOptimizer {
52        LogicalOptimizer::with_rules(vec![Arc::new(EliminateLimit)])
53    }
54
55    #[test]
56    fn eliminate_limit() {
57        let mut db = Database::new_temp().unwrap();
58        db.run("create table t1 (a int)").unwrap();
59
60        let plan = db.create_logical_plan("select a from t1 limit 0").unwrap();
61        let optimized_plan = build_optimizer().optimize(&plan).unwrap();
62        assert!(matches!(optimized_plan, LogicalPlan::EmptyRelation(_)));
63
64        let plan = db.create_logical_plan("select a from t1 offset 0").unwrap();
65        let optimized_plan = build_optimizer().optimize(&plan).unwrap();
66        if let LogicalPlan::Project(p) = optimized_plan {
67            assert!(matches!(p.input.as_ref(), LogicalPlan::TableScan(_)));
68        } else {
69            panic!("the first node should be project");
70        }
71    }
72}