Repository: tobymao/sqlglot

Issue:
TIMESTAMPDIFF is not working property for Spark SQL
TIMESTAMPDIFF is generating incorrect output for output dialects `spark`
Input dialect: `databricks`
Output dialect: `spark`

<img width="526" alt="image" src="https://github.com/user-attachments/assets/4621a61c-bf2d-4ff6-aa81-681eaa8c377f">

Correct Output is `SELECT TIMESTAMPDIFF(DAY, a, b)`

Discussion:

Solution patch:
diff --git a/sqlglot/dialects/databricks.py b/sqlglot/dialects/databricks.py
index d1127f8658..ed5e64235e 100644
--- a/sqlglot/dialects/databricks.py
+++ b/sqlglot/dialects/databricks.py
@@ -7,7 +7,6 @@
     date_delta_sql,
     build_date_delta,
     timestamptrunc_sql,
-    timestampdiff_sql,
 )
 from sqlglot.dialects.spark import Spark
 from sqlglot.tokens import TokenType
@@ -46,7 +45,6 @@ class Parser(Spark.Parser):
             "DATE_ADD": build_date_delta(exp.DateAdd),
             "DATEDIFF": build_date_delta(exp.DateDiff),
             "DATE_DIFF": build_date_delta(exp.DateDiff),
-            "TIMESTAMPDIFF": build_date_delta(exp.TimestampDiff),
             "GET_JSON_OBJECT": _build_json_extract,
         }
 
@@ -75,8 +73,6 @@ class Generator(Spark.Generator):
                 exp.Mul(this=e.expression, expression=exp.Literal.number(-1)),
                 e.this,
             ),
-            exp.DatetimeDiff: timestampdiff_sql,
-            exp.TimestampDiff: timestampdiff_sql,
             exp.DatetimeTrunc: timestamptrunc_sql(),
             exp.Select: transforms.preprocess(
                 [
diff --git a/sqlglot/dialects/spark.py b/sqlglot/dialects/spark.py
index a58b708d29..08d99e4e51 100644
--- a/sqlglot/dialects/spark.py
+++ b/sqlglot/dialects/spark.py
@@ -3,7 +3,7 @@
 import typing as t
 
 from sqlglot import exp
-from sqlglot.dialects.dialect import rename_func, unit_to_var
+from sqlglot.dialects.dialect import rename_func, unit_to_var, timestampdiff_sql, build_date_delta
 from sqlglot.dialects.hive import _build_with_ignore_nulls
 from sqlglot.dialects.spark2 import Spark2, temporary_storage_provider, _build_as_cast
 from sqlglot.helper import ensure_list, seq_get
@@ -108,6 +108,7 @@ class Parser(Spark2.Parser):
             "DATE_ADD": _build_dateadd,
             "DATEADD": _build_dateadd,
             "TIMESTAMPADD": _build_dateadd,
+            "TIMESTAMPDIFF": build_date_delta(exp.TimestampDiff),
             "DATEDIFF": _build_datediff,
             "DATE_DIFF": _build_datediff,
             "TIMESTAMP_LTZ": _build_as_cast("TIMESTAMP_LTZ"),
@@ -167,6 +168,8 @@ class Generator(Spark2.Generator):
             exp.StartsWith: rename_func("STARTSWITH"),
             exp.TsOrDsAdd: _dateadd_sql,
             exp.TimestampAdd: _dateadd_sql,
+            exp.DatetimeDiff: timestampdiff_sql,
+            exp.TimestampDiff: timestampdiff_sql,
             exp.TryCast: lambda self, e: (
                 self.trycast_sql(e) if e.args.get("safe") else self.cast_sql(e)
             ),


Test patch:
diff --git a/tests/dialects/test_spark.py b/tests/dialects/test_spark.py
index 486bf79263..1aa5c2138a 100644
--- a/tests/dialects/test_spark.py
+++ b/tests/dialects/test_spark.py
@@ -754,6 +754,17 @@ def test_spark(self):
             },
         )
 
+        self.validate_all(
+            "SELECT TIMESTAMPDIFF(MONTH, foo, bar)",
+            read={
+                "databricks": "SELECT TIMESTAMPDIFF(MONTH, foo, bar)",
+            },
+            write={
+                "spark": "SELECT TIMESTAMPDIFF(MONTH, foo, bar)",
+                "databricks": "SELECT TIMESTAMPDIFF(MONTH, foo, bar)",
+            },
+        )
+
     def test_bool_or(self):
         self.validate_all(
             "SELECT a, LOGICAL_OR(b) FROM table GROUP BY a",
diff --git a/tests/test_parser.py b/tests/test_parser.py
index 3e5e27e9d8..ba1240c4f7 100644
--- a/tests/test_parser.py
+++ b/tests/test_parser.py
@@ -712,20 +712,20 @@ def test_pivot_columns(self):
                 self.assertEqual(expected_columns, [col.sql(dialect=dialect) for col in columns])
 
     def test_parse_nested(self):
-        now = time.time()
-        query = parse_one("SELECT * FROM a " + ("LEFT JOIN b ON a.id = b.id " * 38))
-        self.assertIsNotNone(query)
-        self.assertLessEqual(time.time() - now, 0.1)
-
-        now = time.time()
-        query = parse_one("SELECT * FROM a " + ("LEFT JOIN UNNEST(ARRAY[]) " * 15))
-        self.assertIsNotNone(query)
-        self.assertLessEqual(time.time() - now, 0.1)
-
-        now = time.time()
-        query = parse_one("SELECT * FROM a " + ("OUTER APPLY (SELECT * FROM b) " * 30))
-        self.assertIsNotNone(query)
-        self.assertLessEqual(time.time() - now, 0.1)
+        def warn_over_threshold(query: str, max_threshold: float = 0.2):
+            now = time.time()
+            ast = parse_one(query)
+            end = time.time() - now
+
+            self.assertIsNotNone(ast)
+            if end >= max_threshold:
+                parser_logger.warning(
+                    f"Query {query[:100]}... surpassed the time threshold of {max_threshold} seconds"
+                )
+
+        warn_over_threshold("SELECT * FROM a " + ("LEFT JOIN b ON a.id = b.id " * 38))
+        warn_over_threshold("SELECT * FROM a " + ("LEFT JOIN UNNEST(ARRAY[]) " * 15))
+        warn_over_threshold("SELECT * FROM a " + ("OUTER APPLY (SELECT * FROM b) " * 30))
 
     def test_parse_properties(self):
         self.assertEqual(