Add nullable and to_column to ExprTrait

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

Hide file tree

Showing 7 changed files with 92 additions and 21 deletions.
diff --git a/bustubx/src/expression/alias.rs b/bustubx/src/expression/alias.rs
@@ -1,9 +1,10 @@
-use crate::catalog::DataType;
 use crate::catalog::Schema;
+use crate::catalog::{Column, DataType};
 use crate::common::ScalarValue;
 use crate::error::BustubxResult;
 use crate::expression::{Expr, ExprTrait};
 use crate::storage::Tuple;
+use crate::BustubxError;
 
 #[derive(Clone, PartialEq, Eq, Debug)]
 pub struct Alias {
@@ -16,7 +17,18 @@ impl ExprTrait for Alias {
         self.expr.data_type(input_schema)
     }
 
+    fn nullable(&self, input_schema: &Schema) -> BustubxResult<bool> {
+        self.expr.nullable(input_schema)
+    }
+
     fn evaluate(&self, tuple: &Tuple) -> BustubxResult<ScalarValue> {
         self.expr.evaluate(tuple)
     }
+
+    fn to_column(&self, input_schema: &Schema) -> BustubxResult<Column> {
+        Err(BustubxError::Plan(format!(
+            "expr {:?} as column not supported",
+            self
+        )))
+    }
 }
diff --git a/bustubx/src/expression/binary.rs b/bustubx/src/expression/binary.rs
@@ -1,5 +1,5 @@
-use crate::catalog::DataType;
 use crate::catalog::Schema;
+use crate::catalog::{Column, DataType};
 use crate::common::ScalarValue;
 use crate::error::BustubxResult;
 use crate::expression::{Expr, ExprTrait};
@@ -37,6 +37,10 @@ impl ExprTrait for BinaryExpr {
         }
     }
 
+    fn nullable(&self, input_schema: &Schema) -> BustubxResult<bool> {
+        Ok(self.left.nullable(input_schema)? || self.right.nullable(input_schema)?)
+    }
+
     fn evaluate(&self, tuple: &Tuple) -> BustubxResult<ScalarValue> {
         let l = self.left.evaluate(tuple)?;
         let r = self.right.evaluate(tuple)?;
@@ -83,6 +87,13 @@ impl ExprTrait for BinaryExpr {
             ))),
         }
     }
+
+    fn to_column(&self, input_schema: &Schema) -> BustubxResult<Column> {
+        Err(BustubxError::Plan(format!(
+            "expr {:?} as column not supported",
+            self
+        )))
+    }
 }
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Hash)]

diff --git a/bustubx/src/expression/cast.rs b/bustubx/src/expression/cast.rs
@@ -1,4 +1,4 @@
-use crate::catalog::{DataType, Schema};
+use crate::catalog::{Column, DataType, Schema};
 use crate::common::ScalarValue;
 use crate::expression::{Expr, ExprTrait};
 use crate::{BustubxError, BustubxResult, Tuple};
@@ -17,6 +17,10 @@ impl ExprTrait for Cast {
         Ok(self.data_type)
     }
 
+    fn nullable(&self, input_schema: &Schema) -> BustubxResult<bool> {
+        self.expr.nullable(input_schema)
+    }
+
     fn evaluate(&self, tuple: &Tuple) -> BustubxResult<ScalarValue> {
         let value = self.expr.evaluate(tuple)?;
         match value {
@@ -66,4 +70,11 @@ impl ExprTrait for Cast {
             },
         }
     }
+
+    fn to_column(&self, input_schema: &Schema) -> BustubxResult<Column> {
+        Err(BustubxError::Plan(format!(
+            "expr {:?} as column not supported",
+            self
+        )))
+    }
 }
diff --git a/bustubx/src/expression/column.rs b/bustubx/src/expression/column.rs
@@ -1,5 +1,5 @@
-use crate::catalog::DataType;
 use crate::catalog::Schema;
+use crate::catalog::{Column, DataType};
 use crate::common::table_ref::TableReference;
 use crate::common::ScalarValue;
 use crate::error::{BustubxError, BustubxResult};
@@ -23,7 +23,25 @@ impl ExprTrait for ColumnExpr {
         )
     }
 
+    fn nullable(&self, input_schema: &Schema) -> BustubxResult<bool> {
+        input_schema.get_col_by_name(&self.name).map_or(
+            Err(BustubxError::Plan(format!(
+                "Not found column {} in input schema {:?}",
+                self.name, input_schema
+            ))),
+            |col| Ok(col.nullable),
+        )
+    }
+
     fn evaluate(&self, tuple: &Tuple) -> BustubxResult<ScalarValue> {
         Ok(tuple.get_value_by_col_name(&tuple.schema, &self.name))
     }
+
+    fn to_column(&self, input_schema: &Schema) -> BustubxResult<Column> {
+        Ok(Column {
+            name: self.name.clone(),
+            data_type: self.data_type(input_schema)?,
+            nullable: self.nullable(input_schema)?,
+        })
+    }
 }
diff --git a/bustubx/src/expression/literal.rs b/bustubx/src/expression/literal.rs
@@ -1,9 +1,10 @@
-use crate::catalog::DataType;
 use crate::catalog::Schema;
+use crate::catalog::{Column, DataType};
 use crate::common::ScalarValue;
 use crate::error::BustubxResult;
 use crate::expression::ExprTrait;
 use crate::storage::Tuple;
+use crate::BustubxError;
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Literal {
@@ -15,7 +16,18 @@ impl ExprTrait for Literal {
         Ok(self.value.data_type())
     }
 
+    fn nullable(&self, _input_schema: &Schema) -> BustubxResult<bool> {
+        Ok(self.value.is_null())
+    }
+
     fn evaluate(&self, _tuple: &Tuple) -> BustubxResult<ScalarValue> {
         Ok(self.value.clone())
     }
+
+    fn to_column(&self, _input_schema: &Schema) -> BustubxResult<Column> {
+        Err(BustubxError::Plan(format!(
+            "expr {:?} as column not supported",
+            self
+        )))
+    }
 }
diff --git a/bustubx/src/expression/mod.rs b/bustubx/src/expression/mod.rs
@@ -20,8 +20,14 @@ pub trait ExprTrait {
     /// Get the data type of this expression, given the schema of the input
     fn data_type(&self, input_schema: &Schema) -> BustubxResult<DataType>;
 
+    /// Determine whether this expression is nullable, given the schema of the input
+    fn nullable(&self, input_schema: &Schema) -> BustubxResult<bool>;
+
     /// Evaluate an expression against a Tuple
     fn evaluate(&self, tuple: &Tuple) -> BustubxResult<ScalarValue>;
+
+    /// convert to a column with respect to a schema
+    fn to_column(&self, input_schema: &Schema) -> BustubxResult<Column>;
 }
 
 #[derive(Clone, PartialEq, Eq, Debug)]
@@ -50,6 +56,16 @@ impl ExprTrait for Expr {
         }
     }
 
+    fn nullable(&self, input_schema: &Schema) -> BustubxResult<bool> {
+        match self {
+            Expr::Alias(alias) => alias.nullable(input_schema),
+            Expr::Column(column) => column.nullable(input_schema),
+            Expr::Literal(literal) => literal.nullable(input_schema),
+            Expr::BinaryExpr(binary) => binary.nullable(input_schema),
+            Expr::Cast(cast) => cast.nullable(input_schema),
+        }
+    }
+
     fn evaluate(&self, tuple: &Tuple) -> BustubxResult<ScalarValue> {
         match self {
             Expr::Alias(alias) => alias.evaluate(tuple),
@@ -59,23 +75,14 @@ impl ExprTrait for Expr {
             Expr::Cast(cast) => cast.evaluate(tuple),
         }
     }
-}
 
-impl Expr {
-    pub fn to_column(&self, input_schema: &Schema) -> BustubxResult<Column> {
+    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
-                )))
-            }
+            Expr::Alias(alias) => alias.to_column(input_schema),
+            Expr::Column(column) => column.to_column(input_schema),
+            Expr::Literal(literal) => literal.to_column(input_schema),
+            Expr::BinaryExpr(binary) => binary.to_column(input_schema),
+            Expr::Cast(cast) => cast.to_column(input_schema),
         }
     }
 }
diff --git a/bustubx/src/planner/logical_plan_v2/util.rs b/bustubx/src/planner/logical_plan_v2/util.rs
@@ -1,5 +1,5 @@
 use crate::catalog::{ColumnRef, Schema};
-use crate::expression::{ColumnExpr, Expr};
+use crate::expression::{ColumnExpr, Expr, ExprTrait};
 use crate::planner::logical_plan_v2::LogicalPlanV2;
 use crate::planner::table_ref::join::JoinType;
 use crate::{BustubxError, BustubxResult};