This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new 4444cb706 feat(compute): Support week0 (PostgreSQL behaviour) for
temporal (#2052)
4444cb706 is described below
commit 4444cb70607a9f0fe4c0f4130fa03b34faffef9c
Author: Dmitry Patsura <[email protected]>
AuthorDate: Thu Jul 14 16:07:06 2022 +0300
feat(compute): Support week0 (PostgreSQL behaviour) for temporal (#2052)
* feat(compute): Support dow for temporal
* update to reflect changes
* Update arrow/src/compute/kernels/temporal.rs
Co-authored-by: Andrew Lamb <[email protected]>
* Update arrow/src/compute/kernels/temporal.rs
Co-authored-by: Andrew Lamb <[email protected]>
* requested changes from review
Co-authored-by: Andrew Lamb <[email protected]>
---
arrow/src/compute/kernels/temporal.rs | 88 ++++++++++++++++++++++++++++++-----
1 file changed, 77 insertions(+), 11 deletions(-)
diff --git a/arrow/src/compute/kernels/temporal.rs
b/arrow/src/compute/kernels/temporal.rs
index 9998649ea..593882423 100644
--- a/arrow/src/compute/kernels/temporal.rs
+++ b/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,15 +112,22 @@ 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
}
@@ -128,15 +135,13 @@ impl ChronoDateQuarter for NaiveDateTime {
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
