From 561f63a232843fb36486f139faa9249652a5053c Mon Sep 17 00:00:00 2001
From: Christian Salvati <81280761+src255@users.noreply.github.com>
Date: Fri, 11 Nov 2022 14:06:14 -0500
Subject: [PATCH] Improved UX of  creating `TimestampNanosecondArray` with
 timezones (#3088)

* Make with_timezone more flexible

Make with_timezone method accept both &str and String values.

* Add alias for UTC

Add a method to PrimitiveArray<T: ArrowTimestampType> for using UTC as
the timezone.

Co-authored-by: Raphael Taylor-Davies <r.taylordavies@googlemail.com>
---
 arrow-array/src/array/primitive_array.rs | 24 ++++++++++++++++++++++--
 1 file changed, 22 insertions(+), 2 deletions(-)

diff --git a/arrow-array/src/array/primitive_array.rs b/arrow-array/src/array/primitive_array.rs
index 195e0009c0cc..34abfeb0a3de 100644
--- a/arrow-array/src/array/primitive_array.rs
+++ b/arrow-array/src/array/primitive_array.rs
@@ -801,8 +801,13 @@ impl<T: ArrowTimestampType> PrimitiveArray<T> {
     }
 
     /// Construct a timestamp array with new timezone
-    pub fn with_timezone(&self, timezone: String) -> Self {
-        self.with_timezone_opt(Some(timezone))
+    pub fn with_timezone(&self, timezone: impl Into<String>) -> Self {
+        self.with_timezone_opt(Some(timezone.into()))
+    }
+
+    /// Construct a timestamp array with UTC
+    pub fn with_timezone_utc(&self) -> Self {
+        self.with_timezone("+00:00")
     }
 
     /// Construct a timestamp array with an optional timezone
@@ -1344,6 +1349,21 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_timestamp_utc_fmt_debug() {
+        let arr: PrimitiveArray<TimestampMillisecondType> =
+            TimestampMillisecondArray::from(vec![
+                1546214400000,
+                1546214400000,
+                -1546214400000,
+            ])
+            .with_timezone_utc();
+        assert_eq!(
+            "PrimitiveArray<Timestamp(Millisecond, Some(\"+00:00\"))>\n[\n  2018-12-31T00:00:00+00:00,\n  2018-12-31T00:00:00+00:00,\n  1921-01-02T00:00:00+00:00,\n]",
+            format!("{:?}", arr)
+        );
+    }
+
     #[test]
     #[cfg(feature = "chrono-tz")]
     fn test_timestamp_with_named_tz_fmt_debug() {