Repository: narwhals-dev/narwhals
Pull request #716 (2024-08-04 15:40:22)

patch: pandas-like and pyarrow scalar reduction fix
## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [ ] ✨ Feature
- [x] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue #694, as mentioned in [comment](https://github.com/narwhals-dev/narwhals/pull/702#discussion_r1703230227) there are cases in which the issue in 694 doesn't match polars behavior for pandas and pyarrow as well.

Compared to dask, this is easier to fix as we know series length due to eagerness

## Checklist

- [x] Code follows style guide (ruff)
- [x] Tests added 
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below.




Patch:
diff --git a/docs/api-reference/expr_str.md b/docs/api-reference/expr_str.md
index f64b360f..3cfd8757 100644
--- a/docs/api-reference/expr_str.md
+++ b/docs/api-reference/expr_str.md
@@ -9,6 +9,7 @@
         - head
         - slice
         - starts_with
+        - strip_chars
         - tail
         - to_datetime
         - to_lowercase
diff --git a/docs/api-reference/series_str.md b/docs/api-reference/series_str.md
index 3405b8a4..96dabcb1 100644
--- a/docs/api-reference/series_str.md
+++ b/docs/api-reference/series_str.md
@@ -9,6 +9,7 @@
         - head
         - slice
         - starts_with
+        - strip_chars
         - tail
       show_source: false
       show_bases: false
diff --git a/narwhals/_arrow/dataframe.py b/narwhals/_arrow/dataframe.py
index 4aee5c2f..738732e4 100644
--- a/narwhals/_arrow/dataframe.py
+++ b/narwhals/_arrow/dataframe.py
@@ -8,6 +8,7 @@ from typing import Literal
 from typing import Sequence
 from typing import overload
 
+from narwhals._arrow.utils import broadcast_series
 from narwhals._arrow.utils import translate_dtype
 from narwhals._arrow.utils import validate_dataframe_comparand
 from narwhals._expression_parsing import evaluate_into_exprs
@@ -192,7 +193,10 @@ class ArrowDataFrame:
             )
         names = [s.name for s in new_series]
         pa = get_pyarrow()
-        df = pa.Table.from_arrays([s._native_series for s in new_series], names=names)
+        df = pa.Table.from_arrays(
+            broadcast_series(new_series),
+            names=names,
+        )
         return self._from_native_dataframe(df)
 
     def with_columns(
diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py
index cef2ca35..844901da 100644
--- a/narwhals/_arrow/expr.py
+++ b/narwhals/_arrow/expr.py
@@ -400,6 +400,14 @@ class ArrowExprStringNamespace:
     def __init__(self, expr: ArrowExpr) -> None:
         self._expr = expr
 
+    def strip_chars(self, characters: str | None = None) -> ArrowExpr:
+        return reuse_series_namespace_implementation(
+            self._expr,
+            "str",
+            "strip_chars",
+            characters,
+        )
+
     def starts_with(self, prefix: str) -> ArrowExpr:
         return reuse_series_namespace_implementation(
             self._expr,
diff --git a/narwhals/_arrow/group_by.py b/narwhals/_arrow/group_by.py
index 4da0356a..b05b41b8 100644
--- a/narwhals/_arrow/group_by.py
+++ b/narwhals/_arrow/group_by.py
@@ -4,6 +4,7 @@ from copy import copy
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Callable
+from typing import Iterator
 
 from narwhals._expression_parsing import is_simple_aggregation
 from narwhals._expression_parsing import parse_into_exprs
@@ -53,6 +54,23 @@ class ArrowGroupBy:
             self._df._from_native_dataframe,
         )
 
+    def __iter__(self) -> Iterator[tuple[Any, ArrowDataFrame]]:
+        key_values = (
+            self._df.select(*self._keys)
+            .unique(subset=self._keys, keep="first")
+            .iter_rows()
+        )
+        nw_namespace = self._df.__narwhals_namespace__()
+        yield from (
+            (
+                key_value,
+                self._df.filter(
+                    *[nw_namespace.col(k) == v for k, v in zip(self._keys, key_value)]
+                ),
+            )
+            for key_value in key_values
+        )
+
 
 def agg_arrow(
     grouped: Any,
diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py
index 731db14a..ff3ee94c 100644
--- a/narwhals/_arrow/series.py
+++ b/narwhals/_arrow/series.py
@@ -791,6 +791,16 @@ class ArrowSeriesStringNamespace:
     def __init__(self: Self, series: ArrowSeries) -> None:
         self._arrow_series = series
 
+    def strip_chars(self: Self, characters: str | None = None) -> ArrowSeries:
+        pc = get_pyarrow_compute()
+        whitespace = " \t\n\r\v\f"
+        return self._arrow_series._from_native_series(
+            pc.utf8_trim(
+                self._arrow_series._native_series,
+                characters or whitespace,
+            )
+        )
+
     def starts_with(self: Self, prefix: str) -> ArrowSeries:
         pc = get_pyarrow_compute()
         return self._arrow_series._from_native_series(
diff --git a/narwhals/_arrow/utils.py b/narwhals/_arrow/utils.py
index f98eb2e5..3fc3ec9d 100644
--- a/narwhals/_arrow/utils.py
+++ b/narwhals/_arrow/utils.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
 from typing import Any
 
 from narwhals import dtypes
@@ -7,6 +8,9 @@ from narwhals.dependencies import get_pyarrow
 from narwhals.dependencies import get_pyarrow_compute
 from narwhals.utils import isinstance_or_issubclass
 
+if TYPE_CHECKING:
+    from narwhals._arrow.series import ArrowSeries
+
 
 def translate_dtype(dtype: Any) -> dtypes.DType:
     pa = get_pyarrow()
@@ -246,3 +250,27 @@ def cast_for_truediv(arrow_array: Any, pa_object: Any) -> tuple[Any, Any]:
         )
 
     return arrow_array, pa_object
+
+
+def broadcast_series(series: list[ArrowSeries]) -> list[Any]:
+    lengths = [len(s) for s in series]
+    max_length = max(lengths)
+    fast_path = all(_len == max_length for _len in lengths)
+
+    if fast_path:
+        return [s._native_series for s in series]
+
+    pa = get_pyarrow()
+
+    reshaped = []
+    for s, length in zip(series, lengths):
+        s_native = s._native_series
+        if max_length > 1 and length == 1:
+            value = s_native[0]
+            if hasattr(value, "as_py"):  # pragma: no cover
+                value = value.as_py()
+            reshaped.append(pa.array([value] * max_length, type=s_native.type))
+        else:
+            reshaped.append(s_native)
+
+    return reshaped
diff --git a/narwhals/_dask/expr.py b/narwhals/_dask/expr.py
index 28db2a5c..aed30b1e 100644
--- a/narwhals/_dask/expr.py
+++ b/narwhals/_dask/expr.py
@@ -283,6 +283,11 @@ class DaskExprStringNamespace:
     def __init__(self, expr: DaskExpr) -> None:
         self._expr = expr
 
+    def strip_chars(self, characters: str | None = None) -> DaskExpr:
+        return self._expr._from_call(
+            lambda _input, characters: _input.str.strip(characters), "strip", characters
+        )
+
     def starts_with(self, prefix: str) -> DaskExpr:
         return self._expr._from_call(
             lambda _input, prefix: _input.str.startswith(prefix), "starts_with", prefix
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
index 21588616..0f64cba9 100644
--- a/narwhals/_pandas_like/dataframe.py
+++ b/narwhals/_pandas_like/dataframe.py
@@ -11,11 +11,11 @@ from typing import overload
 
 from narwhals._expression_parsing import evaluate_into_exprs
 from narwhals._pandas_like.expr import PandasLikeExpr
+from narwhals._pandas_like.utils import broadcast_series
 from narwhals._pandas_like.utils import create_native_series
 from narwhals._pandas_like.utils import horizontal_concat
 from narwhals._pandas_like.utils import translate_dtype
 from narwhals._pandas_like.utils import validate_dataframe_comparand
-from narwhals._pandas_like.utils import validate_indices
 from narwhals.dependencies import get_cudf
 from narwhals.dependencies import get_modin
 from narwhals.dependencies import get_numpy
@@ -225,7 +225,7 @@ class PandasLikeDataFrame:
         if not new_series:
             # return empty dataframe, like Polars does
             return self._from_native_dataframe(self._native_dataframe.__class__())
-        new_series = validate_indices(new_series)
+        new_series = broadcast_series(new_series)
         df = horizontal_concat(
             new_series,
             implementation=self._implementation,
diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py
index cb6211dd..e798bd4d 100644
--- a/narwhals/_pandas_like/expr.py
+++ b/narwhals/_pandas_like/expr.py
@@ -365,6 +365,14 @@ class PandasLikeExprStringNamespace:
     def __init__(self, expr: PandasLikeExpr) -> None:
         self._expr = expr
 
+    def strip_chars(self, characters: str | None = None) -> PandasLikeExpr:
+        return reuse_series_namespace_implementation(
+            self._expr,
+            "str",
+            "strip_chars",
+            characters,
+        )
+
     def starts_with(self, prefix: str) -> PandasLikeExpr:
         return reuse_series_namespace_implementation(
             self._expr,
diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py
index 8cb98a9d..e0746186 100644
--- a/narwhals/_pandas_like/series.py
+++ b/narwhals/_pandas_like/series.py
@@ -659,6 +659,11 @@ class PandasLikeSeriesStringNamespace:
     def __init__(self, series: PandasLikeSeries) -> None:
         self._pandas_series = series
 
+    def strip_chars(self, characters: str | None) -> PandasLikeSeries:
+        return self._pandas_series._from_native_series(
+            self._pandas_series._native_series.str.strip(characters),
+        )
+
     def starts_with(self, prefix: str) -> PandasLikeSeries:
         return self._pandas_series._from_native_series(
             self._pandas_series._native_series.str.startswith(prefix),
diff --git a/narwhals/_pandas_like/utils.py b/narwhals/_pandas_like/utils.py
index f58c1297..9c233ecf 100644
--- a/narwhals/_pandas_like/utils.py
+++ b/narwhals/_pandas_like/utils.py
@@ -416,21 +416,38 @@ def reverse_translate_dtype(  # noqa: PLR0915
     raise AssertionError(msg)
 
 
-def validate_indices(series: list[PandasLikeSeries]) -> list[Any]:
-    idx = series[0]._native_series.index
-    reindexed = [series[0]._native_series]
-    for s in series[1:]:
-        if s._native_series.index is not idx:
+def broadcast_series(series: list[PandasLikeSeries]) -> list[Any]:
+    native_namespace = series[0].__native_namespace__()
+
+    lengths = [len(s) for s in series]
+    max_length = max(lengths)
+
+    idx = series[lengths.index(max_length)]._native_series.index
+    reindexed = []
+
+    for s, length in zip(series, lengths):
+        s_native = s._native_series
+        if max_length > 1 and length == 1:
+            reindexed.append(
+                native_namespace.Series(
+                    [s_native.iloc[0]] * max_length,
+                    index=idx,
+                    name=s_native.name,
+                    dtype=s_native.dtype,
+                )
+            )
+
+        elif s_native.index is not idx:
             reindexed.append(
                 set_axis(
-                    s._native_series,
+                    s_native,
                     idx,
                     implementation=s._implementation,
                     backend_version=s._backend_version,
                 )
             )
         else:
-            reindexed.append(s._native_series)
+            reindexed.append(s_native)
     return reindexed
 
 
diff --git a/narwhals/expr.py b/narwhals/expr.py
index 195c6470..baab0f25 100644
--- a/narwhals/expr.py
+++ b/narwhals/expr.py
@@ -1833,6 +1833,40 @@ class ExprStringNamespace:
     def __init__(self, expr: Expr) -> None:
         self._expr = expr
 
+    def strip_chars(self, characters: str | None = None) -> Expr:
+        r"""
+        Remove leading and trailing characters.
+
+        Arguments:
+            characters: The set of characters to be removed. All combinations of this set of characters will be stripped from the start and end of the string. If set to None (default), all leading and trailing whitespace is removed instead.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> data = {"fruits": ["apple", "\nmango"]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.DataFrame(data)
+
+            We define a dataframe-agnostic function:
+
+            >>> @nw.narwhalify
+            ... def func(df):
+            ...     df = df.with_columns(stripped=nw.col("fruits").str.strip_chars())
+            ...     return df.to_dict(as_series=False)
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(df_pd)
+            {'fruits': ['apple', '\nmango'], 'stripped': ['apple', 'mango']}
+
+            >>> func(df_pl)
+            {'fruits': ['apple', '\nmango'], 'stripped': ['apple', 'mango']}
+        """
+        return self._expr.__class__(
+            lambda plx: self._expr._call(plx).str.strip_chars(characters)
+        )
+
     def starts_with(self, prefix: str) -> Expr:
         r"""
         Check if string values start with a substring.
diff --git a/narwhals/series.py b/narwhals/series.py
index 7caee0c7..2582978c 100644
--- a/narwhals/series.py
+++ b/narwhals/series.py
@@ -2148,6 +2148,40 @@ class SeriesStringNamespace:
     def __init__(self, series: Series) -> None:
         self._narwhals_series = series
 
+    def strip_chars(self, characters: str | None = None) -> Series:
+        r"""
+        Remove leading and trailing characters.
+
+        Arguments:
+            characters: The set of characters to be removed. All combinations of this set of characters will be stripped from the start and end of the string. If set to None (default), all leading and trailing whitespace is removed instead.
+
+        Examples:
+            >>> import pandas as pd
+            >>> import polars as pl
+            >>> import narwhals as nw
+            >>> data = ["apple", "\nmango"]
+            >>> s_pd = pd.Series(data)
+            >>> s_pl = pl.Series(data)
+
+            We define a dataframe-agnostic function:
+
+            >>> @nw.narwhalify
+            ... def func(s):
+            ...     s = s.str.strip_chars()
+            ...     return s.to_list()
+
+            We can then pass either pandas or Polars to `func`:
+
+            >>> func(s_pd)
+            ['apple', 'mango']
+
+            >>> func(s_pl)
+            ['apple', 'mango']
+        """
+        return self._narwhals_series._from_compliant_series(
+            self._narwhals_series._compliant_series.str.strip_chars(characters)
+        )
+
     def starts_with(self, prefix: str) -> Series:
         r"""
         Check if string values start with a substring.


Test patch:
diff --git a/tests/expr_and_series/reduction_test.py b/tests/expr_and_series/reduction_test.py
new file mode 100644
index 00000000..f00356aa
--- /dev/null
+++ b/tests/expr_and_series/reduction_test.py
@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+
+import narwhals.stable.v1 as nw
+from tests.utils import compare_dicts
+
+
+@pytest.mark.parametrize(
+    ("expr", "expected"),
+    [
+        (
+            [nw.col("a").min().alias("min"), nw.col("a", "b").mean()],
+            {"min": [1], "a": [2], "b": [5]},
+        ),
+        ([(nw.col("a") + nw.col("b").max()).alias("x")], {"x": [7, 8, 9]}),
+        ([nw.col("a"), nw.col("b").min()], {"a": [1, 2, 3], "b": [4, 4, 4]}),
+        ([nw.col("a").max(), nw.col("b")], {"a": [3, 3, 3], "b": [4, 5, 6]}),
+        (
+            [nw.col("a"), nw.col("b").min().alias("min")],
+            {"a": [1, 2, 3], "min": [4, 4, 4]},
+        ),
+    ],
+    ids=range(5),
+)
+def test_scalar_reduction_select(
+    request: Any, constructor: Any, expr: list[Any], expected: dict[str, list[Any]]
+) -> None:
+    if "dask" in str(constructor) and int(request.node.callspec.id[-1]) != 1:
+        request.applymarker(pytest.mark.xfail)
+
+    data = {"a": [1, 2, 3], "b": [4, 5, 6]}
+    df = nw.from_native(constructor(data))
+    result = df.select(*expr)
+    compare_dicts(result, expected)
+
+
+@pytest.mark.parametrize(
+    ("expr", "expected"),
+    [
+        (
+            [nw.col("a").min().alias("min"), nw.col("a", "b").mean()],
+            {"min": [1, 1, 1], "a": [2, 2, 2], "b": [5, 5, 5]},
+        ),
+        ([(nw.col("a") + nw.col("b").max()).alias("x")], {"x": [7, 8, 9]}),
+        ([nw.col("a"), nw.col("b").min()], {"a": [1, 2, 3], "b": [4, 4, 4]}),
+        ([nw.col("a").max(), nw.col("b")], {"a": [3, 3, 3], "b": [4, 5, 6]}),
+        (
+            [nw.col("a"), nw.col("b").min().alias("min")],
+            {"a": [1, 2, 3], "min": [4, 4, 4]},
+        ),
+    ],
+    ids=range(5),
+)
+def test_scalar_reduction_with_columns(
+    request: Any, constructor: Any, expr: list[Any], expected: dict[str, list[Any]]
+) -> None:
+    if "dask" in str(constructor) and int(request.node.callspec.id[-1]) != 1:
+        request.applymarker(pytest.mark.xfail)
+
+    data = {"a": [1, 2, 3], "b": [4, 5, 6]}
+    df = nw.from_native(constructor(data))
+    result = df.with_columns(*expr).select(*expected.keys())
+    compare_dicts(result, expected)
diff --git a/tests/expr_and_series/str/strip_chars_test.py b/tests/expr_and_series/str/strip_chars_test.py
new file mode 100644
index 00000000..dffd41fd
--- /dev/null
+++ b/tests/expr_and_series/str/strip_chars_test.py
@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+
+import narwhals.stable.v1 as nw
+from tests.utils import compare_dicts
+
+data = {"a": ["foobar", "bar\n", " baz"]}
+
+
+@pytest.mark.parametrize(
+    ("characters", "expected"),
+    [
+        (None, {"a": ["foobar", "bar", "baz"]}),
+        ("foo", {"a": ["bar", "bar\n", " baz"]}),
+    ],
+)
+def test_str_strip_chars(constructor: Any, characters: str | None, expected: Any) -> None:
+    df = nw.from_native(constructor(data))
+    result_frame = df.select(nw.col("a").str.strip_chars(characters))
+    compare_dicts(result_frame, expected)
+
+
+@pytest.mark.parametrize(
+    ("characters", "expected"),
+    [
+        (None, {"a": ["foobar", "bar", "baz"]}),
+        ("foo", {"a": ["bar", "bar\n", " baz"]}),
+    ],
+)
+def test_str_strip_chars_series(
+    constructor_eager: Any, characters: str | None, expected: Any
+) -> None:
+    df = nw.from_native(constructor_eager(data), eager_only=True)
+
+    result_series = df["a"].str.strip_chars(characters)
+    assert result_series.to_list() == expected["a"]
diff --git a/tests/test_group_by.py b/tests/test_group_by.py
index dc63f0aa..20537185 100644
--- a/tests/test_group_by.py
+++ b/tests/test_group_by.py
@@ -53,10 +53,7 @@ def test_invalid_group_by() -> None:
         )
 
 
-def test_group_by_iter(request: Any, constructor_eager: Any) -> None:
-    if "pyarrow_table" in str(constructor_eager):
-        request.applymarker(pytest.mark.xfail)
-
+def test_group_by_iter(constructor_eager: Any) -> None:
     df = nw.from_native(constructor_eager(data), eager_only=True)
     expected_keys = [(1,), (3,)]
     keys = []