quill_sql/optimizer/rule/
merge_limit.rs

1use crate::error::QuillSQLResult;
2use crate::optimizer::logical_optimizer::ApplyOrder;
3use crate::optimizer::LogicalOptimizerRule;
4use crate::plan::logical_plan::{Limit, LogicalPlan};
5use std::cmp::min;
6use std::sync::Arc;
7
8pub struct MergeLimit;
9
10impl LogicalOptimizerRule for MergeLimit {
11    fn try_optimize(&self, plan: &LogicalPlan) -> QuillSQLResult<Option<LogicalPlan>> {
12        let LogicalPlan::Limit(parent) = plan else {
13            return Ok(None);
14        };
15
16        if let LogicalPlan::Limit(child) = &*parent.input {
17            let new_limit = match (parent.limit, child.limit) {
18                (Some(parent_limit), Some(child_limit)) => {
19                    Some(min(parent_limit, child_limit.saturating_sub(parent.offset)))
20                }
21                (Some(parent_limit), None) => Some(parent_limit),
22                (None, Some(child_limit)) => Some(child_limit.saturating_sub(parent.offset)),
23                (None, None) => None,
24            };
25            let plan = LogicalPlan::Limit(Limit {
26                limit: new_limit,
27                offset: child.offset + parent.offset,
28                input: Arc::new((*child.input).clone()),
29            });
30            self.try_optimize(&plan)
31                .map(|opt_plan| opt_plan.or_else(|| Some(plan)))
32        } else {
33            Ok(None)
34        }
35    }
36
37    fn name(&self) -> &str {
38        "MergeLimit"
39    }
40
41    fn apply_order(&self) -> Option<ApplyOrder> {
42        Some(ApplyOrder::TopDown)
43    }
44}
45
46#[cfg(test)]
47mod tests {
48    use crate::catalog::EMPTY_SCHEMA_REF;
49    use crate::optimizer::rule::MergeLimit;
50    use crate::optimizer::LogicalOptimizer;
51    use crate::plan::logical_plan::{EmptyRelation, Limit, LogicalPlan};
52    use std::sync::Arc;
53
54    fn build_optimizer() -> LogicalOptimizer {
55        LogicalOptimizer::with_rules(vec![Arc::new(MergeLimit)])
56    }
57
58    #[test]
59    fn merge_limit() {
60        let plan = LogicalPlan::Limit(Limit {
61            limit: Some(10),
62            offset: 0,
63            input: Arc::new(LogicalPlan::Limit(Limit {
64                limit: Some(1000),
65                offset: 0,
66                input: Arc::new(LogicalPlan::Limit(Limit {
67                    limit: None,
68                    offset: 10,
69                    input: Arc::new(LogicalPlan::EmptyRelation(EmptyRelation {
70                        produce_one_row: false,
71                        schema: EMPTY_SCHEMA_REF.clone(),
72                    })),
73                })),
74            })),
75        });
76        let optimized_plan = build_optimizer().optimize(&plan).unwrap();
77
78        if let LogicalPlan::Limit(Limit {
79            limit,
80            offset,
81            input,
82        }) = optimized_plan
83        {
84            assert_eq!(limit, Some(10));
85            assert_eq!(offset, 10);
86            assert!(matches!(input.as_ref(), LogicalPlan::EmptyRelation(_)));
87        } else {
88            panic!("the first node should be limit");
89        }
90    }
91}
quill_sql/optimizer/rule/merge_limit.rs

quill_sql/optimizer/rule/
merge_limit.rs
