[arrow-rs] branch master updated: Add API to change timezone for timestamp array (#2347)

[viirya]

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

viirya 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 5fae29927 Add API to change timezone for timestamp array (#2347)
5fae29927 is described below

commit 5fae29927cd808ca25f1e9294085daeb41564617
Author: Liang-Chi Hsieh <vii...@gmail.com>
AuthorDate: Sun Aug 7 11:16:48 2022 -0700

    Add API to change timezone for timestamp array (#2347)
    
    * Add with_timezone API
    
    * Use into_builder()
    
    * Fix clippy
---
 arrow/src/array/array_primitive.rs | 29 +++++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/arrow/src/array/array_primitive.rs 
b/arrow/src/array/array_primitive.rs
index eb731a2b2..a10104d98 100644
--- a/arrow/src/array/array_primitive.rs
+++ b/arrow/src/array/array_primitive.rs
@@ -549,6 +549,18 @@ impl<T: ArrowTimestampType> PrimitiveArray<T> {
         let array_data = unsafe { array_data.build_unchecked() };
         PrimitiveArray::from(array_data)
     }
+
+    /// Construct a timestamp array with new timezone
+    pub fn with_timezone(&self, timezone: String) -> Self {
+        let array_data = unsafe {
+            self.data
+                .clone()
+                .into_builder()
+                .data_type(DataType::Timestamp(T::get_time_unit(), 
Some(timezone)))
+                .build_unchecked()
+        };
+        PrimitiveArray::from(array_data)
+    }
 }
 
 impl<T: ArrowTimestampType> PrimitiveArray<T> {
@@ -1099,4 +1111,21 @@ mod tests {
             BooleanArray::from(vec![true, true, true, true, true])
         );
     }
+
+    #[cfg(feature = "chrono-tz")]
+    #[test]
+    fn test_with_timezone() {
+        use crate::compute::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));
+
+        let a = a.with_timezone(String::from("America/Los_Angeles"));
+
+        let b = hour(&a).unwrap();
+        assert_eq!(2, b.value(0));
+        assert_eq!(15, b.value(1));
+    }
 }