Feat(spark): from_utc_timestamp -> at time zone

tobymao · Oct 5, 2023 · ae27e46 · ae27e46
1 parent e947b9c
commit ae27e46
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 1 deletion.
diff --git a/sqlglot/dataframe/sql/functions.py b/sqlglot/dataframe/sql/functions.py
@@ -656,7 +656,7 @@ def unix_timestamp(
 
 def from_utc_timestamp(timestamp: ColumnOrName, tz: ColumnOrName) -> Column:
     tz_column = tz if isinstance(tz, Column) else lit(tz)
-    return Column.invoke_anonymous_function(timestamp, "FROM_UTC_TIMESTAMP", tz_column)
+    return Column.invoke_expression_over_column(timestamp, expression.AtTimeZone, zone=tz_column)
 
 
 def to_utc_timestamp(timestamp: ColumnOrName, tz: ColumnOrName) -> Column:

diff --git a/sqlglot/dialects/spark2.py b/sqlglot/dialects/spark2.py
@@ -117,6 +117,10 @@ class Parser(Hive.Parser):
             "DAYOFYEAR": lambda args: exp.DayOfYear(this=exp.TsOrDsToDate(this=seq_get(args, 0))),
             "DOUBLE": _parse_as_cast("double"),
             "FLOAT": _parse_as_cast("float"),
+            "FROM_UTC_TIMESTAMP": lambda args: exp.AtTimeZone(
+                this=exp.Cast(this=seq_get(args, 0), to=exp.DataType.build("timestamp")),
+                zone=seq_get(args, 1),
+            ),
             "IIF": exp.If.from_arg_list,
             "INT": _parse_as_cast("int"),
             "MAP_FROM_ARRAYS": exp.Map.from_arg_list,

diff --git a/tests/dialects/test_presto.py b/tests/dialects/test_presto.py
@@ -352,6 +352,12 @@ def test_time(self):
                 "presto": "SELECT CAST('2012-10-31 00:00' AS TIMESTAMP) AT TIME ZONE 'America/Sao_Paulo'",
             },
         )
+        self.validate_all(
+            "CAST('2012-10-31 00:00' AS TIMESTAMP) AT TIME ZONE 'America/Sao_Paulo'",
+            read={
+                "spark": "FROM_UTC_TIMESTAMP('2012-10-31 00:00', 'America/Sao_Paulo')",
+            },
+        )
         self.validate_all(
             "CAST(x AS TIMESTAMP)",
             write={"presto": "CAST(x AS TIMESTAMP)"},