add timestamptz (#3660)

apache · Sep 30, 2022 · bfc5ff0 · bfc5ff0
1 parent 65a5c6b
commit bfc5ff0
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 1 deletion.
diff --git a/datafusion/core/tests/sql/timestamp.rs b/datafusion/core/tests/sql/timestamp.rs
@@ -1530,3 +1530,27 @@ async fn cast_timestamp_before_1970() -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn cast_timestamp_to_timestamptz() -> Result<()> {
+    let ctx = SessionContext::new();
+    let table_a = make_timestamp_table::<TimestampNanosecondType>()?;
+
+    ctx.register_table("table_a", table_a)?;
+
+    let sql = "SELECT ts::timestamptz, arrow_typeof(ts::timestamptz) FROM table_a;";
+    let actual = execute_to_batches(&ctx, sql).await;
+
+    let expected = vec![
+        "+----------------------------+------------------------------------+",
+        "| table_a.ts                 | arrowtypeof(table_a.ts)            |",
+        "+----------------------------+------------------------------------+",
+        "| 2020-09-08 13:42:29.190855 | Timestamp(Nanosecond, Some(\"UTC\")) |",
+        "| 2020-09-08 12:42:29.190855 | Timestamp(Nanosecond, Some(\"UTC\")) |",
+        "| 2020-09-08 11:42:29.190855 | Timestamp(Nanosecond, Some(\"UTC\")) |",
+        "+----------------------------+------------------------------------+",
+    ];
+    assert_batches_eq!(expected, &actual);
+
+    Ok(())
+}
diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs
@@ -2703,6 +2703,10 @@ pub fn convert_simple_data_type(sql_type: &SQLDataType) -> Result<DataType> {
         | SQLDataType::Text
         | SQLDataType::String => Ok(DataType::Utf8),
         SQLDataType::Timestamp => Ok(DataType::Timestamp(TimeUnit::Nanosecond, None)),
+        SQLDataType::TimestampTz => Ok(DataType::Timestamp(
+            TimeUnit::Nanosecond,
+            Some("UTC".into()),
+        )),
         SQLDataType::Date => Ok(DataType::Date32),
         SQLDataType::Time => Ok(DataType::Time64(TimeUnit::Nanosecond)),
         SQLDataType::Decimal(precision, scale) => make_decimal_type(*precision, *scale),
@@ -2716,7 +2720,6 @@ pub fn convert_simple_data_type(sql_type: &SQLDataType) -> Result<DataType> {
         | SQLDataType::Varbinary(_)
         | SQLDataType::Blob(_)
         | SQLDataType::Datetime
-        | SQLDataType::TimestampTz
         | SQLDataType::Interval
         | SQLDataType::Regclass
         | SQLDataType::Custom(_)