Fix(optimizer)!: infer timestamp function types as TIMESTAMPTZ for bigquery (#4914)

Author: georgesittas

sqlglot/dialects/bigquery.py

@@ -8,6 +8,7 @@
 from sqlglot.dialects.dialect import (
     Dialect,
     NormalizationStrategy,
+    annotate_with_type_lambda,
     arg_max_or_min_no_count,
     binary_from_function,
     date_add_interval_sql,
@@ -398,8 +399,20 @@ class BigQuery(Dialect):
     # All set operations require either a DISTINCT or ALL specifier
     SET_OP_DISTINCT_BY_DEFAULT = dict.fromkeys((exp.Except, exp.Intersect, exp.Union), None)
 
+    # BigQuery maps Type.TIMESTAMP to DATETIME, so we need to amend the inferred types
+    TYPE_TO_EXPRESSIONS = {
+        **Dialect.TYPE_TO_EXPRESSIONS,
+        exp.DataType.Type.TIMESTAMPTZ: Dialect.TYPE_TO_EXPRESSIONS[exp.DataType.Type.TIMESTAMP],
+    }
+    TYPE_TO_EXPRESSIONS.pop(exp.DataType.Type.TIMESTAMP)
+
     ANNOTATORS = {
         **Dialect.ANNOTATORS,
+        **{
+            expr_type: annotate_with_type_lambda(data_type)
+            for data_type, expressions in TYPE_TO_EXPRESSIONS.items()
+            for expr_type in expressions
+        },
         **{
             expr_type: lambda self, e: _annotate_math_functions(self, e)
             for expr_type in (exp.Floor, exp.Ceil, exp.Log, exp.Ln, exp.Sqrt, exp.Exp, exp.Round)

sqlglot/dialects/dialect.py

@@ -51,7 +51,7 @@
 }
 
 
-def _annotate_with_type_lambda(data_type: exp.DataType.Type) -> t.Callable[[TypeAnnotator, E], E]:
+def annotate_with_type_lambda(data_type: exp.DataType.Type) -> t.Callable[[TypeAnnotator, E], E]:
     return lambda self, e: self._annotate_with_type(e, data_type)
 
 
@@ -683,15 +683,15 @@ class Dialect(metaclass=_Dialect):
             exp.ParseJSON,
         },
         exp.DataType.Type.TIME: {
+            exp.CurrentTime,
             exp.Time,
+            exp.TimeAdd,
+            exp.TimeSub,
         },
         exp.DataType.Type.TIMESTAMP: {
-            exp.CurrentTime,
             exp.CurrentTimestamp,
             exp.StrToTime,
-            exp.TimeAdd,
             exp.TimeStrToTime,
-            exp.TimeSub,
             exp.TimestampAdd,
             exp.TimestampSub,
             exp.UnixToTime,
@@ -733,7 +733,7 @@ class Dialect(metaclass=_Dialect):
             for expr_type in subclasses(exp.__name__, exp.Binary)
         },
         **{
-            expr_type: _annotate_with_type_lambda(data_type)
+            expr_type: annotate_with_type_lambda(data_type)
             for data_type, expressions in TYPE_TO_EXPRESSIONS.items()
             for expr_type in expressions
         },

tests/dialects/test_bigquery.py

@@ -15,6 +15,7 @@
 from sqlglot.helper import logger as helper_logger
 from sqlglot.parser import logger as parser_logger
 from tests.dialects.test_dialect import Validator
+from sqlglot.optimizer.annotate_types import annotate_types
 
 
 class TestBigQuery(Validator):
@@ -2366,3 +2367,18 @@ def test_string_agg(self):
             "STRING_AGG(DISTINCT a ORDER BY b DESC, c DESC LIMIT 10)",
             "STRING_AGG(DISTINCT a, ',' ORDER BY b DESC, c DESC LIMIT 10)",
         )
+
+    def test_annotate_timestamps(self):
+        sql = """
+        SELECT
+          CURRENT_TIMESTAMP() AS curr_ts,
+          TIMESTAMP_SECONDS(2) AS ts_seconds,
+          PARSE_TIMESTAMP('%c', 'Thu Dec 25 07:30:00 2008', 'UTC') AS parsed_ts,
+          TIMESTAMP_ADD(TIMESTAMP "2008-12-25 15:30:00+00", INTERVAL 10 MINUTE) AS ts_add,
+          TIMESTAMP_SUB(TIMESTAMP "2008-12-25 15:30:00+00", INTERVAL 10 MINUTE) AS ts_sub,
+        """
+
+        annotated = annotate_types(self.parse_one(sql), dialect="bigquery")
+
+        for select in annotated.selects:
+            self.assertEqual(select.type.sql("bigquery"), "TIMESTAMP")

tests/fixtures/optimizer/annotate_functions.sql

@@ -19,6 +19,15 @@ INT;
 LEAST(1, 2.5, 3);
 DOUBLE;
 
+CURRENT_TIME();
+TIME;
+
+TIME_ADD(CAST('09:05:03' AS TIME), INTERVAL 2 HOUR);
+TIME;
+
+TIME_SUB(CAST('09:05:03' AS TIME), INTERVAL 2 HOUR);
+TIME;
+
 --------------------------------------
 -- Spark2 / Spark3 / Databricks
 --------------------------------------