ARROW-11439: [Rust] Add year support to temporal kernels

rust/arrow/src/compute/kernels/temporal.rs

@@ -17,34 +17,82 @@
 
 //! Defines temporal kernels for time and date related functions.
 
-use chrono::Timelike;
+use chrono::{Datelike, Timelike};
 
 use crate::array::*;
 use crate::datatypes::*;
-use crate::error::Result;
-
+use crate::error::{ArrowError, Result};
 /// Extracts the hours of a given temporal array as an array of integers
 pub fn hour<T>(array: &PrimitiveArray<T>) -> Result<Int32Array>
 where
     T: ArrowTemporalType + ArrowNumericType,
     i64: std::convert::From<T::Native>,
 {
     let mut b = Int32Builder::new(array.len());
-    for i in 0..array.len() {
-        if array.is_null(i) {
-            b.append_null()?;
-        } else {
-            match array.data_type() {
-                &DataType::Time32(_) | &DataType::Time64(_) => {
+    match array.data_type() {
+        &DataType::Time32(_) | &DataType::Time64(_) => {
+            for i in 0..array.len() {
+                if array.is_null(i) {
+                    b.append_null()?;
+                } else {
                     match array.value_as_time(i) {
                         Some(time) => b.append_value(time.hour() as i32)?,
                         None => b.append_null()?,
+                    };
+                }
+            }
+        }
+        &DataType::Date32 | &DataType::Date64 | &DataType::Timestamp(_, _) => {
+            for i in 0..array.len() {
+                if array.is_null(i) {
+                    b.append_null()?;
+                } else {
+                    match array.value_as_datetime(i) {
+                        Some(dt) => b.append_value(dt.hour() as i32)?,
+                        None => b.append_null()?,
                     }
                 }
-                _ => match array.value_as_datetime(i) {
-                    Some(dt) => b.append_value(dt.hour() as i32)?,
-                    None => b.append_null()?,
-                },
+            }
+        }
+        dt => {
+            return {
+                Err(ArrowError::ComputeError(format!(
+                    "hour does not support type {:?}",
+                    dt
+                )))
+            }
+        }
+    }
+
+    Ok(b.finish())
+}
+
+/// Extracts the years of a given temporal array as an array of integers
+pub fn year<T>(array: &PrimitiveArray<T>) -> Result<Int32Array>
+where
+    T: ArrowTemporalType + ArrowNumericType,
+    i64: std::convert::From<T::Native>,
+{
+    let mut b = Int32Builder::new(array.len());
+    match array.data_type() {
+        &DataType::Date32 | &DataType::Date64 | &DataType::Timestamp(_, _) => {
+            for i in 0..array.len() {
+                if array.is_null(i) {
+                    b.append_null()?;
+                } else {
+                    match array.value_as_datetime(i) {
+                        Some(dt) => b.append_value(dt.year() as i32)?,
+                        None => b.append_null()?,
+                    }
+                }
+            }
+        }
+        dt => {
+            return {
+                Err(ArrowError::ComputeError(format!(
+                    "year does not support type {:?}",
+                    dt
+                )))
             }
         }
     }
@@ -61,20 +109,79 @@ mod tests {
         let a: PrimitiveArray<Date64Type> =
             vec![Some(1514764800000), None, Some(1550636625000)].into();
 
-        // get hour from temporal
         let b = hour(&a).unwrap();
         assert_eq!(0, b.value(0));
         assert_eq!(false, b.is_valid(1));
         assert_eq!(4, b.value(2));
     }
 
+    #[test]
+    fn test_temporal_array_date32_hour() {
+        let a: PrimitiveArray<Date32Type> = vec![Some(15147), None, Some(15148)].into();
+
+        let b = hour(&a).unwrap();
+        assert_eq!(0, b.value(0));
+        assert_eq!(false, b.is_valid(1));
+        assert_eq!(0, b.value(2));
+    }
+
     #[test]
     fn test_temporal_array_time32_second_hour() {
         let a: PrimitiveArray<Time32SecondType> = vec![37800, 86339].into();
 
-        // get hour from temporal
         let b = hour(&a).unwrap();
         assert_eq!(10, b.value(0));
         assert_eq!(23, b.value(1));
     }
+
+    #[test]
+    fn test_temporal_array_time64_micro_hour() {
+        let a: PrimitiveArray<Time64MicrosecondType> =
+            vec![37800000000, 86339000000].into();
+
+        let b = hour(&a).unwrap();
+        assert_eq!(10, b.value(0));
+        assert_eq!(23, b.value(1));
+    }
+
+    #[test]
+    fn test_temporal_array_timestamp_micro_hour() {
+        let a: TimestampMicrosecondArray = vec![37800000000, 86339000000].into();
+
+        let b = hour(&a).unwrap();
+        assert_eq!(10, b.value(0));
+        assert_eq!(23, b.value(1));
+    }
+
+    #[test]
+    fn test_temporal_array_date64_year() {
+        let a: PrimitiveArray<Date64Type> =
+            vec![Some(1514764800000), None, Some(1550636625000)].into();
+
+        let b = year(&a).unwrap();
+        assert_eq!(2018, b.value(0));
+        assert_eq!(false, b.is_valid(1));
+        assert_eq!(2019, b.value(2));
+    }
+
+    #[test]
+    fn test_temporal_array_date32_year() {
+        let a: PrimitiveArray<Date32Type> = vec![Some(15147), None, Some(15448)].into();
+
+        let b = year(&a).unwrap();
+        assert_eq!(2011, b.value(0));
+        assert_eq!(false, b.is_valid(1));
+        assert_eq!(2012, b.value(2));
+    }
+
+    #[test]
+    fn test_temporal_array_timestamp_micro_year() {
+        let a: TimestampMicrosecondArray =
+            vec![Some(1612025847000000), None, Some(1722015847000000)].into();
+
+        let b = year(&a).unwrap();
+        assert_eq!(2021, b.value(0));
+        assert_eq!(false, b.is_valid(1));
+        assert_eq!(2024, b.value(2));
+    }
 }