Repository: innobi/pantab (python)

Issue:
TypeError: Cannot use .astype to convert from timezone-naive dtype to timezone-aware dtype. Use obj.tz_localize instead or series.dt.tz_localize instead
**Describe the bug**
Pandas has finally deprecated using `#astypes` to convert between timezone aware and timezone naive data types. The postgres datatype is `timestamp without time zone`, yielded from 

```psql
SELECT
    column_name,
    data_type
FROM
    information_schema.columns
WHERE
    table_schema = 'redacted_schema' AND 
    table_name = 'redacted_table' AND
    column_name = 'redacted_column';
```
| column_name     | data_type                   |
|-----------------|-----------------------------|
| redacted_column | timestamp without time zone |

In python we get 
```python3
Timestamp(2023, 4, 17, 23, 3, 39, 0, UTC)
```
**To Reproduce**
Input the previously mentioned data type into a hyperfile and run `pantab#frame_from_hyper_query` with the latest version of pandas. 

**Expected behavior**
Expected to return a data frame after a query is run

**Screenshots**
Here's a stack trace
```python3
TypeError: Cannot use .astype to convert from timezone-naive dtype to timezone-aware dtype. Use obj.tz_localize instead or series.dt.tz_localize instead
```

**Desktop (please complete the following information):**
```bash
$ lsb_release -a
Distributor ID: Ubuntu
Description:    Ubuntu 22.04.2 LTS
Release:        22.04
Codename:       jammy
```

**Additional context**
Every pandas related package in the env I run
```bash
$ pip freeze | grep pandas
dagster-duckdb-pandas==0.19.14
dagster-pandas==0.19.14
pandas==2.0.3
pandas-flavor==0.6.0
```
The problem here is with line 57 of [_readers.py](https://github.com/innobi/pantab/blob/main/pantab/_reader.py#L57)

Pull request description:
add localization to utc timestamps
Add to the date monkey patch and localize all utc date times. Let me know if this is the route you want to go to fix this, or if you'd like a more sustainable way of handling datetime conversions. If this is acceptable, I'll test it, and conform to your CONTRIBUTE.md

closes #186

Patch:
diff --git a/pantab/_reader.py b/pantab/_reader.py
index 1485e11..e26f64f 100644
--- a/pantab/_reader.py
+++ b/pantab/_reader.py
@@ -48,19 +48,15 @@ def _read_query_result(
     df = pd.DataFrame(libpantab.read_hyper_query(result._Result__cdata, dtype_strs))
 
     df.columns = dtypes.keys()
-
     # TODO: remove this hackery...
-    utc_cols = []
     for k, v in dtypes.items():
         if v == "date":
             dtypes[k] = "datetime64[ns]"
-        elif v == "datetime64[ns, UTC]":
-            utc_cols.append(k)
-            dtypes[k] = "datetime64[ns]"
-
-    df = df.astype(dtypes)
-    for utc_col in utc_cols:
-        df[utc_col] = df[utc_col].dt.tz_localize("utc")
+    date_types = ["datetime64[ns, UTC]", "datetime64[ns]"]
+    for col in df.select_dtypes(include=date_types):
+        df[col] = df[col].dt.tz_localize(None)
+    for col in df.select_dtypes(exclude=date_types):
+        df[col] = df[col].astype(dtypes[col])
 
     df = df.fillna(value=np.nan)  # Replace any appearances of None
 
diff --git a/pantab/_writer.py b/pantab/_writer.py
index adbf223..062a390 100644
--- a/pantab/_writer.py
+++ b/pantab/_writer.py
@@ -105,6 +105,16 @@ def _maybe_convert_timedelta(df: pd.DataFrame) -> Tuple[pd.DataFrame, Tuple[str,
     return df, orig_dtypes
 
 
+def _maybe_convert_utctimestamp(df: pd.DataFrame) -> pd.DataFrame:
+    """
+    Hyper implements a subset of postgres and doesn't implement timezone-aware datetimes
+    Thus, we localize to timezone-naive
+    """
+    for utc_col in df.select_dtypes("datetime64[ns, UTC]"):
+        df[utc_col] = df[utc_col].dt.tz_convert(None)
+    return df
+
+
 def _insert_frame(
     df: pd.DataFrame,
     *,
@@ -154,6 +164,7 @@ def _insert_frame(
 
         with tab_api.Inserter(connection, table_def) as inserter:
             if compat.PANDAS_130:
+                df = _maybe_convert_utctimestamp(df)
                 libpantab.write_to_hyper(df, null_mask, inserter._buffer, dtypes)
             else:
                 libpantab.write_to_hyper_legacy(


Test patch:
diff --git a/pantab/tests/conftest.py b/pantab/tests/conftest.py
index 2ec6c77..91f4105 100644
--- a/pantab/tests/conftest.py
+++ b/pantab/tests/conftest.py
@@ -24,7 +24,7 @@ def df():
                 5.0,
                 True,
                 pd.to_datetime("2018-01-01"),
-                # pd.to_datetime("2018-01-01", utc=True),
+                pd.to_datetime("2018-01-01", utc=True),
                 pd.Timedelta("1 days 2 hours 3 minutes 4 seconds"),
                 "foo",
                 np.iinfo(np.int16).min,
@@ -45,7 +45,7 @@ def df():
                 10.0,
                 False,
                 pd.to_datetime("1/1/19"),
-                # pd.to_datetime("2019-01-01", utc=True),
+                pd.to_datetime("2019-01-01", utc=True),
                 pd.Timedelta("-1 days 2 hours 3 minutes 4 seconds"),
                 "bar",
                 np.iinfo(np.int16).max,
@@ -66,7 +66,7 @@ def df():
                 np.nan,
                 False,
                 pd.NaT,
-                # pd.NaT,
+                pd.NaT,
                 pd.NaT,
                 np.nan,
                 0,
@@ -88,7 +88,7 @@ def df():
             "float64",
             "bool",
             "datetime64",
-            # "datetime64_utc",
+            "datetime64_utc",
             "timedelta64",
             "object",
             "int16_limits",
@@ -112,7 +112,7 @@ def df():
             "float64": np.float64,
             "bool": bool,
             "datetime64": "datetime64[ns]",
-            # "datetime64_utc": "datetime64[ns, UTC]",
+            "datetime64_utc": "datetime64[ns, UTC]",
             "timedelta64": "timedelta64[ns]",
             "object": "object",
             "int16_limits": np.int16,
diff --git a/pantab/tests/test_roundtrip.py b/pantab/tests/test_roundtrip.py
index 749b544..1d23c0d 100644
--- a/pantab/tests/test_roundtrip.py
+++ b/pantab/tests/test_roundtrip.py
@@ -14,6 +14,7 @@ def assert_roundtrip_equal(result, expected):
 
     expected["object"] = expected["object"].astype("string")
     expected["non-ascii"] = expected["non-ascii"].astype("string")
+    expected["datetime64_utc"] = expected["datetime64_utc"].dt.tz_localize(None)
 
     tm.assert_frame_equal(result, expected)
 
diff --git a/pantab/tests/test_writer.py b/pantab/tests/test_writer.py
index 88500d5..922322e 100644
--- a/pantab/tests/test_writer.py
+++ b/pantab/tests/test_writer.py
@@ -1,9 +1,11 @@
 import re
+from datetime import datetime, timezone
 
 import numpy as np
 import pandas as pd
 import pytest
 import tableauhyperapi as tab_api
+from tableauhyperapi import Connection, CreateMode, HyperProcess, Telemetry
 
 import pantab
 
@@ -128,3 +130,43 @@ def test_use_parquet_with_timedelta_raises(df, tmp_hyper):
     msg = "Writing timedelta values with use_parquet=True is not yet supported."
     with pytest.raises(ValueError, match=msg):
         pantab.frame_to_hyper(df, tmp_hyper, table="test", use_parquet=True)
+
+
+@pytest.mark.skipif(
+    not pantab._compat.PANDAS_130, reason="bug is specifically with >=pandas 1.3"
+)
+def test_utc_bug(tmp_hyper):
+    """
+    Red-Green for UTC bug
+    """
+    df = pd.DataFrame(
+        {"utc_time": [datetime.now(timezone.utc), pd.Timestamp("today", tz="UTC")]}
+    )
+    pantab.frame_to_hyper(df, tmp_hyper, table="exp")
+    with HyperProcess(Telemetry.DO_NOT_SEND_USAGE_DATA_TO_TABLEAU) as hyper:
+        with Connection(
+            hyper.endpoint, tmp_hyper, CreateMode.CREATE_IF_NOT_EXISTS
+        ) as connection:
+            resp = connection.execute_list_query("select utc_time from exp")
+    assert all(
+        [actual[0].year == expected.year for actual, expected in zip(resp, df.utc_time)]
+    ), f"""
+    expected: {df.utc_time}
+    actual: {[c[0] for c in resp]}
+    """
+
+
+@pytest.mark.skipif(
+    not pantab._compat.PANDAS_130, reason="bug is specifically with >=pandas 1.3"
+)
+def test_maybe_convert_utc(tmp_hyper):
+    """
+    timezone aware is not supported, thus we ensure timezone naive
+    """
+    df = pd.DataFrame(
+        {"utc_time": [datetime.now(timezone.utc), pd.Timestamp("today", tz="UTC")]}
+    )
+    assert not df.select_dtypes("datetime64[ns, UTC]").empty
+    df = pantab._writer._maybe_convert_utctimestamp(df)
+    assert df.select_dtypes("datetime64[ns, UTC]").empty
+    assert not df.select_dtypes("datetime64[ns]").empty