如何在Python Polars中按分组条件添加计数新字段？

要实现按LN、FN、SSN分组并新增每组观测次数的列，Polars提供了两种简洁高效的方法：

方法一：使用窗口函数（推荐） #

窗口函数可以直接在原DataFrame上添加分组计数，无需修改原表结构，所有行都会保留并对应正确的计数值：

import polars as pl

df_names = pl.DataFrame({
    'LN':['Mallesham','Bhavik','Mallesham','Bhavik','Mahesh','Naresh','Sharath','Rakesh','Mallesham'],
    'FN':['Yamulla','Yamulla','Yamulla','Yamulla','Dayala','Burre','Velmala','Uppu','Yamulla'],
    'SSN':['123','456','123','456','893','111','222','333','123'],
    'Address':['A','B','C','D','E','F','G','H','S']
})

# 新增分组计数列
result = df_names.with_columns(
    pl.count().over(['LN', 'FN', 'SSN']).alias('LN_FN_SSN_count')
)

print(result)

说明 #

• pl.count()用于计算行数，over(['LN', 'FN', 'SSN'])指定分组依据，该函数会自动为每个分组的所有行填充对应的计数。
• 这种方法性能更优，代码简洁，适合绝大多数场景。

方法二：分组聚合后连接 #

先通过分组聚合得到每组的计数，再通过左连接将计数合并回原表：

# 分组计算每组计数
count_df = df_names.group_by(['LN', 'FN', 'SSN']).agg(
    pl.count().alias('LN_FN_SSN_count')
)

# 左连接原表，保留所有行并匹配计数
result = df_names.join(count_df, on=['LN', 'FN', 'SSN'], how='left')

预期输出 #

shape: (9, 5)
┌────────────┬─────────┬─────┬─────────┬──────────────────┐
│ LN         ┆ FN      ┆ SSN ┆ Address ┆ LN_FN_SSN_count  │
│ ---        ┆ ---     ┆ --- ┆ ---     ┆ ---              │
│ str        ┆ str     ┆ str ┆ str     ┆ u32              │
╞════════════╪═════════╪═════╪═════════╪══════════════════╡
│ Mallesham  ┆ Yamulla ┆ 123 ┆ A       ┆ 3                │
│ Bhavik     ┆ Yamulla ┆ 456 ┆ B       ┆ 2                │
│ Mallesham  ┆ Yamulla ┆ 123 ┆ C       ┆ 3                │
│ Bhavik     ┆ Yamulla ┆ 456 ┆ D       ┆ 2                │
│ Mahesh     ┆ Dayala  ┆ 893 ┆ E       ┆ 1                │
│ Naresh     ┆ Burre   ┆ 111 ┆ F       ┆ 1                │
│ Sharath    ┆ Velmala ┆ 222 ┆ G       ┆ 1                │
│ Rakesh     ┆ Uppu    ┆ 333 ┆ H       ┆ 1                │
│ Mallesham  ┆ Yamulla ┆ 123 ┆ S       ┆ 3                │
└────────────┴─────────┴─────┴─────────┴──────────────────┘

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