在Polars中实现多组数据合并分组聚合

原始数据 #

import polars as pl

df = pl.from_repr("""
┌──────────┬──────┬───────┐
│ category ┆ year ┆ count │
│ ---      ┆ ---  ┆ ---   │
│ str      ┆ i64  ┆ i64   │
╞══════════╪══════╪═══════╡
│ apple    ┆ 2022 ┆ 5     │
│ apple    ┆ 2021 ┆ 8     │
│ banana   ┆ 2022 ┆ 1     │
│ cold     ┆ 2022 ┆ 9     │
│ cold     ┆ 2021 ┆ 2     │
│ warm     ┆ 2022 ┆ 1     │
│ warm     ┆ 2021 ┆ 3     │
└──────────┴──────┴───────┘
""")

需求说明 #

通过自定义字典将原始category映射到预设分组（fruit、temp），再按新分组和year聚合count字段，得到目标结构的DataFrame。

解决方案 #

利用Polars的矢量化map_dict方法实现高效映射，适配大数据量场景：

# 定义分组映射字典
group_mapping = {
    "apple": "fruit",
    "banana": "fruit",
    "cold": "temp",
    "warm": "temp"
}

# 映射分组 + 聚合计算
result_df = (
    df
    .with_columns(pl.col("category").map_dict(group_mapping).alias("category"))
    .group_by(["category", "year"])
    .agg(pl.col("count").sum().alias("count"))
    .sort(["category", "year"], descending=[False, True])  # 匹配示例输出顺序
)

print(result_df)

输出结果 #

shape: (4, 3)
┌──────────┬──────┬───────┐
│ category ┆ year ┆ count │
│ ---      ┆ ---  ┆ ---   │
│ str      ┆ i64  ┆ i64   │
╞══════════╪══════╪═══════╡
│ fruit    ┆ 2022 ┆ 6     │
│ fruit    ┆ 2021 ┆ 8     │
│ temp     ┆ 2022 ┆ 10    │
│ temp     ┆ 2021 ┆ 5     │
└──────────┴──────┴───────┘

关键说明 #

• map_dict是Polars的矢量化操作，比循环遍历效率高得多，适合处理大规模数据集
• 先完成分类映射再执行分组聚合，逻辑清晰且性能优异

内容的提问来源于stack exchange，提问作者zoned post meridiem