fix(databricks/pyspark): unify timestamp/timestamp_ntz behavior

Author: cpcloud

ibis/backends/pyspark/datatypes.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from functools import partial
+
 import pyspark
 import pyspark.sql.types as pt
 from packaging.version import parse as vparse
@@ -27,9 +29,15 @@
     pt.NullType: dt.Null,
     pt.ShortType: dt.Int16,
     pt.StringType: dt.String,
-    pt.TimestampType: dt.Timestamp,
 }
 
+try:
+    _from_pyspark_dtypes[pt.TimestampNTZType] = dt.Timestamp
+except AttributeError:
+    _from_pyspark_dtypes[pt.TimestampType] = dt.Timestamp
+else:
+    _from_pyspark_dtypes[pt.TimestampType] = partial(dt.Timestamp, timezone="UTC")
+
 _to_pyspark_dtypes = {v: k for k, v in _from_pyspark_dtypes.items()}
 _to_pyspark_dtypes[dt.JSON] = pt.StringType
 _to_pyspark_dtypes[dt.UUID] = pt.StringType
@@ -54,9 +62,7 @@ def to_ibis(cls, typ, nullable=True):
             return dt.Array(cls.to_ibis(typ.elementType), nullable=nullable)
         elif isinstance(typ, pt.MapType):
             return dt.Map(
-                cls.to_ibis(typ.keyType),
-                cls.to_ibis(typ.valueType),
-                nullable=nullable,
+                cls.to_ibis(typ.keyType), cls.to_ibis(typ.valueType), nullable=nullable
             )
         elif isinstance(typ, pt.StructType):
             fields = {f.name: cls.to_ibis(f.dataType) for f in typ.fields}
@@ -102,6 +108,11 @@ def from_ibis(cls, dtype):
                 for n, t in dtype.fields.items()
             ]
             return pt.StructType(fields)
+        elif dtype.is_timestamp():
+            if dtype.timezone is not None:
+                return pt.TimestampType()
+            else:
+                return pt.TimestampNTZType()
         else:
             try:
                 return _to_pyspark_dtypes[type(dtype)]()

ibis/backends/pyspark/tests/conftest.py

@@ -337,8 +337,16 @@ def _load_data(self, **_: Any) -> None:
 
             for name, schema in TEST_TABLES.items():
                 path = str(self.data_dir / "directory" / "parquet" / name)
+                sch = ibis.schema(
+                    {
+                        col: dtype.copy(timezone="UTC")
+                        if dtype.is_timestamp()
+                        else dtype
+                        for col, dtype in schema.items()
+                    }
+                )
                 t = (
-                    s.readStream.schema(PySparkSchema.from_ibis(schema))
+                    s.readStream.schema(PySparkSchema.from_ibis(sch))
                     .parquet(path)
                     .repartition(num_partitions)
                 )

ibis/backends/pyspark/tests/test_window.py

@@ -50,7 +50,7 @@ def test_time_indexed_window(t, spark_table, ibis_windows, spark_range):
         F.mean(spark_table["value"]).over(spark_window),
     ).toPandas()
 
-    tm.assert_frame_equal(result, expected)
+    tm.assert_frame_equal(result, expected, check_dtype=False)
 
 
 @pytest.mark.parametrize(
@@ -90,7 +90,7 @@ def test_multiple_windows(t, spark_table, ibis_windows, spark_range):
         )
         .toPandas()
     )
-    tm.assert_frame_equal(result, expected)
+    tm.assert_frame_equal(result, expected, check_dtype=False)
 
 
 def test_tumble_window_by_grouped_agg(con_streaming, tmp_path):

ibis/backends/tests/test_temporal.py

@@ -42,6 +42,7 @@
 )
 from ibis.common.annotations import ValidationError
 from ibis.conftest import IS_SPARK_REMOTE
+from ibis.util import gen_name
 
 np = pytest.importorskip("numpy")
 pd = pytest.importorskip("pandas")
@@ -2345,3 +2346,17 @@ def test_simple_unix_date_offset(con):
     result = con.execute(expr)
     delta = datetime.date(2023, 4, 7) - datetime.date(1970, 1, 1)
     assert result == delta.days
+
+
+def test_basic_timestamp_with_timezone(con):
+    name = gen_name("tmp_tz")
+    ts = "2023-01-07 13:20:05.561021"
+    dtype = dt.Timestamp(timezone="UTC")
+    colname = "ts"
+    result = con.create_table(
+        name, ibis.timestamp(ts).cast(dtype).name(colname).as_table()
+    )
+    try:
+        assert result.schema() == ibis.schema({colname: dtype})
+    finally:
+        con.drop_table(name, force=True)