feat: add to_string method to SparkLikeExprDateTimeNamespace (#1842)

* feat: add `to_string` method to Spark `Expr.dt`

* refactor(spark): centralize datetime format conversion

* fix: raise on selectors addition (#1854)

* avoid checking datetime objects

* less hacky workaround

* fix get_native_namespace test

---------

Co-authored-by: Francesco Bruzzesi <[email protected]>
Co-authored-by: FBruzzesi <[email protected]>

Author: Dhanunjaya-Elluri

narwhals/_spark_like/expr_dt.py

@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING
 
 from narwhals._duration import parse_interval_string
-from narwhals._spark_like.utils import UNITS_DICT
+from narwhals._spark_like.utils import UNITS_DICT, strptime_to_pyspark_format
 
 if TYPE_CHECKING:
     from sqlframe.base.column import Column
@@ -15,6 +15,40 @@ class SparkLikeExprDateTimeNamespace:
     def __init__(self, expr: SparkLikeExpr) -> None:
         self._compliant_expr = expr
 
+    def to_string(self, format: str) -> SparkLikeExpr:
+        F = self._compliant_expr._F  # noqa: N806
+
+        def _to_string(_input: Column) -> Column:
+            # Handle special formats
+            if format == "%G-W%V":
+                return self._format_iso_week(_input)
+            if format == "%G-W%V-%u":
+                return self._format_iso_week_with_day(_input)
+
+            format_, suffix = self._format_microseconds(_input, format)
+
+            # Convert Python format to PySpark format
+            pyspark_fmt = strptime_to_pyspark_format(format_)
+
+            result = F.date_format(_input, pyspark_fmt)
+            if "T" in format_:
+                # `strptime_to_pyspark_format` replaces "T" with " " since pyspark
+                # does not support the literal "T" in `date_format`.
+                # If no other spaces are in the given format, then we can revert this
+                # operation, otherwise we raise an exception.
+                if " " not in format_:
+                    result = F.replace(result, F.lit(" "), F.lit("T"))
+                else:  # pragma: no cover
+                    msg = (
+                        "`dt.to_string` with a format that contains both spaces and "
+                        " the literal 'T' is not supported for spark-like backends."
+                    )
+                    raise NotImplementedError(msg)
+
+            return F.concat(result, *suffix)
+
+        return self._compliant_expr._with_callable(_to_string)
+
     def date(self) -> SparkLikeExpr:
         return self._compliant_expr._with_callable(self._compliant_expr._F.to_date)
 
@@ -89,3 +123,40 @@ def replace_time_zone(self, time_zone: str | None) -> SparkLikeExpr:
         else:  # pragma: no cover
             msg = "`replace_time_zone` with non-null `time_zone` not yet implemented for spark-like"
             raise NotImplementedError(msg)
+
+    def _format_iso_week_with_day(self, _input: Column) -> Column:
+        """Format datetime as ISO week string with day."""
+        F = self._compliant_expr._F  # noqa: N806
+
+        year = F.date_format(_input, "yyyy")
+        week = F.lpad(F.weekofyear(_input).cast("string"), 2, "0")
+        day = F.dayofweek(_input)
+        # Adjust Sunday from 1 to 7
+        day = F.when(day == 1, 7).otherwise(day - 1)
+        return F.concat(year, F.lit("-W"), week, F.lit("-"), day.cast("string"))
+
+    def _format_iso_week(self, _input: Column) -> Column:
+        """Format datetime as ISO week string."""
+        F = self._compliant_expr._F  # noqa: N806
+
+        year = F.date_format(_input, "yyyy")
+        week = F.lpad(F.weekofyear(_input).cast("string"), 2, "0")
+        return F.concat(year, F.lit("-W"), week)
+
+    def _format_microseconds(
+        self, _input: Column, format: str
+    ) -> tuple[str, tuple[Column, ...]]:
+        """Format microseconds if present in format, else it's a no-op."""
+        F = self._compliant_expr._F  # noqa: N806
+
+        suffix: tuple[Column, ...]
+        if format.endswith((".%f", "%.f")):
+            import re
+
+            micros = F.unix_micros(_input) % 1_000_000
+            micros_str = F.lpad(micros.cast("string"), 6, "0")
+            suffix = (F.lit("."), micros_str)
+            format_ = re.sub(r"(.%|%.)f$", "", format)
+            return format_, suffix
+
+        return format, ()

narwhals/_spark_like/expr_str.py

@@ -3,6 +3,7 @@
 from functools import partial
 from typing import TYPE_CHECKING
 
+from narwhals._spark_like.utils import strptime_to_pyspark_format
 from narwhals.utils import _is_naive_format
 
 if TYPE_CHECKING:
@@ -112,34 +113,3 @@ def to_datetime(self, format: str | None) -> SparkLikeExpr:
         return self._compliant_expr._with_callable(
             lambda expr: function(F.replace(expr, F.lit("T"), F.lit(" ")))
         )
-
-
-def strptime_to_pyspark_format(format: str) -> str:
-    """Converts a Python strptime datetime format string to a PySpark datetime format string."""
-    # Mapping from Python strptime format to PySpark format
-
-    # see https://spark.apache.org/docs/latest/sql-ref-datetime-pattern.html
-    # and https://docs.python.org/3/library/datetime.html#strftime-strptime-behavior
-    format_mapping = {
-        "%Y": "y",  # Year with century
-        "%y": "y",  # Year without century
-        "%m": "M",  # Month
-        "%d": "d",  # Day of the month
-        "%H": "H",  # Hour (24-hour clock) 0-23
-        "%I": "h",  # Hour (12-hour clock) 1-12
-        "%M": "m",  # Minute
-        "%S": "s",  # Second
-        "%f": "S",  # Microseconds -> Milliseconds
-        "%p": "a",  # AM/PM
-        "%a": "E",  # Abbreviated weekday name
-        "%A": "E",  # Full weekday name
-        "%j": "D",  # Day of the year
-        "%z": "Z",  # Timezone offset
-        "%s": "X",  # Unix timestamp
-    }
-
-    # Replace Python format specifiers with PySpark specifiers
-    pyspark_format = format
-    for py_format, spark_format in format_mapping.items():
-        pyspark_format = pyspark_format.replace(py_format, spark_format)
-    return pyspark_format.replace("T", " ")

narwhals/_spark_like/utils.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from importlib import import_module
-from typing import TYPE_CHECKING, Any, Sequence
+from typing import TYPE_CHECKING, Any, Sequence, overload
 
 from narwhals.exceptions import UnsupportedDTypeError
 from narwhals.utils import Implementation, isinstance_or_issubclass
@@ -33,6 +33,26 @@
     "ns": "nanosecond",
 }
 
+# see https://spark.apache.org/docs/latest/sql-ref-datetime-pattern.html
+# and https://docs.python.org/3/library/datetime.html#strftime-strptime-behavior
+DATETIME_PATTERNS_MAPPING = {
+    "%Y": "yyyy",  # Year with century (4 digits)
+    "%y": "yy",  # Year without century (2 digits)
+    "%m": "MM",  # Month (01-12)
+    "%d": "dd",  # Day of the month (01-31)
+    "%H": "HH",  # Hour (24-hour clock) (00-23)
+    "%I": "hh",  # Hour (12-hour clock) (01-12)
+    "%M": "mm",  # Minute (00-59)
+    "%S": "ss",  # Second (00-59)
+    "%f": "S",  # Microseconds -> Milliseconds
+    "%p": "a",  # AM/PM
+    "%a": "E",  # Abbreviated weekday name
+    "%A": "E",  # Full weekday name
+    "%j": "D",  # Day of the year
+    "%z": "Z",  # Timezone offset
+    "%s": "X",  # Unix timestamp
+}
+
 
 class WindowInputs:
     __slots__ = ("expr", "order_by", "partition_by")
@@ -250,3 +270,23 @@ def import_window(implementation: Implementation, /) -> type[Any]:
     return import_module(
         f"sqlframe.{_BaseSession().execution_dialect_name}.window"
     ).Window
+
+
+@overload
+def strptime_to_pyspark_format(format: None) -> None: ...
+
+
+@overload
+def strptime_to_pyspark_format(format: str) -> str: ...
+
+
+def strptime_to_pyspark_format(format: str | None) -> str | None:
+    """Converts a Python strptime datetime format string to a PySpark datetime format string."""
+    if format is None:  # pragma: no cover
+        return None
+
+    # Replace Python format specifiers with PySpark specifiers
+    pyspark_format = format
+    for py_format, spark_format in DATETIME_PATTERNS_MAPPING.items():
+        pyspark_format = pyspark_format.replace(py_format, spark_format)
+    return pyspark_format.replace("T", " ")

tests/expr_and_series/dt/replace_time_zone_test.py

@@ -48,19 +48,14 @@ def test_replace_time_zone(
     assert_equal_data(result_str, expected)
 
 
-def test_replace_time_zone_none(
-    constructor: Constructor, request: pytest.FixtureRequest
-) -> None:
+def test_replace_time_zone_none(constructor: Constructor) -> None:
     if (
         ("pyarrow" in str(constructor) and is_windows())
         or ("pandas_pyarrow" in str(constructor) and PANDAS_VERSION < (2,))
         or ("modin_pyarrow" in str(constructor) and PANDAS_VERSION < (2,))
         or ("pyarrow_table" in str(constructor) and PYARROW_VERSION < (12,))
     ):
         pytest.skip()
-    if any(x in str(constructor) for x in ("pyspark",)):
-        # pyspark: needs `to_string`
-        request.applymarker(pytest.mark.xfail)
     data = {
         "a": [
             datetime(2020, 1, 1, tzinfo=timezone.utc),

tests/expr_and_series/dt/to_string_test.py

@@ -39,11 +39,7 @@ def test_dt_to_string_series(constructor_eager: ConstructorEager, fmt: str) -> N
     "fmt", ["%Y-%m-%d", "%Y-%m-%d %H:%M:%S", "%Y/%m/%d %H:%M:%S", "%G-W%V-%u", "%G-W%V"]
 )
 @pytest.mark.skipif(is_windows(), reason="pyarrow breaking on windows")
-def test_dt_to_string_expr(
-    constructor: Constructor, fmt: str, request: pytest.FixtureRequest
-) -> None:
-    if "pyspark" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
+def test_dt_to_string_expr(constructor: Constructor, fmt: str) -> None:
     input_frame = nw.from_native(constructor(data))
 
     expected_col = [datetime.strftime(d, fmt) for d in data["a"]]
@@ -115,23 +111,19 @@ def test_dt_to_string_iso_local_datetime_expr(
     expected: str,
     request: pytest.FixtureRequest,
 ) -> None:
-    if (
-        ("pyspark" in str(constructor))
-        or "duckdb" in str(constructor)
-        or "ibis" in str(constructor)
-    ):
+    if "duckdb" in str(constructor) or "ibis" in str(constructor):
         request.applymarker(pytest.mark.xfail)
     df = constructor({"a": [data]})
 
-    result = nw.from_native(df).with_columns(
-        _clean_string_expr(nw.col("a").dt.to_string("%Y-%m-%dT%H:%M:%S.%f")).alias("b")
+    result = nw.from_native(df).select(
+        b=_clean_string_expr(nw.col("a").dt.to_string("%Y-%m-%dT%H:%M:%S.%f"))
     )
-    assert_equal_data(result, {"a": [data], "b": [_clean_string(expected)]})
+    assert_equal_data(result, {"b": [_clean_string(expected)]})
 
-    result = nw.from_native(df).with_columns(
-        _clean_string_expr(nw.col("a").dt.to_string("%Y-%m-%dT%H:%M:%S%.f")).alias("b")
+    result = nw.from_native(df).select(
+        b=_clean_string_expr(nw.col("a").dt.to_string("%Y-%m-%dT%H:%M:%S%.f"))
     )
-    assert_equal_data(result, {"a": [data], "b": [_clean_string(expected)]})
+    assert_equal_data(result, {"b": [_clean_string(expected)]})
 
 
 @pytest.mark.parametrize(("data", "expected"), [(datetime(2020, 1, 9), "2020-01-09")])
@@ -147,15 +139,8 @@ def test_dt_to_string_iso_local_date_series(
 @pytest.mark.parametrize(("data", "expected"), [(datetime(2020, 1, 9), "2020-01-09")])
 @pytest.mark.skipif(is_windows(), reason="pyarrow breaking on windows")
 def test_dt_to_string_iso_local_date_expr(
-    constructor: Constructor,
-    data: datetime,
-    expected: str,
-    request: pytest.FixtureRequest,
+    constructor: Constructor, data: datetime, expected: str
 ) -> None:
-    if "pyspark" in str(constructor):
-        request.applymarker(pytest.mark.xfail)
     df = constructor({"a": [data]})
-    result = nw.from_native(df).with_columns(
-        nw.col("a").dt.to_string("%Y-%m-%d").alias("b")
-    )
-    assert_equal_data(result, {"a": [data], "b": [expected]})
+    result = nw.from_native(df).select(b=nw.col("a").dt.to_string("%Y-%m-%d"))
+    assert_equal_data(result, {"b": [expected]})