如何在Python DataFrame中基于客户与交易状态重置连续行号

需求背景 #

• 需生成字段transaction_in_row：按customer分组，当has_transaction字段值发生变化时重置计数（要求连续相同值的行形成独立块，而非将所有同值行归为同一分区）
• 原SQL方案缺陷：使用row_number() over(partition by customer, has_transaction order by month asc)会错误地将同一customer下所有相同has_transaction值的行归为同一分区，导致状态变化后计数不重置（例如目标行返回3而非预期的1）
• 数据规模：约3000万行，需Python DataFrame高效实现，禁止低效循环

高效实现方案 #

核心思路：先识别每个customer内连续相同has_transaction的块，再在每个块内生成行号，全程使用pandas矢量化操作，避免逐行循环。

1. 确保数据按分组和时间排序
   必须保证每个customer的行按month升序排列，否则状态块的识别会出错：

   df = df.sort_values(by=['customer', 'month'], ascending=True)
   

2. 标记状态变化点
   对每个customer，标记当前行与上一行has_transaction值不同的位置：

   df['change_flag'] = df.groupby('customer')['has_transaction'].diff().ne(0).astype(int)
   

   • diff()：计算当前行与上一行has_transaction的差值（非数值类型会返回NaN，首行默认NaN）
   • ne(0)：判断是否发生变化（包括首行，因为NaN不等于0），返回布尔值
   • astype(int)：将布尔值转为0/1的标记

3. 生成连续状态块的分组ID
   对每个customer内的change_flag做累加，得到每个连续状态块的唯一ID：

   df['group_id'] = df.groupby('customer')['change_flag'].cumsum()
   

   同一连续状态块内的行group_id相同，状态变化后group_id递增。

4. 生成目标行号
   按customer和group_id分组，使用cumcount()生成块内的行号：

   df['transaction_in_row'] = df.groupby(['customer', 'group_id']).cumcount() + 1
   

   • cumcount()默认从0开始计数，加1后得到从1开始的行号

5. 清理中间字段（可选）
   如果不需要中间标记字段，可以删除：

   df = df.drop(['change_flag', 'group_id'], axis=1)
   

方案优势 #

所有操作均为pandas内置的矢量化分组运算，避免了自定义循环或apply逐行处理，能够高效支撑千万级数据的计算，性能远优于循环方案。

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