Maintain time_zone in new_list (#7899)

apache · Oct 23, 2023 · f72ee5e · f72ee5e
1 parent ca5dc8c
commit f72ee5e
Showing 1 changed file with 31 additions and 7 deletions.
diff --git a/datafusion/common/src/scalar.rs b/datafusion/common/src/scalar.rs
@@ -620,26 +620,30 @@ macro_rules! build_timestamp_list {
                         TimestampSecondBuilder,
                         TimestampSecond,
                         values,
-                        $SIZE
+                        $SIZE,
+                        $TIME_ZONE
                     )
                 }
                 TimeUnit::Millisecond => build_values_list_tz!(
                     TimestampMillisecondBuilder,
                     TimestampMillisecond,
                     values,
-                    $SIZE
+                    $SIZE,
+                    $TIME_ZONE
                 ),
                 TimeUnit::Microsecond => build_values_list_tz!(
                     TimestampMicrosecondBuilder,
                     TimestampMicrosecond,
                     values,
-                    $SIZE
+                    $SIZE,
+                    $TIME_ZONE
                 ),
                 TimeUnit::Nanosecond => build_values_list_tz!(
                     TimestampNanosecondBuilder,
                     TimestampNanosecond,
                     values,
-                    $SIZE
+                    $SIZE,
+                    $TIME_ZONE
                 ),
             },
         }
@@ -683,9 +687,10 @@ macro_rules! build_values_list {
 }
 
 macro_rules! build_values_list_tz {
-    ($VALUE_BUILDER_TY:ident, $SCALAR_TY:ident, $VALUES:expr, $SIZE:expr) => {{
-        let mut builder =
-            ListBuilder::new($VALUE_BUILDER_TY::with_capacity($VALUES.len()));
+    ($VALUE_BUILDER_TY:ident, $SCALAR_TY:ident, $VALUES:expr, $SIZE:expr, $TIME_ZONE:expr) => {{
+        let mut builder = ListBuilder::new(
+            $VALUE_BUILDER_TY::with_capacity($VALUES.len()).with_timezone_opt($TIME_ZONE),
+        );
 
         for _ in 0..$SIZE {
             for scalar_value in $VALUES {
@@ -5185,6 +5190,25 @@ mod tests {
         assert_eq!(1, arr.len());
     }
 
+    #[test]
+    fn test_newlist_timestamp_zone() {
+        let s: &'static str = "UTC";
+        let values = vec![ScalarValue::TimestampMillisecond(Some(1), Some(s.into()))];
+        let arr = ScalarValue::new_list(
+            &values,
+            &DataType::Timestamp(TimeUnit::Millisecond, Some(s.into())),
+        );
+        assert_eq!(1, arr.len());
+        assert_eq!(
+            arr.data_type(),
+            &DataType::List(Arc::new(Field::new(
+                "item",
+                DataType::Timestamp(TimeUnit::Millisecond, Some(s.into())),
+                true
+            )))
+        );
+    }
+
     fn get_random_timestamps(sample_size: u64) -> Vec<ScalarValue> {
         let vector_size = sample_size;
         let mut timestamp = vec![];