From 92b1aca1fe3f4f883fe9815d2e61c26c97b9b425 Mon Sep 17 00:00:00 2001
From: Dmitry Patsura <talk@dmtry.me>
Date: Thu, 14 Jul 2022 16:21:49 +0300
Subject: [PATCH] feat(compute): Support doy (day of year) for temporal

---
 arrow/src/compute/kernels/temporal.rs | 51 ++++++++++++++++++++++++++-
 1 file changed, 50 insertions(+), 1 deletion(-)

diff --git a/arrow/src/compute/kernels/temporal.rs b/arrow/src/compute/kernels/temporal.rs
index 59388242354a..527b5be259eb 100644
--- a/arrow/src/compute/kernels/temporal.rs
+++ b/arrow/src/compute/kernels/temporal.rs
@@ -338,7 +338,7 @@ where
                 scratch
             )
         }
-        dt => return_compute_error_with!("weekday does not support", dt),
+        dt => return_compute_error_with!("weekday0 does not support", dt),
     }
 
     Ok(b.finish())
@@ -372,6 +372,35 @@ where
     Ok(b.finish())
 }
 
+/// Extracts the day of year of a given temporal array as an array of integers
+/// The day of year that ranges from 1 to 366
+pub fn doy<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(_, None) => {
+            extract_component_from_array!(array, b, ordinal, value_as_datetime)
+        }
+        &DataType::Timestamp(_, Some(ref tz)) => {
+            let mut scratch = Parsed::new();
+            extract_component_from_array!(
+                array,
+                b,
+                ordinal,
+                value_as_datetime_with_tz,
+                tz,
+                scratch
+            )
+        }
+        dt => return_compute_error_with!("doy does not support", dt),
+    }
+
+    Ok(b.finish())
+}
+
 /// Extracts the minutes of a given temporal array as an array of integers
 pub fn minute<T>(array: &PrimitiveArray<T>) -> Result<Int32Array>
 where
@@ -681,6 +710,26 @@ mod tests {
         assert_eq!(1, b.value(2));
     }
 
+    #[test]
+    fn test_temporal_array_date64_doy() {
+        //1483228800000 -> 2017-01-01 (Sunday)
+        //1514764800000 -> 2018-01-01
+        //1550636625000 -> 2019-02-20
+        let a: PrimitiveArray<Date64Type> = vec![
+            Some(1483228800000),
+            Some(1514764800000),
+            None,
+            Some(1550636625000),
+        ]
+        .into();
+
+        let b = doy(&a).unwrap();
+        assert_eq!(1, b.value(0));
+        assert_eq!(1, b.value(1));
+        assert!(!b.is_valid(2));
+        assert_eq!(51, b.value(3));
+    }
+
     #[test]
     fn test_temporal_array_timestamp_micro_year() {
         let a: TimestampMicrosecondArray =