quill_sql/optimizer/rule/
push_limit_to_scan.rs

1use crate::error::QuillSQLResult;
2use crate::optimizer::logical_optimizer::ApplyOrder;
3use crate::optimizer::LogicalOptimizerRule;
4use crate::plan::logical_plan::{Limit, LogicalPlan, Project, TableScan};
5use std::sync::Arc;
6
7/// Propagate LIMIT/OFFSET requirements into the underlying table scan so that
8/// execution can stop scanning early.
9pub struct PushLimitIntoScan;
10
11impl LogicalOptimizerRule for PushLimitIntoScan {
12    fn try_optimize(&self, plan: &LogicalPlan) -> QuillSQLResult<Option<LogicalPlan>> {
13        let LogicalPlan::Limit(limit) = plan else {
14            return Ok(None);
15        };
16
17        let Some(limit_value) = limit.limit else {
18            return Ok(None);
19        };
20
21        let required_rows = limit.offset.saturating_add(limit_value);
22        match limit.input.as_ref() {
23            LogicalPlan::TableScan(scan) => {
24                maybe_attach_limit(scan, required_rows).map_or(Ok(None), |new_scan| {
25                    Ok(Some(LogicalPlan::Limit(Limit {
26                        limit: limit.limit,
27                        offset: limit.offset,
28                        input: Arc::new(LogicalPlan::TableScan(new_scan)),
29                    })))
30                })
31            }
32            LogicalPlan::Project(project) => {
33                if let LogicalPlan::TableScan(scan) = project.input.as_ref() {
34                    maybe_attach_limit(scan, required_rows).map_or(Ok(None), |new_scan| {
35                        let new_project = LogicalPlan::Project(Project {
36                            exprs: project.exprs.clone(),
37                            input: Arc::new(LogicalPlan::TableScan(new_scan)),
38                            schema: project.schema.clone(),
39                        });
40                        Ok(Some(LogicalPlan::Limit(Limit {
41                            limit: limit.limit,
42                            offset: limit.offset,
43                            input: Arc::new(new_project),
44                        })))
45                    })
46                } else {
47                    Ok(None)
48                }
49            }
50            _ => Ok(None),
51        }
52    }
53
54    fn name(&self) -> &str {
55        "PushLimitIntoScan"
56    }
57
58    fn apply_order(&self) -> Option<ApplyOrder> {
59        Some(ApplyOrder::TopDown)
60    }
61}
62
63fn maybe_attach_limit(scan: &TableScan, required_rows: usize) -> Option<TableScan> {
64    let mut new_scan = scan.clone();
65    let new_limit = match new_scan.limit {
66        Some(existing) => existing.min(required_rows),
67        None => required_rows,
68    };
69
70    if new_scan.limit == Some(new_limit) {
71        return None;
72    }
73
74    new_scan.limit = Some(new_limit);
75    Some(new_scan)
76}
77
78#[cfg(test)]
79mod tests {
80    use crate::database::Database;
81    use crate::optimizer::rule::PushLimitIntoScan;
82    use crate::optimizer::LogicalOptimizer;
83    use crate::plan::logical_plan::{Limit, LogicalPlan};
84    use std::sync::Arc;
85
86    fn build_optimizer() -> LogicalOptimizer {
87        LogicalOptimizer::with_rules(vec![Arc::new(PushLimitIntoScan)])
88    }
89
90    #[test]
91    fn pushes_limit_into_scan() {
92        let mut db = Database::new_temp().unwrap();
93        db.run("create table t1 (a int)").unwrap();
94
95        let plan = db
96            .create_logical_plan("select * from t1 limit 5 offset 2")
97            .unwrap();
98        let optimized_plan = build_optimizer().optimize(&plan).unwrap();
99
100        match optimized_plan {
101            LogicalPlan::Limit(Limit { input, .. }) => match input.as_ref() {
102                LogicalPlan::Project(project) => match project.input.as_ref() {
103                    LogicalPlan::TableScan(scan) => assert_eq!(scan.limit, Some(7)),
104                    other => panic!("expected TableScan under project, got {other:?}"),
105                },
106                other => panic!("expected Project inside limit, got {other:?}"),
107            },
108            other => panic!("expected Limit after optimization, got {other:?}"),
109        }
110    }
111}
quill_sql/optimizer/rule/push_limit_to_scan.rs

quill_sql/optimizer/rule/
push_limit_to_scan.rs
