fix: support various reductions in pyspark

narwhals/_spark_like/dataframe.py

@@ -6,8 +6,12 @@
 from typing import Literal
 from typing import Sequence
 
+from pyspark.sql import Window
+from pyspark.sql import functions as F  # noqa: N812
+
 from narwhals._spark_like.utils import native_to_narwhals_dtype
 from narwhals._spark_like.utils import parse_exprs_and_named_exprs
+from narwhals.typing import CompliantLazyFrame
 from narwhals.utils import Implementation
 from narwhals.utils import check_column_exists
 from narwhals.utils import parse_columns_to_drop
@@ -26,7 +30,6 @@
     from narwhals._spark_like.namespace import SparkLikeNamespace
     from narwhals.dtypes import DType
     from narwhals.utils import Version
-from narwhals.typing import CompliantLazyFrame
 
 
 class SparkLikeLazyFrame(CompliantLazyFrame):
@@ -94,7 +97,9 @@ def select(
         *exprs: SparkLikeExpr,
         **named_exprs: SparkLikeExpr,
     ) -> Self:
-        new_columns = parse_exprs_and_named_exprs(self)(*exprs, **named_exprs)
+        new_columns, returns_scalar = parse_exprs_and_named_exprs(self)(
+            *exprs, **named_exprs
+        )
 
         if not new_columns:
             # return empty dataframe, like Polars does
@@ -105,8 +110,38 @@ def select(
 
             return self._from_native_frame(spark_df)
 
-        new_columns_list = [col.alias(col_name) for col_name, col in new_columns.items()]
-        return self._from_native_frame(self._native_frame.select(*new_columns_list))
+        if all(returns_scalar):
+            new_columns_list = [
+                col.alias(col_name) for col_name, col in new_columns.items()
+            ]
+            return self._from_native_frame(self._native_frame.agg(*new_columns_list))
+        else:
+            new_columns_list = [
+                col.over(Window.partitionBy(F.lit(1))).alias(col_name)
+                if _returns_scalar
+                else col.alias(col_name)
+                for (col_name, col), _returns_scalar in zip(
+                    new_columns.items(), returns_scalar
+                )
+            ]
+            return self._from_native_frame(self._native_frame.select(*new_columns_list))
+
+    def with_columns(
+        self: Self,
+        *exprs: SparkLikeExpr,
+        **named_exprs: SparkLikeExpr,
+    ) -> Self:
+        new_columns, returns_scalar = parse_exprs_and_named_exprs(self)(
+            *exprs, **named_exprs
+        )
+
+        new_columns_map = {
+            col_name: col.over(Window.partitionBy(F.lit(1))) if _returns_scalar else col
+            for (col_name, col), _returns_scalar in zip(
+                new_columns.items(), returns_scalar
+            )
+        }
+        return self._from_native_frame(self._native_frame.withColumns(new_columns_map))
 
     def filter(self: Self, *predicates: SparkLikeExpr, **constraints: Any) -> Self:
         plx = self.__narwhals_namespace__()
@@ -130,14 +165,6 @@ def schema(self: Self) -> dict[str, DType]:
     def collect_schema(self: Self) -> dict[str, DType]:
         return self.schema
 
-    def with_columns(
-        self: Self,
-        *exprs: SparkLikeExpr,
-        **named_exprs: SparkLikeExpr,
-    ) -> Self:
-        new_columns_map = parse_exprs_and_named_exprs(self)(*exprs, **named_exprs)
-        return self._from_native_frame(self._native_frame.withColumns(new_columns_map))
-
     def drop(self: Self, columns: list[str], strict: bool) -> Self:  # noqa: FBT001
         columns_to_drop = parse_columns_to_drop(
             compliant_frame=self, columns=columns, strict=strict

narwhals/_spark_like/expr.py

@@ -126,7 +126,7 @@ def _from_call(
         def func(df: SparkLikeLazyFrame) -> list[Column]:
             native_series_list = self._call(df)
             other_native_series = {
-                key: maybe_evaluate(df, value)
+                key: maybe_evaluate(df, value, returns_scalar=returns_scalar)
                 for key, value in expressifiable_args.items()
             }
             return [
@@ -140,7 +140,7 @@ def func(df: SparkLikeLazyFrame) -> list[Column]:
             function_name=f"{self._function_name}->{expr_name}",
             evaluate_output_names=self._evaluate_output_names,
             alias_output_names=self._alias_output_names,
-            returns_scalar=self._returns_scalar or returns_scalar,
+            returns_scalar=returns_scalar,
             backend_version=self._backend_version,
             version=self._version,
             kwargs=expressifiable_args,
@@ -349,26 +349,30 @@ def sum(self: Self) -> Self:
         return self._from_call(F.sum, "sum", returns_scalar=True)
 
     def std(self: Self, ddof: int) -> Self:
-        from functools import partial
-
         import numpy as np  # ignore-banned-import
 
         from narwhals._spark_like.utils import _std
 
-        func = partial(_std, ddof=ddof, np_version=parse_version(np.__version__))
-
-        return self._from_call(func, "std", returns_scalar=True, ddof=ddof)
+        return self._from_call(
+            _std,
+            "std",
+            returns_scalar=True,
+            ddof=ddof,
+            np_version=parse_version(np.__version__),
+        )
 
     def var(self: Self, ddof: int) -> Self:
-        from functools import partial
-
         import numpy as np  # ignore-banned-import
 
         from narwhals._spark_like.utils import _var
 
-        func = partial(_var, ddof=ddof, np_version=parse_version(np.__version__))
-
-        return self._from_call(func, "var", returns_scalar=True, ddof=ddof)
+        return self._from_call(
+            _var,
+            "var",
+            returns_scalar=True,
+            ddof=ddof,
+            np_version=parse_version(np.__version__),
+        )
 
     def clip(
         self: Self,

narwhals/_spark_like/utils.py

@@ -5,14 +5,15 @@
 from typing import Any
 from typing import Callable
 
+from pyspark.sql import Column
+from pyspark.sql import Window
 from pyspark.sql import functions as F  # noqa: N812
 
 from narwhals.exceptions import UnsupportedDTypeError
 from narwhals.utils import import_dtypes_module
 from narwhals.utils import isinstance_or_issubclass
 
 if TYPE_CHECKING:
-    from pyspark.sql import Column
     from pyspark.sql import types as pyspark_types
 
     from narwhals._spark_like.dataframe import SparkLikeLazyFrame
@@ -112,9 +113,16 @@ def narwhals_to_native_dtype(
 
 def parse_exprs_and_named_exprs(
     df: SparkLikeLazyFrame,
-) -> Callable[..., dict[str, Column]]:
-    def func(*exprs: SparkLikeExpr, **named_exprs: SparkLikeExpr) -> dict[str, Column]:
+) -> Callable[..., tuple[dict[str, Column], list[bool]]]:
+    def func(
+        *exprs: SparkLikeExpr, **named_exprs: SparkLikeExpr
+    ) -> tuple[dict[str, Column], list[bool]]:
         native_results: dict[str, list[Column]] = {}
+
+        # `returns_scalar` keeps track if an expression returns a scalar and is not lit.
+        # Notice that lit is quite special case, since it gets broadcasted by pyspark
+        # without the need of adding `.over(Window.partitionBy(F.lit(1)))`
+        returns_scalar: list[bool] = []
         for expr in exprs:
             native_series_list = expr._call(df)
             output_names = expr._evaluate_output_names(df)
@@ -124,18 +132,23 @@ def func(*exprs: SparkLikeExpr, **named_exprs: SparkLikeExpr) -> dict[str, Colum
                 msg = f"Internal error: got output names {output_names}, but only got {len(native_series_list)} results"
                 raise AssertionError(msg)
             native_results.update(zip(output_names, native_series_list))
+            returns_scalar.extend(
+                [expr._returns_scalar and expr._function_name != "lit"]
+                * len(output_names)
+            )
         for col_alias, expr in named_exprs.items():
             native_series_list = expr._call(df)
             if len(native_series_list) != 1:  # pragma: no cover
                 msg = "Named expressions must return a single column"
                 raise ValueError(msg)
             native_results[col_alias] = native_series_list[0]
-        return native_results
+            returns_scalar.append(expr._returns_scalar and expr._function_name != "lit")
+        return native_results, returns_scalar
 
     return func
 
 
-def maybe_evaluate(df: SparkLikeLazyFrame, obj: Any) -> Any:
+def maybe_evaluate(df: SparkLikeLazyFrame, obj: Any, *, returns_scalar: bool) -> Column:
     from narwhals._spark_like.expr import SparkLikeExpr
 
     if isinstance(obj, SparkLikeExpr):
@@ -144,10 +157,9 @@ def maybe_evaluate(df: SparkLikeLazyFrame, obj: Any) -> Any:
             msg = "Multi-output expressions (e.g. `nw.all()` or `nw.col('a', 'b')`) not supported in this context"
             raise NotImplementedError(msg)
         column_result = column_results[0]
-        if obj._returns_scalar:
-            # Return scalar, let PySpark do its broadcasting
-            from pyspark.sql.window import Window
-
+        if obj._returns_scalar and obj._function_name != "lit" and not returns_scalar:
+            # Returns scalar, but overall expression doesn't.
+            # Let PySpark do its broadcasting
             return column_result.over(Window.partitionBy(F.lit(1)))
         return column_result
     return obj

pyproject.toml

@@ -176,7 +176,8 @@ xfail_strict = true
 markers = ["slow: marks tests as slow (deselect with '-m \"not slow\"')"]
 env = [
   "MODIN_ENGINE=python",
-  "PYARROW_IGNORE_TIMEZONE=1"
+  "PYARROW_IGNORE_TIMEZONE=1",
+  "TZ=UTC",
 ]
 
 [tool.coverage.run]

tests/expr_and_series/binary_test.py

@@ -9,9 +9,7 @@
 
 
 def test_expr_binary(constructor: Constructor, request: pytest.FixtureRequest) -> None:
-    if ("dask" in str(constructor) and DASK_VERSION < (2024, 10)) or "pyspark" in str(
-        constructor
-    ):
+    if "dask" in str(constructor) and DASK_VERSION < (2024, 10):
         request.applymarker(pytest.mark.xfail)
     data = {"a": [1, 3, 2], "b": [4, 4, 6], "z": [7.0, 8.0, 9.0]}
     df_raw = constructor(data)

tests/expr_and_series/lit_test.py

@@ -105,13 +105,6 @@ def test_lit_operation(
         and DASK_VERSION < (2024, 10)
     ):
         request.applymarker(pytest.mark.xfail)
-    if "pyspark" in str(constructor) and col_name in {
-        "left_lit_with_agg",
-        "left_scalar_with_agg",
-        "right_lit_with_agg",
-        "right_lit",
-    }:
-        request.applymarker(pytest.mark.xfail)
 
     data = {"a": [1, 3, 2]}
     df_raw = constructor(data)

tests/expr_and_series/reduction_test.py

@@ -33,13 +33,6 @@ def test_scalar_reduction_select(
     expected: dict[str, list[Any]],
     request: pytest.FixtureRequest,
 ) -> None:
-    if "pyspark" in str(constructor) and request.node.callspec.id in {
-        "pyspark-2",
-        "pyspark-3",
-        "pyspark-4",
-    }:
-        request.applymarker(pytest.mark.xfail)
-
     if "duckdb" in str(constructor) and request.node.callspec.id not in {"duckdb-0"}:
         request.applymarker(pytest.mark.xfail)
 
@@ -72,9 +65,7 @@ def test_scalar_reduction_with_columns(
     expected: dict[str, list[Any]],
     request: pytest.FixtureRequest,
 ) -> None:
-    if "duckdb" in str(constructor) or (
-        "pyspark" in str(constructor) and request.node.callspec.id != "pyspark-1"
-    ):
+    if "duckdb" in str(constructor):
         request.applymarker(pytest.mark.xfail)
     data = {"a": [1, 2, 3], "b": [4, 5, 6]}
     df = nw.from_native(constructor(data))
@@ -85,6 +76,7 @@ def test_scalar_reduction_with_columns(
 def test_empty_scalar_reduction_select(
     constructor: Constructor, request: pytest.FixtureRequest
 ) -> None:
+    # pyspark doesn't necessarely fails, but returns all None's
     if "pyspark" in str(constructor) or "duckdb" in str(constructor):
         request.applymarker(pytest.mark.xfail)
     data = {
@@ -118,7 +110,7 @@ def test_empty_scalar_reduction_select(
 def test_empty_scalar_reduction_with_columns(
     constructor: Constructor, request: pytest.FixtureRequest
 ) -> None:
-    if "pyspark" in str(constructor) or "duckdb" in str(constructor):
+    if "duckdb" in str(constructor):
         request.applymarker(pytest.mark.xfail)
     from itertools import chain
 

tests/stable_api_test.py

@@ -22,7 +22,7 @@ def remove_docstring_examples(doc: str) -> str:
 def test_renamed_taxicab_norm(
     constructor: Constructor, request: pytest.FixtureRequest
 ) -> None:
-    if ("pyspark" in str(constructor)) or "duckdb" in str(constructor):
+    if "duckdb" in str(constructor):
         request.applymarker(pytest.mark.xfail)
     # Suppose we need to rename `_l1_norm` to `_taxicab_norm`.
     # We need `narwhals.stable.v1` to stay stable. So, we