Repository: narwhals-dev/narwhals
Pull request #1269 (2024-10-28 09:25:39)

feat: Adding support for Expr.map_batches
<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [x] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues 

- Related issue # 
- Closes # https://github.com/narwhals-dev/narwhals/issues/1226

## Checklist

- [ ] Code follows style guide (ruff)
- [ ] Tests added 
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below.
I don't know if this is what was expected and if I should add the expression to dask as well. 


DeaMariaLeon: Before I try to fix this for older versions, I would like to know if I took the right approach.
MarcoGorelli: hey, thanks for working on this! from an initial review, this looks really solid! will do a full review shortly
DeaMariaLeon: I fixed the issue with the names @MarcoGorelli. (I'm not sure if I have to ping you, no rush of course). 
Again, older versions are complaining, but maybe you can tell me about the rest? thanks.
DeaMariaLeon: It seems that `map_batches` has a bug in old polars (0.20.3). Doing the same operation than the test that was failing (with just polars) there is the same error.

```
import polars as pl
data = {"a": [1, 2, 3], "b": [4, 5, 6], "z": [7.0, 8.0, 9.0]}
df_pl = pl.DataFrame(data)
print(pl.Float64())
print(df_pl.select(pl.all().map_batches(lambda s: s.to_numpy() + 1, return_dtype=pl.Float64())))
```
MarcoGorelli: thanks for updating! this is probably good to go, will do a full review soon-ish
MarcoGorelli: i just tried this out, and it looks like we're getting different outputs across different versions:
```
# ruff: noqa
from typing import Any
import pandas as pd
import polars as pl
import narwhals as nw
import pyarrow as pa
data = {"statistic": [1, 2, -3, 4]}
df_pd = pd.DataFrame(data)
df_pl = pl.DataFrame(data)
df_pa = pa.table(data)

@nw.narwhalify
def p_value(df: nw.DataFrame[Any]) -> nw.DataFrame[Any]:
    return df.select(
        nw.all().map_batches(lambda x: x.to_numpy())
    )

print(p_value(df_pl))
print(p_value(df_pd))
print(p_value(df_pa))
```

I think it might be OK to raise an error if we end up in `            if is_numpy_array(result[0]):` and `self._output_names` is `None`

The other thing I noticed is that this currently raises with the output from the docstring in https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.map_batches.html (i.e. `pl.all().map_batches(lambda x: x.to_numpy().argmax())`)

```
AttributeError: 'numpy.int64' object has no attribute '__native_namespace__'. Did you mean: '__array_namespace__'?
```
DeaMariaLeon: >i just tried this out, and it looks like we're getting different outputs across different versions:

Which versions, do you know? It's passing the old versions in the tests. And locally I'm testing with an old version too and it works.. You didn't test with an older version of my PR, did you? (🤞) 
```
shape: (1, 1)
┌─────────────┐
│ statistic   │
│ ---         │
│ list[i64]   │
╞═════════════╡
│ [1, 2, … 4] │
└─────────────┘
   statistic
0          1
1          2
2         -3
3          4
pyarrow.Table
statistic: int64
----
statistic: [[1,2,-3,4]]
```

For the second part of your comment, ("The other thing I noticed is... ") . I'll work on it. 
MarcoGorelli: sorry i meant different packages - for polars it outpus one column, for pandas and pyarrow the output is empty
DeaMariaLeon: So, I did the same example you posted, with the normal/recent versions, and I still have the right answers. (Just like on my previous comment). I'm lost. 🫣
MarcoGorelli: sure, just to clarify, this is what I mean:
```python
In [5]: from typing import Any
   ...: import pandas as pd
   ...: import polars as pl
   ...: import narwhals as nw
   ...: import pyarrow as pa
   ...: data = {"statistic": [1, 2, -3, 4]}
   ...: df_pd = pd.DataFrame(data)
   ...: df_pl = pl.DataFrame(data)
   ...: df_pa = pa.table(data)
   ...: 
   ...: @nw.narwhalify
   ...: def func(df: nw.DataFrame[Any]) -> nw.DataFrame[Any]:
   ...:     return df.select(
   ...:         nw.all().map_batches(lambda x: x.to_numpy())
   ...:     )
   ...: 

In [6]: func(pd.DataFrame(data))  # empty output
Out[6]: 
Empty DataFrame
Columns: []
Index: []

In [7]: func(pl.DataFrame(data))  # one column
Out[7]: 
shape: (4, 1)
┌───────────┐
│ statistic │
│ ---       │
│ i64       │
╞═══════════╡
│ 1         │
│ 2         │
│ -3        │
│ 4         │
└───────────┘
```
MarcoGorelli: ah sorry my bad 🤦 i had indeed to updated my branch because of a conflict, so sorry
MarcoGorelli: after all these years i'm still making git mistakes
DeaMariaLeon: Hey, that's OK! No need to apologise.

>after all these years i'm still making git mistakes

We are only humans. 🤗

DeaMariaLeon: Fixed this part, it seems: 
>The other thing I noticed is that this currently raises with the output from the docstring in https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.map_batches.html (i.e. pl.all().map_batches(lambda x: x.to_numpy().argmax()))

Patch:
diff --git a/docs/api-reference/expr.md b/docs/api-reference/expr.md
index bf68689a..e99ffd4f 100644
--- a/docs/api-reference/expr.md
+++ b/docs/api-reference/expr.md
@@ -27,6 +27,7 @@
         - is_null
         - is_unique
         - len
+        - map_batches
         - max
         - mean
         - median
diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py
index 802a900b..b8aba0ae 100644
--- a/narwhals/_arrow/expr.py
+++ b/narwhals/_arrow/expr.py
@@ -8,6 +8,8 @@ from typing import Sequence
 
 from narwhals._expression_parsing import reuse_series_implementation
 from narwhals._expression_parsing import reuse_series_namespace_implementation
+from narwhals.dependencies import get_numpy
+from narwhals.dependencies import is_numpy_array
 from narwhals.utils import Implementation
 
 if TYPE_CHECKING:
@@ -383,6 +385,44 @@ class ArrowExpr:
     def mode(self: Self) -> Self:
         return reuse_series_implementation(self, "mode")
 
+    def map_batches(
+        self: Self,
+        function: Callable[[Any], Any],
+        return_dtype: DType | None = None,
+    ) -> Self:
+        def func(df: ArrowDataFrame) -> list[ArrowSeries]:
+            input_series_list = self._call(df)
+            output_names = [input_series.name for input_series in input_series_list]
+            result = [function(series) for series in input_series_list]
+
+            if is_numpy_array(result[0]):
+                result = [
+                    df.__narwhals_namespace__()
+                    ._create_compliant_series(array)
+                    .alias(output_name)
+                    for array, output_name in zip(result, output_names)
+                ]
+            elif (np := get_numpy()) is not None and np.isscalar(result[0]):
+                result = [
+                    df.__narwhals_namespace__()
+                    ._create_compliant_series([array])
+                    .alias(output_name)
+                    for array, output_name in zip(result, output_names)
+                ]
+            if return_dtype is not None:
+                result = [series.cast(return_dtype) for series in result]
+            return result
+
+        return self.__class__(
+            func,
+            depth=self._depth + 1,
+            function_name=self._function_name + "->map_batches",
+            root_names=self._root_names,
+            output_names=self._output_names,
+            backend_version=self._backend_version,
+            dtypes=self._dtypes,
+        )
+
     @property
     def dt(self: Self) -> ArrowExprDateTimeNamespace:
         return ArrowExprDateTimeNamespace(self)
diff --git a/narwhals/_dask/expr.py b/narwhals/_dask/expr.py
index 169adc43..1f863f20 100644
--- a/narwhals/_dask/expr.py
+++ b/narwhals/_dask/expr.py
@@ -376,6 +376,14 @@ class DaskExpr:
             returns_scalar=False,
         )
 
+    def map_batches(
+        self: Self,
+        function: Callable[[Any], Any],
+        return_dtype: DType | None = None,
+    ) -> NoReturn:
+        msg = "`Expr.map_batches` is not implemented for Dask yet"
+        raise NotImplementedError(msg)
+
     def mean(self) -> Self:
         return self._from_call(
             lambda _input: _input.mean(),
diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py
index 9ebbc315..6b2487dc 100644
--- a/narwhals/_pandas_like/expr.py
+++ b/narwhals/_pandas_like/expr.py
@@ -9,6 +9,8 @@ from typing import Sequence
 from narwhals._expression_parsing import reuse_series_implementation
 from narwhals._expression_parsing import reuse_series_namespace_implementation
 from narwhals._pandas_like.series import PandasLikeSeries
+from narwhals.dependencies import get_numpy
+from narwhals.dependencies import is_numpy_array
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -399,6 +401,39 @@ class PandasLikeExpr:
     def mode(self: Self) -> Self:
         return reuse_series_implementation(self, "mode")
 
+    def map_batches(
+        self: Self,
+        function: Callable[[Any], Any],
+        return_dtype: DType | None = None,
+    ) -> Self:
+        def func(df: PandasLikeDataFrame) -> list[PandasLikeSeries]:
+            input_series_list = self._call(df)
+            output_names = [input_series.name for input_series in input_series_list]
+            result = [function(series) for series in input_series_list]
+            if is_numpy_array(result[0]) or (
+                (np := get_numpy()) is not None and np.isscalar(result[0])
+            ):
+                result = [
+                    df.__narwhals_namespace__()
+                    ._create_compliant_series(array)
+                    .alias(output_name)
+                    for array, output_name in zip(result, output_names)
+                ]
+            if return_dtype is not None:
+                result = [series.cast(return_dtype) for series in result]
+            return result
+
+        return self.__class__(
+            func,
+            depth=self._depth + 1,
+            function_name=self._function_name + "->map_batches",
+            root_names=self._root_names,
+            output_names=self._output_names,
+            implementation=self._implementation,
+            backend_version=self._backend_version,
+            dtypes=self._dtypes,
+        )
+
     @property
     def str(self: Self) -> PandasLikeExprStringNamespace:
         return PandasLikeExprStringNamespace(self)
diff --git a/narwhals/_polars/expr.py b/narwhals/_polars/expr.py
index 46b415dd..7aa0fb29 100644
--- a/narwhals/_polars/expr.py
+++ b/narwhals/_polars/expr.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 from typing import Any
+from typing import Callable
 from typing import Sequence
 
 from narwhals._polars.utils import extract_args_kwargs
@@ -47,6 +48,19 @@ class PolarsExpr:
         dtype = narwhals_to_native_dtype(dtype, self._dtypes)
         return self._from_native_expr(expr.cast(dtype))
 
+    def map_batches(
+        self,
+        function: Callable[[Any], Self],
+        return_dtype: DType | None = None,
+    ) -> Self:
+        if return_dtype is not None:
+            return_dtype = narwhals_to_native_dtype(return_dtype, self._dtypes)
+            return self._from_native_expr(
+                self._native_expr.map_batches(function, return_dtype)
+            )
+        else:
+            return self._from_native_expr(self._native_expr.map_batches(function))
+
     def replace_strict(
         self, old: Sequence[Any], new: Sequence[Any], *, return_dtype: DType | None
     ) -> Self:
diff --git a/narwhals/expr.py b/narwhals/expr.py
index 9349dd7d..4c7af41d 100644
--- a/narwhals/expr.py
+++ b/narwhals/expr.py
@@ -539,6 +539,75 @@ class Expr:
         """
         return self.__class__(lambda plx: self._call(plx).std(ddof=ddof))
 
+    def map_batches(
+        self,
+        function: Callable[[Any], Self],
+        return_dtype: DType | None = None,
+    ) -> Self:
+        """
+        Apply a custom python function to a whole Series or sequence of Series.
+
+        The output of this custom function is presumed to be either a Series,
+        or a NumPy array (in which case it will be automatically converted into
+        a Series).
+
+        Arguments:
+            return_dtype: Dtype of the output Series.
+                          If not set, the dtype will be inferred based on the first non-null value
+                          that is returned by the function.
+
+        Examples:
+            >>> import polars as pl
+            >>> import pandas as pd
+            >>> import pyarrow as pa
+            >>> import narwhals as nw
+            >>> data = {"a": [1, 2, 3], "b": [4, 5, 6]}
+            >>> df_pd = pd.DataFrame(data)
+            >>> df_pl = pl.DataFrame(data)
+            >>> df_pa = pa.table(data)
+
+            Let's define a dataframe-agnostic function:
+
+            >>> @nw.narwhalify
+            ... def func(df):
+            ...     return df.select(
+            ...         nw.col("a", "b").map_batches(
+            ...             lambda s: s.to_numpy() + 1, return_dtype=nw.Float64
+            ...         )
+            ...     )
+
+            We can then pass any supported library such as Pandas, Polars, or PyArrow to `func`:
+
+            >>> func(df_pd)
+                 a    b
+            0  2.0  5.0
+            1  3.0  6.0
+            2  4.0  7.0
+            >>> func(df_pl)
+            shape: (3, 2)
+            ┌─────┬─────┐
+            │ a   ┆ b   │
+            │ --- ┆ --- │
+            │ f64 ┆ f64 │
+            ╞═════╪═════╡
+            │ 2.0 ┆ 5.0 │
+            │ 3.0 ┆ 6.0 │
+            │ 4.0 ┆ 7.0 │
+            └─────┴─────┘
+            >>> func(df_pa)
+            pyarrow.Table
+            a: double
+            b: double
+            ----
+            a: [[2,3,4]]
+            b: [[5,6,7]]
+        """
+        return self.__class__(
+            lambda plx: self._call(plx).map_batches(
+                function=function, return_dtype=return_dtype
+            )
+        )
+
     def sum(self) -> Expr:
         """
         Return the sum value.
diff --git a/utils/check_api_reference.py b/utils/check_api_reference.py
index b7d8595a..199b60ed 100644
--- a/utils/check_api_reference.py
+++ b/utils/check_api_reference.py
@@ -12,7 +12,7 @@ from narwhals.utils import remove_suffix
 ret = 0
 
 NAMESPACES = {"dt", "str", "cat", "name"}
-EXPR_ONLY_METHODS = {"over"}
+EXPR_ONLY_METHODS = {"over", "map_batches"}
 SERIES_ONLY_METHODS = {
     "dtype",
     "is_empty",


Test patch:
diff --git a/tests/expr_and_series/map_batches_test.py b/tests/expr_and_series/map_batches_test.py
new file mode 100644
index 00000000..f2720e20
--- /dev/null
+++ b/tests/expr_and_series/map_batches_test.py
@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+import pandas as pd
+import pytest
+
+import narwhals.stable.v1 as nw
+from tests.utils import POLARS_VERSION
+from tests.utils import Constructor
+from tests.utils import assert_equal_data
+
+data = {"a": [1, 2, 3], "b": [4, 5, 6], "z": [7.0, 8.0, 9.0]}
+
+
+def test_map_batches_expr(
+    request: pytest.FixtureRequest, constructor: Constructor
+) -> None:
+    if "dask" in str(constructor):
+        request.applymarker(pytest.mark.xfail)
+
+    df = nw.from_native(constructor(data))
+    expected = df.select(nw.col("a", "b").map_batches(lambda s: s + 1))
+    assert_equal_data(expected, {"a": [2, 3, 4], "b": [5, 6, 7]})
+
+
+def test_map_batches_expr_numpy(
+    request: pytest.FixtureRequest, constructor: Constructor
+) -> None:
+    if "dask" in str(constructor) or (
+        "polars" in str(constructor) and POLARS_VERSION <= (0, 20, 3)
+    ):
+        request.applymarker(pytest.mark.xfail)
+
+    df = nw.from_native(constructor(data))
+    expected = df.select(
+        nw.col("a")
+        .map_batches(lambda s: s.to_numpy() + 1, return_dtype=nw.Float64())
+        .sum()
+    )
+    assert_equal_data(expected, {"a": [9.0]})
+
+    expected = df.select(nw.all().map_batches(lambda s: s.to_numpy().argmax()))
+    assert_equal_data(expected, {"a": [2], "b": [2], "z": [2]})
+
+
+def test_map_batches_expr_names(
+    request: pytest.FixtureRequest, constructor: Constructor
+) -> None:
+    if "dask" in str(constructor) or (
+        "polars" in str(constructor) and POLARS_VERSION <= (0, 20, 3)
+    ):
+        request.applymarker(pytest.mark.xfail)
+
+    df = nw.from_native(constructor(data))
+    expected = nw.from_native(df.select(nw.all().map_batches(lambda x: x.to_numpy())))
+    assert_equal_data(expected, {"a": [1, 2, 3], "b": [4, 5, 6], "z": [7.0, 8.0, 9.0]})
+
+
+def test_map_batches_raise() -> None:
+    pytest.importorskip("dask")
+    pytest.importorskip("dask_expr", exc_type=ImportError)
+    import dask.dataframe as dd
+
+    df = nw.from_native(dd.from_pandas(pd.DataFrame({"a": [1, 2, 3]})))
+    with pytest.raises(
+        NotImplementedError, match="`Expr.map_batches` is not implemented for Dask yet"
+    ):
+        df.select(nw.col("a").map_batches(lambda x: x.to_numpy()))