polars.Expr.map_batches#

Expr.map_batches(
function: Callable[[Series], Series | Any],
return_dtype: PolarsDataType | None = None,
*,
agg_list: bool = False,
is_elementwise: bool = False,
returns_scalar: bool = False,
) Expr[source]#

对整个Series或Series序列应用自定义的python函数。

此自定义函数的输出假定为Series,或NumPy数组(在这种情况下将自动转换为Series),或将被转换为Series的标量。如果结果是标量并且您希望它保持为标量,请传入returns_scalar=True。如果您希望对单个值逐元素应用自定义函数,请参见map_elements()map函数的一个合理用例是使用第三方库转换表达式表示的值。

Parameters:
function

要应用的Lambda/函数。

return_dtype

输出Series的数据类型。 如果未设置,将根据函数返回的第一个非空值推断数据类型。

agg_list

在应用函数之前,将表达式的值聚合到一个列表中。此参数仅在分组上下文中有效。 该函数将在分组列表上仅调用一次,而不是每个组调用一次。

is_elementwise

如果设置为true,这可以在流式引擎中运行,但可能在分组时产生不正确的结果。请确保你知道你在做什么!

returns_scalar

如果函数返回一个标量,默认情况下它会被包装在输出中的一个列表中,因为假设是函数总是返回类似Series的内容。如果你想保持结果为标量,请将此参数设置为True。

警告

如果未提供return_dtype,可能会导致意外结果。 我们允许这种情况,但这被视为用户查询中的错误。

另请参阅

map_elements
replace

示例

>>> df = pl.DataFrame(
...     {
...         "sine": [0.0, 1.0, 0.0, -1.0],
...         "cosine": [1.0, 0.0, -1.0, 0.0],
...     }
... )
>>> df.select(pl.all().map_batches(lambda x: x.to_numpy().argmax()))
shape: (1, 2)
┌──────┬────────┐
│ sine ┆ cosine │
│ ---  ┆ ---    │
│ i64  ┆ i64    │
╞══════╪════════╡
│ 1    ┆ 0      │
└──────┴────────┘

在分组上下文中,如果正确使用agg_list参数,可以提高性能。以下示例将agg_list设置为False,这会导致函数在每个组中应用一次。函数的输入是一个类型为Int64的Series。这样效率较低。

>>> df = pl.DataFrame(
...     {
...         "a": [0, 1, 0, 1],
...         "b": [1, 2, 3, 4],
...     }
... )
>>> df.group_by("a").agg(
...     pl.col("b").map_batches(lambda x: x + 2, agg_list=False)
... )  
shape: (2, 2)
┌─────┬───────────┐
│ a   ┆ b         │
│ --- ┆ ---       │
│ i64 ┆ list[i64] │
╞═════╪═══════════╡
│ 1   ┆ [4, 6]    │
│ 0   ┆ [3, 5]    │
└─────┴───────────┘

使用 agg_list=True 会更高效。在这个例子中,函数的输入是一个类型为 List(Int64) 的 Series。

>>> df.group_by("a").agg(
...     pl.col("b").map_batches(
...         lambda x: x.list.eval(pl.element() + 2), agg_list=True
...     )
... )  
shape: (2, 2)
┌─────┬───────────┐
│ a   ┆ b         │
│ --- ┆ ---       │
│ i64 ┆ list[i64] │
╞═════╪═══════════╡
│ 0   ┆ [3, 5]    │
│ 1   ┆ [4, 6]    │
└─────┴───────────┘

这是一个返回标量的函数示例,我们希望它保持为标量:

>>> df = pl.DataFrame(
...     {
...         "a": [0, 1, 0, 1],
...         "b": [1, 2, 3, 4],
...     }
... )
>>> df.group_by("a").agg(
...     pl.col("b").map_batches(lambda x: x.max(), returns_scalar=True)
... )  
shape: (2, 2)
┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 1   ┆ 4   │
│ 0   ┆ 3   │
└─────┴─────┘

通过创建一个struct并在函数调用中引用其字段来调用一个接受多个参数的函数。

>>> df = pl.DataFrame(
...     {
...         "a": [5, 1, 0, 3],
...         "b": [4, 2, 3, 4],
...     }
... )
>>> df.with_columns(
...     a_times_b=pl.struct("a", "b").map_batches(
...         lambda x: np.multiply(x.struct.field("a"), x.struct.field("b"))
...     )
... )
shape: (4, 3)
┌─────┬─────┬───────────┐
│ a   ┆ b   ┆ a_times_b │
│ --- ┆ --- ┆ ---       │
│ i64 ┆ i64 ┆ i64       │
╞═════╪═════╪═══════════╡
│ 5   ┆ 4   ┆ 20        │
│ 1   ┆ 2   ┆ 2         │
│ 0   ┆ 3   ┆ 0         │
│ 3   ┆ 4   ┆ 12        │
└─────┴─────┴───────────┘