Plan query v2

systemxlabs · Feb 1, 2024 · 8c31f27 · 8c31f27
1 parent afcc2ed
commit 8c31f27
Show file tree

Hide file tree

Showing 5 changed files with 65 additions and 11 deletions.
diff --git a/bustubx/src/planner/logical_plan_v2/project.rs b/bustubx/src/planner/logical_plan_v2/project.rs
@@ -4,6 +4,6 @@ use std::sync::Arc;
 
 #[derive(derive_new::new, Debug, Clone)]
 pub struct Project {
-    pub expressions: Vec<Expr>,
+    pub exprs: Vec<Expr>,
     pub input: Arc<LogicalPlanV2>,
 }
diff --git a/bustubx/src/planner/logical_planner/logical_planner.rs b/bustubx/src/planner/logical_planner/logical_planner.rs
@@ -198,7 +198,7 @@ impl<'a> LogicalPlanner<'a> {
         }
     }
 
-    pub fn plan_order_by_v2(
+    pub fn plan_order_by_expr(
         &self,
         order_by: &sqlparser::ast::OrderByExpr,
     ) -> BustubxResult<OrderByExpr> {

diff --git a/bustubx/src/planner/logical_planner/plan_create_index.rs b/bustubx/src/planner/logical_planner/plan_create_index.rs
@@ -54,7 +54,7 @@ impl<'a> LogicalPlanner<'a> {
         let table = self.plan_table_name(table_name)?;
         let mut columns_expr = vec![];
         for col in columns.iter() {
-            let col_expr = self.plan_order_by_v2(&col)?;
+            let col_expr = self.plan_order_by_expr(&col)?;
             columns_expr.push(col_expr);
         }
         Ok(LogicalPlanV2::CreateIndex(CreateIndex {

diff --git a/bustubx/src/planner/logical_planner/plan_query.rs b/bustubx/src/planner/logical_planner/plan_query.rs
@@ -5,7 +5,7 @@ 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::logical_plan_v2::{Limit, LogicalPlanV2, Sort};
 use crate::planner::operator::LogicalOperator;
 use crate::planner::order_by::BoundOrderBy;
 
@@ -151,11 +151,38 @@ impl<'a> LogicalPlanner<'a> {
             .collect::<Vec<BoundOrderBy>>()
     }
 
+    pub fn plan_query_v2(&self, query: &sqlparser::ast::Query) -> BustubxResult<LogicalPlanV2> {
+        let plan = self.plan_set_expr(&query.body)?;
+        let plan = self.plan_order_by(plan, &query.order_by)?;
+        self.plan_limit_v2(plan, &query.limit, &query.offset)
+    }
+
+    pub fn plan_order_by(
+        &self,
+        input: LogicalPlanV2,
+        order_by: &Vec<sqlparser::ast::OrderByExpr>,
+    ) -> BustubxResult<LogicalPlanV2> {
+        if order_by.is_empty() {
+            return Ok(input);
+        }
+
+        let mut order_by_exprs = vec![];
+        for order in order_by {
+            order_by_exprs.push(self.plan_order_by_expr(order)?);
+        }
+
+        Ok(LogicalPlanV2::Sort(Sort {
+            expr: order_by_exprs,
+            input: Arc::new(input),
+            limit: None,
+        }))
+    }
+
     pub fn plan_limit_v2(
         &self,
         input: LogicalPlanV2,
-        limit: Option<sqlparser::ast::Expr>,
-        offset: Option<sqlparser::ast::Offset>,
+        limit: &Option<sqlparser::ast::Expr>,
+        offset: &Option<sqlparser::ast::Offset>,
     ) -> BustubxResult<LogicalPlanV2> {
         if limit.is_none() && offset.is_none() {
             return Ok(input);

diff --git a/bustubx/src/planner/logical_planner/plan_set_expr.rs b/bustubx/src/planner/logical_planner/plan_set_expr.rs
@@ -1,7 +1,7 @@
 use crate::catalog::{Column, Schema};
-use crate::expression::ExprTrait;
+use crate::expression::{Alias, Expr, ExprTrait};
 use crate::planner::logical_plan_v2::{
-    EmptyRelation, Filter, Join, LogicalPlanV2, TableScan, Values,
+    EmptyRelation, Filter, Join, LogicalPlanV2, Project, TableScan, Values,
 };
 use crate::planner::table_ref::join::JoinType;
 use crate::planner::LogicalPlanner;
@@ -14,8 +14,7 @@ impl LogicalPlanner<'_> {
         set_expr: &sqlparser::ast::SetExpr,
     ) -> BustubxResult<LogicalPlanV2> {
         match set_expr {
-            sqlparser::ast::SetExpr::Select(select) => todo!(),
-            sqlparser::ast::SetExpr::Query(_) => todo!(),
+            sqlparser::ast::SetExpr::Select(select) => self.plan_select(select),
             sqlparser::ast::SetExpr::Values(values) => self.plan_values(values),
             _ => Err(BustubxError::Plan(format!(
                 "Failed to plan set expr: {}",
@@ -27,8 +26,36 @@ impl LogicalPlanner<'_> {
     pub fn plan_select(&self, select: &sqlparser::ast::Select) -> BustubxResult<LogicalPlanV2> {
         let table_scan = self.plan_from_tables(&select.from)?;
         let selection = self.plan_selection(table_scan, &select.selection)?;
+        self.plan_project(selection, &select.projection)
+    }
 
-        todo!()
+    pub fn plan_project(
+        &self,
+        input: LogicalPlanV2,
+        project: &Vec<sqlparser::ast::SelectItem>,
+    ) -> BustubxResult<LogicalPlanV2> {
+        let mut exprs = vec![];
+        for select_item in project {
+            match select_item {
+                sqlparser::ast::SelectItem::UnnamedExpr(expr) => exprs.push(self.plan_expr(expr)?),
+                sqlparser::ast::SelectItem::ExprWithAlias { expr, alias } => {
+                    exprs.push(Expr::Alias(Alias {
+                        name: alias.value.clone(),
+                        expr: Box::new(self.plan_expr(expr)?),
+                    }))
+                }
+                _ => {
+                    return Err(BustubxError::Plan(format!(
+                        "sqlparser select item {} not supported",
+                        select_item
+                    )));
+                }
+            }
+        }
+        Ok(LogicalPlanV2::Project(Project {
+            exprs,
+            input: Arc::new(input),
+        }))
     }
 
     pub fn plan_selection(