feat(compute): Support week0 (PostgreSQL behaviour) for temporal

arrow/src/compute/kernels/temporal.rs

@@ -17,7 +17,7 @@
 
 //! Defines temporal kernels for time and date related functions.
 
-use chrono::{Datelike, NaiveDate, NaiveDateTime, Timelike};
+use chrono::{Datelike, Timelike};
 
 use crate::array::*;
 use crate::datatypes::*;
@@ -112,31 +112,36 @@ macro_rules! return_compute_error_with {
     };
 }
 
-trait ChronoDateQuarter {
+// Internal trait, which is used for mapping values from DateLike structures
+trait ChronoDateExt {
     /// Returns a value in range `1..=4` indicating the quarter this date falls into
     fn quarter(&self) -> u32;
 
     /// Returns a value in range `0..=3` indicating the quarter (zero-based) this date falls into
     fn quarter0(&self) -> u32;
+
+    /// Returns the day of week; Monday is encoded as `0`, Tuesday as `1`, etc.
+    fn num_days_from_monday(&self) -> i32;
+
+    /// Returns the day of week; Sunday is encoded as `0`, Monday as `1`, etc.
+    fn num_days_from_sunday(&self) -> i32;
 }
 
-impl ChronoDateQuarter for NaiveDateTime {
+impl<T: Datelike> ChronoDateExt for T {
     fn quarter(&self) -> u32 {
         self.quarter0() + 1
     }
 
     fn quarter0(&self) -> u32 {
         self.month0() / 3
     }
-}
 
-impl ChronoDateQuarter for NaiveDate {
-    fn quarter(&self) -> u32 {
-        self.quarter0() + 1
+    fn num_days_from_monday(&self) -> i32 {
+        self.weekday().num_days_from_monday() as i32
     }
 
-    fn quarter0(&self) -> u32 {
-        self.month0() / 3
+    fn num_days_from_sunday(&self) -> i32 {
+        self.weekday().num_days_from_sunday() as i32
     }
 }
 
@@ -279,14 +284,55 @@ where
     let mut b = Int32Builder::new(array.len());
     match array.data_type() {
         &DataType::Date32 | &DataType::Date64 | &DataType::Timestamp(_, None) => {
-            extract_component_from_array!(array, b, weekday, value_as_datetime)
+            extract_component_from_array!(
+                array,
+                b,
+                num_days_from_monday,
+                value_as_datetime
+            )
+        }
+        &DataType::Timestamp(_, Some(ref tz)) => {
+            let mut scratch = Parsed::new();
+            extract_component_from_array!(
+                array,
+                b,
+                num_days_from_monday,
+                value_as_datetime_with_tz,
+                tz,
+                scratch
+            )
+        }
+        dt => return_compute_error_with!("weekday does not support", dt),
+    }
+
+    Ok(b.finish())
+}
+
+/// Extracts the day of week of a given temporal array as an array of
+/// integers, starting at Sunday. This is different than [`weekday`] which starts at Monday.
+///
+/// Sunday is encoded as `0`, Monday as `1`, etc.
+pub fn weekday0<T>(array: &PrimitiveArray<T>) -> Result<Int32Array>
+where
+    T: ArrowTemporalType + ArrowNumericType,
+    i64: std::convert::From<T::Native>,
+{
+    let mut b = Int32Builder::new(array.len());
+    match array.data_type() {
+        &DataType::Date32 | &DataType::Date64 | &DataType::Timestamp(_, None) => {
+            extract_component_from_array!(
+                array,
+                b,
+                num_days_from_sunday,
+                value_as_datetime
+            )
         }
         &DataType::Timestamp(_, Some(ref tz)) => {
             let mut scratch = Parsed::new();
             extract_component_from_array!(
                 array,
                 b,
-                weekday,
+                num_days_from_sunday,
                 value_as_datetime_with_tz,
                 tz,
                 scratch
@@ -592,6 +638,26 @@ mod tests {
         assert_eq!(2, b.value(2));
     }
 
+    #[test]
+    fn test_temporal_array_date64_weekday0() {
+        //1483228800000 -> 2017-01-01 (Sunday)
+        //1514764800000 -> 2018-01-01 (Monday)
+        //1550636625000 -> 2019-02-20 (Wednesday)
+        let a: PrimitiveArray<Date64Type> = vec![
+            Some(1483228800000),
+            None,
+            Some(1514764800000),
+            Some(1550636625000),
+        ]
+        .into();
+
+        let b = weekday0(&a).unwrap();
+        assert_eq!(0, b.value(0));
+        assert!(!b.is_valid(1));
+        assert_eq!(1, b.value(2));
+        assert_eq!(3, b.value(3));
+    }
+
     #[test]
     fn test_temporal_array_date64_day() {
         //1514764800000 -> 2018-01-01