Polars 自定义函数指南

原创

东方佑 2024-10-31 00:27:43 ©著作权

文章标签 python Python 自定义函数代码示例 文章分类 Python 后端开发

©著作权归作者所有：来自51CTO博客作者东方佑的原创作品，请联系作者获取转载授权，否则将追究法律责任

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)

输出结果正确。