From fbea0b664ce81765ade0ebd32433ea0285143418 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Wed, 2 Nov 2022 15:12:06 -0700 Subject: [PATCH 1/2] Replace hour_generic with hour_dyn --- arrow/src/compute/kernels/temporal.rs | 83 +++++++++++++++++++++------ 1 file changed, 66 insertions(+), 17 deletions(-) diff --git a/arrow/src/compute/kernels/temporal.rs b/arrow/src/compute/kernels/temporal.rs index 412adb9a9fe9..6158bb6562c0 100644 --- a/arrow/src/compute/kernels/temporal.rs +++ b/arrow/src/compute/kernels/temporal.rs @@ -17,7 +17,9 @@ //! Defines temporal kernels for time and date related functions. +use arrow_array::downcast_dictionary_array; use chrono::{DateTime, Datelike, NaiveDateTime, NaiveTime, Offset, Timelike}; +use std::sync::Arc; use crate::array::*; use crate::datatypes::*; @@ -180,21 +182,69 @@ where T: ArrowTemporalType + ArrowNumericType, i64: From, { - hour_generic::(array) + hour_internal::(array, array.data_type()) } -/// Extracts the hours of a given temporal array as an array of integers within -/// the range of [0, 23]. -pub fn hour_generic>(array: A) -> Result -where - T: ArrowTemporalType + ArrowNumericType, - i64: From, -{ +/// Extracts the hours of a given array as an array of integers within +/// the range of [0, 23]. If the given array isn't temporal primitive or dictionary array, +/// an `Err` will be returned. +pub fn hour_dyn(array: &dyn Array) -> Result { match array.data_type().clone() { - DataType::Dictionary(_, value_type) => { - hour_internal::(array, value_type.as_ref()) + DataType::Dictionary(_, _) => { + downcast_dictionary_array!( + array => { + let hour_values = hour_dyn(array.values())?; + Ok(Arc::new(array.with_values(&hour_values))) + } + dt => return_compute_error_with!("hour does not support", dt), + ) + } + DataType::Time32(_) => { + let array = as_primitive_array::(array); + hour_internal::(array, array.data_type()) + .map(|a| Arc::new(a) as ArrayRef) + } + DataType::Time64(TimeUnit::Microsecond) => { + let array = as_primitive_array::(array); + hour_internal::(array, array.data_type()) + .map(|a| Arc::new(a) as ArrayRef) + } + DataType::Time64(TimeUnit::Nanosecond) => { + let array = as_primitive_array::(array); + hour_internal::(array, array.data_type()) + .map(|a| Arc::new(a) as ArrayRef) } - dt => hour_internal::(array, &dt), + DataType::Date32 => { + let array = as_primitive_array::(array); + hour_internal::(array, array.data_type()) + .map(|a| Arc::new(a) as ArrayRef) + } + DataType::Date64 => { + let array = as_primitive_array::(array); + hour_internal::(array, array.data_type()) + .map(|a| Arc::new(a) as ArrayRef) + } + DataType::Timestamp(TimeUnit::Second, _) => { + let array = as_primitive_array::(array); + hour_internal::(array, array.data_type()) + .map(|a| Arc::new(a) as ArrayRef) + } + DataType::Timestamp(TimeUnit::Millisecond, _) => { + let array = as_primitive_array::(array); + hour_internal::(array, array.data_type()) + .map(|a| Arc::new(a) as ArrayRef) + } + DataType::Timestamp(TimeUnit::Microsecond, _) => { + let array = as_primitive_array::(array); + hour_internal::(array, array.data_type()) + .map(|a| Arc::new(a) as ArrayRef) + } + DataType::Timestamp(TimeUnit::Nanosecond, _) => { + let array = as_primitive_array::(array); + hour_internal::(array, array.data_type()) + .map(|a| Arc::new(a) as ArrayRef) + } + dt => return_compute_error_with!("hour does not support", dt), } } @@ -1204,13 +1254,12 @@ mod tests { let keys = Int8Array::from_iter_values([0_i8, 0, 1, 2, 1]); let dict = DictionaryArray::try_new(&keys, &a).unwrap(); - let b = hour_generic::( - dict.downcast_dict::().unwrap(), - ) - .unwrap(); + let b = hour_dyn(&dict).unwrap(); - let expected = Int32Array::from(vec![11, 11, 21, 7, 21]); - assert_eq!(expected, b); + let expected_dict = + DictionaryArray::try_new(&keys, &Int32Array::from(vec![11, 21, 7])).unwrap(); + let expected = Arc::new(expected_dict) as ArrayRef; + assert_eq!(&expected, &b); let b = minute_generic::( dict.downcast_dict::().unwrap(), From f9e45ac14dd3330f2ac74c96aa9e96df69963744 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Wed, 2 Nov 2022 15:53:12 -0700 Subject: [PATCH 2/2] Add Time32MillisecondType --- arrow/src/compute/kernels/temporal.rs | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/arrow/src/compute/kernels/temporal.rs b/arrow/src/compute/kernels/temporal.rs index 6158bb6562c0..7e94d1da0601 100644 --- a/arrow/src/compute/kernels/temporal.rs +++ b/arrow/src/compute/kernels/temporal.rs @@ -199,11 +199,16 @@ pub fn hour_dyn(array: &dyn Array) -> Result { dt => return_compute_error_with!("hour does not support", dt), ) } - DataType::Time32(_) => { + DataType::Time32(TimeUnit::Second) => { let array = as_primitive_array::(array); hour_internal::(array, array.data_type()) .map(|a| Arc::new(a) as ArrayRef) } + DataType::Time32(TimeUnit::Microsecond) => { + let array = as_primitive_array::(array); + hour_internal::(array, array.data_type()) + .map(|a| Arc::new(a) as ArrayRef) + } DataType::Time64(TimeUnit::Microsecond) => { let array = as_primitive_array::(array); hour_internal::(array, array.data_type())