From 5aad8137265698885ab5751a8ce140acbf4f82c4 Mon Sep 17 00:00:00 2001 From: Dmitry Patsura Date: Tue, 12 Jul 2022 19:13:50 +0300 Subject: [PATCH 1/5] feat(compute): Support dow for temporal --- arrow/src/compute/kernels/temporal.rs | 80 ++++++++++++++++++++++++--- 1 file changed, 71 insertions(+), 9 deletions(-) diff --git a/arrow/src/compute/kernels/temporal.rs b/arrow/src/compute/kernels/temporal.rs index 9998649ead30..91841e28fee9 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,18 @@ 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; + + fn weekday_from_sunday(&self) -> i32; } -impl ChronoDateQuarter for NaiveDateTime { +impl ChronoDateExt for T { fn quarter(&self) -> u32 { self.quarter0() + 1 } @@ -128,16 +131,43 @@ impl ChronoDateQuarter for NaiveDateTime { fn quarter0(&self) -> u32 { self.month0() / 3 } -} -impl ChronoDateQuarter for NaiveDate { - fn quarter(&self) -> u32 { - self.quarter0() + 1 + fn weekday_from_sunday(&self) -> i32 { + self.weekday().num_days_from_sunday() as i32 } +} - fn quarter0(&self) -> u32 { - self.month0() / 3 +/// Extracts the day of week of a given temporal array as an array of integers +pub fn dow(array: &PrimitiveArray) -> Result +where + T: ArrowTemporalType + ArrowNumericType, + i64: std::convert::From, +{ + 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_from_sunday, + value_as_datetime + ) + } + &DataType::Timestamp(_, Some(ref tz)) => { + let mut scratch = Parsed::new(); + extract_component_from_array!( + array, + b, + weekday_from_sunday, + value_as_datetime_with_tz, + tz, + scratch + ) + } + dt => return_compute_error_with!("dow does not support", dt), } + + Ok(b.finish()) } #[cfg(not(feature = "chrono-tz"))] @@ -518,6 +548,38 @@ mod tests { assert_eq!(1, b.value(0)); } + #[test] + fn test_temporal_array_date64_dow() { + //1657486800 -> 2022-07-11 + //1657573200 -> 2022-07-12 + //1658005200 -> 2022-07-17 + let a: PrimitiveArray = vec![ + Some(1657486800000), + Some(1657573200000), + Some(1658005200000), + None, + ] + .into(); + + let b = dow(&a).unwrap(); + assert_eq!(0, b.value(0)); + assert_eq!(1, b.value(1)); + assert_eq!(6, b.value(2)); + assert!(!b.is_valid(3)); + } + + #[test] + fn test_temporal_array_date32_dow() { + let a: PrimitiveArray = + vec![Some(1657497600), Some(1657584000), Some(1658016000), None].into(); + + let b = dow(&a).unwrap(); + assert_eq!(0, b.value(0)); + assert_eq!(1, b.value(1)); + assert_eq!(6, b.value(2)); + assert!(!b.is_valid(3)); + } + #[test] fn test_temporal_array_date64_month() { //1514764800000 -> 2018-01-01 From a610ccb59fbd16921175e0cbbb07495eef810848 Mon Sep 17 00:00:00 2001 From: Dmitry Patsura Date: Tue, 12 Jul 2022 20:05:25 +0300 Subject: [PATCH 2/5] update to reflect changes --- arrow/src/compute/kernels/temporal.rs | 120 ++++++++++++-------------- 1 file changed, 53 insertions(+), 67 deletions(-) diff --git a/arrow/src/compute/kernels/temporal.rs b/arrow/src/compute/kernels/temporal.rs index 91841e28fee9..9764ddb6823d 100644 --- a/arrow/src/compute/kernels/temporal.rs +++ b/arrow/src/compute/kernels/temporal.rs @@ -120,7 +120,7 @@ trait ChronoDateExt { /// Returns a value in range `0..=3` indicating the quarter (zero-based) this date falls into fn quarter0(&self) -> u32; - fn weekday_from_sunday(&self) -> i32; + fn weekday0(&self) -> i32; } impl ChronoDateExt for T { @@ -132,44 +132,11 @@ impl ChronoDateExt for T { self.month0() / 3 } - fn weekday_from_sunday(&self) -> i32 { + fn weekday0(&self) -> i32 { self.weekday().num_days_from_sunday() as i32 } } -/// Extracts the day of week of a given temporal array as an array of integers -pub fn dow(array: &PrimitiveArray) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: std::convert::From, -{ - 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_from_sunday, - value_as_datetime - ) - } - &DataType::Timestamp(_, Some(ref tz)) => { - let mut scratch = Parsed::new(); - extract_component_from_array!( - array, - b, - weekday_from_sunday, - value_as_datetime_with_tz, - tz, - scratch - ) - } - dt => return_compute_error_with!("dow does not support", dt), - } - - Ok(b.finish()) -} - #[cfg(not(feature = "chrono-tz"))] pub fn using_chrono_tz_and_utc_naive_date_time( _tz: &str, @@ -328,6 +295,37 @@ where Ok(b.finish()) } +/// Extracts the day of week of a given temporal array as an array of +/// integers. +/// +/// Sunday is encoded as `0`, Monday as `1`, etc. +pub fn weekday0(array: &PrimitiveArray) -> Result +where + T: ArrowTemporalType + ArrowNumericType, + i64: std::convert::From, +{ + let mut b = Int32Builder::new(array.len()); + match array.data_type() { + &DataType::Date32 | &DataType::Date64 | &DataType::Timestamp(_, None) => { + extract_component_from_array!(array, b, weekday0, value_as_datetime) + } + &DataType::Timestamp(_, Some(ref tz)) => { + let mut scratch = Parsed::new(); + extract_component_from_array!( + array, + b, + weekday0, + value_as_datetime_with_tz, + tz, + scratch + ) + } + dt => return_compute_error_with!("weekday does not support", dt), + } + + Ok(b.finish()) +} + /// Extracts the day of a given temporal array as an array of integers pub fn day(array: &PrimitiveArray) -> Result where @@ -548,38 +546,6 @@ mod tests { assert_eq!(1, b.value(0)); } - #[test] - fn test_temporal_array_date64_dow() { - //1657486800 -> 2022-07-11 - //1657573200 -> 2022-07-12 - //1658005200 -> 2022-07-17 - let a: PrimitiveArray = vec![ - Some(1657486800000), - Some(1657573200000), - Some(1658005200000), - None, - ] - .into(); - - let b = dow(&a).unwrap(); - assert_eq!(0, b.value(0)); - assert_eq!(1, b.value(1)); - assert_eq!(6, b.value(2)); - assert!(!b.is_valid(3)); - } - - #[test] - fn test_temporal_array_date32_dow() { - let a: PrimitiveArray = - vec![Some(1657497600), Some(1657584000), Some(1658016000), None].into(); - - let b = dow(&a).unwrap(); - assert_eq!(0, b.value(0)); - assert_eq!(1, b.value(1)); - assert_eq!(6, b.value(2)); - assert!(!b.is_valid(3)); - } - #[test] fn test_temporal_array_date64_month() { //1514764800000 -> 2018-01-01 @@ -654,6 +620,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 = 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 From c349773c005e8bb407775bd82c0f9100c8da24bc Mon Sep 17 00:00:00 2001 From: Dmitry Patsura Date: Wed, 13 Jul 2022 13:48:15 +0300 Subject: [PATCH 3/5] Update arrow/src/compute/kernels/temporal.rs Co-authored-by: Andrew Lamb --- arrow/src/compute/kernels/temporal.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/arrow/src/compute/kernels/temporal.rs b/arrow/src/compute/kernels/temporal.rs index 9764ddb6823d..404f38adb3bb 100644 --- a/arrow/src/compute/kernels/temporal.rs +++ b/arrow/src/compute/kernels/temporal.rs @@ -120,6 +120,7 @@ trait ChronoDateExt { /// 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; Sunday is encoded as `0`, Monday as `1`, etc. fn weekday0(&self) -> i32; } From 2d0235c7bf4daeffcb09590a1935bacb154dc3b8 Mon Sep 17 00:00:00 2001 From: Dmitry Patsura Date: Wed, 13 Jul 2022 13:48:41 +0300 Subject: [PATCH 4/5] Update arrow/src/compute/kernels/temporal.rs Co-authored-by: Andrew Lamb --- arrow/src/compute/kernels/temporal.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arrow/src/compute/kernels/temporal.rs b/arrow/src/compute/kernels/temporal.rs index 404f38adb3bb..b7005ad5d235 100644 --- a/arrow/src/compute/kernels/temporal.rs +++ b/arrow/src/compute/kernels/temporal.rs @@ -297,7 +297,7 @@ where } /// Extracts the day of week of a given temporal array as an array of -/// integers. +/// 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(array: &PrimitiveArray) -> Result From 6e8070226678939bf9ccd44f03fe5768933b5e24 Mon Sep 17 00:00:00 2001 From: Dmitry Patsura Date: Thu, 14 Jul 2022 00:05:44 +0300 Subject: [PATCH 5/5] requested changes from review --- arrow/src/compute/kernels/temporal.rs | 29 +++++++++++++++++++++------ 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/arrow/src/compute/kernels/temporal.rs b/arrow/src/compute/kernels/temporal.rs index b7005ad5d235..59388242354a 100644 --- a/arrow/src/compute/kernels/temporal.rs +++ b/arrow/src/compute/kernels/temporal.rs @@ -120,8 +120,11 @@ trait ChronoDateExt { /// 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 weekday0(&self) -> i32; + fn num_days_from_sunday(&self) -> i32; } impl ChronoDateExt for T { @@ -133,7 +136,11 @@ impl ChronoDateExt for T { self.month0() / 3 } - fn weekday0(&self) -> i32 { + fn num_days_from_monday(&self) -> i32 { + self.weekday().num_days_from_monday() as i32 + } + + fn num_days_from_sunday(&self) -> i32 { self.weekday().num_days_from_sunday() as i32 } } @@ -277,14 +284,19 @@ 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, - weekday, + num_days_from_monday, value_as_datetime_with_tz, tz, scratch @@ -308,14 +320,19 @@ 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, weekday0, value_as_datetime) + 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, - weekday0, + num_days_from_sunday, value_as_datetime_with_tz, tz, scratch