Support more functions and scalar types (#7)

Author: Dandandan

src/consumer.rs

@@ -16,6 +16,37 @@ use substrait::protobuf::{
     Expression, Rel,
 };
 
+pub fn reference_to_op(reference: u32) -> Result<Operator> {
+    match reference {
+        1 => Ok(Operator::Eq),
+        2 => Ok(Operator::NotEq),
+        3 => Ok(Operator::Lt),
+        4 => Ok(Operator::LtEq),
+        5 => Ok(Operator::Gt),
+        6 => Ok(Operator::GtEq),
+        7 => Ok(Operator::Plus),
+        8 => Ok(Operator::Minus),
+        9 => Ok(Operator::Multiply),
+        10 => Ok(Operator::Divide),
+        11 => Ok(Operator::Modulo),
+        12 => Ok(Operator::And),
+        13 => Ok(Operator::Or),
+        14 => Ok(Operator::Like),
+        15 => Ok(Operator::NotLike),
+        16 => Ok(Operator::IsDistinctFrom),
+        17 => Ok(Operator::IsNotDistinctFrom),
+        18 => Ok(Operator::RegexMatch),
+        19 => Ok(Operator::RegexIMatch),
+        20 => Ok(Operator::RegexNotMatch),
+        21 => Ok(Operator::RegexNotIMatch),
+        22 => Ok(Operator::BitwiseAnd),
+        _ => Err(DataFusionError::NotImplemented(format!(
+            "Unsupported function_reference: {:?}",
+            reference
+        ))),
+    }
+}
+
 /// Convert Substrait Rel to DataFusion DataFrame
 #[async_recursion]
 pub async fn from_substrait_rel(
@@ -95,19 +126,7 @@ pub async fn from_substrait_rex(e: &Expression, input: &dyn DataFrame) -> Result
         },
         Some(RexType::ScalarFunction(f)) => {
             assert!(f.args.len() == 2);
-            let op = match f.function_reference {
-                1 => Operator::Eq,
-                2 => Operator::Lt,
-                3 => Operator::LtEq,
-                4 => Operator::Gt,
-                5 => Operator::GtEq,
-                _ => {
-                    return Err(DataFusionError::NotImplemented(format!(
-                        "Unsupported function_reference: {:?}",
-                        f.function_reference
-                    )))
-                }
-            };
+            let op = reference_to_op(f.function_reference)?;
             Ok(Arc::new(Expr::BinaryExpr {
                 left: Box::new(
                     from_substrait_rex(&f.args[0], input)
@@ -124,19 +143,37 @@ pub async fn from_substrait_rex(e: &Expression, input: &dyn DataFrame) -> Result
                 ),
             }))
         }
-        Some(RexType::Literal(lit)) => match lit.literal_type {
+        Some(RexType::Literal(lit)) => match &lit.literal_type {
             Some(LiteralType::I8(n)) => {
-                Ok(Arc::new(Expr::Literal(ScalarValue::Int8(Some(n as i8)))))
+                Ok(Arc::new(Expr::Literal(ScalarValue::Int8(Some(*n as i8)))))
             }
             Some(LiteralType::I16(n)) => {
-                Ok(Arc::new(Expr::Literal(ScalarValue::Int16(Some(n as i16)))))
+                Ok(Arc::new(Expr::Literal(ScalarValue::Int16(Some(*n as i16)))))
             }
             Some(LiteralType::I32(n)) => {
-                Ok(Arc::new(Expr::Literal(ScalarValue::Int32(Some(n as i32)))))
+                Ok(Arc::new(Expr::Literal(ScalarValue::Int32(Some(*n as i32)))))
             }
             Some(LiteralType::I64(n)) => {
-                Ok(Arc::new(Expr::Literal(ScalarValue::Int64(Some(n as i64)))))
+                Ok(Arc::new(Expr::Literal(ScalarValue::Int64(Some(*n as i64)))))
+            }
+            Some(LiteralType::Boolean(b)) => {
+                Ok(Arc::new(Expr::Literal(ScalarValue::Boolean(Some(*b)))))
+            }
+            Some(LiteralType::Date(d)) => {
+                Ok(Arc::new(Expr::Literal(ScalarValue::Date32(Some(*d)))))
+            }
+            Some(LiteralType::Fp32(f)) => {
+                Ok(Arc::new(Expr::Literal(ScalarValue::Float32(Some(*f)))))
+            }
+            Some(LiteralType::Fp64(f)) => {
+                Ok(Arc::new(Expr::Literal(ScalarValue::Float64(Some(*f)))))
             }
+            Some(LiteralType::String(s)) => Ok(Arc::new(Expr::Literal(ScalarValue::LargeUtf8(
+                Some(s.clone()),
+            )))),
+            Some(LiteralType::Binary(b)) => Ok(Arc::new(Expr::Literal(ScalarValue::Binary(Some(
+                b.clone(),
+            ))))),
             _ => {
                 return Err(DataFusionError::NotImplemented(format!(
                     "Unsupported literal_type: {:?}",

src/lib.rs

@@ -23,6 +23,16 @@ mod tests {
         roundtrip("SELECT * FROM data WHERE a > 1").await
     }
 
+    #[tokio::test]
+    async fn select_with_filter_date() -> Result<()> {
+        roundtrip("SELECT * FROM data WHERE c > CAST('2020-01-01' AS DATE)").await
+    }
+
+    #[tokio::test]
+    async fn select_with_filter_bool_expr() -> Result<()> {
+        roundtrip("SELECT * FROM data WHERE d AND a > 1").await
+    }
+
     async fn roundtrip(sql: &str) -> Result<()> {
         let mut ctx = ExecutionContext::new();
         ctx.register_csv("data", "testdata/data.csv", CsvReadOptions::new())

src/producer.rs

@@ -94,6 +94,33 @@ pub fn to_substrait_rel(plan: &LogicalPlan) -> Result<Box<Rel>> {
     }
 }
 
+pub fn operator_to_reference(op: Operator) -> u32 {
+    match op {
+        Operator::Eq => 1,
+        Operator::NotEq => 2,
+        Operator::Lt => 3,
+        Operator::LtEq => 4,
+        Operator::Gt => 5,
+        Operator::GtEq => 6,
+        Operator::Plus => 7,
+        Operator::Minus => 8,
+        Operator::Multiply => 9,
+        Operator::Divide => 10,
+        Operator::Modulo => 11,
+        Operator::And => 12,
+        Operator::Or => 13,
+        Operator::Like => 14,
+        Operator::NotLike => 15,
+        Operator::IsDistinctFrom => 16,
+        Operator::IsNotDistinctFrom => 17,
+        Operator::RegexMatch => 18,
+        Operator::RegexIMatch => 19,
+        Operator::RegexNotMatch => 20,
+        Operator::RegexNotIMatch => 21,
+        Operator::BitwiseAnd => 22,
+    }
+}
+
 /// Convert DataFusion Expr to Substrait Rex
 pub fn to_substrait_rex(expr: &Expr, schema: &DFSchemaRef) -> Result<Expression> {
     match expr {
@@ -114,19 +141,7 @@ pub fn to_substrait_rex(expr: &Expr, schema: &DFSchemaRef) -> Result<Expression>
         Expr::BinaryExpr { left, op, right } => {
             let l = to_substrait_rex(left, schema)?;
             let r = to_substrait_rex(right, schema)?;
-            let function_reference: u32 = match op {
-                Operator::Eq => 1,
-                Operator::Lt => 2,
-                Operator::LtEq => 3,
-                Operator::Gt => 4,
-                Operator::GtEq => 5,
-                _ => {
-                    return Err(DataFusionError::NotImplemented(format!(
-                        "Unsupported operator: {:?}",
-                        op
-                    )))
-                }
-            };
+            let function_reference: u32 = operator_to_reference(*op);
             Ok(Expression {
                 rex_type: Some(RexType::ScalarFunction(ScalarFunction {
                     function_reference,
@@ -141,6 +156,14 @@ pub fn to_substrait_rex(expr: &Expr, schema: &DFSchemaRef) -> Result<Expression>
                 ScalarValue::Int16(Some(n)) => Some(LiteralType::I16(*n as i32)),
                 ScalarValue::Int32(Some(n)) => Some(LiteralType::I32(*n)),
                 ScalarValue::Int64(Some(n)) => Some(LiteralType::I64(*n)),
+                ScalarValue::Boolean(Some(b)) => Some(LiteralType::Boolean(*b)),
+                ScalarValue::Float32(Some(f)) => Some(LiteralType::Fp32(*f)),
+                ScalarValue::Float64(Some(f)) => Some(LiteralType::Fp64(*f)),
+                ScalarValue::Utf8(Some(s)) => Some(LiteralType::String(s.clone())),
+                ScalarValue::LargeUtf8(Some(s)) => Some(LiteralType::String(s.clone())),
+                ScalarValue::Binary(Some(b)) => Some(LiteralType::Binary(b.clone())),
+                ScalarValue::LargeBinary(Some(b)) => Some(LiteralType::Binary(b.clone())),
+                ScalarValue::Date32(Some(d)) => Some(LiteralType::Date(*d)),
                 _ => {
                     return Err(DataFusionError::NotImplemented(format!(
                         "Unsupported literal: {:?}",

testdata/data.csv

@@ -1,3 +1,3 @@
-a,b
-1,2
-3,4
+a,b,c,d
+1,2,2020-01-01,false
+3,4,2020-01-01,true