diff --git a/arrow/src/compute/kernels/temporal.rs b/arrow/src/compute/kernels/temporal.rs index 412adb9a9fe9..7e94d1da0601 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,74 @@ 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(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()) + .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 +1259,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(),