如何基于DATE、CALL_ID两列实现行级分组计数（复刻Alteryx MultiRow工具）

需求说明 #

按DATE、CALL_ID两列对数据分组，为每个分组内的每一行生成逐行递增的计数，需保留原数据的所有行。

示例数据 #

DATE        CALL_ID   
2023-11-21  29933702  
2023-11-21  29933703    
2023-11-21  29933703    
2023-11-21  29933704    
2023-11-21  29933704
2023-11-22  29933704  

期望输出 #

DATE        CALL_ID    COUNT
2023-11-21  29933702    1
2023-11-21  29933703    1
2023-11-21  29933703    2
2023-11-21  29933704    1
2023-11-21  29933704    2
2023-11-22  29933704    1

当前问题代码及输出 #

当前使用groupby().size()只能得到聚合后的分组总数量，无法保留原行：

# 当前代码
g = df.groupby(['DATE','CALL_ID']).size()

输出结果：

DATE        CALL_ID 
2023-11-21  29933702    1
            29933703    2
            29933704    2
2023-11-22  29933704    1

解决方案 #

使用groupby结合cumcount()方法，该方法会为每个分组内的行生成从0开始的递增序号，加1后即可得到从1起始的计数，同时保留原数据所有行：

# 正确代码
df['COUNT'] = df.groupby(['DATE', 'CALL_ID']).cumcount() + 1

效果验证 #

执行上述代码后，原DataFrame会新增COUNT列，结果与期望输出完全一致。

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