polars.map_batches#

polars.map_batches( exprs: Sequence[str] | Sequence[Expr], function: Callable[[Sequence[Series]], Series], return_dtype: PolarsDataType | None = None, ) → Expr[source]#

Map a custom function over multiple columns/expressions.

Produces a single Series result.

Parameters:

exprs: Expression(s) representing the input Series to the function.
function: Function to apply over the input.
return_dtype: dtype of the output Series.

Returns:

Expr: Expression with the data type given by return_dtype.

Examples

>>> def test_func(a, b, c):
...     return a + b + c
...
>>> df = pl.DataFrame(
...     {
...         "a": [1, 2, 3, 4],
...         "b": [4, 5, 6, 7],
...     }
... )
>>>
>>> df.with_columns(
...     (
...         pl.struct(["a", "b"]).map_batches(
...             lambda x: test_func(x.struct.field("a"), x.struct.field("b"), 1)
...         )
...     ).alias("a+b+c")
... )
shape: (4, 3)
┌─────┬─────┬───────┐
│ a   ┆ b   ┆ a+b+c │
│ --- ┆ --- ┆ ---   │
│ i64 ┆ i64 ┆ i64   │
╞═════╪═════╪═══════╡
│ 1   ┆ 4   ┆ 6     │
│ 2   ┆ 5   ┆ 8     │
│ 3   ┆ 6   ┆ 10    │
│ 4   ┆ 7   ┆ 12    │
└─────┴─────┴───────┘