implement day

apache · Mar 3, 2022 · b5d97a5 · b5d97a5
1 parent 11c3515
commit b5d97a5
Showing 1 changed file with 98 additions and 28 deletions.
diff --git a/arrow/src/compute/kernels/temporal.rs b/arrow/src/compute/kernels/temporal.rs
@@ -121,66 +121,94 @@ pub fn using_chrono_tz_and_utc_naive_date_time(
         .ok()
 }
 
-/// Extracts the hours of a given temporal array as an array of integers
-pub fn hour<T>(array: &PrimitiveArray<T>) -> Result<Int32Array>
+/// 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::Time32(_) | &DataType::Time64(_) => {
-            extract_component_from_array!(array, b, hour, value_as_time)
-        }
-        &DataType::Date32 | &DataType::Date64 | &DataType::Timestamp(_, None) => {
-            extract_component_from_array!(array, b, hour, value_as_datetime)
-        }
-        &DataType::Timestamp(_, Some(ref tz)) => {
-            let mut scratch = Parsed::new();
-            extract_component_from_array!(
-                array,
-                b,
-                hour,
-                value_as_datetime_with_tz,
-                tz,
-                scratch
-            )
+        &DataType::Date32 | &DataType::Date64 | &DataType::Timestamp(_, _) => {
+            extract_component_from_array!(array, b, year, value_as_datetime)
         }
-        dt => return_compute_error_with!("hour does not support", dt),
+        dt => return_compute_error_with!("year does not support", 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>
+/// Extracts the month of a given temporal array as an array of integers
+pub fn month<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(_, _) => {
-            extract_component_from_array!(array, b, year, value_as_datetime)
+            extract_component_from_array!(array, b, month, value_as_datetime)
         }
-        dt => return_compute_error_with!("year does not support", dt),
+        dt => return_compute_error_with!("month does not support", dt),
     }
 
     Ok(b.finish())
 }
 
-/// Extracts the month of a given temporal array as an array of integers
-pub fn month<T>(array: &PrimitiveArray<T>) -> Result<Int32Array>
+/// Extracts the day of a given temporal array as an array of integers
+pub fn day<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(_, _) => {
-            extract_component_from_array!(array, b, month, value_as_datetime)
+        &DataType::Date32 | &DataType::Date64 | &DataType::Timestamp(_, None) => {
+            extract_component_from_array!(array, b, day, value_as_datetime)
         }
-        dt => return_compute_error_with!("year does not support", dt),
+        &DataType::Timestamp(_, Some(ref tz)) => {
+            let mut scratch = Parsed::new();
+            extract_component_from_array!(
+                array,
+                b,
+                day,
+                value_as_datetime_with_tz,
+                tz,
+                scratch
+            )
+        }
+        dt => return_compute_error_with!("day does not support", dt),
+    }
+
+    Ok(b.finish())
+}
+
+/// 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());
+    match array.data_type() {
+        &DataType::Time32(_) | &DataType::Time64(_) => {
+            extract_component_from_array!(array, b, hour, value_as_time)
+        }
+        &DataType::Date32 | &DataType::Date64 | &DataType::Timestamp(_, None) => {
+            extract_component_from_array!(array, b, hour, value_as_datetime)
+        }
+        &DataType::Timestamp(_, Some(ref tz)) => {
+            let mut scratch = Parsed::new();
+            extract_component_from_array!(
+                array,
+                b,
+                hour,
+                value_as_datetime_with_tz,
+                tz,
+                scratch
+            )
+        }
+        dt => return_compute_error_with!("hour does not support", dt),
     }
 
     Ok(b.finish())
@@ -341,6 +369,48 @@ mod tests {
         assert_eq!(2, b.value(2));
     }
 
+    #[test]
+    fn test_temporal_array_timestamp_day_with_timezone() {
+        use std::sync::Arc;
+
+        // 24 * 60 * 60 = 8640
+        let a = Arc::new(TimestampSecondArray::from_vec(
+            vec![86400],
+            Some("+00:00".to_string()),
+        ));
+        let b = day(&a).unwrap();
+        assert_eq!(2, b.value(0));
+        let a = Arc::new(TimestampSecondArray::from_vec(
+            vec![86400],
+            Some("-10:00".to_string()),
+        ));
+        let b = day(&a).unwrap();
+        assert_eq!(1, b.value(0));
+    }
+
+    #[test]
+    fn test_temporal_array_date64_day() {
+        //1514764800000 -> 2018-01-01
+        //1550636625000 -> 2019-02-20
+        let a: PrimitiveArray<Date64Type> =
+            vec![Some(1514764800000), None, Some(1550636625000)].into();
+
+        let b = day(&a).unwrap();
+        assert_eq!(1, b.value(0));
+        assert!(!b.is_valid(1));
+        assert_eq!(20, b.value(2));
+    }
+
+    #[test]
+    fn test_temporal_array_date32_day() {
+        let a: PrimitiveArray<Date32Type> = vec![Some(0), None, Some(31)].into();
+
+        let b = day(&a).unwrap();
+        assert_eq!(1, b.value(0));
+        assert!(!b.is_valid(1));
+        assert_eq!(1, b.value(2));
+    }
+
     #[test]
     fn test_temporal_array_timestamp_micro_year() {
         let a: TimestampMicrosecondArray =