add constructors for timestamps with timezones

Author: nevi-me

rust/arrow/src/array/array.rs

@@ -621,6 +621,51 @@ def_numeric_from_vec!(
     DataType::Time64(TimeUnit::Nanosecond)
 );
 
+impl<T: ArrowTimestampType> PrimitiveArray<T> {
+    pub fn from_vec(data: Vec<i64>, timezone: Option<Arc<String>>) -> Self {
+        let array_data =
+            ArrayData::builder(DataType::Timestamp(T::get_time_unit(), timezone))
+                .len(data.len())
+                .add_buffer(Buffer::from(data.to_byte_slice()))
+                .build();
+        PrimitiveArray::from(array_data)
+    }
+}
+
+impl<T: ArrowTimestampType> PrimitiveArray<T> {
+    /// Construct a timestamp array from a vec of Option<i64> values and an optional timezone
+    pub fn from_opt_vec(data: Vec<Option<i64>>, timezone: Option<Arc<String>>) -> Self {
+        // TODO: duplicated from def_numeric_from_vec! macro, it looks possible to convert to generic
+        let data_len = data.len();
+        let num_bytes = bit_util::ceil(data_len, 8);
+        let mut null_buf = MutableBuffer::new(num_bytes).with_bitset(num_bytes, false);
+        let mut val_buf = MutableBuffer::new(data_len * mem::size_of::<i64>());
+
+        {
+            let null = vec![0; mem::size_of::<i64>()];
+            let null_slice = null_buf.data_mut();
+            for (i, v) in data.iter().enumerate() {
+                if let Some(n) = v {
+                    bit_util::set_bit(null_slice, i);
+                    // unwrap() in the following should be safe here since we've
+                    // made sure enough space is allocated for the values.
+                    val_buf.write(&n.to_byte_slice()).unwrap();
+                } else {
+                    val_buf.write(&null).unwrap();
+                }
+            }
+        }
+
+        let array_data =
+            ArrayData::builder(DataType::Timestamp(T::get_time_unit(), timezone))
+                .len(data_len)
+                .add_buffer(val_buf.freeze())
+                .null_bit_buffer(null_buf.freeze())
+                .build();
+        PrimitiveArray::from(array_data)
+    }
+}
+
 /// Constructs a boolean array from a vector. Should only be used for testing.
 impl From<Vec<bool>> for BooleanArray {
     fn from(data: Vec<bool>) -> Self {

rust/arrow/src/compute/kernels/cast.rs

@@ -1062,11 +1062,10 @@ mod tests {
 
     #[test]
     fn test_cast_timestamp_to_date32() {
-        let a = TimestampMillisecondArray::from(vec![
-            Some(864000000005),
-            Some(1545696000001),
-            None,
-        ]);
+        let a = TimestampMillisecondArray::from_opt_vec(
+            vec![Some(864000000005), Some(1545696000001), None],
+            Some(Arc::new(String::from("UTC"))),
+        );
         let array = Arc::new(a) as ArrayRef;
         let b = cast(&array, &DataType::Date32(DateUnit::Day)).unwrap();
         let c = b.as_any().downcast_ref::<Date32Array>().unwrap();

rust/arrow/src/datatypes.rs

@@ -28,6 +28,7 @@ use std::mem::size_of;
 use std::ops::{Add, Div, Mul, Sub};
 use std::slice::from_raw_parts;
 use std::str::FromStr;
+use std::sync::Arc;
 
 #[cfg(feature = "simd")]
 use packed_simd::*;
@@ -37,7 +38,6 @@ use serde_json::{
 };
 
 use crate::error::{ArrowError, Result};
-use std::sync::Arc;
 
 /// The possible relative types that are supported.
 ///
@@ -65,7 +65,7 @@ pub enum DataType {
     Float32,
     Float64,
     /// A timestamp with an optional timezone
-    Timestamp(TimeUnit, Option<String>),
+    Timestamp(TimeUnit, Option<Arc<String>>),
     Date32(DateUnit),
     Date64(DateUnit),
     Time32(TimeUnit),
@@ -516,6 +516,32 @@ impl ArrowTemporalType for Time64NanosecondType {}
 impl ArrowTemporalType for IntervalYearMonthType {}
 impl ArrowTemporalType for IntervalDayTimeType {}
 
+/// A timestamp type allows us to create array builders that take a timestamp
+pub trait ArrowTimestampType: ArrowTemporalType {
+    fn get_time_unit() -> TimeUnit;
+}
+
+impl ArrowTimestampType for TimestampSecondType {
+    fn get_time_unit() -> TimeUnit {
+        TimeUnit::Second
+    }
+}
+impl ArrowTimestampType for TimestampMillisecondType {
+    fn get_time_unit() -> TimeUnit {
+        TimeUnit::Millisecond
+    }
+}
+impl ArrowTimestampType for TimestampMicrosecondType {
+    fn get_time_unit() -> TimeUnit {
+        TimeUnit::Microsecond
+    }
+}
+impl ArrowTimestampType for TimestampNanosecondType {
+    fn get_time_unit() -> TimeUnit {
+        TimeUnit::Nanosecond
+    }
+}
+
 /// Allows conversion from supported Arrow types to a byte slice.
 pub trait ToByteSlice {
     /// Converts this instance into a byte slice
@@ -574,7 +600,7 @@ impl DataType {
                     };
                     let tz = match map.get("timezone") {
                         None => Ok(None),
-                        Some(VString(tz)) => Ok(Some(tz.to_string())),
+                        Some(VString(tz)) => Ok(Some(Arc::new(tz.to_string()))),
                         _ => Err(ArrowError::ParseError(
                             "timezone must be a string".to_string(),
                         )),
@@ -1280,14 +1306,17 @@ mod tests {
                 Field::new("c15", DataType::Timestamp(TimeUnit::Second, None), false),
                 Field::new(
                     "c16",
-                    DataType::Timestamp(TimeUnit::Millisecond, Some("UTC".to_string())),
+                    DataType::Timestamp(
+                        TimeUnit::Millisecond,
+                        Some(Arc::new("UTC".to_string())),
+                    ),
                     false,
                 ),
                 Field::new(
                     "c17",
                     DataType::Timestamp(
                         TimeUnit::Microsecond,
-                        Some("Africa/Johannesburg".to_string()),
+                        Some(Arc::new("Africa/Johannesburg".to_string())),
                     ),
                     false,
                 ),

rust/arrow/src/ipc/convert.rs

@@ -24,6 +24,7 @@ use flatbuffers::{
     FlatBufferBuilder, ForwardsUOffset, UnionWIPOffset, Vector, WIPOffset,
 };
 use std::collections::HashMap;
+use std::sync::Arc;
 
 /// Serialize a schema in IPC format
 fn schema_to_fb(schema: &Schema) -> FlatBufferBuilder {
@@ -166,7 +167,8 @@ fn get_data_type(field: ipc::Field) -> DataType {
         }
         ipc::Type::Timestamp => {
             let timestamp = field.type_as_timestamp().unwrap();
-            let timezone: Option<String> = timestamp.timezone().map(|tz| tz.to_string());
+            let timezone: Option<Arc<String>> =
+                timestamp.timezone().map(|tz| Arc::new(tz.to_string()));
             match timestamp.unit() {
                 ipc::TimeUnit::SECOND => DataType::Timestamp(TimeUnit::Second, timezone),
                 ipc::TimeUnit::MILLISECOND => {
@@ -308,7 +310,7 @@ fn get_fb_field_type<'a: 'b, 'b>(
             (ipc::Type::Time, builder.finish().as_union_value(), None)
         }
         Timestamp(unit, tz) => {
-            let tz = tz.clone().unwrap_or(String::new());
+            let tz = tz.clone().unwrap_or(Arc::new(String::new()));
             let tz_str = fbb.create_string(tz.as_str());
             let mut builder = ipc::TimestampBuilder::new(&mut fbb);
             let time_unit = match unit {

rust/arrow/src/util/integration_util.rs

@@ -378,6 +378,10 @@ mod tests {
 
     #[test]
     fn test_arrow_data_equality() {
+        let secs_tz = Some(Arc::new("Europe/Budapest".to_string()));
+        let millis_tz = Some(Arc::new("America/New_York".to_string()));
+        let micros_tz = Some(Arc::new("UTC".to_string()));
+        let nanos_tz = Some(Arc::new("Africa/Johannesburg".to_string()));
         let schema = Schema::new(vec![
             Field::new("bools", DataType::Boolean, true),
             Field::new("int8s", DataType::Int8, true),
@@ -412,6 +416,26 @@ mod tests {
                 DataType::Timestamp(TimeUnit::Nanosecond, None),
                 true,
             ),
+            Field::new(
+                "ts_secs_tz",
+                DataType::Timestamp(TimeUnit::Second, secs_tz.clone()),
+                true,
+            ),
+            Field::new(
+                "ts_millis_tz",
+                DataType::Timestamp(TimeUnit::Millisecond, millis_tz.clone()),
+                true,
+            ),
+            Field::new(
+                "ts_micros_tz",
+                DataType::Timestamp(TimeUnit::Microsecond, micros_tz.clone()),
+                true,
+            ),
+            Field::new(
+                "ts_nanos_tz",
+                DataType::Timestamp(TimeUnit::Nanosecond, nanos_tz.clone()),
+                true,
+            ),
             Field::new("utf8s", DataType::Utf8, true),
             Field::new("lists", DataType::List(Box::new(DataType::Int32)), true),
             Field::new(
@@ -464,6 +488,20 @@ mod tests {
         let ts_micros = TimestampMicrosecondArray::from(vec![None, None, None]);
         let ts_nanos =
             TimestampNanosecondArray::from(vec![None, None, Some(-6473623571954960143)]);
+        let ts_secs_tz = TimestampSecondArray::from_opt_vec(
+            vec![None, Some(193438817552), None],
+            secs_tz,
+        );
+        let ts_millis_tz = TimestampMillisecondArray::from_opt_vec(
+            vec![None, Some(38606916383008), Some(58113709376587)],
+            millis_tz,
+        );
+        let ts_micros_tz =
+            TimestampMicrosecondArray::from_opt_vec(vec![None, None, None], micros_tz);
+        let ts_nanos_tz = TimestampNanosecondArray::from_opt_vec(
+            vec![None, None, Some(-6473623571954960143)],
+            nanos_tz,
+        );
         let utf8s = StringArray::try_from(vec![Some("aa"), None, Some("bbb")]).unwrap();
 
         let value_data = Int32Array::from(vec![None, Some(2), None, None]);
@@ -514,6 +552,10 @@ mod tests {
                 Arc::new(ts_millis),
                 Arc::new(ts_micros),
                 Arc::new(ts_nanos),
+                Arc::new(ts_secs_tz),
+                Arc::new(ts_millis_tz),
+                Arc::new(ts_micros_tz),
+                Arc::new(ts_nanos_tz),
                 Arc::new(utf8s),
                 Arc::new(lists),
                 Arc::new(structs),

rust/arrow/test/data/integration.json

@@ -201,6 +201,46 @@
         "nullable": true,
         "children": []
       },
+      {
+        "name": "ts_secs_tz",
+        "type": {
+          "name": "timestamp",
+          "unit": "SECOND",
+          "timezone": "Europe/Budapest"
+        },
+        "nullable": true,
+        "children": []
+      },
+      {
+        "name": "ts_millis_tz",
+        "type": {
+          "name": "timestamp",
+          "unit": "MILLISECOND",
+          "timezone": "America/New_York"
+        },
+        "nullable": true,
+        "children": []
+      },
+      {
+        "name": "ts_micros_tz",
+        "type": {
+          "name": "timestamp",
+          "unit": "MICROSECOND",
+          "timezone": "UTC"
+        },
+        "nullable": true,
+        "children": []
+      },
+      {
+        "name": "ts_nanos_tz",
+        "type": {
+          "name": "timestamp",
+          "unit": "NANOSECOND",
+          "timezone": "Africa/Johannesburg"
+        },
+        "nullable": true,
+        "children": []
+      },
       {
         "name": "utf8s",
         "type": {
@@ -555,6 +595,62 @@
             -6473623571954960143
           ]
         },
+        {
+          "name": "ts_secs_tz",
+          "count": 3,
+          "VALIDITY": [
+            0,
+            1,
+            0
+          ],
+          "DATA": [
+            209869064422,
+            193438817552,
+            51757838205
+          ]
+        },
+        {
+          "name": "ts_millis_tz",
+          "count": 3,
+          "VALIDITY": [
+            0,
+            1,
+            1
+          ],
+          "DATA": [
+            228315043570185,
+            38606916383008,
+            58113709376587
+          ]
+        },
+        {
+          "name": "ts_micros_tz",
+          "count": 3,
+          "VALIDITY": [
+            0,
+            0,
+            0
+          ],
+          "DATA": [
+            133457416537791415,
+            129522736067409280,
+            177110451066832967
+          ]
+        },
+        {
+          "name": "ts_nanos_tz",
+          "count": 3,
+          "VALIDITY": [
+            0,
+            0,
+            1
+          ],
+          "DATA": [
+            -804525722984600007,
+            8166038652634779458,
+            -6473623571954960143
+          ]
+        },
         {
           "name": "utf8s",
           "count": 3,