feat: support tpch q5 and q6 (#69)

- Support gt and lt for bin op type - Support `LikeExpr` (it's for Q2, but it turns out Q2 requires more work) - Support for cast types
cmu-db · Feb 17, 2024 · 214848c · 214848c
1 parent d34c22c
commit 214848c
Show file tree

Hide file tree

Showing 6 changed files with 365 additions and 8 deletions.
diff --git a/optd-datafusion-bridge/src/from_optd.rs b/optd-datafusion-bridge/src/from_optd.rs
@@ -23,7 +23,7 @@ use datafusion::{
 use optd_datafusion_repr::{
     plan_nodes::{
         BetweenExpr, BinOpExpr, BinOpType, CastExpr, ColumnRefExpr, ConstantExpr, ConstantType,
-        Expr, ExprList, FuncExpr, FuncType, JoinType, LogOpExpr, LogOpType, OptRelNode,
+        Expr, ExprList, FuncExpr, FuncType, JoinType, LikeExpr, LogOpExpr, LogOpType, OptRelNode,
         OptRelNodeRef, OptRelNodeTyp, PhysicalAgg, PhysicalEmptyRelation, PhysicalFilter,
         PhysicalHashJoin, PhysicalLimit, PhysicalNestedLoopJoin, PhysicalProjection, PhysicalScan,
         PhysicalSort, PlanNode, SortOrderExpr, SortOrderType,
@@ -49,6 +49,7 @@ fn from_optd_schema(optd_schema: OptdSchema) -> Schema {
         ConstantType::Int16 => DataType::Int16,
         ConstantType::Int32 => DataType::Int32,
         ConstantType::Int64 => DataType::Int64,
+        ConstantType::Float64 => DataType::Float64,
         ConstantType::Date => DataType::Date32,
         ConstantType::Decimal => DataType::Float64,
         ConstantType::Utf8String => DataType::Utf8,
@@ -145,6 +146,7 @@ impl OptdPlanContext<'_> {
                     ConstantType::Int16 => ScalarValue::Int16(Some(value.as_i16())),
                     ConstantType::Int32 => ScalarValue::Int32(Some(value.as_i32())),
                     ConstantType::Int64 => ScalarValue::Int64(Some(value.as_i64())),
+                    ConstantType::Float64 => ScalarValue::Float64(Some(value.as_f64())),
                     ConstantType::Decimal => {
                         ScalarValue::Decimal128(Some(value.as_f64() as i128), 20, 0)
                         // TODO(chi): no hard code decimal
@@ -214,7 +216,9 @@ impl OptdPlanContext<'_> {
                     BinOpType::Eq => Operator::Eq,
                     BinOpType::Neq => Operator::NotEq,
                     BinOpType::Leq => Operator::LtEq,
+                    BinOpType::Lt => Operator::Lt,
                     BinOpType::Geq => Operator::GtEq,
+                    BinOpType::Gt => Operator::Gt,
                     BinOpType::And => Operator::And,
                     BinOpType::Add => Operator::Plus,
                     BinOpType::Sub => Operator::Minus,
@@ -256,6 +260,19 @@ impl OptdPlanContext<'_> {
                     datafusion::physical_plan::expressions::CastExpr::new(child, data_type, None),
                 ))
             }
+            OptRelNodeTyp::Like => {
+                let expr = LikeExpr::from_rel_node(expr.into_rel_node()).unwrap();
+                let child = Self::conv_from_optd_expr(expr.child(), context)?;
+                let pattern = Self::conv_from_optd_expr(expr.pattern(), context)?;
+                Ok(Arc::new(
+                    datafusion::physical_plan::expressions::LikeExpr::new(
+                        expr.negated(),
+                        expr.case_insensitive(),
+                        child,
+                        pattern,
+                    ),
+                ))
+            }
             _ => unimplemented!("{}", expr.into_rel_node()),
         }
     }

diff --git a/optd-datafusion-bridge/src/into_optd.rs b/optd-datafusion-bridge/src/into_optd.rs
@@ -9,9 +9,10 @@ use optd_core::rel_node::RelNode;
 use optd_datafusion_repr::{
     plan_nodes::{
         BetweenExpr, BinOpExpr, BinOpType, CastExpr, ColumnRefExpr, ConstantExpr, Expr, ExprList,
-        FuncExpr, FuncType, JoinType, LogOpExpr, LogOpType, LogicalAgg, LogicalEmptyRelation,
-        LogicalFilter, LogicalJoin, LogicalLimit, LogicalProjection, LogicalScan, LogicalSort,
-        OptRelNode, OptRelNodeRef, OptRelNodeTyp, PlanNode, SortOrderExpr, SortOrderType,
+        FuncExpr, FuncType, JoinType, LikeExpr, LogOpExpr, LogOpType, LogicalAgg,
+        LogicalEmptyRelation, LogicalFilter, LogicalJoin, LogicalLimit, LogicalProjection,
+        LogicalScan, LogicalSort, OptRelNode, OptRelNodeRef, OptRelNodeTyp, PlanNode,
+        SortOrderExpr, SortOrderType,
     },
     Value,
 };
@@ -54,7 +55,9 @@ impl OptdPlanContext<'_> {
                     Operator::Eq => BinOpType::Eq,
                     Operator::NotEq => BinOpType::Neq,
                     Operator::LtEq => BinOpType::Leq,
+                    Operator::Lt => BinOpType::Lt,
                     Operator::GtEq => BinOpType::Geq,
+                    Operator::Gt => BinOpType::Gt,
                     Operator::And => BinOpType::And,
                     Operator::Plus => BinOpType::Add,
                     Operator::Minus => BinOpType::Sub,
@@ -101,6 +104,10 @@ impl OptdPlanContext<'_> {
                     let x = x.as_ref().unwrap();
                     Ok(ConstantExpr::int64(*x).into_expr())
                 }
+                ScalarValue::Float64(x) => {
+                    let x = x.as_ref().unwrap();
+                    Ok(ConstantExpr::float64(*x).into_expr())
+                }
                 ScalarValue::Utf8(x) => {
                     let x = x.as_ref().unwrap();
                     Ok(ConstantExpr::string(x).into_expr())
@@ -166,12 +173,24 @@ impl OptdPlanContext<'_> {
                 let data_type = x.data_type.clone();
                 let val = match data_type {
                     arrow_schema::DataType::Int8 => Value::Int8(0),
+                    arrow_schema::DataType::Int16 => Value::Int16(0),
+                    arrow_schema::DataType::Int32 => Value::Int32(0),
+                    arrow_schema::DataType::Int64 => Value::Int64(0),
+                    arrow_schema::DataType::UInt8 => Value::UInt8(0),
+                    arrow_schema::DataType::UInt16 => Value::UInt16(0),
+                    arrow_schema::DataType::UInt32 => Value::UInt32(0),
+                    arrow_schema::DataType::UInt64 => Value::UInt64(0),
                     arrow_schema::DataType::Date32 => Value::Date32(0),
                     arrow_schema::DataType::Decimal128(_, _) => Value::Decimal128(0),
                     other => unimplemented!("unimplemented datatype {:?}", other),
                 };
                 Ok(CastExpr::new(expr, val).into_expr())
             }
+            Expr::Like(x) => {
+                let expr = self.conv_into_optd_expr(x.expr.as_ref(), context)?;
+                let pattern = self.conv_into_optd_expr(x.pattern.as_ref(), context)?;
+                Ok(LikeExpr::new(x.negated, x.case_insensitive, expr, pattern).into_expr())
+            }
             _ => bail!("Unsupported expression: {:?}", expr),
         }
     }

diff --git a/optd-datafusion-repr/src/plan_nodes.rs b/optd-datafusion-repr/src/plan_nodes.rs
@@ -24,8 +24,8 @@ pub use apply::{ApplyType, LogicalApply};
 pub use empty_relation::{LogicalEmptyRelation, PhysicalEmptyRelation};
 pub use expr::{
     BetweenExpr, BinOpExpr, BinOpType, CastExpr, ColumnRefExpr, ConstantExpr, ConstantType,
-    ExprList, FuncExpr, FuncType, LogOpExpr, LogOpType, SortOrderExpr, SortOrderType, UnOpExpr,
-    UnOpType,
+    ExprList, FuncExpr, FuncType, LikeExpr, LogOpExpr, LogOpType, SortOrderExpr, SortOrderType,
+    UnOpExpr, UnOpType,
 };
 pub use filter::{LogicalFilter, PhysicalFilter};
 pub use join::{JoinType, LogicalJoin, PhysicalHashJoin, PhysicalNestedLoopJoin};
@@ -76,6 +76,7 @@ pub enum OptRelNodeTyp {
     SortOrder(SortOrderType),
     Between,
     Cast,
+    Like,
 }
 
 impl OptRelNodeTyp {
@@ -116,6 +117,7 @@ impl OptRelNodeTyp {
                 | Self::LogOp(_)
                 | Self::Between
                 | Self::Cast
+                | Self::Like
         )
     }
 }
@@ -382,6 +384,9 @@ pub fn explain(rel_node: OptRelNodeRef) -> Pretty<'static> {
         OptRelNodeTyp::Cast => CastExpr::from_rel_node(rel_node)
             .unwrap()
             .dispatch_explain(),
+        OptRelNodeTyp::Like => LikeExpr::from_rel_node(rel_node)
+            .unwrap()
+            .dispatch_explain(),
     }
 }
 

diff --git a/optd-datafusion-repr/src/plan_nodes/expr.rs b/optd-datafusion-repr/src/plan_nodes/expr.rs
@@ -1,4 +1,4 @@
-use std::fmt::Display;
+use std::{fmt::Display, sync::Arc};
 
 use itertools::Itertools;
 use pretty_xmlish::Pretty;
@@ -76,6 +76,7 @@ pub enum ConstantType {
     Int16,
     Int32,
     Int64,
+    Float64,
     Date,
     Decimal,
     Any,
@@ -97,7 +98,7 @@ impl ConstantExpr {
             Value::Int16(_) => ConstantType::Int16,
             Value::Int32(_) => ConstantType::Int32,
             Value::Int64(_) => ConstantType::Int64,
-            Value::Float(_) => ConstantType::Decimal,
+            Value::Float(_) => ConstantType::Float64,
             _ => unimplemented!(),
         };
         Self::new_with_type(value, typ)
@@ -157,6 +158,10 @@ impl ConstantExpr {
         Self::new_with_type(Value::Int64(value), ConstantType::Int64)
     }
 
+    pub fn float64(value: f64) -> Self {
+        Self::new_with_type(Value::Float(value.into()), ConstantType::Float64)
+    }
+
     pub fn date(value: i64) -> Self {
         Self::new_with_type(Value::Int64(value), ConstantType::Date)
     }
@@ -681,3 +686,68 @@ impl OptRelNode for CastExpr {
         )
     }
 }
+
+#[derive(Clone, Debug)]
+pub struct LikeExpr(pub Expr);
+
+impl LikeExpr {
+    pub fn new(negated: bool, case_insensitive: bool, expr: Expr, pattern: Expr) -> Self {
+        // TODO: support multiple values in data.
+        let negated = if negated { 1 } else { 0 };
+        let case_insensitive = if case_insensitive { 1 } else { 0 };
+        LikeExpr(Expr(
+            RelNode {
+                typ: OptRelNodeTyp::Like,
+                children: vec![expr.into_rel_node(), pattern.into_rel_node()],
+                data: Some(Value::Serialized(Arc::new([negated, case_insensitive]))),
+            }
+            .into(),
+        ))
+    }
+
+    pub fn child(&self) -> Expr {
+        Expr(self.0.child(0))
+    }
+
+    pub fn pattern(&self) -> Expr {
+        Expr(self.0.child(1))
+    }
+
+    pub fn negated(&self) -> bool {
+        match self.0 .0.data.as_ref().unwrap() {
+            Value::Serialized(data) => data[0] != 0,
+            _ => panic!("not a serialized value"),
+        }
+    }
+
+    pub fn case_insensitive(&self) -> bool {
+        match self.0 .0.data.as_ref().unwrap() {
+            Value::Serialized(data) => data[1] != 0,
+            _ => panic!("not a serialized value"),
+        }
+    }
+}
+
+impl OptRelNode for LikeExpr {
+    fn into_rel_node(self) -> OptRelNodeRef {
+        self.0.into_rel_node()
+    }
+
+    fn from_rel_node(rel_node: OptRelNodeRef) -> Option<Self> {
+        if !matches!(rel_node.typ, OptRelNodeTyp::Like) {
+            return None;
+        }
+        Expr::from_rel_node(rel_node).map(Self)
+    }
+
+    fn dispatch_explain(&self) -> Pretty<'static> {
+        Pretty::simple_record(
+            "Like",
+            vec![
+                ("expr", self.child().explain()),
+                ("pattern", self.pattern().explain()),
+            ],
+            vec![],
+        )
+    }
+}