From 56d34ece5719d537889857e01e6f3d6e03631da8 Mon Sep 17 00:00:00 2001
From: "Jorge C. Leitao" <jorgecarleitao@gmail.com>
Date: Fri, 3 Sep 2021 18:34:06 +0000
Subject: [PATCH] Fixed display of timestamp with tz.

---
 src/array/display.rs            | 24 ++++++++++++------------
 src/temporal_conversions.rs     | 33 +++++++++++++++++++++++++++++++--
 tests/it/array/primitive/mod.rs | 12 ++++++++++++
 tests/it/io/print.rs            | 12 ++++++------
 4 files changed, 61 insertions(+), 20 deletions(-)

diff --git a/src/array/display.rs b/src/array/display.rs
index ed403d33cba..059c5e56feb 100644
--- a/src/array/display.rs
+++ b/src/array/display.rs
@@ -68,11 +68,11 @@ pub fn get_value_display<'a>(array: &'a dyn Array) -> Box<dyn Fn(usize) -> Strin
         Time64(_) => unreachable!(), // remaining are not valid
         Timestamp(TimeUnit::Second, tz) => {
             if let Some(tz) = tz {
+                let offset = temporal_conversions::parse_offset(tz).unwrap();
                 dyn_primitive!(array, i64, |x| {
-                    format!(
-                        "{} {}",
+                    chrono::DateTime::<chrono::FixedOffset>::from_utc(
                         temporal_conversions::timestamp_s_to_datetime(x),
-                        tz
+                        offset,
                     )
                 })
             } else {
@@ -81,11 +81,11 @@ pub fn get_value_display<'a>(array: &'a dyn Array) -> Box<dyn Fn(usize) -> Strin
         }
         Timestamp(TimeUnit::Millisecond, tz) => {
             if let Some(tz) = tz {
+                let offset = temporal_conversions::parse_offset(tz).unwrap();
                 dyn_primitive!(array, i64, |x| {
-                    format!(
-                        "{} {}",
+                    chrono::DateTime::<chrono::FixedOffset>::from_utc(
                         temporal_conversions::timestamp_ms_to_datetime(x),
-                        tz
+                        offset,
                     )
                 })
             } else {
@@ -94,11 +94,11 @@ pub fn get_value_display<'a>(array: &'a dyn Array) -> Box<dyn Fn(usize) -> Strin
         }
         Timestamp(TimeUnit::Microsecond, tz) => {
             if let Some(tz) = tz {
+                let offset = temporal_conversions::parse_offset(tz).unwrap();
                 dyn_primitive!(array, i64, |x| {
-                    format!(
-                        "{} {}",
+                    chrono::DateTime::<chrono::FixedOffset>::from_utc(
                         temporal_conversions::timestamp_us_to_datetime(x),
-                        tz
+                        offset,
                     )
                 })
             } else {
@@ -107,11 +107,11 @@ pub fn get_value_display<'a>(array: &'a dyn Array) -> Box<dyn Fn(usize) -> Strin
         }
         Timestamp(TimeUnit::Nanosecond, tz) => {
             if let Some(tz) = tz {
+                let offset = temporal_conversions::parse_offset(tz).unwrap();
                 dyn_primitive!(array, i64, |x| {
-                    format!(
-                        "{} {}",
+                    chrono::DateTime::<chrono::FixedOffset>::from_utc(
                         temporal_conversions::timestamp_ns_to_datetime(x),
-                        tz
+                        offset,
                     )
                 })
             } else {
diff --git a/src/temporal_conversions.rs b/src/temporal_conversions.rs
index 47b25cfea5f..ae83469fd8f 100644
--- a/src/temporal_conversions.rs
+++ b/src/temporal_conversions.rs
@@ -1,7 +1,11 @@
 //! Conversion methods for dates and times.
 
-use crate::datatypes::TimeUnit;
-use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
+use chrono::{FixedOffset, NaiveDate, NaiveDateTime, NaiveTime};
+
+use crate::{
+    datatypes::TimeUnit,
+    error::{ArrowError, Result},
+};
 
 /// Number of seconds in a day
 pub const SECONDS_IN_DAY: i64 = 86_400;
@@ -149,3 +153,28 @@ pub fn timeunit_scale(a: &TimeUnit, b: &TimeUnit) -> f64 {
         (TimeUnit::Nanosecond, TimeUnit::Nanosecond) => 1.0,
     }
 }
+
+pub(crate) fn parse_offset(offset: &str) -> Result<FixedOffset> {
+    if offset == "UTC" {
+        return Ok(FixedOffset::east(0));
+    }
+    let mut a = offset.split(':');
+    let first = a.next().map(Ok).unwrap_or_else(|| {
+        Err(ArrowError::InvalidArgumentError(
+            "timezone offset must be of the form [-]00:00".to_string(),
+        ))
+    })?;
+    let last = a.next().map(Ok).unwrap_or_else(|| {
+        Err(ArrowError::InvalidArgumentError(
+            "timezone offset must be of the form [-]00:00".to_string(),
+        ))
+    })?;
+    let hours: i32 = first.parse().map_err(|_| {
+        ArrowError::InvalidArgumentError("timezone offset must be of the form [-]00:00".to_string())
+    })?;
+    let minutes: i32 = last.parse().map_err(|_| {
+        ArrowError::InvalidArgumentError("timezone offset must be of the form [-]00:00".to_string())
+    })?;
+
+    Ok(FixedOffset::east(hours * 60 * 60 + minutes * 60))
+}
diff --git a/tests/it/array/primitive/mod.rs b/tests/it/array/primitive/mod.rs
index 0dc9cab6ffc..344cdf90b91 100644
--- a/tests/it/array/primitive/mod.rs
+++ b/tests/it/array/primitive/mod.rs
@@ -181,6 +181,18 @@ fn display_timestamp_ns() {
     );
 }
 
+#[test]
+fn display_timestamp_tz_ns() {
+    let array = Int64Array::from(&[Some(1), None, Some(2)]).to(DataType::Timestamp(
+        TimeUnit::Nanosecond,
+        Some("+02:00".to_string()),
+    ));
+    assert_eq!(
+        format!("{}", array),
+        "Timestamp(Nanosecond, Some(\"+02:00\"))[1970-01-01 02:00:00.000000001 +02:00, , 1970-01-01 02:00:00.000000002 +02:00]"
+    );
+}
+
 #[test]
 fn display_duration_ms() {
     let array =
diff --git a/tests/it/io/print.rs b/tests/it/io/print.rs
index 040e7c4b555..bce6750ee7f 100644
--- a/tests/it/io/print.rs
+++ b/tests/it/io/print.rs
@@ -185,12 +185,12 @@ fn write_timestamp_second() {
 #[test]
 fn write_timestamp_second_with_tz() {
     let expected = vec![
-        "+-------------------------+",
-        "| f                       |",
-        "+-------------------------+",
-        "| 1970-05-09 14:25:11 UTC |",
-        "|                         |",
-        "+-------------------------+",
+        "+----------------------------+",
+        "| f                          |",
+        "+----------------------------+",
+        "| 1970-05-09 14:25:11 +00:00 |",
+        "|                            |",
+        "+----------------------------+",
     ];
     check_datetime!(
         i64,