Support for extract(epoch from date) for Date32 and Date64 (#8695)

apache · Jan 2, 2024 · 9a6cc88 · 9a6cc88
1 parent 67baf10
commit 9a6cc88
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 20 deletions.
diff --git a/datafusion/core/tests/sql/expr.rs b/datafusion/core/tests/sql/expr.rs
@@ -741,6 +741,7 @@ async fn test_extract_date_part() -> Result<()> {
 
 #[tokio::test]
 async fn test_extract_epoch() -> Result<()> {
+    // timestamp
     test_expression!(
         "extract(epoch from '1870-01-01T07:29:10.256'::timestamp)",
         "-3155646649.744"
@@ -754,6 +755,39 @@ async fn test_extract_epoch() -> Result<()> {
         "946684800.0"
     );
     test_expression!("extract(epoch from NULL::timestamp)", "NULL");
+    // date
+    test_expression!(
+        "extract(epoch from arrow_cast('1970-01-01', 'Date32'))",
+        "0.0"
+    );
+    test_expression!(
+        "extract(epoch from arrow_cast('1970-01-02', 'Date32'))",
+        "86400.0"
+    );
+    test_expression!(
+        "extract(epoch from arrow_cast('1970-01-11', 'Date32'))",
+        "864000.0"
+    );
+    test_expression!(
+        "extract(epoch from arrow_cast('1969-12-31', 'Date32'))",
+        "-86400.0"
+    );
+    test_expression!(
+        "extract(epoch from arrow_cast('1970-01-01', 'Date64'))",
+        "0.0"
+    );
+    test_expression!(
+        "extract(epoch from arrow_cast('1970-01-02', 'Date64'))",
+        "86400.0"
+    );
+    test_expression!(
+        "extract(epoch from arrow_cast('1970-01-11', 'Date64'))",
+        "864000.0"
+    );
+    test_expression!(
+        "extract(epoch from arrow_cast('1969-12-31', 'Date64'))",
+        "-86400.0"
+    );
     Ok(())
 }
 

diff --git a/datafusion/physical-expr/src/datetime_expressions.rs b/datafusion/physical-expr/src/datetime_expressions.rs
@@ -19,7 +19,6 @@
 
 use crate::datetime_expressions;
 use crate::expressions::cast_column;
-use arrow::array::Float64Builder;
 use arrow::compute::cast;
 use arrow::{
     array::{Array, ArrayRef, Float64Array, OffsetSizeTrait, PrimitiveArray},
@@ -887,28 +886,33 @@ where
     T: ArrowTemporalType + ArrowNumericType,
     i64: From<T::Native>,
 {
-    let mut b = Float64Builder::with_capacity(array.len());
-    match array.data_type() {
+    let b = match array.data_type() {
         DataType::Timestamp(tu, _) => {
-            for i in 0..array.len() {
-                if array.is_null(i) {
-                    b.append_null();
-                } else {
-                    let scale = match tu {
-                        TimeUnit::Second => 1,
-                        TimeUnit::Millisecond => 1_000,
-                        TimeUnit::Microsecond => 1_000_000,
-                        TimeUnit::Nanosecond => 1_000_000_000,
-                    };
-
-                    let n: i64 = array.value(i).into();
-                    b.append_value(n as f64 / scale as f64);
-                }
-            }
+            let scale = match tu {
+                TimeUnit::Second => 1,
+                TimeUnit::Millisecond => 1_000,
+                TimeUnit::Microsecond => 1_000_000,
+                TimeUnit::Nanosecond => 1_000_000_000,
+            } as f64;
+            array.unary(|n| {
+                let n: i64 = n.into();
+                n as f64 / scale
+            })
         }
+        DataType::Date32 => {
+            let seconds_in_a_day = 86400_f64;
+            array.unary(|n| {
+                let n: i64 = n.into();
+                n as f64 * seconds_in_a_day
+            })
+        }
+        DataType::Date64 => array.unary(|n| {
+            let n: i64 = n.into();
+            n as f64 / 1_000_f64
+        }),
         _ => return internal_err!("Can not convert {:?} to epoch", array.data_type()),
-    }
-    Ok(b.finish())
+    };
+    Ok(b)
 }
 
 /// to_timestammp() SQL function implementation