如何在Polars中实现分组按位OR聚合操作？

在Polars中，直接用or_()做聚合会返回列表，因为它是元素级的逻辑或操作，不是针对分组的聚合函数。要实现分组后对auth字段做按位OR的聚合，直接用Polars提供的原生聚合函数bitwise_or_agg()即可，无需借助map_batches。

完整实现代码 #

import polars as pl

df_in = pl.DataFrame(
    {
        "k": ["a", "a", "b", "b", "c"],
        "auth": [1, 3, 1, 0, 0],
    }
)

# 使用bitwise_or_agg做按位OR聚合
df_wanted_or = df_in.group_by("k", maintain_order=True).agg(
    pl.col("auth").bitwise_or_agg()
)

print(df_wanted_or)

输出结果 #

shape: (3, 2)
┌─────┬──────┐
│ k   ┆ auth │
│ --- ┆ ---  │
│ str ┆ i64  │
╞═════╪══════╡
│ a   ┆ 3    │
│ b   ┆ 1    │
│ c   ┆ 0    │
└─────┴──────┘

这个方法完全是Polars原生实现，效率比map_batches调用Python函数高很多，适合处理大数据量的场景。

内容的提问来源于stack exchange，提问作者JL Peyret