Plan limit

systemxlabs · Jan 31, 2024 · 988bcd3 · 988bcd3
1 parent 5c031d0
commit 988bcd3
Show file tree

Hide file tree

Showing 5 changed files with 72 additions and 7 deletions.
diff --git a/bustubx/src/error.rs b/bustubx/src/error.rs
@@ -10,6 +10,9 @@ pub enum BustubxError {
     #[error("Internal error: {0}")]
     Internal(String),
 
+    #[error("Plan error: {0}")]
+    Plan(String),
+
     #[error("IO error: {0}")]
     Io(#[from] std::io::Error),
 

diff --git a/bustubx/src/expression/cast.rs b/bustubx/src/expression/cast.rs
@@ -13,7 +13,7 @@ pub struct Cast {
 }
 
 impl ExprTrait for Cast {
-    fn data_type(&self, input_schema: &Schema) -> BustubxResult<DataType> {
+    fn data_type(&self, _input_schema: &Schema) -> BustubxResult<DataType> {
         Ok(self.data_type)
     }
 
@@ -49,6 +49,7 @@ impl ExprTrait for Cast {
                 ))),
             },
             ScalarValue::Int64(v) => match self.data_type {
+                DataType::Int32 => Ok(v.map(|v| v as i32).into()),
                 DataType::Int64 => Ok(value),
                 _ => Err(BustubxError::Internal(format!(
                     "Failed to cast {} as {}",

diff --git a/bustubx/src/planner/logical_plan_v2/limit.rs b/bustubx/src/planner/logical_plan_v2/limit.rs
@@ -4,6 +4,6 @@ use std::sync::Arc;
 #[derive(derive_new::new, Debug, Clone)]
 pub struct Limit {
     pub limit: Option<usize>,
-    pub offset: Option<usize>,
+    pub offset: usize,
     pub input: Arc<LogicalPlanV2>,
 }
diff --git a/bustubx/src/planner/logical_planner/plan_expr.rs b/bustubx/src/planner/logical_planner/plan_expr.rs
@@ -59,8 +59,8 @@ impl LogicalPlanner<'_> {
     pub fn plan_value(&self, value: &sqlparser::ast::Value) -> BustubxResult<Expr> {
         match value {
             sqlparser::ast::Value::Number(s, _) => {
-                let num: u64 = s.parse::<u64>().map_err(|e| {
-                    BustubxError::Internal("Failed to parse literal as u64".to_string())
+                let num: i64 = s.parse::<i64>().map_err(|e| {
+                    BustubxError::Internal("Failed to parse literal as i64".to_string())
                 })?;
                 Ok(Expr::Literal(Literal { value: num.into() }))
             }

diff --git a/bustubx/src/planner/logical_planner/plan_select.rs b/bustubx/src/planner/logical_planner/plan_select.rs
@@ -1,9 +1,11 @@
 use crate::common::ScalarValue;
-use crate::expression::{Alias, Expr, Literal};
-use sqlparser::ast::{Offset, OrderByExpr, Query, SelectItem, SetExpr};
+use crate::expression::{Alias, Expr};
+use crate::{BustubxError, BustubxResult};
+use sqlparser::ast::{OrderByExpr, Query, SelectItem, SetExpr};
 use std::sync::Arc;
 
 use crate::planner::logical_plan::LogicalPlan;
+use crate::planner::logical_plan_v2::{Limit, LogicalPlanV2};
 use crate::planner::operator::LogicalOperator;
 use crate::planner::order_by::BoundOrderBy;
 
@@ -130,7 +132,7 @@ impl<'a> LogicalPlanner<'a> {
     pub fn bind_limit(
         &self,
         limit: &Option<sqlparser::ast::Expr>,
-        offset: &Option<Offset>,
+        offset: &Option<sqlparser::ast::Offset>,
     ) -> (Option<Expr>, Option<Expr>) {
         let limit = limit.as_ref().map(|expr| self.plan_expr(&expr).unwrap());
         let offset = offset
@@ -148,4 +150,63 @@ impl<'a> LogicalPlanner<'a> {
             })
             .collect::<Vec<BoundOrderBy>>()
     }
+
+    pub fn plan_limit_v2(
+        &self,
+        input: LogicalPlanV2,
+        limit: Option<sqlparser::ast::Expr>,
+        offset: Option<sqlparser::ast::Offset>,
+    ) -> BustubxResult<LogicalPlanV2> {
+        if limit.is_none() && offset.is_none() {
+            return Ok(input);
+        }
+
+        let limit = match limit {
+            None => None,
+            Some(limit_expr) => {
+                let n = match self.plan_expr(&limit_expr)? {
+                    Expr::Literal(lit) => match lit.value {
+                        ScalarValue::Int64(Some(v)) if v >= 0 => Ok(v as usize),
+                        _ => Err(BustubxError::Plan(format!(
+                            "LIMIT must not be negative, {}",
+                            lit.value
+                        ))),
+                    },
+                    _ => Err(BustubxError::Plan(format!(
+                        "LIMIT must be literal, {}",
+                        limit_expr
+                    ))),
+                }?;
+                Some(n)
+            }
+        };
+
+        let offset = match offset {
+            None => 0,
+            Some(offset_expr) => match self.plan_expr(&offset_expr.value)? {
+                Expr::Literal(lit) => match lit.value {
+                    ScalarValue::Int64(Some(v)) => {
+                        if v < 0 {
+                            return Err(BustubxError::Plan(format!("Offset must be >= 0, {}", v)));
+                        }
+                        Ok(v as usize)
+                    }
+                    _ => Err(BustubxError::Plan(format!(
+                        "Offset value not int64, {}",
+                        lit.value
+                    ))),
+                },
+                _ => Err(BustubxError::Plan(format!(
+                    "Offset expression not expected, {}",
+                    offset_expr
+                ))),
+            }?,
+        };
+
+        Ok(LogicalPlanV2::Limit(Limit {
+            limit,
+            offset,
+            input: Arc::new(input),
+        }))
+    }
 }