如何在Polars DataFrame的with_columns方法中按分组列统计每个ID的唯一年份数量

嘿，这问题我熟！你想要给原DataFrame的每一行都添加上对应ID的唯一年份数量，用Polars的with_columns配合分组窗口函数就能轻松实现，我给你一步步拆解：

首先先确认下你的原始数据（和你给出的一致）：

import polars as pl

df = pl.DataFrame({
    'ID': [1, 1, 5, 5, 7, 7, 7],
    'YEAR': [2025, 2025, 2023, 2024, 2020, 2021, 2021]
})

接下来就是核心的实现代码，直接用with_columns加上窗口统计就行：

result_df = df.with_columns(
    # 按ID分组统计YEAR的唯一值数量，并映射回每一行
    pl.col("YEAR").n_unique().over("ID").alias("UNIQUE_YEAR_COUNT")
)

print(result_df)

运行后你就能得到想要的结果，每一行都保留原始数据，同时新增了对应ID的唯一年份统计列：

shape: (7, 3)
┌─────┬──────┬──────────────────┐
│ ID  ┆ YEAR ┆ UNIQUE_YEAR_COUNT│
│ --- ┆ ---  ┆ ---              │
│ i64 ┆ i64  ┆ u32              │
╞═════╪══════╪══════════════════╡
│ 1   ┆ 2025 ┆ 1                │
│ 1   ┆ 2025 ┆ 1                │
│ 5   ┆ 2023 ┆ 2                │
│ 5   ┆ 2024 ┆ 2                │
│ 7   ┆ 2020 ┆ 2                │
│ 7   ┆ 2021 ┆ 2                │
│ 7   ┆ 2021 ┆ 2                │
└─────┴──────┴──────────────────┘

我再给你解释下这段代码的逻辑：

• pl.col("YEAR").n_unique()：这部分是用来计算某组内YEAR列的唯一值数量
• .over("ID")：这个窗口函数指定了按ID分组进行统计，关键是它会把分组统计的结果映射回原DataFrame的每一行，而不是只返回分组后的聚合结果，正好符合你要保留所有原始行的需求
• .alias("UNIQUE_YEAR_COUNT")：给新生成的统计列起个清晰好懂的名字
• 最后用with_columns把这个新列直接添加到原DataFrame中，完美达成你的要求

备注：内容来源于stack exchange，提问作者Phil-ZXX