Feat(teradata): handle transpile of quarter function (#3303)

Add an expression type for `QUARTER` instead of treating it as an anonymous function. Teradata doesn't have a `QUARTER` function to extract the quarter from a date, so transform it to an `EXTRACT(QUARTER ...)` expression. Teradata also doesn't support `EXTRACT(QUARTER ...)`, but this is already handled in `extract_sql`.
tobymao · Apr 11, 2024 · 1bc51df · 1bc51df
1 parent 3c97d34
commit 1bc51df
Show file tree

Hide file tree

Showing 7 changed files with 20 additions and 1 deletion.
diff --git a/sqlglot/dataframe/sql/functions.py b/sqlglot/dataframe/sql/functions.py
@@ -536,7 +536,7 @@ def year(col: ColumnOrName) -> Column:
 
 
 def quarter(col: ColumnOrName) -> Column:
-    return Column.invoke_anonymous_function(col, "QUARTER")
+    return Column.invoke_expression_over_column(col, expression.Quarter)
 
 
 def month(col: ColumnOrName) -> Column:

diff --git a/sqlglot/dialects/teradata.py b/sqlglot/dialects/teradata.py
@@ -36,6 +36,10 @@ def func(self: Teradata.Generator, expression: exp.DateAdd | exp.DateSub) -> str
     return func
 
 
+def _quarter_sql(self: Teradata.Generator, expression: exp.Quarter) -> str:
+    return self.sql(exp.Extract(this="QUARTER", expression=expression.this))
+
+
 class Teradata(Dialect):
     SUPPORTS_SEMI_ANTI_JOIN = False
     TYPED_DIVISION = True
@@ -241,6 +245,7 @@ class Generator(generator.Generator):
             exp.CurrentTimestamp: lambda *_: "CURRENT_TIMESTAMP",
             exp.DateAdd: _date_add_sql("+"),
             exp.DateSub: _date_add_sql("-"),
+            exp.Quarter: _quarter_sql,
         }
 
         def cast_sql(self, expression: exp.Cast, safe_prefix: t.Optional[str] = None) -> str:

diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -5486,6 +5486,10 @@ class ApproxQuantile(Quantile):
     arg_types = {"this": True, "quantile": True, "accuracy": False, "weight": False}
 
 
+class Quarter(Func):
+    pass
+
+
 class Rand(Func):
     _sql_names = ["RAND", "RANDOM"]
     arg_types = {"this": False}

diff --git a/sqlglot/optimizer/annotate_types.py b/sqlglot/optimizer/annotate_types.py
@@ -212,6 +212,7 @@ class TypeAnnotator(metaclass=_TypeAnnotator):
             exp.Month,
             exp.Week,
             exp.Year,
+            exp.Quarter,
         },
         exp.DataType.Type.VARCHAR: {
             exp.ArrayConcat,

diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
@@ -66,6 +66,7 @@ def test_snowflake(self):
         self.validate_identity("SELECT DAYOFYEAR(CURRENT_TIMESTAMP())")
         self.validate_identity("LISTAGG(data['some_field'], ',')")
         self.validate_identity("WEEKOFYEAR(tstamp)")
+        self.validate_identity("SELECT QUARTER(CURRENT_TIMESTAMP())")
         self.validate_identity("SELECT SUM(amount) FROM mytable GROUP BY ALL")
         self.validate_identity("WITH x AS (SELECT 1 AS foo) SELECT foo FROM IDENTIFIER('x')")
         self.validate_identity("WITH x AS (SELECT 1 AS foo) SELECT IDENTIFIER('foo') FROM x")

diff --git a/tests/dialects/test_teradata.py b/tests/dialects/test_teradata.py
@@ -292,3 +292,10 @@ def test_time(self):
                 "bigquery": "EXTRACT(MONTH FROM x)",
             },
         )
+        self.validate_all(
+            "CAST(TO_CHAR(x, 'Q') AS INT)",
+            read={
+                "snowflake": "quarter(x)",
+                "teradata": "CAST(TO_CHAR(x, 'Q') AS INT)",
+            },
+        )
diff --git a/tests/test_expressions.py b/tests/test_expressions.py
@@ -634,6 +634,7 @@ def test_functions(self):
         self.assertIsInstance(parse_one("MAX(a)"), exp.Max)
         self.assertIsInstance(parse_one("MIN(a)"), exp.Min)
         self.assertIsInstance(parse_one("MONTH(a)"), exp.Month)
+        self.assertIsInstance(parse_one("QUARTER(a)"), exp.Quarter)
         self.assertIsInstance(parse_one("POSITION(' ' IN a)"), exp.StrPosition)
         self.assertIsInstance(parse_one("POW(a, 2)"), exp.Pow)
         self.assertIsInstance(parse_one("POWER(a, 2)"), exp.Pow)