From b758df9df0c30cfdf07b2ecb303f6d2609287fef Mon Sep 17 00:00:00 2001
From: VasanthakumarV <vasanth260m12@gmail.com>
Date: Thu, 23 Sep 2021 11:37:32 +0530
Subject: [PATCH] Rearrange functions in temporal compute

---
 src/compute/temporal.rs | 438 ++++++++++++++++++++--------------------
 1 file changed, 219 insertions(+), 219 deletions(-)

diff --git a/src/compute/temporal.rs b/src/compute/temporal.rs
index 1866875751a..fa1210f4ebc 100644
--- a/src/compute/temporal.rs
+++ b/src/compute/temporal.rs
@@ -28,84 +28,82 @@ use crate::types::NaturalDataType;
 
 use super::arity::unary;
 
-fn extract_impl<T, A, F>(
-    array: &PrimitiveArray<i64>,
-    time_unit: TimeUnit,
-    timezone: T,
-    extract: F,
-) -> PrimitiveArray<A>
-where
-    T: chrono::TimeZone,
-    A: NativeType + NaturalDataType,
-    F: Fn(chrono::DateTime<T>) -> A,
-{
-    match time_unit {
-        TimeUnit::Second => {
-            let op = |x| {
-                let datetime = timestamp_s_to_datetime(x);
-                let offset = timezone.offset_from_utc_datetime(&datetime);
-                extract(chrono::DateTime::<T>::from_utc(datetime, offset))
-            };
-            unary(array, op, DataType::UInt32)
+/// Extracts the years of a temporal array as [`PrimitiveArray<i32>`].
+/// Use [`can_year`] to check if this operation is supported for the target [`DataType`].
+pub fn year(array: &dyn Array) -> Result<PrimitiveArray<i32>> {
+    match array.data_type() {
+        DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, None) => {
+            date_like(array, DataType::Int32, |x| x.year())
         }
-        TimeUnit::Millisecond => {
-            let op = |x| {
-                let datetime = timestamp_ms_to_datetime(x);
-                let offset = timezone.offset_from_utc_datetime(&datetime);
-                extract(chrono::DateTime::<T>::from_utc(datetime, offset))
-            };
-            unary(array, op, A::DATA_TYPE)
+        DataType::Timestamp(time_unit, Some(timezone_str)) => {
+            let time_unit = *time_unit;
+            let timezone = parse_offset(timezone_str);
+
+            let array = array.as_any().downcast_ref().unwrap();
+
+            if let Ok(timezone) = timezone {
+                Ok(extract_impl(array, time_unit, timezone, |x| x.year()))
+            } else {
+                chrono_tz(array, time_unit, timezone_str, |x| x.year())
+            }
         }
-        TimeUnit::Microsecond => {
-            let op = |x| {
-                let datetime = timestamp_us_to_datetime(x);
-                let offset = timezone.offset_from_utc_datetime(&datetime);
-                extract(chrono::DateTime::<T>::from_utc(datetime, offset))
-            };
-            unary(array, op, A::DATA_TYPE)
+        dt => Err(ArrowError::NotYetImplemented(format!(
+            "\"year\" does not support type {:?}",
+            dt
+        ))),
+    }
+}
+
+/// Extracts the months of a temporal array as [`PrimitiveArray<u32>`].
+/// Use [`can_month`] to check if this operation is supported for the target [`DataType`].
+pub fn month(array: &dyn Array) -> Result<PrimitiveArray<u32>> {
+    match array.data_type() {
+        DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, None) => {
+            date_like(array, DataType::UInt32, |x| x.month())
         }
-        TimeUnit::Nanosecond => {
-            let op = |x| {
-                let datetime = timestamp_ns_to_datetime(x);
-                let offset = timezone.offset_from_utc_datetime(&datetime);
-                extract(chrono::DateTime::<T>::from_utc(datetime, offset))
-            };
-            unary(array, op, A::DATA_TYPE)
+        DataType::Timestamp(time_unit, Some(timezone_str)) => {
+            let time_unit = *time_unit;
+            let timezone = parse_offset(timezone_str);
+
+            let array = array.as_any().downcast_ref().unwrap();
+
+            if let Ok(timezone) = timezone {
+                Ok(extract_impl(array, time_unit, timezone, |x| x.month()))
+            } else {
+                chrono_tz(array, time_unit, timezone_str, |x| x.month())
+            }
         }
+        dt => Err(ArrowError::NotYetImplemented(format!(
+            "\"month\" does not support type {:?}",
+            dt
+        ))),
     }
 }
 
-#[cfg(feature = "chrono-tz")]
-#[cfg_attr(docsrs, doc(cfg(feature = "chrono-tz")))]
-fn chrono_tz<F, O>(
-    array: &PrimitiveArray<i64>,
-    time_unit: TimeUnit,
-    timezone_str: &str,
-    op: F,
-) -> Result<PrimitiveArray<O>>
-where
-    O: NativeType,
-    F: Fn(chrono::DateTime<chrono_tz::Tz>) -> O,
-{
-    let timezone = parse_offset_tz(timezone_str)?;
-    Ok(extract_impl(array, time_unit, timezone, op))
-}
+/// Extracts the days of a temporal array as [`PrimitiveArray<u32>`].
+/// Use [`can_day`] to check if this operation is supported for the target [`DataType`].
+pub fn day(array: &dyn Array) -> Result<PrimitiveArray<u32>> {
+    match array.data_type() {
+        DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, None) => {
+            date_like(array, DataType::UInt32, |x| x.day())
+        }
+        DataType::Timestamp(time_unit, Some(timezone_str)) => {
+            let time_unit = *time_unit;
+            let timezone = parse_offset(timezone_str);
 
-#[cfg(not(feature = "chrono-tz"))]
-fn chrono_tz<F, O>(
-    _: &PrimitiveArray<i64>,
-    _: TimeUnit,
-    timezone_str: &str,
-    _: F,
-) -> Result<PrimitiveArray<O>>
-where
-    O: NativeType,
-    F: Fn(chrono::DateTime<chrono_tz::Tz>) -> O,
-{
-    Err(ArrowError::InvalidArgumentError(format!(
-        "timezone \"{}\" cannot be parsed (feature chrono-tz is not active)",
-        timezone_str
-    )))
+            let array = array.as_any().downcast_ref().unwrap();
+
+            if let Ok(timezone) = timezone {
+                Ok(extract_impl(array, time_unit, timezone, |x| x.day()))
+            } else {
+                chrono_tz(array, time_unit, timezone_str, |x| x.day())
+            }
+        }
+        dt => Err(ArrowError::NotYetImplemented(format!(
+            "\"day\" does not support type {:?}",
+            dt
+        ))),
+    }
 }
 
 /// Extracts the hours of a temporal array as [`PrimitiveArray<u32>`].
@@ -195,6 +193,43 @@ pub fn second(array: &dyn Array) -> Result<PrimitiveArray<u32>> {
     }
 }
 
+pub fn date_like<F, O>(array: &dyn Array, data_type: DataType, op: F) -> Result<PrimitiveArray<O>>
+where
+    O: NativeType,
+    F: Fn(chrono::NaiveDateTime) -> O,
+{
+    match array.data_type() {
+        DataType::Date32 => {
+            let array = array
+                .as_any()
+                .downcast_ref::<PrimitiveArray<i32>>()
+                .unwrap();
+            Ok(unary(array, |x| op(date32_to_datetime(x)), data_type))
+        }
+        DataType::Date64 => {
+            let array = array
+                .as_any()
+                .downcast_ref::<PrimitiveArray<i64>>()
+                .unwrap();
+            Ok(unary(array, |x| op(date64_to_datetime(x)), data_type))
+        }
+        DataType::Timestamp(time_unit, None) => {
+            let array = array
+                .as_any()
+                .downcast_ref::<PrimitiveArray<i64>>()
+                .unwrap();
+            let func = match time_unit {
+                TimeUnit::Second => timestamp_s_to_datetime,
+                TimeUnit::Millisecond => timestamp_ms_to_datetime,
+                TimeUnit::Microsecond => timestamp_us_to_datetime,
+                TimeUnit::Nanosecond => timestamp_ns_to_datetime,
+            };
+            Ok(unary(array, |x| op(func(x)), data_type))
+        }
+        _ => unreachable!(),
+    }
+}
+
 fn time_like<F, O>(array: &dyn Array, data_type: DataType, op: F) -> Result<PrimitiveArray<O>>
 where
     O: NativeType,
@@ -233,6 +268,120 @@ where
     }
 }
 
+#[cfg(feature = "chrono-tz")]
+#[cfg_attr(docsrs, doc(cfg(feature = "chrono-tz")))]
+fn chrono_tz<F, O>(
+    array: &PrimitiveArray<i64>,
+    time_unit: TimeUnit,
+    timezone_str: &str,
+    op: F,
+) -> Result<PrimitiveArray<O>>
+where
+    O: NativeType,
+    F: Fn(chrono::DateTime<chrono_tz::Tz>) -> O,
+{
+    let timezone = parse_offset_tz(timezone_str)?;
+    Ok(extract_impl(array, time_unit, timezone, op))
+}
+
+#[cfg(not(feature = "chrono-tz"))]
+fn chrono_tz<F, O>(
+    _: &PrimitiveArray<i64>,
+    _: TimeUnit,
+    timezone_str: &str,
+    _: F,
+) -> Result<PrimitiveArray<O>>
+where
+    O: NativeType,
+    F: Fn(chrono::DateTime<chrono_tz::Tz>) -> O,
+{
+    Err(ArrowError::InvalidArgumentError(format!(
+        "timezone \"{}\" cannot be parsed (feature chrono-tz is not active)",
+        timezone_str
+    )))
+}
+
+fn extract_impl<T, A, F>(
+    array: &PrimitiveArray<i64>,
+    time_unit: TimeUnit,
+    timezone: T,
+    extract: F,
+) -> PrimitiveArray<A>
+where
+    T: chrono::TimeZone,
+    A: NativeType + NaturalDataType,
+    F: Fn(chrono::DateTime<T>) -> A,
+{
+    match time_unit {
+        TimeUnit::Second => {
+            let op = |x| {
+                let datetime = timestamp_s_to_datetime(x);
+                let offset = timezone.offset_from_utc_datetime(&datetime);
+                extract(chrono::DateTime::<T>::from_utc(datetime, offset))
+            };
+            unary(array, op, DataType::UInt32)
+        }
+        TimeUnit::Millisecond => {
+            let op = |x| {
+                let datetime = timestamp_ms_to_datetime(x);
+                let offset = timezone.offset_from_utc_datetime(&datetime);
+                extract(chrono::DateTime::<T>::from_utc(datetime, offset))
+            };
+            unary(array, op, A::DATA_TYPE)
+        }
+        TimeUnit::Microsecond => {
+            let op = |x| {
+                let datetime = timestamp_us_to_datetime(x);
+                let offset = timezone.offset_from_utc_datetime(&datetime);
+                extract(chrono::DateTime::<T>::from_utc(datetime, offset))
+            };
+            unary(array, op, A::DATA_TYPE)
+        }
+        TimeUnit::Nanosecond => {
+            let op = |x| {
+                let datetime = timestamp_ns_to_datetime(x);
+                let offset = timezone.offset_from_utc_datetime(&datetime);
+                extract(chrono::DateTime::<T>::from_utc(datetime, offset))
+            };
+            unary(array, op, A::DATA_TYPE)
+        }
+    }
+}
+
+/// Checks if an array of type `datatype` can perform year operation
+///
+/// # Examples
+/// ```
+/// use arrow2::compute::temporal::can_year;
+/// use arrow2::datatypes::{DataType};
+///
+/// let data_type = DataType::Date32;
+/// assert_eq!(can_year(&data_type), true);
+
+/// let data_type = DataType::Int8;
+/// assert_eq!(can_year(&data_type), false);
+/// ```
+pub fn can_year(data_type: &DataType) -> bool {
+    matches!(
+        data_type,
+        DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _)
+    )
+}
+
+pub fn can_month(data_type: &DataType) -> bool {
+    matches!(
+        data_type,
+        DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _)
+    )
+}
+
+pub fn can_day(data_type: &DataType) -> bool {
+    matches!(
+        data_type,
+        DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _)
+    )
+}
+
 /// Checks if an array of type `datatype` can perform hour operation
 ///
 /// # Examples
@@ -284,152 +433,3 @@ pub fn can_second(data_type: &DataType) -> bool {
             | DataType::Timestamp(_, _)
     )
 }
-
-/// Extracts the years of a temporal array as [`PrimitiveArray<i32>`].
-/// Use [`can_year`] to check if this operation is supported for the target [`DataType`].
-pub fn year(array: &dyn Array) -> Result<PrimitiveArray<i32>> {
-    match array.data_type() {
-        DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, None) => {
-            date_like(array, DataType::Int32, |x| x.year())
-        }
-        DataType::Timestamp(time_unit, Some(timezone_str)) => {
-            let time_unit = *time_unit;
-            let timezone = parse_offset(timezone_str);
-
-            let array = array.as_any().downcast_ref().unwrap();
-
-            if let Ok(timezone) = timezone {
-                Ok(extract_impl(array, time_unit, timezone, |x| x.year()))
-            } else {
-                chrono_tz(array, time_unit, timezone_str, |x| x.year())
-            }
-        }
-        dt => Err(ArrowError::NotYetImplemented(format!(
-            "\"year\" does not support type {:?}",
-            dt
-        ))),
-    }
-}
-
-/// Extracts the months of a temporal array as [`PrimitiveArray<u32>`].
-/// Use [`can_month`] to check if this operation is supported for the target [`DataType`].
-pub fn month(array: &dyn Array) -> Result<PrimitiveArray<u32>> {
-    match array.data_type() {
-        DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, None) => {
-            date_like(array, DataType::UInt32, |x| x.month())
-        }
-        DataType::Timestamp(time_unit, Some(timezone_str)) => {
-            let time_unit = *time_unit;
-            let timezone = parse_offset(timezone_str);
-
-            let array = array.as_any().downcast_ref().unwrap();
-
-            if let Ok(timezone) = timezone {
-                Ok(extract_impl(array, time_unit, timezone, |x| x.month()))
-            } else {
-                chrono_tz(array, time_unit, timezone_str, |x| x.month())
-            }
-        }
-        dt => Err(ArrowError::NotYetImplemented(format!(
-            "\"month\" does not support type {:?}",
-            dt
-        ))),
-    }
-}
-
-/// Extracts the days of a temporal array as [`PrimitiveArray<u32>`].
-/// Use [`can_day`] to check if this operation is supported for the target [`DataType`].
-pub fn day(array: &dyn Array) -> Result<PrimitiveArray<u32>> {
-    match array.data_type() {
-        DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, None) => {
-            date_like(array, DataType::UInt32, |x| x.day())
-        }
-        DataType::Timestamp(time_unit, Some(timezone_str)) => {
-            let time_unit = *time_unit;
-            let timezone = parse_offset(timezone_str);
-
-            let array = array.as_any().downcast_ref().unwrap();
-
-            if let Ok(timezone) = timezone {
-                Ok(extract_impl(array, time_unit, timezone, |x| x.day()))
-            } else {
-                chrono_tz(array, time_unit, timezone_str, |x| x.day())
-            }
-        }
-        dt => Err(ArrowError::NotYetImplemented(format!(
-            "\"day\" does not support type {:?}",
-            dt
-        ))),
-    }
-}
-
-pub fn date_like<F, O>(array: &dyn Array, data_type: DataType, op: F) -> Result<PrimitiveArray<O>>
-where
-    O: NativeType,
-    F: Fn(chrono::NaiveDateTime) -> O,
-{
-    match array.data_type() {
-        DataType::Date32 => {
-            let array = array
-                .as_any()
-                .downcast_ref::<PrimitiveArray<i32>>()
-                .unwrap();
-            Ok(unary(array, |x| op(date32_to_datetime(x)), data_type))
-        }
-        DataType::Date64 => {
-            let array = array
-                .as_any()
-                .downcast_ref::<PrimitiveArray<i64>>()
-                .unwrap();
-            Ok(unary(array, |x| op(date64_to_datetime(x)), data_type))
-        }
-        DataType::Timestamp(time_unit, None) => {
-            let array = array
-                .as_any()
-                .downcast_ref::<PrimitiveArray<i64>>()
-                .unwrap();
-            let func = match time_unit {
-                TimeUnit::Second => timestamp_s_to_datetime,
-                TimeUnit::Millisecond => timestamp_ms_to_datetime,
-                TimeUnit::Microsecond => timestamp_us_to_datetime,
-                TimeUnit::Nanosecond => timestamp_ns_to_datetime,
-            };
-            Ok(unary(array, |x| op(func(x)), data_type))
-        }
-        _ => unreachable!(),
-    }
-}
-
-/// Checks if an array of type `datatype` can perform year operation
-///
-/// # Examples
-/// ```
-/// use arrow2::compute::temporal::can_year;
-/// use arrow2::datatypes::{DataType};
-///
-/// let data_type = DataType::Date32;
-/// assert_eq!(can_year(&data_type), true);
-
-/// let data_type = DataType::Int8;
-/// assert_eq!(can_year(&data_type), false);
-/// ```
-pub fn can_year(data_type: &DataType) -> bool {
-    matches!(
-        data_type,
-        DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _)
-    )
-}
-
-pub fn can_month(data_type: &DataType) -> bool {
-    matches!(
-        data_type,
-        DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _)
-    )
-}
-
-pub fn can_day(data_type: &DataType) -> bool {
-    matches!(
-        data_type,
-        DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _)
-    )
-}