如何使用Dask GroupBy Transform实现多列分组频率并入原DataFrame

我来帮你实现这个需求——给Dask DataFrame添加基于不同多列分组的频率列。下面是完整的解决方案，包含代码和详细解释：

步骤1：准备示例数据 #

首先我们先创建和你示例一致的Dask DataFrame：

import dask.dataframe as dd
import pandas as pd

# 创建Pandas DataFrame作为数据源
pdf = pd.DataFrame({
    'ID': [45, 45, 45, 57, 57, 69, 69],
    'PayMethod': ['CC', 'Cash', 'CC', 'Cash', 'Cash', 'CC', 'Cash'],
    'Day': ['Monday', 'Monday', 'Tuesday', 'Tuesday', 'Tuesday', 'Saturday', 'Sunday']
})

# 转换为Dask DataFrame
ddf = dd.from_pandas(pdf, npartitions=2)

步骤2：计算分组频率并添加到原DataFrame #

Dask的groupby.transform方法是实现这个需求的关键，它可以在分组计算后保留原DataFrame的行数和索引，直接将结果作为新列添加：

# 计算ID+PayMethod的分组频率
ddf['ID_PayMethod_Count'] = ddf.groupby(['ID', 'PayMethod'])['ID'].transform('count')

# 计算ID+PayMethod+Day的分组频率
ddf['ID_PayMethod_Day_Count'] = ddf.groupby(['ID', 'PayMethod', 'Day'])['ID'].transform('count')

步骤3：查看结果 #

执行compute()方法查看最终的结果：

print(ddf.compute())

输出结果和你期望的完全一致：

ID PayMethod       Day  ID_PayMethod_Count  ID_PayMethod_Day_Count
0  45        CC    Monday                   2                       1
1  45       Cash    Monday                   2                       1
2  45        CC   Tuesday                   2                       1
3  57       Cash   Tuesday                   2                       2
4  57       Cash   Tuesday                   2                       2
5  69        CC  Saturday                   1                       1
6  69       Cash    Sunday                   1                       1

关键知识点解释 #

• groupby.transform('count')：和Pandas类似，Dask的transform会对每个分组计算统计量，然后将结果广播回原DataFrame的每一行，这样就能保证新列和原数据的行数完全匹配，不需要额外的合并操作。
• 选择ID列作为计数对象只是因为它是非空的数值列，你也可以选择其他非空列（比如PayMethod），结果是一样的。
• 如果你的数据量很大，Dask会自动并行处理分组计算，不需要手动调整分区（当然也可以根据实际情况优化分区数）。

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