From 49d90c85c15adfa2705581f9d77d2dce64cc9195 Mon Sep 17 00:00:00 2001
From: Arttu Voutilainen <avoutilainen@palantir.com>
Date: Mon, 22 Jul 2024 17:01:29 +0200
Subject: [PATCH] more cleanup

---
 .../substrait/src/logical_plan/producer.rs    | 41 ++++++++-----------
 1 file changed, 16 insertions(+), 25 deletions(-)

diff --git a/datafusion/substrait/src/logical_plan/producer.rs b/datafusion/substrait/src/logical_plan/producer.rs
index d5a5052cbadf..ae48768312a2 100644
--- a/datafusion/substrait/src/logical_plan/producer.rs
+++ b/datafusion/substrait/src/logical_plan/producer.rs
@@ -1410,41 +1410,31 @@ fn to_substrait_type(
                 nullability,
             })),
         }),
-        DataType::Timestamp(unit, None) => {
+        DataType::Timestamp(unit, tz) => {
             let precision = match unit {
                 TimeUnit::Second => 0,
                 TimeUnit::Millisecond => 3,
                 TimeUnit::Microsecond => 6,
                 TimeUnit::Nanosecond => 9,
             };
-            Ok(substrait::proto::Type {
-                kind: Some(r#type::Kind::PrecisionTimestamp(
-                    r#type::PrecisionTimestamp {
+            let kind = match tz {
+                None => r#type::Kind::PrecisionTimestamp(r#type::PrecisionTimestamp {
+                    type_variation_reference: DEFAULT_TYPE_VARIATION_REF,
+                    nullability,
+                    precision,
+                }),
+                Some(_) => {
+                    // If timezone is present, no matter what the actual tz value is, it indicates the
+                    // value of the timestamp is tied to UTC epoch. That's all that Substrait cares about.
+                    // As the timezone is lost, this conversion may be lossy for downstream use of the value.
+                    r#type::Kind::PrecisionTimestampTz(r#type::PrecisionTimestampTz {
                         type_variation_reference: DEFAULT_TYPE_VARIATION_REF,
                         nullability,
                         precision,
-                    },
-                )),
-            })
-        }
-        DataType::Timestamp(unit, Some(_)) => {
-            // If timezone is present, no matter what the actual tz value is, it indicates the
-            // value of the timestamp is tied to UTC epoch. That's all that Substrait cares about.
-            let precision = match unit {
-                TimeUnit::Second => 0,
-                TimeUnit::Millisecond => 3,
-                TimeUnit::Microsecond => 6,
-                TimeUnit::Nanosecond => 9,
+                    })
+                }
             };
-            Ok(substrait::proto::Type {
-                kind: Some(r#type::Kind::PrecisionTimestampTz(
-                    r#type::PrecisionTimestampTz {
-                        type_variation_reference: DEFAULT_TYPE_VARIATION_REF,
-                        nullability,
-                        precision,
-                    },
-                )),
-            })
+            Ok(substrait::proto::Type { kind: Some(kind) })
         }
         DataType::Date32 => Ok(substrait::proto::Type {
             kind: Some(r#type::Kind::Date(r#type::Date {
@@ -1883,6 +1873,7 @@ fn to_substrait_literal(
         ),
         // If timezone is present, no matter what the actual tz value is, it indicates the
         // value of the timestamp is tied to UTC epoch. That's all that Substrait cares about.
+        // As the timezone is lost, this conversion may be lossy for downstream use of the value.
         ScalarValue::TimestampSecond(Some(t), Some(_)) => (
             LiteralType::PrecisionTimestampTz(PrecisionTimestamp {
                 precision: 0,