Pandas按分组执行Transform日期计算的性能优化方案

需求说明 #

现有数据集需按group_id分组计算first_date字段，计算规则：

• 同组内当前行start_date与上一行end_date的差值大于等于20天时，当前行first_date取自身start_date
• 不满足上述条件时，当前行继承上一行的first_date值

原始数据集样例 #

id         group_id        start_date   end_date
45891167   00000dc8-b87e   2017-09-14   2017-12-15
45891167   00000dc8-b87e   2017-12-15   2018-03-18
50692679   000012e9-e137   2018-03-29   2019-03-05
50692679   000012e9-e137   2019-03-05   2019-03-27
50692679   000012e9-e137   2019-03-27   2019-04-05
51622620   000011e5-e107   2021-03-15   NaT
81652621   000019e1-e111   2020-06-15   2020-06-19
81652621   000019e1-e111   2020-07-16   NaT

期望输出样例 #

id         group_id        start_date   end_date    first_date    
45891167   00000dc8-b87e   2017-09-14   2017-12-15  2017-09-14
45891167   00000dc8-b87e   2017-12-15   2018-03-18  2017-09-14
50692679   000012e9-e137   2018-03-29   2019-03-05  2018-03-29 
50692679   000012e9-e137   2019-03-07   2019-03-27  2018-03-29
50692679   000012e9-e137   2019-03-27   2019-04-05  2018-03-29
51622620   000011e5-e107   2021-03-15   NaT         2021-03-15
81652621   000019e1-e111   2020-06-15   2020-06-19  2020-06-15
81652621   000019e1-e111   2021-07-16   NaT         2021-07-16

原有实现问题 #

原有代码采用逐group_id遍历+组内iterrows逐行计算的逻辑，在数据集包含超100万个独立group_id时，Python层循环开销极大，性能无法满足要求。原有代码如下：

for group in df['group_id'].unique():
    indexes = (df.loc[df['group_id'] == group]).index.tolist()

    df['end_date_previous_row'] = (df.loc[df['group_id'] == group])['end_date'].shift(1)

    if len((df.loc[df['group_id'] == group]).index) == 1:
        df['first_date'] = (df.loc[df['group_id'] == group])['start_date']
    else:
        for idx, row in (df.loc[df['group_id'] == group]).iterrows():
            if idx == indexes[0]:
                df.loc[idx, 'first_date'] = df.loc[idx, 'start_date']
            else:
                difference = row['start_date'].to_pydatetime() - row['end_date_previous_row'].to_pydatetime()
                df.loc[idx, 'difference'] = difference
                if difference.days >= 20:
                    df.loc[idx, 'first_date'] = df.loc[idx, 'start_date']
                else:
                    df['first_date'].fillna(method='ffill', limit=1, inplace=True)

向量化高性能实现 #

全程使用Pandas底层C实现的向量化接口，完全抛弃Python层显式循环，在百万级分组场景下性能可提升数十至上百倍：

1. 首先确保日期字段格式正确

import pandas as pd
import numpy as np

df['start_date'] = pd.to_datetime(df['start_date'])
df['end_date'] = pd.to_datetime(df['end_date'])

2. 分组计算同组上一行的结束日期，向量化计算日期间隔

# 分组shift取上一行end_date，无需逐组循环
df['prev_end_date'] = df.groupby('group_id')['end_date'].shift(1)
# 向量化计算两行日期间隔天差
df['day_gap'] = (df['start_date'] - df['prev_end_date']).dt.days

3. 标记分段起点，按组向前填充得到first_date

# 分段起点判定：组内第一行（上一行end_date为空）、间隔>=20天
reset_mask = df['prev_end_date'].isna() | (df['day_gap'] >= 20)
# 起点位置赋值当前start_date，其余位置设为空值
df['first_date'] = np.where(reset_mask, df['start_date'], pd.NaT)
# 按组向前填充空值，继承上一个分段的首日期
df['first_date'] = df.groupby('group_id')['first_date'].ffill()

4. 可选：删除临时计算列

df = df.drop(columns=['prev_end_date', 'day_gap'])

注意：执行上述操作前需要保证每个group_id内的行已经按照时间先后顺序排序，否则shift取值会出错。如果未排序，可先执行df = df.sort_values(by=['group_id', 'start_date']).reset_index(drop=True)预处理。

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