fix: convert datetime64[us] to datetime64[ns] for TIMESTAMP and DATETIME columns

Author: kitagry

pandas_gbq/gbq.py

@@ -630,6 +630,7 @@ def _finalize_dtypes(
     """
     import db_dtypes
     import pandas.api.types
+    import pandas as pd
 
     # If you update this mapping, also update the table at
     # `docs/reading.rst`.
@@ -652,6 +653,15 @@ def _finalize_dtypes(
         if dtype and pandas.api.types.is_object_dtype(df[name]):
             df[name] = df[name].astype(dtype, errors="ignore")
 
+        # From pandas 2.0, returns datetime64[us]. Convert to datetime64[ns].
+        if pd.__version__ >= "2.0.0":
+            if (
+                field["type"].upper() in ("DATETIME", "TIMESTAMP")
+                and pandas.api.types.is_datetime64_dtype(df[name])
+                and not pandas.api.types.is_datetime64_ns_dtype(df[name])
+            ):
+                df[name] = df[name].dt.as_unit("ns")
+
     # Ensure any TIMESTAMP columns are tz-aware.
     df = pandas_gbq.timestamp.localize_df(df, schema_fields)
 

tests/unit/test_gbq.py

@@ -113,6 +113,29 @@ def test__bqschema_to_nullsafe_dtypes(type_, expected):
         assert result == {"x": expected}
 
 
+@pytest.mark.parametrize(
+    ("data", "schema_type", "expected"),
+    [
+        (
+            pandas.Timestamp("2017-01-01T12:00:00Z").as_unit("us"),
+            "TIMESTAMP",
+            pandas.DatetimeTZDtype(unit="ns", tz="UTC"),
+        ),
+        (
+            pandas.Timestamp("2017-01-01T12:00:00").as_unit("us"),
+            "DATETIME",
+            numpy.dtype("datetime64[ns]"),
+        ),
+    ],
+)
+def test__finalize_dtypes(data, schema_type, expected):
+    result = gbq._finalize_dtypes(
+        pandas.DataFrame([dict(x=data)]),
+        [dict(name="x", type=schema_type, mode="NULLABLE")],
+    )
+    assert result["x"].dtype == expected
+
+
 @pytest.mark.parametrize(
     ["query_or_table", "expected"],
     [