Polars 自定义函数指南
Polars 提供了 map
和 apply
两种表达式来处理自定义函数。
1. map 函数
map
函数将表达式所支持的 Series
数据原封不动地传递。
例子
在 groupby
中使用 map
时需要小心,因为它可能会导致不正确的结果。例如,对于一个分组的数据帧,map
会在聚合之前应用函数。
df = pl.DataFrame({
"keys": ["a", "a", "b"],
"values": [10, 7, 1],
})
out = df.groupby("keys", maintain_order=True).agg(
[
pl.col("values").map(lambda s: s.shift()).alias("shift_map"),
pl.col("values").shift().alias("shift_expression"),
]
)
print(out)
输出结果可能不是预期的,因为 map
在聚合之前应用了函数。
2. apply 函数
apply
函数可以对该操作的最小逻辑元素起作用。在 select
中,它作用于单个元素;在 groupby
中,它作用于单个分组。
例子
使用 apply
来解决上述问题:
out = df.groupby("keys", maintain_order=True).agg(
[
pl.col("values").apply(lambda s: s.shift()).alias("shift_map"),
pl.col("values").shift().alias("shift_expression"),
]
)
print(out)
3. select 中的 apply
在 select
中,apply
表达式将列的元素传递给 Python 函数。
例子
创建一个全局的 counter
并在每次处理一个元素时将其值添加到全局状态中:
counter = 0
def add_counter(val: int) -> int:
global counter
counter += 1
return counter + val
out = df.select(
[
pl.col("values").apply(add_counter).alias("solution_apply"),
(pl.col("values") + pl.arange(1, pl.count() + 1)).alias("solution_expr"),
]
)
print(out)
输出结果正确。
4. 返回类型
自定义 Python 函数对 Polars 而言是黑箱。数据类型是自动推断出来的,通过等待第一个非空值来确定 Series
的类型。
这个Markdown总结包含了原网页中的主要概念和代码示例,以便于快速理解和参考。更多详细信息和高级用法,请访问原网页。