ARROW-8117 [Datafusion] [Rust] allow cast SQLTimestamp to Timestamp

rust/datafusion/src/execution/physical_plan/expressions.rs

@@ -27,8 +27,8 @@ use crate::execution::physical_plan::{Accumulator, AggregateExpr, PhysicalExpr};
 use crate::logicalplan::{Operator, ScalarValue};
 use arrow::array::{
     ArrayRef, BooleanArray, Float32Array, Float64Array, Int16Array, Int32Array,
-    Int64Array, Int8Array, StringArray, UInt16Array, UInt32Array, UInt64Array,
-    UInt8Array,
+    Int64Array, Int8Array, StringArray, TimestampNanosecondArray, UInt16Array,
+    UInt32Array, UInt64Array, UInt8Array,
 };
 use arrow::array::{
     Float32Builder, Float64Builder, Int16Builder, Int32Builder, Int64Builder,
@@ -43,7 +43,7 @@ use arrow::compute::kernels::comparison::{eq, gt, gt_eq, lt, lt_eq, neq};
 use arrow::compute::kernels::comparison::{
     eq_utf8, gt_eq_utf8, gt_utf8, like_utf8, lt_eq_utf8, lt_utf8, neq_utf8, nlike_utf8,
 };
-use arrow::datatypes::{DataType, Schema};
+use arrow::datatypes::{DataType, Schema, TimeUnit};
 use arrow::record_batch::RecordBatch;
 
 /// Represents the column at a given index in a RecordBatch
@@ -936,6 +936,9 @@ macro_rules! binary_array_op {
             DataType::Float32 => compute_op!($LEFT, $RIGHT, $OP, Float32Array),
             DataType::Float64 => compute_op!($LEFT, $RIGHT, $OP, Float64Array),
             DataType::Utf8 => compute_utf8_op!($LEFT, $RIGHT, $OP, StringArray),
+            DataType::Timestamp(TimeUnit::Nanosecond, None) => {
+                compute_op!($LEFT, $RIGHT, $OP, TimestampNanosecondArray)
+            }
             other => Err(ExecutionError::General(format!(
                 "Unsupported data type {:?}",
                 other
@@ -1122,6 +1125,10 @@ impl CastExpr {
             Ok(Self { expr, cast_type })
         } else if expr_type == DataType::Binary && cast_type == DataType::Utf8 {
             Ok(Self { expr, cast_type })
+        } else if is_numeric(&expr_type)
+            && cast_type == DataType::Timestamp(TimeUnit::Nanosecond, None)
+        {
+            Ok(Self { expr, cast_type })
         } else {
             Err(ExecutionError::General(format!(
                 "Invalid CAST from {:?} to {:?}",
@@ -1230,7 +1237,7 @@ mod tests {
     use super::*;
     use crate::error::Result;
     use crate::execution::physical_plan::common::get_scalar_value;
-    use arrow::array::{PrimitiveArray, StringArray};
+    use arrow::array::{PrimitiveArray, StringArray, Time64NanosecondArray};
     use arrow::datatypes::*;
 
     #[test]
@@ -1357,6 +1364,29 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn cast_i64_to_timestamp_nanoseconds() -> Result<()> {
+        let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]);
+        let a = Int64Array::from(vec![1, 2, 3, 4, 5]);
+        let batch = RecordBatch::try_new(Arc::new(schema.clone()), vec![Arc::new(a)])?;
+
+        let cast = CastExpr::try_new(
+            col(0),
+            &schema,
+            DataType::Timestamp(TimeUnit::Nanosecond, None),
+        )?;
+        let result = cast.evaluate(&batch)?;
+        assert_eq!(result.len(), 5);
+        let expected_result = Time64NanosecondArray::from(vec![1, 2, 3, 4]);
+        let result = result
+            .as_any()
+            .downcast_ref::<TimestampNanosecondArray>()
+            .expect("failed to downcast to TimestampNanosecondArray");
+        assert_eq!(result.value(0), expected_result.value(0));
+
+        Ok(())
+    }
+
     #[test]
     fn invalid_cast() -> Result<()> {
         let schema = Schema::new(vec![Field::new("a", DataType::Utf8, false)]);

rust/datafusion/src/sql/planner.rs

@@ -470,6 +470,7 @@ pub fn convert_data_type(sql: &SQLType) -> Result<DataType> {
         SQLType::Float(_) | SQLType::Real => Ok(DataType::Float64),
         SQLType::Double => Ok(DataType::Float64),
         SQLType::Char(_) | SQLType::Varchar(_) => Ok(DataType::Utf8),
+        SQLType::Timestamp => Ok(DataType::Timestamp(TimeUnit::Nanosecond, None)),
         other => Err(ExecutionError::NotImplemented(format!(
             "Unsupported SQL type {:?}",
             other
@@ -532,6 +533,17 @@ mod tests {
         quick_test(sql, expected);
     }
 
+    #[test]
+    fn test_timestamp_selection() {
+        let sql = "SELECT state FROM person WHERE birth_date < CAST (158412331400600000 as timestamp)";
+
+        let expected = "Projection: #4\
+            \n  Selection: #6 Lt CAST(Int64(158412331400600000) AS Timestamp(Nanosecond, None))\
+            \n    TableScan: person projection=None";
+
+        quick_test(sql, expected);
+    }
+
     #[test]
     fn select_all_boolean_operators() {
         let sql = "SELECT age, first_name, last_name \
@@ -658,6 +670,11 @@ mod tests {
                     Field::new("age", DataType::Int32, false),
                     Field::new("state", DataType::Utf8, false),
                     Field::new("salary", DataType::Float64, false),
+                    Field::new(
+                        "birth_date",
+                        DataType::Timestamp(TimeUnit::Nanosecond, None),
+                        false,
+                    ),
                 ]))),
                 _ => None,
             }