diff --git a/narwhals/_spark_like/expr.py b/narwhals/_spark_like/expr.py index 7c48731443..965627ff82 100644 --- a/narwhals/_spark_like/expr.py +++ b/narwhals/_spark_like/expr.py @@ -24,6 +24,7 @@ if TYPE_CHECKING: from pyspark.sql import Column + from pyspark.sql import Window from typing_extensions import Self from narwhals._spark_like.dataframe import SparkLikeLazyFrame @@ -78,7 +79,11 @@ def func(df: SparkLikeLazyFrame) -> Sequence[Column]: ) @property - def _F(self: Self) -> Any: # noqa: N802 + def _F(self: Self): # type: ignore[no-untyped-def] # noqa: ANN202, N802 + if TYPE_CHECKING: + from pyspark.sql import functions + + return functions if self._implementation is Implementation.SQLFRAME: from sqlframe.base.session import _BaseSession @@ -91,7 +96,11 @@ def _F(self: Self) -> Any: # noqa: N802 return functions @property - def _native_dtypes(self: Self) -> Any: + def _native_dtypes(self: Self): # type: ignore[no-untyped-def] # noqa: ANN202 + if TYPE_CHECKING: + from pyspark.sql import types + + return types if self._implementation is Implementation.SQLFRAME: from sqlframe.base.session import _BaseSession @@ -104,7 +113,7 @@ def _native_dtypes(self: Self) -> Any: return types @property - def _Window(self: Self) -> Any: # noqa: N802 + def _Window(self: Self) -> type[Window]: # noqa: N802 if self._implementation is Implementation.SQLFRAME: from sqlframe.base.session import _BaseSession diff --git a/narwhals/_spark_like/expr_str.py b/narwhals/_spark_like/expr_str.py index c84eefecd1..d8646d4dee 100644 --- a/narwhals/_spark_like/expr_str.py +++ b/narwhals/_spark_like/expr_str.py @@ -1,7 +1,7 @@ from __future__ import annotations +from functools import partial from typing import TYPE_CHECKING -from typing import overload if TYPE_CHECKING: from pyspark.sql import Column @@ -107,47 +107,29 @@ def to_lowercase(self: Self) -> SparkLikeExpr: ) def to_datetime(self: Self, format: str | None) -> SparkLikeExpr: # noqa: A002 - is_naive = ( - format is not None - and "%s" not in format - and "%z" not in format - and "Z" not in format - ) - function = ( - self._compliant_expr._F.to_timestamp_ntz - if is_naive - else self._compliant_expr._F.to_timestamp - ) - pyspark_format = strptime_to_pyspark_format(format) - format = ( - self._compliant_expr._F.lit(pyspark_format) if is_naive else pyspark_format - ) + F = self._compliant_expr._F # noqa: N806 + if not format: + function = F.to_timestamp + elif is_naive_format(format): + function = partial( + F.to_timestamp_ntz, format=F.lit(strptime_to_pyspark_format(format)) + ) + else: + format = strptime_to_pyspark_format(format) + function = partial(F.to_timestamp, format=format) return self._compliant_expr._from_call( - lambda _input: function( - self._compliant_expr._F.replace( - _input, - self._compliant_expr._F.lit("T"), - self._compliant_expr._F.lit(" "), - ), - format=format, - ), + lambda _input: function(F.replace(_input, F.lit("T"), F.lit(" "))), "to_datetime", ) -@overload -def strptime_to_pyspark_format(format: None) -> None: ... - - -@overload -def strptime_to_pyspark_format(format: str) -> str: ... +def is_naive_format(format: str) -> bool: # noqa: A002 + return not any(x in format for x in ("%s", "%z", "Z")) -def strptime_to_pyspark_format(format: str | None) -> str | None: # noqa: A002 +def strptime_to_pyspark_format(format: str) -> str: # noqa: A002 """Converts a Python strptime datetime format string to a PySpark datetime format string.""" # Mapping from Python strptime format to PySpark format - if format is None: - return None # see https://spark.apache.org/docs/latest/sql-ref-datetime-pattern.html # and https://docs.python.org/3/library/datetime.html#strftime-strptime-behavior