Add with_timezone API

apache · Aug 7, 2022 · c4d393d · c4d393d
1 parent 30c94db
commit c4d393d
Showing 1 changed file with 30 additions and 1 deletion.
diff --git a/arrow/src/array/array_primitive.rs b/arrow/src/array/array_primitive.rs
@@ -549,6 +549,19 @@ 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 =
+            ArrayData::builder(DataType::Timestamp(T::get_time_unit(), Some(timezone)))
+                .len(self.len())
+                .null_count(self.data.null_count())
+                .offset(self.data.offset())
+                .add_buffer(self.data.buffers()[0].clone())
+                .null_bit_buffer(self.data.null_buffer().map(|b| b.clone()));
+        let array_data = unsafe { array_data.build_unchecked() };
+        PrimitiveArray::from(array_data)
+    }
 }
 
 impl<T: ArrowTimestampType> PrimitiveArray<T> {
@@ -605,7 +618,7 @@ mod tests {
     use std::thread;
 
     use crate::buffer::Buffer;
-    use crate::compute::eq_dyn;
+    use crate::compute::{eq_dyn, hour};
     use crate::datatypes::DataType;
 
     #[test]
@@ -1099,4 +1112,20 @@ mod tests {
             BooleanArray::from(vec![true, true, true, true, true])
         );
     }
+
+    #[cfg(feature = "chrono-tz")]
+    #[test]
+    fn test_with_timezone() {
+        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));
+    }
 }