[arrow-datafusion] branch main updated: feat: extract (epoch from col) (#5555)

[alamb]

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/main by this push:
     new 146a94921 feat: extract (epoch from col) (#5555)
146a94921 is described below

commit 146a949218ec970784974137277cde3b4e547d0a
Author: Alex Huang <huangweijun1...@gmail.com>
AuthorDate: Mon Mar 13 18:00:39 2023 +0100

    feat: extract (epoch from col) (#5555)
    
    * feat: extract (epoch from col)
    
    * add more tests and refine error info
---
 datafusion/core/tests/sql/expr.rs                  | 18 +++++++++++
 .../physical-expr/src/datetime_expressions.rs      | 36 ++++++++++++++++++++++
 2 files changed, 54 insertions(+)

diff --git a/datafusion/core/tests/sql/expr.rs 
b/datafusion/core/tests/sql/expr.rs
index 24017f9cd..3d5f134a8 100644
--- a/datafusion/core/tests/sql/expr.rs
+++ b/datafusion/core/tests/sql/expr.rs
@@ -1313,6 +1313,24 @@ async fn test_extract_date_part() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn test_extract_epoch() -> Result<()> {
+    test_expression!(
+        "extract(epoch from '1870-01-01T07:29:10.256'::timestamp)",
+        "-3155646649.744"
+    );
+    test_expression!(
+        "extract(epoch from '2000-01-01T00:00:00.000'::timestamp)",
+        "946684800.0"
+    );
+    test_expression!(
+        "extract(epoch from to_timestamp('2000-01-01T00:00:00+00:00'))",
+        "946684800.0"
+    );
+    test_expression!("extract(epoch from NULL::timestamp)", "NULL");
+    Ok(())
+}
+
 #[tokio::test]
 async fn test_extract_date_part_func() -> Result<()> {
     test_expression!(
diff --git a/datafusion/physical-expr/src/datetime_expressions.rs 
b/datafusion/physical-expr/src/datetime_expressions.rs
index 4b3eeb442..4ce1f18e8 100644
--- a/datafusion/physical-expr/src/datetime_expressions.rs
+++ b/datafusion/physical-expr/src/datetime_expressions.rs
@@ -17,6 +17,7 @@
 
 //! DateTime expressions
 
+use arrow::array::Float64Builder;
 use arrow::compute::cast;
 use arrow::{
     array::TimestampNanosecondArray, compute::kernels::temporal, 
datatypes::TimeUnit,
@@ -478,6 +479,7 @@ pub fn date_part(args: &[ColumnarValue]) -> 
Result<ColumnarValue> {
         "millisecond" => extract_date_part!(&array, millis),
         "microsecond" => extract_date_part!(&array, micros),
         "nanosecond" => extract_date_part!(&array, nanos),
+        "epoch" => extract_date_part!(&array, epoch),
         _ => Err(DataFusionError::Execution(format!(
             "Date part '{date_part}' not supported"
         ))),
@@ -537,6 +539,40 @@ where
     to_ticks(array, 1_000_000_000)
 }
 
+fn epoch<T>(array: &PrimitiveArray<T>) -> Result<Float64Array>
+where
+    T: ArrowTemporalType + ArrowNumericType,
+    i64: From<T::Native>,
+{
+    let mut b = Float64Builder::with_capacity(array.len());
+    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);
+                }
+            }
+        }
+        _ => {
+            return Err(DataFusionError::Internal(format!(
+                "Can not convert {:?} to epoch",
+                array.data_type()
+            )))
+        }
+    }
+    Ok(b.finish())
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;