[arrow-rs] branch master updated: feat: microsecond and millisecond (#4375)

[al...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

alamb pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new e121d5058 feat: microsecond and millisecond (#4375)
e121d5058 is described below

commit e121d5058a8f49caf51b22a9b5e93b788140715f
Author: Igor Izvekov <iz...@gmail.com>
AuthorDate: Thu Jun 8 15:56:07 2023 +0300

    feat: microsecond and millisecond (#4375)
---
 arrow-arith/src/temporal.rs | 71 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 71 insertions(+)

diff --git a/arrow-arith/src/temporal.rs b/arrow-arith/src/temporal.rs
index f62e7e9a6..0a313718c 100644
--- a/arrow-arith/src/temporal.rs
+++ b/arrow-arith/src/temporal.rs
@@ -432,6 +432,41 @@ pub fn nanosecond_dyn(array: &dyn Array) -> Result<ArrayRef, ArrowError> {
     time_fraction_dyn(array, "nanosecond", |t| t.nanosecond() as i32)
 }
 
+/// Extracts the microseconds of a given temporal primitive array as an array of integers
+pub fn microsecond<T>(array: &PrimitiveArray<T>) -> Result<Int32Array, ArrowError>
+where
+    T: ArrowTemporalType + ArrowNumericType,
+    i64: From<T::Native>,
+{
+    time_fraction_internal(array, "microsecond", |t| (t.nanosecond() / 1_000) as i32)
+}
+
+/// Extracts the microseconds of a given temporal primitive array as an array of integers.
+/// If the given array isn't temporal primitive or dictionary array,
+/// an `Err` will be returned.
+pub fn microsecond_dyn(array: &dyn Array) -> Result<ArrayRef, ArrowError> {
+    time_fraction_dyn(array, "microsecond", |t| (t.nanosecond() / 1_000) as i32)
+}
+
+/// Extracts the milliseconds of a given temporal primitive array as an array of integers
+pub fn millisecond<T>(array: &PrimitiveArray<T>) -> Result<Int32Array, ArrowError>
+where
+    T: ArrowTemporalType + ArrowNumericType,
+    i64: From<T::Native>,
+{
+    time_fraction_internal(array, "millisecond", |t| {
+        (t.nanosecond() / 1_000_000) as i32
+    })
+}
+/// Extracts the milliseconds of a given temporal primitive array as an array of integers.
+/// If the given array isn't temporal primitive or dictionary array,
+/// an `Err` will be returned.
+pub fn millisecond_dyn(array: &dyn Array) -> Result<ArrayRef, ArrowError> {
+    time_fraction_dyn(array, "millisecond", |t| {
+        (t.nanosecond() / 1_000_000) as i32
+    })
+}
+
 /// Extracts the time fraction of a given temporal array as an array of integers
 fn time_fraction_dyn<F>(
     array: &dyn Array,
@@ -1118,4 +1153,40 @@ mod tests {
         let expected = Arc::new(expected_dict) as ArrayRef;
         assert_eq!(&expected, &b);
     }
+
+    #[test]
+    fn test_temporal_array_date64_microsecond() {
+        let a: PrimitiveArray<Date64Type> = vec![None, Some(1667328721453)].into();
+
+        let b = microsecond(&a).unwrap();
+        assert!(!b.is_valid(0));
+        assert_eq!(453_000, b.value(1));
+
+        let keys = Int8Array::from(vec![Some(0_i8), Some(1), Some(1)]);
+        let dict = DictionaryArray::new(keys.clone(), Arc::new(a));
+        let b = microsecond_dyn(&dict).unwrap();
+
+        let a = Int32Array::from(vec![None, Some(453_000)]);
+        let expected_dict = DictionaryArray::new(keys, Arc::new(a));
+        let expected = Arc::new(expected_dict) as ArrayRef;
+        assert_eq!(&expected, &b);
+    }
+
+    #[test]
+    fn test_temporal_array_date64_millisecond() {
+        let a: PrimitiveArray<Date64Type> = vec![None, Some(1667328721453)].into();
+
+        let b = millisecond(&a).unwrap();
+        assert!(!b.is_valid(0));
+        assert_eq!(453, b.value(1));
+
+        let keys = Int8Array::from(vec![Some(0_i8), Some(1), Some(1)]);
+        let dict = DictionaryArray::new(keys.clone(), Arc::new(a));
+        let b = millisecond_dyn(&dict).unwrap();
+
+        let a = Int32Array::from(vec![None, Some(453)]);
+        let expected_dict = DictionaryArray::new(keys, Arc::new(a));
+        let expected = Arc::new(expected_dict) as ArrayRef;
+        assert_eq!(&expected, &b);
+    }
 }