Replace complicated temporal macro with generic functions

arrow/src/compute/kernels/cast.rs

@@ -37,14 +37,13 @@
 
 use chrono::format::strftime::StrftimeItems;
 use chrono::format::{parse, Parsed};
-use chrono::Timelike;
+use chrono::{NaiveDateTime, Timelike};
 use std::ops::{Div, Mul};
 use std::str;
 use std::sync::Arc;
 
 use crate::buffer::MutableBuffer;
 use crate::compute::kernels::cast_utils::string_to_timestamp_nanos;
-use crate::compute::kernels::temporal::extract_component_from_array;
 use crate::compute::kernels::temporal::return_compute_error_with;
 use crate::compute::{divide_scalar, multiply_scalar};
 use crate::compute::{try_unary, using_chrono_tz_and_utc_naive_date_time};
@@ -1638,6 +1637,98 @@ where
     unsafe { PrimitiveArray::<R>::from_trusted_len_iter(iter) }
 }
 
+fn as_time_with_string_op<
+    A: ArrayAccessor<Item = T::Native>,
+    OffsetSize,
+    T: ArrowTemporalType,
+    F,
+>(
+    iter: ArrayIter<A>,
+    mut builder: GenericStringBuilder<OffsetSize>,
+    op: F,
+) -> ArrayRef
+where
+    OffsetSize: OffsetSizeTrait,
+    F: Fn(NaiveDateTime) -> String,
+    i64: From<T::Native>,
+{
+    iter.into_iter().for_each(|value| {
+        if let Some(value) = value {
+            match as_datetime::<T>(<i64 as From<_>>::from(value)) {
+                Some(dt) => builder.append_value(op(dt)),
+                None => builder.append_null(),
+            }
+        } else {
+            builder.append_null();
+        }
+    });
+
+    Arc::new(builder.finish())
+}
+
+fn extract_component_from_datatime_array<
+    A: ArrayAccessor<Item = T::Native>,
+    OffsetSize,
+    T: ArrowTemporalType,
+    F,
+>(
+    iter: ArrayIter<A>,
+    mut builder: GenericStringBuilder<OffsetSize>,
+    tz: &str,
+    mut parsed: Parsed,
+    op: F,
+) -> Result<ArrayRef>
+where
+    OffsetSize: OffsetSizeTrait,
+    F: Fn(NaiveDateTime) -> String,
+    i64: From<T::Native>,
+{
+    if (tz.starts_with('+') || tz.starts_with('-')) && !tz.contains(':') {
+        return_compute_error_with!(
+            "Invalid timezone",
+            "Expected format [+-]XX:XX".to_string()
+        )
+    } else {
+        let tz_parse_result = parse(&mut parsed, tz, StrftimeItems::new("%z"));
+        let fixed_offset_from_parsed = match tz_parse_result {
+            Ok(_) => match parsed.to_fixed_offset() {
+                Ok(fo) => Some(fo),
+                err => return_compute_error_with!("Invalid timezone", err),
+            },
+            _ => None,
+        };
+
+        for value in iter {
+            if let Some(value) = value {
+                match as_datetime::<T>(<i64 as From<_>>::from(value)) {
+                    Some(utc) => {
+                        let fixed_offset = match fixed_offset_from_parsed {
+                            Some(fo) => fo,
+                            None => {
+                                match using_chrono_tz_and_utc_naive_date_time(tz, utc) {
+                                    Some(fo) => fo,
+                                    err => return_compute_error_with!(
+                                        "Unable to parse timezone",
+                                        err
+                                    ),
+                                }
+                            }
+                        };
+                        builder.append_value(op(utc + fixed_offset));
+                    }
+                    err => return_compute_error_with!(
+                        "Unable to read value as datetime",
+                        err
+                    ),
+                }
+            } else {
+                builder.append_null();
+            }
+        }
+    }
+    Ok(Arc::new(builder.finish()))
+}
+
 /// Cast timestamp types to Utf8/LargeUtf8
 fn cast_timestamp_to_string<T, OffsetSize>(
     array: &ArrayRef,
@@ -1650,38 +1741,30 @@ where
 {
     let array = array.as_any().downcast_ref::<PrimitiveArray<T>>().unwrap();
 
-    let mut builder = GenericStringBuilder::<OffsetSize>::new();
+    let builder = GenericStringBuilder::<OffsetSize>::new();
 
     if let Some(tz) = tz {
-        let mut scratch = Parsed::new();
+        let scratch = Parsed::new();
         // 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!(
+        extract_component_from_datatime_array::<_, OffsetSize, T, _>(
             iter,
             builder,
-            to_string,
-            |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
+            |t| t.to_string(),
         )
     } else {
         // No timezone available. Calling `to_string` on the datatime value simply.
         let iter = ArrayIter::new(array);
-        extract_component_from_array!(
+        Ok(as_time_with_string_op::<_, OffsetSize, T, _>(
             iter,
             builder,
-            to_string,
-            |value| as_datetime::<T>(<i64 as From<_>>::from(value)),
-            |h| h
-        )
+            |t| t.to_string(),
+        ))
     }
-
-    Ok(Arc::new(builder.finish()) as ArrayRef)
 }
 
 /// Cast date32 types to Utf8/LargeUtf8