Return correct scalar types for date_trunc (#6638)

* Return correct scalar for date_trunc * Fix test
apache · Jun 12, 2023 · 9846da6 · 9846da6
1 parent 6df158e
commit 9846da6
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 14 deletions.
diff --git a/datafusion/core/tests/sql/timestamp.rs b/datafusion/core/tests/sql/timestamp.rs
@@ -712,7 +712,7 @@ async fn test_arrow_typeof() -> Result<()> {
         "+--------------------------------------------------------------------------+",
         "| arrow_typeof(date_trunc(Utf8(\"minute\"),to_timestamp_seconds(Int64(61)))) |",
         "+--------------------------------------------------------------------------+",
-        "| Timestamp(Second, None)                                                  |",
+        "| Timestamp(Nanosecond, None)                                              |",
         "+--------------------------------------------------------------------------+",
     ];
     assert_batches_eq!(expected, &actual);
@@ -723,7 +723,7 @@ async fn test_arrow_typeof() -> Result<()> {
         "+-------------------------------------------------------------------------+",
         "| arrow_typeof(date_trunc(Utf8(\"second\"),to_timestamp_millis(Int64(61)))) |",
         "+-------------------------------------------------------------------------+",
-        "| Timestamp(Millisecond, None)                                            |",
+        "| Timestamp(Nanosecond, None)                                             |",
         "+-------------------------------------------------------------------------+",
     ];
     assert_batches_eq!(expected, &actual);
@@ -734,7 +734,7 @@ async fn test_arrow_typeof() -> Result<()> {
         "+------------------------------------------------------------------------------+",
         "| arrow_typeof(date_trunc(Utf8(\"millisecond\"),to_timestamp_micros(Int64(61)))) |",
         "+------------------------------------------------------------------------------+",
-        "| Timestamp(Microsecond, None)                                                 |",
+        "| Timestamp(Nanosecond, None)                                                  |",
         "+------------------------------------------------------------------------------+",
     ];
     assert_batches_eq!(expected, &actual);

diff --git a/datafusion/physical-expr/src/datetime_expressions.rs b/datafusion/physical-expr/src/datetime_expressions.rs
@@ -285,31 +285,31 @@ pub fn date_trunc(args: &[ColumnarValue]) -> Result<ColumnarValue> {
             let nano = (f)(*v)?;
             match granularity.as_str() {
                 "minute" => {
-                    // cast to second
-                    let second = ScalarValue::TimestampSecond(
-                        Some(nano.unwrap() / 1_000_000_000),
+                    // trunc to minute
+                    let second = ScalarValue::TimestampNanosecond(
+                        Some(nano.unwrap() / 1_000_000_000 * 1_000_000_000),
                         tz_opt.clone(),
                     );
                     ColumnarValue::Scalar(second)
                 }
                 "second" => {
-                    // cast to millisecond
-                    let mill = ScalarValue::TimestampMillisecond(
-                        Some(nano.unwrap() / 1_000_000),
+                    // trunc to second
+                    let mill = ScalarValue::TimestampNanosecond(
+                        Some(nano.unwrap() / 1_000_000 * 1_000_000),
                         tz_opt.clone(),
                     );
                     ColumnarValue::Scalar(mill)
                 }
                 "millisecond" => {
-                    // cast to microsecond
-                    let micro = ScalarValue::TimestampMicrosecond(
-                        Some(nano.unwrap() / 1_000),
+                    // trunc to microsecond
+                    let micro = ScalarValue::TimestampNanosecond(
+                        Some(nano.unwrap() / 1_000 * 1_000),
                         tz_opt.clone(),
                     );
                     ColumnarValue::Scalar(micro)
                 }
                 _ => {
-                    // cast to nanosecond
+                    // trunc to nanosecond
                     let nano = ScalarValue::TimestampNanosecond(
                         Some(nano.unwrap()),
                         tz_opt.clone(),
@@ -329,7 +329,7 @@ pub fn date_trunc(args: &[ColumnarValue]) -> Result<ColumnarValue> {
         }
         _ => {
             return Err(DataFusionError::Execution(
-                "array of `date_trunc` must be non-null scalar Utf8".to_string(),
+                "second argument of `date_trunc` must be nanosecond timestamp scalar or array".to_string(),
             ));
         }
     })