Project schema

systemxlabs · Feb 1, 2024 · 9f506c2 · 9f506c2
1 parent ad4c63d
commit 9f506c2
Show file tree

Hide file tree

Showing 5 changed files with 40 additions and 5 deletions.
diff --git a/bustubx/src/expression/mod.rs b/bustubx/src/expression/mod.rs
@@ -10,11 +10,11 @@ pub use cast::Cast;
 pub use column::ColumnExpr;
 pub use literal::Literal;
 
-use crate::catalog::DataType;
 use crate::catalog::Schema;
+use crate::catalog::{Column, DataType};
 use crate::common::ScalarValue;
 use crate::storage::Tuple;
-use crate::BustubxResult;
+use crate::{BustubxError, BustubxResult};
 
 pub trait ExprTrait {
     /// Get the data type of this expression, given the schema of the input
@@ -60,3 +60,22 @@ impl ExprTrait for Expr {
         }
     }
 }
+
+impl Expr {
+    pub fn to_column(&self, input_schema: &Schema) -> BustubxResult<Column> {
+        match self {
+            Expr::Column(ColumnExpr { relation, name }) => Ok(Column {
+                name: name.clone(),
+                data_type: self.data_type(input_schema)?,
+                // TODO fix
+                nullable: false,
+            }),
+            _ => {
+                return Err(BustubxError::Plan(format!(
+                    "expr {:?} as column not supported",
+                    self
+                )))
+            }
+        }
+    }
+}
diff --git a/bustubx/src/planner/logical_plan_v2/mod.rs b/bustubx/src/planner/logical_plan_v2/mod.rs
@@ -49,7 +49,7 @@ impl LogicalPlanV2 {
             LogicalPlanV2::Insert(_) => todo!(),
             LogicalPlanV2::Join(Join { schema, .. }) => schema,
             LogicalPlanV2::Limit(Limit { input, .. }) => input.schema(),
-            LogicalPlanV2::Project(_) => todo!(),
+            LogicalPlanV2::Project(Project { schema, .. }) => schema,
             LogicalPlanV2::TableScan(TableScan { schema, .. }) => schema,
             LogicalPlanV2::Sort(Sort { input, .. }) => input.schema(),
             LogicalPlanV2::Values(Values { schema, .. }) => schema,

diff --git a/bustubx/src/planner/logical_plan_v2/project.rs b/bustubx/src/planner/logical_plan_v2/project.rs
@@ -1,3 +1,4 @@
+use crate::catalog::SchemaRef;
 use crate::expression::Expr;
 use crate::planner::logical_plan_v2::LogicalPlanV2;
 use std::sync::Arc;
@@ -6,4 +7,5 @@ use std::sync::Arc;
 pub struct Project {
     pub exprs: Vec<Expr>,
     pub input: Arc<LogicalPlanV2>,
+    pub schema: SchemaRef,
 }
diff --git a/bustubx/src/planner/logical_plan_v2/util.rs b/bustubx/src/planner/logical_plan_v2/util.rs
@@ -1,6 +1,8 @@
 use crate::catalog::{ColumnRef, Schema};
+use crate::expression::{ColumnExpr, Expr};
+use crate::planner::logical_plan_v2::LogicalPlanV2;
 use crate::planner::table_ref::join::JoinType;
-use crate::BustubxResult;
+use crate::{BustubxError, BustubxResult};
 use std::sync::Arc;
 
 pub fn build_join_schema(
@@ -40,3 +42,12 @@ pub fn build_join_schema(
     };
     Ok(Schema { columns })
 }
+
+pub fn project_schema(input: &LogicalPlanV2, exprs: &[Expr]) -> BustubxResult<Schema> {
+    let input_schema = &input.schema();
+    let mut columns = vec![];
+    for expr in exprs {
+        columns.push(expr.to_column(input_schema)?)
+    }
+    Ok(Schema::new(columns))
+}
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,8 @@
 use crate::catalog::{Column, Schema};
 use crate::expression::{Alias, Expr, ExprTrait};
 use crate::planner::logical_plan_v2::{
-    build_join_schema, EmptyRelation, Filter, Join, LogicalPlanV2, Project, TableScan, Values,
+    build_join_schema, project_schema, EmptyRelation, Filter, Join, LogicalPlanV2, Project,
+    TableScan, Values,
 };
 use crate::planner::table_ref::join::JoinType;
 use crate::planner::LogicalPlanner;
@@ -52,9 +53,11 @@ impl LogicalPlanner<'_> {
                 }
             }
         }
+        let schema = Arc::new(project_schema(&input, &exprs)?);
         Ok(LogicalPlanV2::Project(Project {
             exprs,
             input: Arc::new(input),
+            schema,
         }))
     }