diff --git a/bustubx/src/planner/logical_plan_v2/insert.rs b/bustubx/src/planner/logical_plan_v2/insert.rs index 4d391e5..7953cef 100644 --- a/bustubx/src/planner/logical_plan_v2/insert.rs +++ b/bustubx/src/planner/logical_plan_v2/insert.rs @@ -6,6 +6,6 @@ use std::sync::Arc; #[derive(derive_new::new, Debug, Clone)] pub struct Insert { pub table: TableReference, - pub schema: SchemaRef, + pub columns: Vec, pub input: Arc, } diff --git a/bustubx/src/planner/logical_planner/mod.rs b/bustubx/src/planner/logical_planner/mod.rs index 7e16d47..dc76dac 100644 --- a/bustubx/src/planner/logical_planner/mod.rs +++ b/bustubx/src/planner/logical_planner/mod.rs @@ -4,5 +4,6 @@ mod plan_create_table; mod plan_expr; mod plan_insert; mod plan_select; +mod plan_set_expr; pub use logical_planner::{LogicalPlanner, PlannerContext}; diff --git a/bustubx/src/planner/logical_planner/plan_insert.rs b/bustubx/src/planner/logical_planner/plan_insert.rs index 5cccbc2..57d4f9a 100644 --- a/bustubx/src/planner/logical_planner/plan_insert.rs +++ b/bustubx/src/planner/logical_planner/plan_insert.rs @@ -1,8 +1,10 @@ use crate::expression::{Cast, Expr}; +use crate::BustubxResult; use sqlparser::ast::{Ident, ObjectName, Query, SetExpr}; use std::sync::Arc; use crate::planner::logical_plan::LogicalPlan; +use crate::planner::logical_plan_v2::{Insert, LogicalPlanV2}; use crate::planner::operator::LogicalOperator; use super::LogicalPlanner; @@ -66,4 +68,23 @@ impl<'a> LogicalPlanner<'a> { unimplemented!() } } + + pub fn plan_insert_v2( + &self, + table_name: &ObjectName, + columns_ident: &Vec, + source: &Query, + ) -> BustubxResult { + let values = self.plan_set_expr(source.body.as_ref())?; + let table = self.plan_table_name(table_name)?; + let columns = columns_ident + .iter() + .map(|ident| ident.value.clone()) + .collect(); + Ok(LogicalPlanV2::Insert(Insert { + table, + columns, + input: Arc::new(values), + })) + } } diff --git a/bustubx/src/planner/logical_planner/plan_set_expr.rs b/bustubx/src/planner/logical_planner/plan_set_expr.rs new file mode 100644 index 0000000..77331d2 --- /dev/null +++ b/bustubx/src/planner/logical_planner/plan_set_expr.rs @@ -0,0 +1,55 @@ +use crate::catalog::{Column, Schema}; +use crate::expression::ExprTrait; +use crate::planner::logical_plan_v2::{LogicalPlanV2, Values}; +use crate::planner::LogicalPlanner; +use crate::{BustubxError, BustubxResult}; +use std::sync::Arc; + +impl LogicalPlanner<'_> { + pub fn plan_set_expr( + &self, + set_expr: &sqlparser::ast::SetExpr, + ) -> BustubxResult { + match set_expr { + sqlparser::ast::SetExpr::Select(select) => todo!(), + sqlparser::ast::SetExpr::Query(_) => todo!(), + sqlparser::ast::SetExpr::Values(values) => self.plan_values(values), + _ => Err(BustubxError::Plan(format!( + "Failed to plan set expr: {}", + set_expr + ))), + } + } + + pub fn plan_values(&self, values: &sqlparser::ast::Values) -> BustubxResult { + let mut result = vec![]; + for row in values.rows.iter() { + let mut record = vec![]; + for item in row { + record.push(self.plan_expr(item)?); + } + result.push(record); + } + if result.is_empty() { + return Ok(LogicalPlanV2::Values(Values { + schema: Arc::new(Schema::empty()), + values: vec![], + })); + } + + // parse schema + let first_row = &result[0]; + let mut columns = vec![]; + for (idx, item) in first_row.iter().enumerate() { + columns.push(Column::new( + idx.to_string(), + item.data_type(&Schema::empty())?, + )) + } + + Ok(LogicalPlanV2::Values(Values { + schema: Arc::new(Schema::new(columns)), + values: result, + })) + } +}