fix root names in pyspark reduction with nw.all()

narwhals-dev · Jan 10, 2025 · fd24a5e · fd24a5e
1 parent 58dd1d2
commit fd24a5e
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/narwhals/_spark_like/expr.py b/narwhals/_spark_like/expr.py
@@ -102,7 +102,9 @@ def func(df: SparkLikeLazyFrame) -> list[Column]:
             for _input in inputs:
                 input_col_name = get_column_name(df, _input)
                 column_result = call(_input, **_kwargs)
-                if not returns_scalar:
+                if not returns_scalar or (
+                    (self._depth == 0) and (self._function_name == "all")
+                ):
                     column_result = column_result.alias(input_col_name)
                 results.append(column_result)
             return results

diff --git a/tests/expr_and_series/n_unique_test.py b/tests/expr_and_series/n_unique_test.py
@@ -13,7 +13,7 @@
 
 def test_n_unique(constructor: Constructor) -> None:
     df = nw.from_native(constructor(data))
-    result = df.select(nw.col("a", "b").n_unique())
+    result = df.select(nw.all().n_unique())
     expected = {"a": [3], "b": [4]}
     assert_equal_data(result, expected)