fix(postgres): fix insertion of NaT/None into timestamp columns (#…

…10748)

ibis/backends/postgres/__init__.py

@@ -8,6 +8,7 @@
 from typing import TYPE_CHECKING, Any
 from urllib.parse import unquote_plus
 
+import psycopg
 import sqlglot as sg
 import sqlglot.expressions as sge
 from pandas.api.types import is_float_dtype
@@ -31,10 +32,14 @@
 
     import pandas as pd
     import polars as pl
-    import psycopg
     import pyarrow as pa
 
 
+class NatDumper(psycopg.adapt.Dumper):
+    def dump(self, obj, context: Any | None = None) -> str | None:
+        return None
+
+
 class Backend(SQLBackend, CanListCatalog, CanCreateDatabase):
     name = "postgres"
     compiler = sc.postgres.compiler
@@ -233,6 +238,7 @@ def do_connect(
           year            int32
           month           int32
         """
+        import pandas as pd
         import psycopg
         import psycopg.types.json
 
@@ -248,6 +254,8 @@ def do_connect(
             **kwargs,
         )
 
+        self.con.adapters.register_dumper(type(pd.NaT), NatDumper)
+
         self._post_connect()
 
     @util.experimental

ibis/backends/postgres/tests/test_functions.py

@@ -15,6 +15,7 @@
 import ibis.expr.datatypes as dt
 import ibis.expr.types as ir
 from ibis import literal as L
+from ibis.util import gen_name
 
 pytest.importorskip("psycopg")
 
@@ -1230,3 +1231,18 @@ def test_array_discovery(con):
         )
     )
     assert t.schema() == expected
+
+
+@pytest.mark.parametrize("tz", [None, "UTC", "America/New_York"])
+def test_insert_null_timestamp(con, tz):
+    name = gen_name("test_insert_nat")
+    ts = pd.Timestamp(datetime(2025, 1, 3, 12, 37, 38, 234236), tz=tz)
+    df = pd.DataFrame({"ts": [ts, None]})
+
+    # check that timezones match the input
+    t = con.create_table(name, obj=df, temp=True)
+    assert t.ts.type().timezone == tz
+
+    # check that the NaT went into the database as NULL
+    res = t.ts.count().execute()
+    assert res == 1

ibis/backends/risingwave/__init__.py

@@ -122,7 +122,7 @@ def begin(self):
     def _fetch_from_cursor(self, cursor, schema: sch.Schema) -> pd.DataFrame:
         import pandas as pd
 
-        from ibis.backends.postgres.converter import PostgresPandasData
+        from ibis.backends.risingwave.converter import RisingWavePandasData
 
         try:
             df = pd.DataFrame.from_records(
@@ -135,7 +135,7 @@ def _fetch_from_cursor(self, cursor, schema: sch.Schema) -> pd.DataFrame:
             # artificially locked tables
             cursor.close()
             raise
-        df = PostgresPandasData.convert_table(df, schema)
+        df = RisingWavePandasData.convert_table(df, schema)
         return df
 
     @property

ibis/backends/risingwave/converter.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from ibis.formats.pandas import PandasData
+
+
+class RisingWavePandasData(PandasData):
+    @classmethod
+    def convert_Binary(cls, s, dtype, pandas_type):
+        return s.map(bytes, na_action="ignore")