enh: reflect connection time zone in Datetime for pyspark, support convert_time_zone and replace_time_zone (#2592)

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: MarcoGorelli

narwhals/_spark_like/dataframe.py

@@ -281,9 +281,10 @@ def schema(self) -> dict[str, DType]:
         if self._cached_schema is None:
             self._cached_schema = {
                 field.name: native_to_narwhals_dtype(
-                    dtype=field.dataType,
-                    version=self._version,
-                    spark_types=self._native_dtypes,
+                    field.dataType,
+                    self._version,
+                    self._native_dtypes,
+                    self.native.sparkSession,
                 )
                 for field in self.native.schema
             }

narwhals/_spark_like/expr_dt.py

@@ -1,13 +1,18 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Sequence
 
 from narwhals._duration import parse_interval_string
-from narwhals._spark_like.utils import UNITS_DICT, strptime_to_pyspark_format
+from narwhals._spark_like.utils import (
+    UNITS_DICT,
+    fetch_session_time_zone,
+    strptime_to_pyspark_format,
+)
 
 if TYPE_CHECKING:
     from sqlframe.base.column import Column
 
+    from narwhals._spark_like.dataframe import SparkLikeLazyFrame
     from narwhals._spark_like.expr import SparkLikeExpr
 
 
@@ -115,14 +120,40 @@ def _truncate(expr: Column) -> Column:
 
         return self._compliant_expr._with_callable(_truncate)
 
-    def replace_time_zone(self, time_zone: str | None) -> SparkLikeExpr:
+    def _no_op_time_zone(self, time_zone: str) -> SparkLikeExpr:  # pragma: no cover
+        def func(df: SparkLikeLazyFrame) -> Sequence[Column]:
+            native_series_list = self._compliant_expr(df)
+            conn_time_zone = fetch_session_time_zone(df.native.sparkSession)
+            if conn_time_zone != time_zone:
+                msg = (
+                    "PySpark stores the time zone in the session, rather than in the "
+                    f"data type, so changing the timezone to anything other than {conn_time_zone} "
+                    " (the current session time zone) is not supported."
+                )
+                raise NotImplementedError(msg)
+            return native_series_list
+
+        return self._compliant_expr.__class__(
+            func,
+            evaluate_output_names=self._compliant_expr._evaluate_output_names,
+            alias_output_names=self._compliant_expr._alias_output_names,
+            backend_version=self._compliant_expr._backend_version,
+            version=self._compliant_expr._version,
+            implementation=self._compliant_expr._implementation,
+        )
+
+    def convert_time_zone(self, time_zone: str) -> SparkLikeExpr:  # pragma: no cover
+        return self._no_op_time_zone(time_zone)
+
+    def replace_time_zone(
+        self, time_zone: str | None
+    ) -> SparkLikeExpr:  # pragma: no cover
         if time_zone is None:
             return self._compliant_expr._with_callable(
                 lambda _input: _input.cast("timestamp_ntz")
             )
-        else:  # pragma: no cover
-            msg = "`replace_time_zone` with non-null `time_zone` not yet implemented for spark-like"
-            raise NotImplementedError(msg)
+        else:
+            return self._no_op_time_zone(time_zone)
 
     def _format_iso_week_with_day(self, _input: Column) -> Column:
         """Format datetime as ISO week string with day."""

narwhals/_spark_like/utils.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from functools import lru_cache
 from importlib import import_module
 from typing import TYPE_CHECKING, Any, Sequence, overload
 
@@ -11,6 +12,7 @@
 
     import sqlframe.base.types as sqlframe_types
     from sqlframe.base.column import Column
+    from sqlframe.base.session import _BaseSession as Session
     from typing_extensions import TypeAlias
 
     from narwhals._spark_like.dataframe import SparkLikeLazyFrame
@@ -19,6 +21,7 @@
     from narwhals.utils import Version
 
     _NativeDType: TypeAlias = sqlframe_types.DataType
+    SparkSession = Session[Any, Any, Any, Any, Any, Any, Any]
 
 UNITS_DICT = {
     "y": "year",
@@ -75,7 +78,7 @@ def __init__(self, expr: Column, partition_by: Sequence[str | Column]) -> None:
 
 # NOTE: don't lru_cache this as `ModuleType` isn't hashable
 def native_to_narwhals_dtype(  # noqa: C901, PLR0912
-    dtype: _NativeDType, version: Version, spark_types: ModuleType
+    dtype: _NativeDType, version: Version, spark_types: ModuleType, session: SparkSession
 ) -> DType:
     dtypes = version.dtypes
     if TYPE_CHECKING:
@@ -105,16 +108,14 @@ def native_to_narwhals_dtype(  # noqa: C901, PLR0912
         # TODO(marco): cover this
         return dtypes.Datetime()  # pragma: no cover
     if isinstance(dtype, native.TimestampType):
-        # TODO(marco): is UTC correct, or should we be getting the connection timezone?
-        # https://github.com/narwhals-dev/narwhals/issues/2165
-        return dtypes.Datetime(time_zone="UTC")
+        return dtypes.Datetime(time_zone=fetch_session_time_zone(session))
     if isinstance(dtype, native.DecimalType):
         # TODO(marco): cover this
         return dtypes.Decimal()  # pragma: no cover
     if isinstance(dtype, native.ArrayType):
         return dtypes.List(
             inner=native_to_narwhals_dtype(
-                dtype.elementType, version=version, spark_types=spark_types
+                dtype.elementType, version, spark_types, session
             )
         )
     if isinstance(dtype, native.StructType):
@@ -123,7 +124,7 @@ def native_to_narwhals_dtype(  # noqa: C901, PLR0912
                 dtypes.Field(
                     name=field.name,
                     dtype=native_to_narwhals_dtype(
-                        field.dataType, version=version, spark_types=spark_types
+                        field.dataType, version, spark_types, session
                     ),
                 )
                 for field in dtype
@@ -134,6 +135,16 @@ def native_to_narwhals_dtype(  # noqa: C901, PLR0912
     return dtypes.Unknown()  # pragma: no cover
 
 
+@lru_cache(maxsize=4)
+def fetch_session_time_zone(session: SparkSession) -> str:
+    # Timezone can't be changed in PySpark session, so this can be cached.
+    try:
+        return session.conf.get("spark.sql.session.timeZone")  # type: ignore[attr-defined]
+    except Exception:  # noqa: BLE001
+        # https://github.com/eakmanrq/sqlframe/issues/406
+        return "<unknown>"
+
+
 def narwhals_to_native_dtype(  # noqa: C901, PLR0912
     dtype: DType | type[DType], version: Version, spark_types: ModuleType
 ) -> _NativeDType:

tests/dtypes_test.py

@@ -482,3 +482,26 @@ def test_datetime_w_tz_duckdb() -> None:
     result = df.collect_schema()
     assert result["a"] == nw.Datetime("us", "Asia/Kathmandu")
     assert result["b"] == nw.List(nw.List(nw.Datetime("us", "Asia/Kathmandu")))
+
+
+def test_datetime_w_tz_pyspark(constructor: Constructor) -> None:  # pragma: no cover
+    if "pyspark" not in str(constructor):
+        pytest.skip()
+    pytest.importorskip("pyspark")
+    pytest.importorskip("zoneinfo")
+    from pyspark.sql import SparkSession
+
+    session = SparkSession.builder.config(
+        "spark.sql.session.timeZone", "UTC"
+    ).getOrCreate()
+
+    df = nw.from_native(
+        session.createDataFrame([(datetime(2020, 1, 1, tzinfo=timezone.utc),)], ["a"])
+    )
+    result = df.collect_schema()
+    assert result["a"] == nw.Datetime("us", "UTC")
+    df = nw.from_native(
+        session.createDataFrame([([datetime(2020, 1, 1, tzinfo=timezone.utc)],)], ["a"])
+    )
+    result = df.collect_schema()
+    assert result["a"] == nw.List(nw.Datetime("us", "UTC"))

tests/expr_and_series/dt/convert_time_zone_test.py

@@ -153,3 +153,27 @@ def test_convert_time_zone_to_connection_tz_duckdb() -> None:
         result = nw.from_native(rel).with_columns(
             nw.col("a").dt.convert_time_zone("Asia/Kathmandu")
         )
+
+
+def test_convert_time_zone_to_connection_tz_pyspark(
+    constructor: Constructor,
+) -> None:  # pragma: no cover
+    if "pyspark" not in str(constructor):
+        pytest.skip()
+    pytest.importorskip("pyspark")
+    pytest.importorskip("zoneinfo")
+    from pyspark.sql import SparkSession
+
+    session = SparkSession.builder.config(
+        "spark.sql.session.timeZone", "UTC"
+    ).getOrCreate()
+    df = nw.from_native(
+        session.createDataFrame([(datetime(2020, 1, 1, tzinfo=timezone.utc),)], ["a"])
+    )
+    result = nw.from_native(df).with_columns(nw.col("a").dt.convert_time_zone("UTC"))
+    expected = {"a": [datetime(2020, 1, 1, tzinfo=timezone.utc)]}
+    assert_equal_data(result, expected)
+    with pytest.raises(NotImplementedError):
+        result = nw.from_native(df).with_columns(
+            nw.col("a").dt.convert_time_zone("Asia/Kathmandu")
+        )

tests/expr_and_series/dt/replace_time_zone_test.py

@@ -144,3 +144,27 @@ def test_replace_time_zone_to_connection_tz_duckdb() -> None:
         result = nw.from_native(rel).with_columns(
             nw.col("a").dt.replace_time_zone("Asia/Kathmandu")
         )
+
+
+def test_replace_time_zone_to_connection_tz_pyspark(
+    constructor: Constructor,
+) -> None:  # pragma: no cover
+    if "pyspark" not in str(constructor):
+        pytest.skip()
+    pytest.importorskip("pyspark")
+    pytest.importorskip("zoneinfo")
+    from pyspark.sql import SparkSession
+
+    session = SparkSession.builder.config(
+        "spark.sql.session.timeZone", "UTC"
+    ).getOrCreate()
+    df = nw.from_native(
+        session.createDataFrame([(datetime(2020, 1, 1, tzinfo=timezone.utc),)], ["a"])
+    )
+    result = nw.from_native(df).with_columns(nw.col("a").dt.replace_time_zone("UTC"))
+    expected = {"a": [datetime(2020, 1, 1, tzinfo=timezone.utc)]}
+    assert_equal_data(result, expected)
+    with pytest.raises(NotImplementedError):
+        result = nw.from_native(df).with_columns(
+            nw.col("a").dt.replace_time_zone("Asia/Kathmandu")
+        )