diff --git a/rust/arrow/src/compute/kernels/temporal.rs b/rust/arrow/src/compute/kernels/temporal.rs index 4319294e2de..63e412990fd 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::error::{ArrowError, Result}; /// 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!( + "year 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(15448)].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)); + } }