fix incorrect transpilation to tz types, add tests

Author: alistairewj

src/mimic_utils/compare_concepts.py

@@ -55,19 +55,34 @@ def _list_tables_duckdb(con, schema):
 
 def _as_numeric(series):
     """Return a float Series if every non-null value coerces, else None."""
+    if _is_datetime_like_series(series):
+        return None
     coerced = pd.to_numeric(series, errors="coerce")
     if coerced.isna().sum() == series.isna().sum():
         return coerced.astype(float)
     return None
 
 
+def _is_datetime_like_series(series):
+    """Return True when values are already datetimes, even if stored as object dtype."""
+    if pd.api.types.is_datetime64_any_dtype(series.dtype) or pd.api.types.is_datetime64tz_dtype(series.dtype):
+        return True
+    non_null = series.dropna()
+    if non_null.empty:
+        return False
+    sample = non_null.iloc[0]
+    return isinstance(sample, (pd.Timestamp, np.datetime64)) or hasattr(sample, "tzinfo")
+
+
 def _as_datetime(series):
     """Return a datetime Series if every non-null value parses, else None."""
     import warnings
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")  # silence per-element parse fallback notices
         coerced = pd.to_datetime(series, errors="coerce")
     if coerced.isna().sum() == series.isna().sum():
+        if pd.api.types.is_datetime64tz_dtype(coerced.dtype):
+            coerced = coerced.dt.tz_localize(None)
         return coerced
     return None
 

src/mimic_utils/sqlglot_dialects/duckdb.py

@@ -11,6 +11,8 @@
 class MimicDuckDB(DuckDB):
     class Generator(DuckDB.Generator):
         def datatype_sql(self, expression: exp.DataType) -> str:
+            if expression.this == exp.DataType.Type.TIMESTAMPTZ:
+                return "TIMESTAMP"
             if expression.this == exp.DataType.Type.DECIMAL and not expression.expressions:
                 return "DECIMAL(38, 9)"
             return super().datatype_sql(expression)

src/mimic_utils/sqlglot_dialects/postgres.py

@@ -80,6 +80,11 @@ def _generate_array_sql(self: Postgres.Generator, expression: exp.Expression) ->
 
 class MimicPostgres(Postgres):
     class Generator(Postgres.Generator):
+        def datatype_sql(self, expression: exp.DataType) -> str:
+            if expression.this == exp.DataType.Type.TIMESTAMPTZ:
+                return "TIMESTAMP"
+            return super().datatype_sql(expression)
+
         TRANSFORMS = {
             **Postgres.Generator.TRANSFORMS,
             exp.DatetimeDiff: _datetime_diff_sql,

src/mimic_utils/transpile.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import re
 from pathlib import Path
 from typing import Union
 
@@ -25,6 +26,12 @@
 _CATALOG_TO_REMOVE = "physionet-data"
 
 
+def _strip_timezone_types(sql: str) -> str:
+    """Normalize any timezone-aware timestamp types back to naive TIMESTAMP."""
+    sql = re.sub(r"\bTIMESTAMP\s+WITH\s+TIME\s+ZONE\b", "TIMESTAMP", sql)
+    return re.sub(r"\bTIMESTAMPTZ\b", "TIMESTAMP", sql)
+
+
 def _strip_catalog(parsed: exp.Expression) -> None:
     """Remove the ``physionet-data`` catalog from every table reference in place."""
     for table in parsed.find_all(exp.Table):
@@ -50,11 +57,12 @@ def transpile_query(query: str, source_dialect: str = "bigquery", destination_di
     # drop comments when targeting it.
     keep_comments = destination_dialect != "duckdb"
 
-    return parsed.sql(
+    sql = parsed.sql(
         dialect=_DESTINATION_DIALECTS[destination_dialect],
         pretty=True,
         comments=keep_comments,
     )
+    return _strip_timezone_types(sql)
 
 
 def transpile_file(

tests/test_compare_concepts.py

@@ -0,0 +1,22 @@
+from datetime import datetime, timedelta, timezone
+
+import pandas as pd
+
+from mimic_utils.compare_concepts import compare_table
+
+
+def test_compare_table_treats_same_wall_time_as_equal_across_timezones():
+    pg = pd.DataFrame({"charttime": [datetime(2150, 1, 1, 3, 30, tzinfo=timezone(timedelta(hours=2)))]})
+    duck = pd.DataFrame({"charttime": [datetime(2150, 1, 1, 3, 30)]})
+
+    assert compare_table(pg, duck, 1e-6, 1e-9) == (True, "1 rows OK")
+
+
+def test_compare_table_flags_different_wall_times_even_if_instants_match():
+    pg = pd.DataFrame({"charttime": [datetime(2150, 1, 1, 3, 30, tzinfo=timezone(timedelta(hours=2)))]})
+    duck = pd.DataFrame({"charttime": [datetime(2150, 1, 1, 1, 30)]})
+
+    assert compare_table(pg, duck, 1e-6, 1e-9) == (
+        False,
+        "1 differing cells in columns: charttime(1)",
+    )

tests/test_transpile.py

@@ -102,6 +102,18 @@ def test_duckdb_comments_stripped():
     assert "/*" in transpile_query(bq, "bigquery", "postgres")
 
 
+@pytest.mark.parametrize(
+    "bq",
+    [
+        "SELECT TIMESTAMP('2150-01-01 00:00:00') AS ts",
+        "SELECT CAST(DATETIME '2150-01-01 00:00:00' AS TIMESTAMP) AS ts",
+    ],
+)
+def test_timestamp_outputs_stay_timezone_naive(bq):
+    for dialect in ("postgres", "duckdb"):
+        assert "TIMESTAMPTZ" not in transpile_query(bq, "bigquery", dialect)
+
+
 def test_unsupported_dialect_raises():
     with pytest.raises(ValueError):
         transpile_query("SELECT 1", "bigquery", "mysql")