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

Author: cpcloud

ibis/backends/polars/__init__.py

@@ -559,7 +559,7 @@ def _read_in_memory(source: Any, table_name: str, _conn: Backend, **kwargs: Any)
 
 @_read_in_memory.register("ibis.expr.types.Table")
 def _table(source, table_name, _conn, **kwargs: Any):
-    _conn._add_table(table_name, source.to_polars())
+    _conn._add_table(table_name, _conn.to_polars(source))
 
 
 @_read_in_memory.register("polars.DataFrame")

ibis/backends/pyspark/datatypes.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+from functools import partial
+from inspect import isclass
+
 import pyspark
 import pyspark.sql.types as pt
 from packaging.version import parse as vparse
@@ -27,10 +30,20 @@
     pt.NullType: dt.Null,
     pt.ShortType: dt.Int16,
     pt.StringType: dt.String,
-    pt.TimestampType: dt.Timestamp,
 }
 
-_to_pyspark_dtypes = {v: k for k, v in _from_pyspark_dtypes.items()}
+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()
+    if isclass(v) and not issubclass(v, dt.Timestamp) and not isinstance(v, partial)
+}
 _to_pyspark_dtypes[dt.JSON] = pt.StringType
 _to_pyspark_dtypes[dt.UUID] = pt.StringType
 
@@ -54,9 +67,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}
@@ -97,11 +108,17 @@ def from_ibis(cls, dtype):
             value_contains_null = dtype.value_type.nullable
             return pt.MapType(key_type, value_type, value_contains_null)
         elif dtype.is_struct():
-            fields = [
-                pt.StructField(n, cls.from_ibis(t), t.nullable)
-                for n, t in dtype.fields.items()
-            ]
-            return pt.StructType(fields)
+            return pt.StructType(
+                [
+                    pt.StructField(field, cls.from_ibis(dtype), dtype.nullable)
+                    for field, dtype in dtype.fields.items()
+                ]
+            )
+        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)]()
@@ -114,11 +131,7 @@ def from_ibis(cls, dtype):
 class PySparkSchema(SchemaMapper):
     @classmethod
     def from_ibis(cls, schema):
-        fields = [
-            pt.StructField(name, PySparkType.from_ibis(dtype), dtype.nullable)
-            for name, dtype in schema.items()
-        ]
-        return pt.StructType(fields)
+        return PySparkType.from_ibis(schema.as_struct())
 
     @classmethod
     def to_ibis(cls, schema):

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/sql/datatypes.py

@@ -1355,6 +1355,15 @@ def _from_sqlglot_VARIANT(cls, nullable: bool | None = None) -> sge.DataType:
 
     _from_sqlglot_JSON = _from_sqlglot_VARIANT
 
+    @classmethod
+    def _from_ibis_Timestamp(cls, dtype: dt.Timestamp) -> sge.DataType:
+        code = typecode.TIMESTAMPNTZ if dtype.timezone is None else typecode.TIMESTAMPTZ
+        if dtype.scale is not None:
+            scale = sge.DataTypeParam(this=sge.Literal.number(dtype.scale))
+            return sge.DataType(this=code, expressions=[scale])
+        else:
+            return sge.DataType(this=code)
+
 
 class AthenaType(SqlglotType):
     dialect = "athena"

ibis/backends/sql/tests/test_datatypes.py

@@ -10,6 +10,7 @@
 import ibis.tests.strategies as its
 from ibis.backends.sql.datatypes import (
     ClickHouseType,
+    DatabricksType,
     DuckDBType,
     PostgresType,
     SqlglotType,
@@ -85,3 +86,22 @@ def test_interval_without_unit():
 )
 def test_unsupported_dtypes_are_unknown(typengine, typ):
     assert typengine.to_ibis(sge.DataType(this=typ)) == dt.unknown
+
+
+@pytest.mark.parametrize(
+    ("dtype", "expected"),
+    [
+        (
+            dt.Timestamp(timezone="UTC"),
+            sge.DataType(this=sge.DataType.Type.TIMESTAMPTZ),
+        ),
+        (
+            dt.Timestamp(timezone=None),
+            sge.DataType(this=sge.DataType.Type.TIMESTAMPNTZ),
+        ),
+    ],
+    ids=["timezone", "no_timezone"],
+)
+def test_timestamp_with_timezone(dtype, expected):
+    assert DatabricksType.from_ibis(dtype) == expected
+    assert DatabricksType.to_ibis(expected) == dtype

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,35 @@ 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
+
+
+timestamp_with_timezone_params = {
+    "clickhouse": ("UTC", 0),
+    "datafusion": ("+00:00", 9),
+    "duckdb": ("UTC", 6),
+    "impala": (None, None),
+    "oracle": ("UTC", 6),
+    "trino": ("UTC", 3),
+}
+
+
+@pytest.mark.notyet(
+    ["druid"],
+    raises=NotImplementedError,
+    reason="druid doesn't implement `create_table`",
+)
+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"
+    timezone, scale = timestamp_with_timezone_params.get(con.name, ("UTC", None))
+    result = con.create_table(
+        name, ibis.timestamp(ts).cast(dtype).name(colname).as_table()
+    )
+    try:
+        assert result.schema() == ibis.schema(
+            {colname: dtype.copy(timezone=timezone, scale=scale)}
+        )
+    finally:
+        con.drop_table(name, force=True)