From 60142825bc727473eb1715276fda75810c5b42ff Mon Sep 17 00:00:00 2001 From: "Heres, Daniel" Date: Sat, 30 Jan 2021 18:02:01 +0100 Subject: [PATCH 1/5] Add year to temporal kernel, additional tests --- rust/arrow/src/compute/kernels/temporal.rs | 137 ++++++++++++++++++--- 1 file changed, 122 insertions(+), 15 deletions(-) diff --git a/rust/arrow/src/compute/kernels/temporal.rs b/rust/arrow/src/compute/kernels/temporal.rs index 4319294e2de..ac1d66bac82 100644 --- a/rust/arrow/src/compute/kernels/temporal.rs +++ b/rust/arrow/src/compute/kernels/temporal.rs @@ -17,12 +17,11 @@ //! Defines temporal kernels for time and date related functions. -use chrono::Timelike; +use chrono::{Datelike, Timelike}; -use crate::array::*; use crate::datatypes::*; use crate::error::Result; - +use crate::{array::*, error::ArrowError}; /// Extracts the hours of a given temporal array as an array of integers pub fn hour(array: &PrimitiveArray) -> Result where @@ -30,21 +29,70 @@ where i64: std::convert::From, { let mut b = Int32Builder::new(array.len()); - for i in 0..array.len() { - if array.is_null(i) { - b.append_null()?; - } else { - match array.data_type() { - &DataType::Time32(_) | &DataType::Time64(_) => { + match array.data_type() { + &DataType::Time32(_) | &DataType::Time64(_) => { + for i in 0..array.len() { + if array.is_null(i) { + b.append_null()?; + } else { match array.value_as_time(i) { Some(time) => b.append_value(time.hour() as i32)?, None => b.append_null()?, + }; + } + } + } + &DataType::Date32 | &DataType::Date64 | &DataType::Timestamp(_, _) => { + for i in 0..array.len() { + if array.is_null(i) { + b.append_null()?; + } else { + match array.value_as_datetime(i) { + Some(dt) => b.append_value(dt.hour() as i32)?, + None => b.append_null()?, } } - _ => match array.value_as_datetime(i) { - Some(dt) => b.append_value(dt.hour() as i32)?, - None => b.append_null()?, - }, + } + } + dt => { + return { + Err(ArrowError::ComputeError(format!( + "hour does not support type {:?}", + dt + ))) + } + } + } + + Ok(b.finish()) +} + +/// Extracts the years of a given temporal array as an array of integers +pub fn year(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(_, _) => { + for i in 0..array.len() { + if array.is_null(i) { + b.append_null()?; + } else { + match array.value_as_datetime(i) { + Some(dt) => b.append_value(dt.year() as i32)?, + None => b.append_null()?, + } + } + } + } + dt => { + return { + Err(ArrowError::ComputeError(format!( + "hour does not support type {:?}", + dt + ))) } } } @@ -61,20 +109,79 @@ mod tests { let a: PrimitiveArray = vec![Some(1514764800000), None, Some(1550636625000)].into(); - // get hour from temporal let b = hour(&a).unwrap(); assert_eq!(0, b.value(0)); assert_eq!(false, b.is_valid(1)); assert_eq!(4, b.value(2)); } + #[test] + fn test_temporal_array_date32_hour() { + let a: PrimitiveArray = vec![Some(15147), None, Some(15148)].into(); + + let b = hour(&a).unwrap(); + assert_eq!(0, b.value(0)); + assert_eq!(false, b.is_valid(1)); + assert_eq!(0, b.value(2)); + } + #[test] fn test_temporal_array_time32_second_hour() { let a: PrimitiveArray = vec![37800, 86339].into(); - // get hour from temporal let b = hour(&a).unwrap(); assert_eq!(10, b.value(0)); assert_eq!(23, b.value(1)); } + + #[test] + fn test_temporal_array_time64_micro_hour() { + let a: PrimitiveArray = + vec![37800000000, 86339000000].into(); + + let b = hour(&a).unwrap(); + assert_eq!(10, b.value(0)); + assert_eq!(23, b.value(1)); + } + + #[test] + fn test_temporal_array_timestamp_micro_hour() { + let a: TimestampMicrosecondArray = + vec![37800000000, 86339000000].into(); + + let b = hour(&a).unwrap(); + assert_eq!(10, b.value(0)); + assert_eq!(23, b.value(1)); + } + + #[test] + fn test_temporal_array_date64_year() { + let a: PrimitiveArray = + vec![Some(1514764800000), None, Some(1550636625000)].into(); + + let b = year(&a).unwrap(); + assert_eq!(2018, b.value(0)); + assert_eq!(false, b.is_valid(1)); + assert_eq!(2019, b.value(2)); + } + + #[test] + fn test_temporal_array_date32_year() { + let a: PrimitiveArray = vec![Some(15147), None, Some(15248)].into(); + + let b = year(&a).unwrap(); + assert_eq!(2011, b.value(0)); + assert_eq!(false, b.is_valid(1)); + assert_eq!(2012, b.value(2)); + } + + #[test] + fn test_temporal_array_timestamp_micro_year() { + let a: TimestampMicrosecondArray = vec![Some(1612025847000000), None, Some(1722015847000000)].into(); + + let b = year(&a).unwrap(); + assert_eq!(2021, b.value(0)); + assert_eq!(false, b.is_valid(1)); + assert_eq!(2024, b.value(2)); + } } From 3ec352d5b2bcd35f3b011d90e9023f91b6245571 Mon Sep 17 00:00:00 2001 From: "Heres, Daniel" Date: Sat, 30 Jan 2021 18:04:14 +0100 Subject: [PATCH 2/5] Fix error message --- rust/arrow/src/compute/kernels/temporal.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/arrow/src/compute/kernels/temporal.rs b/rust/arrow/src/compute/kernels/temporal.rs index ac1d66bac82..325a0f79778 100644 --- a/rust/arrow/src/compute/kernels/temporal.rs +++ b/rust/arrow/src/compute/kernels/temporal.rs @@ -90,7 +90,7 @@ where dt => { return { Err(ArrowError::ComputeError(format!( - "hour does not support type {:?}", + "year does not support type {:?}", dt ))) } From 22cdd7fca75d073c0e2cdbf8a6f438ef7e082cc5 Mon Sep 17 00:00:00 2001 From: "Heres, Daniel" Date: Sat, 30 Jan 2021 18:20:21 +0100 Subject: [PATCH 3/5] Test fix --- rust/arrow/src/compute/kernels/temporal.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/arrow/src/compute/kernels/temporal.rs b/rust/arrow/src/compute/kernels/temporal.rs index 325a0f79778..6d6c33769b9 100644 --- a/rust/arrow/src/compute/kernels/temporal.rs +++ b/rust/arrow/src/compute/kernels/temporal.rs @@ -167,7 +167,7 @@ mod tests { #[test] fn test_temporal_array_date32_year() { - let a: PrimitiveArray = vec![Some(15147), None, Some(15248)].into(); + let a: PrimitiveArray = vec![Some(15147), None, Some(15448)].into(); let b = year(&a).unwrap(); assert_eq!(2011, b.value(0)); From c6bfe2f49eda8afd743ebdaa47b1399044ca89d2 Mon Sep 17 00:00:00 2001 From: "Heres, Daniel" Date: Mon, 1 Feb 2021 16:52:07 +0100 Subject: [PATCH 4/5] Fmt --- rust/arrow/src/compute/kernels/temporal.rs | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/rust/arrow/src/compute/kernels/temporal.rs b/rust/arrow/src/compute/kernels/temporal.rs index 6d6c33769b9..0c17e10583c 100644 --- a/rust/arrow/src/compute/kernels/temporal.rs +++ b/rust/arrow/src/compute/kernels/temporal.rs @@ -19,9 +19,9 @@ use chrono::{Datelike, Timelike}; +use crate::array::*; use crate::datatypes::*; -use crate::error::Result; -use crate::{array::*, error::ArrowError}; +use crate::error::{ArrowError, Result}; /// Extracts the hours of a given temporal array as an array of integers pub fn hour(array: &PrimitiveArray) -> Result where @@ -136,8 +136,7 @@ mod tests { #[test] fn test_temporal_array_time64_micro_hour() { - let a: PrimitiveArray = - vec![37800000000, 86339000000].into(); + let a: PrimitiveArray = vec![37800000000, 86339000000].into(); let b = hour(&a).unwrap(); assert_eq!(10, b.value(0)); @@ -146,8 +145,7 @@ mod tests { #[test] fn test_temporal_array_timestamp_micro_hour() { - let a: TimestampMicrosecondArray = - vec![37800000000, 86339000000].into(); + let a: TimestampMicrosecondArray = vec![37800000000, 86339000000].into(); let b = hour(&a).unwrap(); assert_eq!(10, b.value(0)); @@ -177,7 +175,8 @@ mod tests { #[test] fn test_temporal_array_timestamp_micro_year() { - let a: TimestampMicrosecondArray = vec![Some(1612025847000000), None, Some(1722015847000000)].into(); + let a: TimestampMicrosecondArray = + vec![Some(1612025847000000), None, Some(1722015847000000)].into(); let b = year(&a).unwrap(); assert_eq!(2021, b.value(0)); From c2571fd98fc93271b4500da05289e09fd0b5fbdb Mon Sep 17 00:00:00 2001 From: "Heres, Daniel" Date: Tue, 2 Feb 2021 08:34:42 +0100 Subject: [PATCH 5/5] Fmt --- rust/arrow/src/compute/kernels/temporal.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/rust/arrow/src/compute/kernels/temporal.rs b/rust/arrow/src/compute/kernels/temporal.rs index 0c17e10583c..63e412990fd 100644 --- a/rust/arrow/src/compute/kernels/temporal.rs +++ b/rust/arrow/src/compute/kernels/temporal.rs @@ -136,7 +136,8 @@ mod tests { #[test] fn test_temporal_array_time64_micro_hour() { - let a: PrimitiveArray = vec![37800000000, 86339000000].into(); + let a: PrimitiveArray = + vec![37800000000, 86339000000].into(); let b = hour(&a).unwrap(); assert_eq!(10, b.value(0));