基于Pandas实现分组后指定行增量及计算编号添加（大数据场景）

现有一个通过嵌套循环计算并与客户数据合并得到的大型DataFrame（df1，1000万+行），简化示例如下：

import pandas as pd

df1 = pd.DataFrame({
    "id": ['z111','z111','z111','z111','z112','z112','z112','z112'],
    "calc_amt": [1000,500,200,300,100,50,30,200],
    "month_end":['28-02-2023','28-02-2023','28-02-2023','28-02-2023','28-02-2023','28-02-2023','28-02-2023','28-02-2023'],
    "period":[2,2,2,2,6,6,6,6],
})

需要按id分组，以最高效的方式完成以下操作：

• 组内第3、4行的period列值加1；
• 组内第3、4行的month_end日期调整为下月月末日期；
• 添加calc_num列，组内按顺序循环标记1、2。

期望输出示例如下：

df2 = pd.DataFrame({
    "id": ['z111','z111','z111','z111','z112','z112','z112','z112'],
    "calc_amt": [1000,500,200,300,100,50,30,200],
    "month_end":['28-02-2023','28-02-2023','31-03-2023','31-03-2023','28-02-2023','28-02-2023','31-03-2023','31-03-2023'],
    "period":[2,2,3,3,6,6,7,7],
    "calc_num":[1,2,1,2,1,2,1,2],
})

高效实现方案 #

针对千万级数据，必须避免循环，采用pandas矢量化分组操作保证性能：

1. 转换日期格式：先将month_end转为datetime类型，方便后续日期计算

df1['month_end'] = pd.to_datetime(df1['month_end'], format='%d-%m-%Y')

2. 生成组内索引与调整标记：按id分组后，用cumcount()生成组内连续索引，以此定位需要调整的行

# 生成组内行号（从0开始计数）
df1['group_idx'] = df1.groupby('id').cumcount()
# 标记组内第3、4行（对应索引2、3）
adjust_mask = df1['group_idx'] >= 2

3. 调整period列：对标记行的period值加1

df1.loc[adjust_mask, 'period'] += 1

4. 调整月末日期：对标记行的日期添加一个月月末偏移量

from pandas.tseries.offsets import MonthEnd
df1.loc[adjust_mask, 'month_end'] = df1.loc[adjust_mask, 'month_end'] + MonthEnd(1)

5. 添加calc_num标记：通过组内行号取余实现1、2循环标记

df1['calc_num'] = (df1['group_idx'] % 2) + 1

6. 清理临时列：删除用于计算的辅助列

df1 = df1.drop('group_idx', axis=1)

7. 还原日期格式（可选）：如果需要将日期转回原字符串格式

df1['month_end'] = df1['month_end'].dt.strftime('%d-%m-%Y')

该方案全程采用矢量化操作，分组计算效率远高于循环，完美适配千万级大型数据集的处理需求。

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