[arrow-datafusion] branch main updated: Maintain time_zone in new_list (#7899)

[mneumann]

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

mneumann pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/main by this push:
     new f72ee5ee50 Maintain time_zone in new_list (#7899)
f72ee5ee50 is described below

commit f72ee5ee506423cccee8b5e57d79ecb84bbf1163
Author: Daniël Heres <danielhe...@gmail.com>
AuthorDate: Mon Oct 23 14:07:37 2023 +0200

    Maintain time_zone in new_list (#7899)
---
 datafusion/common/src/scalar.rs | 38 +++++++++++++++++++++++++++++++-------
 1 file changed, 31 insertions(+), 7 deletions(-)

diff --git a/datafusion/common/src/scalar.rs b/datafusion/common/src/scalar.rs
index 2c3dd4c5ca..be24e2b933 100644
--- 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![];