quill_sql/plan/logical_planner/
plan_query.rs

1use crate::error::{QuillSQLError, QuillSQLResult};
2use crate::expression::Expr;
3use crate::utils::scalar::ScalarValue;
4use std::sync::Arc;
5
6use crate::plan::logical_plan::{Limit, LogicalPlan, Sort};
7
8use super::LogicalPlanner;
9
10impl<'a> LogicalPlanner<'a> {
11    pub fn plan_query(&self, query: &sqlparser::ast::Query) -> QuillSQLResult<LogicalPlan> {
12        let plan = self.plan_set_expr(&query.body)?;
13        let plan = self.plan_order_by(plan, &query.order_by)?;
14        self.plan_limit(plan, &query.limit, &query.offset)
15    }
16
17    pub fn plan_order_by(
18        &self,
19        input: LogicalPlan,
20        order_by: &Vec<sqlparser::ast::OrderByExpr>,
21    ) -> QuillSQLResult<LogicalPlan> {
22        if order_by.is_empty() {
23            return Ok(input);
24        }
25
26        let mut order_by_exprs = vec![];
27        for order in order_by {
28            order_by_exprs.push(self.bind_order_by_expr(order)?);
29        }
30
31        Ok(LogicalPlan::Sort(Sort {
32            order_by: order_by_exprs,
33            input: Arc::new(input),
34            limit: None,
35        }))
36    }
37
38    pub fn plan_limit(
39        &self,
40        input: LogicalPlan,
41        limit: &Option<sqlparser::ast::Expr>,
42        offset: &Option<sqlparser::ast::Offset>,
43    ) -> QuillSQLResult<LogicalPlan> {
44        if limit.is_none() && offset.is_none() {
45            return Ok(input);
46        }
47
48        let limit = match limit {
49            None => None,
50            Some(limit_expr) => {
51                let n = match self.bind_expr(limit_expr)? {
52                    Expr::Literal(lit) => match lit.value {
53                        ScalarValue::Int64(Some(v)) if v >= 0 => Ok(v as usize),
54                        _ => Err(QuillSQLError::Plan(format!(
55                            "LIMIT must not be negative, {}",
56                            lit.value
57                        ))),
58                    },
59                    _ => Err(QuillSQLError::Plan(format!(
60                        "LIMIT must be literal, {}",
61                        limit_expr
62                    ))),
63                }?;
64                Some(n)
65            }
66        };
67
68        let offset = match offset {
69            None => 0,
70            Some(offset_expr) => match self.bind_expr(&offset_expr.value)? {
71                Expr::Literal(lit) => match lit.value {
72                    ScalarValue::Int64(Some(v)) => {
73                        if v < 0 {
74                            return Err(QuillSQLError::Plan(format!("Offset must be >= 0, {}", v)));
75                        }
76                        Ok(v as usize)
77                    }
78                    _ => Err(QuillSQLError::Plan(format!(
79                        "Offset value not int64, {}",
80                        lit.value
81                    ))),
82                },
83                _ => Err(QuillSQLError::Plan(format!(
84                    "Offset expression not expected, {}",
85                    offset_expr
86                ))),
87            }?,
88        };
89
90        Ok(LogicalPlan::Limit(Limit {
91            limit,
92            offset,
93            input: Arc::new(input),
94        }))
95    }
96}
quill_sql/plan/logical_planner/plan_query.rs

quill_sql/plan/logical_planner/
plan_query.rs
