Support DictionaryArray in temporal kernels (#2623)

* Support dictionary array in temporal kernels * Support dictionary array in temporal kernels * Prepare for merging conflicts * Keep same kernel signature for primitive array * For review * Add doc
apache · Sep 7, 2022 · c8bf1ca · c8bf1ca
1 parent d73d78f
commit c8bf1ca
Show file tree

Hide file tree

Showing 4 changed files with 697 additions and 199 deletions.
diff --git a/arrow/src/array/array_primitive.rs b/arrow/src/array/array_primitive.rs
@@ -206,7 +206,7 @@ impl<'a, T: ArrowPrimitiveType> ArrayAccessor for &'a PrimitiveArray<T> {
     }
 }
 
-fn as_datetime<T: ArrowPrimitiveType>(v: i64) -> Option<NaiveDateTime> {
+pub(crate) fn as_datetime<T: ArrowPrimitiveType>(v: i64) -> Option<NaiveDateTime> {
     match T::DATA_TYPE {
         DataType::Date32 => Some(temporal_conversions::date32_to_datetime(v as i32)),
         DataType::Date64 => Some(temporal_conversions::date64_to_datetime(v)),
@@ -233,7 +233,7 @@ fn as_date<T: ArrowPrimitiveType>(v: i64) -> Option<NaiveDate> {
     as_datetime::<T>(v).map(|datetime| datetime.date())
 }
 
-fn as_time<T: ArrowPrimitiveType>(v: i64) -> Option<NaiveTime> {
+pub(crate) fn as_time<T: ArrowPrimitiveType>(v: i64) -> Option<NaiveTime> {
     match T::DATA_TYPE {
         DataType::Time32(unit) => {
             // safe to immediately cast to u32 as `self.value(i)` is positive i32

diff --git a/arrow/src/array/mod.rs b/arrow/src/array/mod.rs
@@ -223,6 +223,8 @@ pub use self::array::make_array;
 pub use self::array::new_empty_array;
 pub use self::array::new_null_array;
 
+pub(crate) use self::array_primitive::{as_datetime, as_time};
+
 ///
 /// # Example: Using `collect`
 /// ```

diff --git a/arrow/src/compute/kernels/cast.rs b/arrow/src/compute/kernels/cast.rs
@@ -42,6 +42,7 @@ use std::ops::{Div, Mul};
 use std::str;
 use std::sync::Arc;
 
+use crate::array::as_datetime;
 use crate::buffer::MutableBuffer;
 use crate::compute::divide_scalar;
 use crate::compute::kernels::arithmetic::{divide, multiply};
@@ -1708,21 +1709,31 @@ where
 
     if let Some(tz) = tz {
         let mut scratch = Parsed::new();
-        // The macro calls `value_as_datetime_with_tz` on timestamp values of the array.
+        // The macro calls `as_datetime` on timestamp values of the array.
         // After applying timezone offset on the datatime, calling `to_string` to get
         // the strings.
+        let iter = ArrayIter::new(array);
         extract_component_from_array!(
-            array,
+            iter,
             builder,
             to_string,
-            value_as_datetime_with_tz,
+            |value, tz| as_datetime::<T>(<i64 as From<_>>::from(value))
+                .map(|datetime| datetime + tz),
             tz,
             scratch,
+            |value| as_datetime::<T>(<i64 as From<_>>::from(value)),
             |h| h
         )
     } else {
         // No timezone available. Calling `to_string` on the datatime value simply.
-        extract_component_from_array!(array, builder, to_string, value_as_datetime, |h| h)
+        let iter = ArrayIter::new(array);
+        extract_component_from_array!(
+            iter,
+            builder,
+            to_string,
+            |value| as_datetime::<T>(<i64 as From<_>>::from(value)),
+            |h| h
+        )
     }
 
     Ok(Arc::new(builder.finish()) as ArrayRef)