Repository: tobymao/sqlglot

Issue:
TRUNC(expr, 'DAY') returns null in SparkSQL
This is an issue while transpiling Trino to SparkSQL:

A valid Trino date_trunc by day

`DATE_TRUNC('day' ,col);`

is transpiled into SparkSQL as

`TRUNC('day', col)`.

However, `TRUNC('day', col)` always returns NULL because `day` is not a valid argument for TRUNC.  It's potentially tricky as DATE_TRUNC in Trino is type parameterized on the input type.

Discussion:

Patch:
diff --git a/sqlglot/dialects/dialect.py b/sqlglot/dialects/dialect.py
index 940775d1..b267521f 100644
--- a/sqlglot/dialects/dialect.py
+++ b/sqlglot/dialects/dialect.py
@@ -388,6 +388,21 @@ def parse_date_delta(
     return inner_func
 
 
+def date_trunc_to_time(args: t.Sequence) -> exp.DateTrunc | exp.TimestampTrunc:
+    unit = seq_get(args, 0)
+    this = seq_get(args, 1)
+
+    if isinstance(this, exp.Cast) and this.is_type(exp.DataType.Type.DATE):
+        return exp.DateTrunc(unit=unit, this=this)
+    return exp.TimestampTrunc(this=this, unit=unit)
+
+
+def timestamptrunc_sql(self: Generator, expression: exp.TimestampTrunc) -> str:
+    return self.func(
+        "DATE_TRUNC", exp.Literal.string(expression.text("unit") or "day"), expression.this
+    )
+
+
 def locate_to_strposition(args: t.Sequence) -> exp.Expression:
     return exp.StrPosition(
         this=seq_get(args, 1),
diff --git a/sqlglot/dialects/duckdb.py b/sqlglot/dialects/duckdb.py
index e04cc2a8..f1d2266f 100644
--- a/sqlglot/dialects/duckdb.py
+++ b/sqlglot/dialects/duckdb.py
@@ -14,6 +14,7 @@ from sqlglot.dialects.dialect import (
     rename_func,
     str_position_sql,
     str_to_time_sql,
+    timestamptrunc_sql,
     timestrtotime_sql,
     ts_or_ds_to_date_sql,
 )
@@ -179,6 +180,7 @@ class DuckDB(Dialect):
             exp.StrToTime: str_to_time_sql,
             exp.StrToUnix: lambda self, e: f"EPOCH(STRPTIME({self.sql(e, 'this')}, {self.format_time(e)}))",
             exp.Struct: _struct_sql,
+            exp.TimestampTrunc: timestamptrunc_sql,
             exp.TimeStrToDate: lambda self, e: f"CAST({self.sql(e, 'this')} AS DATE)",
             exp.TimeStrToTime: timestrtotime_sql,
             exp.TimeStrToUnix: lambda self, e: f"EPOCH(CAST({self.sql(e, 'this')} AS TIMESTAMP))",
diff --git a/sqlglot/dialects/postgres.py b/sqlglot/dialects/postgres.py
index 6304714c..5f556a57 100644
--- a/sqlglot/dialects/postgres.py
+++ b/sqlglot/dialects/postgres.py
@@ -12,6 +12,7 @@ from sqlglot.dialects.dialect import (
     no_trycast_sql,
     rename_func,
     str_position_sql,
+    timestamptrunc_sql,
     trim_sql,
 )
 from sqlglot.helper import seq_get
@@ -34,7 +35,7 @@ def _date_add_sql(kind):
         from sqlglot.optimizer.simplify import simplify
 
         this = self.sql(expression, "this")
-        unit = self.sql(expression, "unit")
+        unit = expression.args.get("unit")
         expression = simplify(expression.args["expression"])
 
         if not isinstance(expression, exp.Literal):
@@ -255,9 +256,8 @@ class Postgres(Dialect):
             "TO_TIMESTAMP": _to_timestamp,
             "TO_CHAR": format_time_lambda(exp.TimeToStr, "postgres"),
             "GENERATE_SERIES": _generate_series,
-            "DATE_TRUNC": lambda args: exp.DateTrunc(
-                unit=exp.Literal.string(seq_get(args, 0).name),  # type: ignore
-                this=seq_get(args, 1),
+            "DATE_TRUNC": lambda args: exp.TimestampTrunc(
+                this=seq_get(args, 1), unit=seq_get(args, 0)
             ),
         }
 
@@ -324,6 +324,7 @@ class Postgres(Dialect):
             exp.StrPosition: str_position_sql,
             exp.StrToTime: lambda self, e: f"TO_TIMESTAMP({self.sql(e, 'this')}, {self.format_time(e)})",
             exp.Substring: _substring_sql,
+            exp.TimestampTrunc: timestamptrunc_sql,
             exp.TimeStrToTime: lambda self, e: f"CAST({self.sql(e, 'this')} AS TIMESTAMP)",
             exp.TimeToStr: lambda self, e: f"TO_CHAR({self.sql(e, 'this')}, {self.format_time(e)})",
             exp.TableSample: no_tablesample_sql,
diff --git a/sqlglot/dialects/presto.py b/sqlglot/dialects/presto.py
index a90990f0..07e8f43f 100644
--- a/sqlglot/dialects/presto.py
+++ b/sqlglot/dialects/presto.py
@@ -3,12 +3,14 @@ from __future__ import annotations
 from sqlglot import exp, generator, parser, tokens, transforms
 from sqlglot.dialects.dialect import (
     Dialect,
+    date_trunc_to_time,
     format_time_lambda,
     if_sql,
     no_ilike_sql,
     no_safe_divide_sql,
     rename_func,
     struct_extract_sql,
+    timestamptrunc_sql,
     timestrtotime_sql,
 )
 from sqlglot.dialects.mysql import MySQL
@@ -98,10 +100,16 @@ def _ts_or_ds_to_date_sql(self, expression):
 
 
 def _ts_or_ds_add_sql(self, expression):
-    this = self.sql(expression, "this")
-    e = self.sql(expression, "expression")
-    unit = self.sql(expression, "unit") or "'day'"
-    return f"DATE_ADD({unit}, {e}, DATE_PARSE(SUBSTR({this}, 1, 10), {Presto.date_format}))"
+    return self.func(
+        "DATE_ADD",
+        exp.Literal.string(expression.text("unit") or "day"),
+        expression.expression,
+        self.func(
+            "DATE_PARSE",
+            self.func("SUBSTR", expression.this, exp.Literal.number(1), exp.Literal.number(10)),
+            Presto.date_format,
+        ),
+    )
 
 
 def _sequence_sql(self, expression):
@@ -195,6 +203,7 @@ class Presto(Dialect):
             ),
             "DATE_FORMAT": format_time_lambda(exp.TimeToStr, "presto"),
             "DATE_PARSE": format_time_lambda(exp.StrToTime, "presto"),
+            "DATE_TRUNC": date_trunc_to_time,
             "FROM_UNIXTIME": _from_unixtime,
             "NOW": exp.CurrentTimestamp.from_arg_list,
             "STRPOS": lambda args: exp.StrPosition(
@@ -251,8 +260,12 @@ class Presto(Dialect):
             exp.BitwiseXor: lambda self, e: f"BITWISE_XOR({self.sql(e, 'this')}, {self.sql(e, 'expression')})",
             exp.CurrentTimestamp: lambda *_: "CURRENT_TIMESTAMP",
             exp.DataType: _datatype_sql,
-            exp.DateAdd: lambda self, e: f"""DATE_ADD({self.sql(e, 'unit') or "'day'"}, {self.sql(e, 'expression')}, {self.sql(e, 'this')})""",
-            exp.DateDiff: lambda self, e: f"""DATE_DIFF({self.sql(e, 'unit') or "'day'"}, {self.sql(e, 'expression')}, {self.sql(e, 'this')})""",
+            exp.DateAdd: lambda self, e: self.func(
+                "DATE_ADD", exp.Literal.string(e.text("unit") or "day"), e.expression, e.this
+            ),
+            exp.DateDiff: lambda self, e: self.func(
+                "DATE_DIFF", exp.Literal.string(e.text("unit") or "day"), e.expression, e.this
+            ),
             exp.DateStrToDate: lambda self, e: f"CAST(DATE_PARSE({self.sql(e, 'this')}, {Presto.date_format}) AS DATE)",
             exp.DateToDi: lambda self, e: f"CAST(DATE_FORMAT({self.sql(e, 'this')}, {Presto.dateint_format}) AS INT)",
             exp.Decode: _decode_sql,
@@ -279,6 +292,7 @@ class Presto(Dialect):
             exp.StructExtract: struct_extract_sql,
             exp.TableFormatProperty: lambda self, e: f"TABLE_FORMAT='{e.name.upper()}'",
             exp.FileFormatProperty: lambda self, e: f"FORMAT='{e.name.upper()}'",
+            exp.TimestampTrunc: timestamptrunc_sql,
             exp.TimeStrToDate: timestrtotime_sql,
             exp.TimeStrToTime: timestrtotime_sql,
             exp.TimeStrToUnix: lambda self, e: f"TO_UNIXTIME(DATE_PARSE({self.sql(e, 'this')}, {Presto.time_format}))",
diff --git a/sqlglot/dialects/redshift.py b/sqlglot/dialects/redshift.py
index a73ae41c..ebd52162 100644
--- a/sqlglot/dialects/redshift.py
+++ b/sqlglot/dialects/redshift.py
@@ -20,7 +20,11 @@ class Redshift(Postgres):
     class Parser(Postgres.Parser):
         FUNCTIONS = {
             **Postgres.Parser.FUNCTIONS,  # type: ignore
-            "DATEADD": exp.DateAdd.from_arg_list,
+            "DATEADD": lambda args: exp.DateAdd(
+                this=seq_get(args, 2),
+                expression=seq_get(args, 1),
+                unit=seq_get(args, 0),
+            ),
             "DATEDIFF": lambda args: exp.DateDiff(
                 this=seq_get(args, 2),
                 expression=seq_get(args, 1),
@@ -77,9 +81,11 @@ class Redshift(Postgres):
         TRANSFORMS = {
             **Postgres.Generator.TRANSFORMS,  # type: ignore
             **transforms.ELIMINATE_DISTINCT_ON,  # type: ignore
-            exp.DateAdd: rename_func("DATEADD"),  # redshift supports DATEADD, postgres doesn't
+            exp.DateAdd: lambda self, e: self.func(
+                "DATEADD", exp.var(e.text("unit") or "day"), e.expression, e.this
+            ),
             exp.DateDiff: lambda self, e: self.func(
-                "DATEDIFF", e.args.get("unit") or "day", e.expression, e.this
+                "DATEDIFF", exp.var(e.text("unit") or "day"), e.expression, e.this
             ),
             exp.DistKeyProperty: lambda self, e: f"DISTKEY({e.name})",
             exp.DistStyleProperty: lambda self, e: self.naked_property(e),
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
index bdcc8de8..799e9a60 100644
--- a/sqlglot/dialects/snowflake.py
+++ b/sqlglot/dialects/snowflake.py
@@ -5,11 +5,13 @@ import typing as t
 from sqlglot import exp, generator, parser, tokens
 from sqlglot.dialects.dialect import (
     Dialect,
+    date_trunc_to_time,
     datestrtodate_sql,
     format_time_lambda,
     inline_array_sql,
     min_or_least,
     rename_func,
+    timestamptrunc_sql,
     timestrtotime_sql,
     ts_or_ds_to_date_sql,
     var_map_sql,
@@ -176,6 +178,7 @@ class Snowflake(Dialect):
             "ARRAYAGG": exp.ArrayAgg.from_arg_list,
             "ARRAY_CONSTRUCT": exp.Array.from_arg_list,
             "ARRAY_TO_STRING": exp.ArrayJoin.from_arg_list,
+            "DATE_TRUNC": date_trunc_to_time,
             "DATEADD": lambda args: exp.DateAdd(
                 this=seq_get(args, 2),
                 expression=seq_get(args, 1),
@@ -186,10 +189,6 @@ class Snowflake(Dialect):
                 expression=seq_get(args, 1),
                 unit=seq_get(args, 0),
             ),
-            "DATE_TRUNC": lambda args: exp.DateTrunc(
-                unit=exp.Literal.string(seq_get(args, 0).name),  # type: ignore
-                this=seq_get(args, 1),
-            ),
             "DECODE": exp.Matches.from_arg_list,
             "DIV0": _div0_to_if,
             "IFF": exp.If.from_arg_list,
@@ -289,6 +288,7 @@ class Snowflake(Dialect):
                 "POSITION", e.args.get("substr"), e.this, e.args.get("position")
             ),
             exp.StrToTime: lambda self, e: f"TO_TIMESTAMP({self.sql(e, 'this')}, {self.format_time(e)})",
+            exp.TimestampTrunc: timestamptrunc_sql,
             exp.TimeStrToTime: timestrtotime_sql,
             exp.TimeToUnix: lambda self, e: f"EXTRACT(epoch_second FROM {self.sql(e, 'this')})",
             exp.Trim: lambda self, e: self.func("TRIM", e.this, e.expression),
diff --git a/sqlglot/dialects/starrocks.py b/sqlglot/dialects/starrocks.py
index 01e6357b..2ba1a923 100644
--- a/sqlglot/dialects/starrocks.py
+++ b/sqlglot/dialects/starrocks.py
@@ -3,9 +3,18 @@ from __future__ import annotations
 from sqlglot import exp
 from sqlglot.dialects.dialect import arrow_json_extract_sql, rename_func
 from sqlglot.dialects.mysql import MySQL
+from sqlglot.helper import seq_get
 
 
 class StarRocks(MySQL):
+    class Parser(MySQL.Parser):  # type: ignore
+        FUNCTIONS = {
+            **MySQL.Parser.FUNCTIONS,
+            "DATE_TRUNC": lambda args: exp.TimestampTrunc(
+                this=seq_get(args, 1), unit=seq_get(args, 0)
+            ),
+        }
+
     class Generator(MySQL.Generator):  # type: ignore
         TYPE_MAPPING = {
             **MySQL.Generator.TYPE_MAPPING,  # type: ignore
@@ -20,6 +29,9 @@ class StarRocks(MySQL):
             exp.JSONExtract: arrow_json_extract_sql,
             exp.DateDiff: rename_func("DATEDIFF"),
             exp.StrToUnix: lambda self, e: f"UNIX_TIMESTAMP({self.sql(e, 'this')}, {self.format_time(e)})",
+            exp.TimestampTrunc: lambda self, e: self.func(
+                "DATE_TRUNC", exp.Literal.string(e.text("unit")), e.this
+            ),
             exp.TimeStrToDate: rename_func("TO_DATE"),
             exp.UnixToStr: lambda self, e: f"FROM_UNIXTIME({self.sql(e, 'this')}, {self.format_time(e)})",
             exp.UnixToTime: rename_func("FROM_UNIXTIME"),
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 72415418..3123fb42 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -3161,7 +3161,7 @@ class TimeUnit(Expression):
 
     def __init__(self, **args):
         unit = args.get("unit")
-        if isinstance(unit, Column):
+        if isinstance(unit, (Column, Literal)):
             args["unit"] = Var(this=unit.name)
         elif isinstance(unit, Week):
             unit.set("this", Var(this=unit.this.name))
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index dd99cb60..e8279854 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -57,10 +57,10 @@ class Generator:
 
     TRANSFORMS = {
         exp.DateAdd: lambda self, e: self.func(
-            "DATE_ADD", e.this, e.expression, e.args.get("unit")
+            "DATE_ADD", e.this, e.expression, exp.Literal.string(e.text("unit"))
         ),
         exp.TsOrDsAdd: lambda self, e: self.func(
-            "TS_OR_DS_ADD", e.this, e.expression, e.args.get("unit")
+            "TS_OR_DS_ADD", e.this, e.expression, exp.Literal.string(e.text("unit"))
         ),
         exp.VarMap: lambda self, e: self.func("MAP", e.args["keys"], e.args["values"]),
         exp.CharacterSetProperty: lambda self, e: f"{'DEFAULT ' if e.args.get('default') else ''}CHARACTER SET={self.sql(e, 'this')}",


Test patch:
diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
index 26eeaee1..e731b50f 100644
--- a/tests/dialects/test_bigquery.py
+++ b/tests/dialects/test_bigquery.py
@@ -283,7 +283,7 @@ class TestBigQuery(Validator):
                 "duckdb": "CURRENT_DATE + INTERVAL 1 DAY",
                 "mysql": "DATE_ADD(CURRENT_DATE, INTERVAL 1 DAY)",
                 "postgres": "CURRENT_DATE + INTERVAL '1' DAY",
-                "presto": "DATE_ADD(DAY, 1, CURRENT_DATE)",
+                "presto": "DATE_ADD('DAY', 1, CURRENT_DATE)",
                 "hive": "DATE_ADD(CURRENT_DATE, 1)",
                 "spark": "DATE_ADD(CURRENT_DATE, 1)",
             },
diff --git a/tests/dialects/test_dialect.py b/tests/dialects/test_dialect.py
index f45b2658..6214c436 100644
--- a/tests/dialects/test_dialect.py
+++ b/tests/dialects/test_dialect.py
@@ -519,7 +519,7 @@ class TestDialect(Validator):
                 "duckdb": "x + INTERVAL 1 day",
                 "hive": "DATE_ADD(x, 1)",
                 "mysql": "DATE_ADD(x, INTERVAL 1 DAY)",
-                "postgres": "x + INTERVAL '1' 'day'",
+                "postgres": "x + INTERVAL '1' day",
                 "presto": "DATE_ADD('day', 1, x)",
                 "snowflake": "DATEADD(day, 1, x)",
                 "spark": "DATE_ADD(x, 1)",
@@ -543,12 +543,48 @@ class TestDialect(Validator):
         )
         self.validate_all(
             "DATE_TRUNC('day', x)",
+            read={
+                "bigquery": "DATE_TRUNC(x, day)",
+                "duckdb": "DATE_TRUNC('day', x)",
+                "spark": "TRUNC(x, 'day')",
+            },
             write={
+                "bigquery": "DATE_TRUNC(x, day)",
+                "duckdb": "DATE_TRUNC('day', x)",
                 "mysql": "DATE(x)",
+                "presto": "DATE_TRUNC('day', x)",
+                "postgres": "DATE_TRUNC('day', x)",
+                "snowflake": "DATE_TRUNC('day', x)",
+                "starrocks": "DATE_TRUNC('day', x)",
+                "spark": "TRUNC(x, 'day')",
+            },
+        )
+        self.validate_all(
+            "TIMESTAMP_TRUNC(x, day)",
+            read={
+                "bigquery": "TIMESTAMP_TRUNC(x, day)",
+                "presto": "DATE_TRUNC('day', x)",
                 "postgres": "DATE_TRUNC('day', x)",
                 "snowflake": "DATE_TRUNC('day', x)",
+                "starrocks": "DATE_TRUNC('day', x)",
+                "spark": "DATE_TRUNC('day', x)",
+            },
+        )
+        self.validate_all(
+            "DATE_TRUNC('day', CAST(x AS DATE))",
+            read={
+                "presto": "DATE_TRUNC('day', x::DATE)",
+                "snowflake": "DATE_TRUNC('day', x::DATE)",
             },
         )
+        self.validate_all(
+            "TIMESTAMP_TRUNC(CAST(x AS DATE), day)",
+            read={
+                "postgres": "DATE_TRUNC('day', x::DATE)",
+                "starrocks": "DATE_TRUNC('day', x::DATE)",
+            },
+        )
+
         self.validate_all(
             "DATE_TRUNC('week', x)",
             write={
@@ -583,9 +619,6 @@ class TestDialect(Validator):
             "DATE_TRUNC('year', x)",
             read={
                 "bigquery": "DATE_TRUNC(x, year)",
-                "postgres": "DATE_TRUNC(year, x)",
-                "snowflake": "DATE_TRUNC(year, x)",
-                "starrocks": "DATE_TRUNC('year', x)",
                 "spark": "TRUNC(x, 'year')",
             },
             write={
@@ -601,7 +634,10 @@ class TestDialect(Validator):
             "TIMESTAMP_TRUNC(x, year)",
             read={
                 "bigquery": "TIMESTAMP_TRUNC(x, year)",
+                "postgres": "DATE_TRUNC(year, x)",
                 "spark": "DATE_TRUNC('year', x)",
+                "snowflake": "DATE_TRUNC(year, x)",
+                "starrocks": "DATE_TRUNC('year', x)",
             },
             write={
                 "bigquery": "TIMESTAMP_TRUNC(x, year)",
diff --git a/tests/dialects/test_redshift.py b/tests/dialects/test_redshift.py
index b83bde49..ff730f8a 100644
--- a/tests/dialects/test_redshift.py
+++ b/tests/dialects/test_redshift.py
@@ -7,7 +7,7 @@ class TestRedshift(Validator):
     def test_redshift(self):
         self.validate_all("CONVERT(INTEGER, x)", write={"redshift": "CAST(x AS INTEGER)"})
         self.validate_all(
-            "DATE_ADD('day', ndays, caldate)", write={"redshift": "DATEADD('day', ndays, caldate)"}
+            "DATEADD('day', ndays, caldate)", write={"redshift": "DATEADD(day, ndays, caldate)"}
         )
         self.validate_all(
             'create table "group" ("col" char(10))',
@@ -83,10 +83,10 @@ class TestRedshift(Validator):
             },
         )
         self.validate_all(
-            "DATEDIFF(d, a, b)",
+            "DATEDIFF('day', a, b)",
             write={
-                "redshift": "DATEDIFF(d, a, b)",
-                "presto": "DATE_DIFF(d, a, b)",
+                "redshift": "DATEDIFF(day, a, b)",
+                "presto": "DATE_DIFF('day', a, b)",
             },
         )
         self.validate_all(