add test for take preseving timezone

apache · Dec 17, 2021 · 01f5c54 · 01f5c54
1 parent 117759c
commit 01f5c54
Showing 1 changed file with 17 additions and 0 deletions.
diff --git a/arrow/src/compute/kernels/take.rs b/arrow/src/compute/kernels/take.rs
@@ -1241,6 +1241,23 @@ mod tests {
         .unwrap();
     }
 
+    #[test]
+    fn test_take_preserve_timezone() {
+        let index = Int64Array::from(vec![Some(0), None]);
+
+        let input = TimestampNanosecondArray::from_vec(
+            vec![1_639_715_368_000_000_000, 1_639_715_368_000_000_000],
+            Some("UTC".to_owned()),
+        );
+        let result = take_impl(&input, &index, None).unwrap();
+        match result.data_type() {
+            DataType::Timestamp(TimeUnit::Nanosecond, tz) => {
+                assert_eq!(tz.clone(), Some("UTC".to_owned()))
+            }
+            _ => panic!(),
+        }
+    }
+
     #[test]
     fn test_take_impl_primitive_with_int64_indices() {
         let index = Int64Array::from(vec![Some(3), None, Some(1), Some(3), Some(2)]);