基于交付日期条件计算Pandas分组的过去12个月平均实际提前期

问题背景 #

现有如下结构的Pandas DataFrame（25万+行）：

<class 'pandas.core.frame.DataFrame'>
Int64Index: 250409 entries, 0 to 350928
Data columns (total 17 columns):
 #   Column               Non-Null Count   Dtype         
---  ------               --------------   -----         
 0   Location             250409 non-null  object        
 1   Code                 250409 non-null  object        
 2   PO                   250409 non-null  float64       
 3   Delivered_Date       250409 non-null  datetime64[ns]
 4   Supplier_Number      250409 non-null  object        
 5   Supplier_Name        250409 non-null  object        
 6   Created_Date         250409 non-null  datetime64[ns]
 7   Original_Quantity    250409 non-null  float64       
 8   Delivered_Quantity   250409 non-null  float64       
 9   Article_description  250409 non-null  object        
10  Brand                250409 non-null  object        
11  Category             250409 non-null  object        
12  Sub_Category         250409 non-null  object        
13  actual_lead_time     250409 non-null  int64         
14  OrderedDateMonth     250409 non-null  int64         
15  OrderedateYear       250409 non-null  int64         
16  PoDeliveredSequence  250409 non-null  int64         
dtypes: datetime64[ns](2), float64(3), int64(4), object(8)
memory usage: 34.4+ MB

需要新增三列，分别按以下分组规则，计算分组内Delivered_Date在当前行Created_Date过去12个月范围内的actual_lead_time平均值：

• Supplier_Number, Brand
• Supplier_Number, Sub_Category
• Supplier_Number, Category

计算规则示例：某行Supplier_Number=101010、Brand=Coca_cola、Created_Date=2024-02-12时，需计算该分组中Delivered_Date在2023-02-12至2024-02-11区间内的actual_lead_time平均值。

原代码使用iterrows逐行循环+全表过滤，导致Python内核卡顿无法完成计算，代码如下：

import pandas as pd

# Assuming hpo_df is your DataFrame with columns Created_Date, Delivered_Date, actual_lead_time, Supplier_Number, Brand, Category, Sub_Category

# Convert Created_Date and Delivered_Date to datetime objects if they are not already
hpo_df['Created_Date'] = pd.to_datetime(hpo_df['Created_Date'])
hpo_df['Delivered_Date'] = pd.to_datetime(hpo_df['Delivered_Date'])

# Function to calculate mean actual_lead_time for last 12 months prior to a given created_date
def mean_lead_time_last_12_months(created_date, supplier_number, additional_col, df):
    # Filter rows based on Supplier_Number and additional_col
    filtered_df = df[(df['Supplier_Number'] == supplier_number) & (df[col] == additional_col)]
    
    # Filter rows where Delivered_Date falls within last 12 months prior to created_date
    filtered_df = filtered_df[(filtered_df['Delivered_Date'] >= created_date - pd.DateOffset(months=12)) & (filtered_df['Delivered_Date'] < created_date)]
    
    # Calculate mean of actual_lead_time for the filtered rows
    mean_lead_time = filtered_df['actual_lead_time'].mean()
    
    return mean_lead_time

# Iterate over rows and calculate mean actual_lead_time for each row and each additional column
mean_lead_times = []
cols = ['Brand', 'Category', 'Sub_Category']

for index, row in hpo_df.iterrows():
    created_date = row['Created_Date']
    supplier_number = row['Supplier_Number']
    for col in cols:
        additional_col = row[col]
        mean_lead_time = mean_lead_time_last_12_months(created_date, supplier_number, additional_col, hpo_df)
        mean_lead_times.append(mean_lead_time)

    # Add mean_lead_times as a new column in the DataFrame
    hpo_df['mean_lead_time_last_12_months'] = mean_lead_times

问题分析 #

原代码的核心性能瓶颈：

• 使用iterrows逐行遍历25万+行，属于低效的循环操作
• 每次循环都对全表进行两次过滤（分组过滤+时间过滤），时间复杂度为O(n²)，数据量较大时完全无法承受

优化解决方案 #

采用分组排序+矢量化时间范围查询的方式，利用Pandas的分组和二分查找能力大幅提升效率，步骤如下：

优化思路 #

1. 对每个目标分组，预先将该分组的记录按Delivered_Date排序，并计算累积均值和累积计数，方便快速计算时间窗口内的均值
2. 对主表的每一行，在对应分组中通过二分查找定位时间窗口的边界，利用累积统计值快速计算区间均值

优化代码 #

import pandas as pd
import numpy as np

# 确保日期列是datetime类型（如果还没转换的话）
hpo_df['Created_Date'] = pd.to_datetime(hpo_df['Created_Date'])
hpo_df['Delivered_Date'] = pd.to_datetime(hpo_df['Delivered_Date'])

# 定义需要计算的分组组合
group_combinations = [
    ('Supplier_Number', 'Brand'),
    ('Supplier_Number', 'Sub_Category'),
    ('Supplier_Number', 'Category')
]

for group_cols in group_combinations:
    # 1. 按分组列+Delivered_Date排序，准备构建累积统计
    grouped_df = hpo_df.sort_values(by=list(group_cols) + ['Delivered_Date']).copy()
    
    # 2. 对每个分组计算累积求和、累积计数，以及对应的日期
    grouped_df['cum_sum'] = grouped_df.groupby(list(group_cols))['actual_lead_time'].cumsum()
    grouped_df['cum_count'] = grouped_df.groupby(list(group_cols))['actual_lead_time'].cumcount() + 1  # 从1开始计数
    grouped_dates = grouped_df.groupby(list(group_cols))['Delivered_Date'].apply(list).to_dict()
    grouped_cum_sum = grouped_df.groupby(list(group_cols))['cum_sum'].apply(list).to_dict()
    grouped_cum_count = grouped_df.groupby(list(group_cols))['cum_count'].apply(list).to_dict()
    
    # 3. 定义函数，根据当前行的分组值和Created_Date计算区间均值
    def calculate_window_mean(row):
        group_key = tuple(row[col] for col in group_cols)
        if group_key not in grouped_dates:
            return np.nan
        
        dates = grouped_dates[group_key]
        cum_sum = grouped_cum_sum[group_key]
        cum_count = grouped_cum_count[group_key]
        
        # 计算时间窗口的上下界
        lower_bound = row['Created_Date'] - pd.DateOffset(months=12)
        upper_bound = row['Created_Date']
        
        # 用二分查找定位上下界的索引
        left_idx = np.searchsorted(dates, lower_bound, side='left')
        right_idx = np.searchsorted(dates, upper_bound, side='left')
        
        if left_idx == right_idx:
            return np.nan
        
        # 计算区间内的总和和数量，得到均值
        total_sum = cum_sum[right_idx - 1] - (cum_sum[left_idx - 1] if left_idx > 0 else 0)
        total_count = cum_count[right_idx - 1] - (cum_count[left_idx - 1] if left_idx > 0 else 0)
        
        return total_sum / total_count if total_count > 0 else np.nan
    
    # 4. 应用函数生成新列，列名根据分组组合命名
    new_col_name = f'mean_lead_time_{"_".join(group_cols)}'
    hpo_df[new_col_name] = hpo_df.apply(calculate_window_mean, axis=1)

进一步优化（可选） #

如果性能还需要提升，可以改用merge_asof结合分组统计的方式，完全避免apply操作：

for group_cols in group_combinations:
    # 1. 准备分组后的时间序列数据
    stats_df = hpo_df.sort_values(by=list(group_cols) + ['Delivered_Date']).copy()
    stats_df['cum_sum'] = stats_df.groupby(list(group_cols))['actual_lead_time'].cumsum()
    stats_df['cum_count'] = stats_df.groupby(list(group_cols))['actual_lead_time'].cumcount() + 1
    
    # 2. 为每个主表行生成时间窗口下界
    hpo_df['lower_bound'] = hpo_df['Created_Date'] - pd.DateOffset(months=12)
    
    # 3. 合并找到每个Created_Date对应的最近Delivered_Date的累积统计（上界）
    upper_merge = pd.merge_asof(
        hpo_df.sort_values('Created_Date'),
        stats_df.sort_values('Delivered_Date'),
        left_on='Created_Date',
        right_on='Delivered_Date',
        by=list(group_cols),
        direction='backward'
    )
    
    # 4. 合并找到每个lower_bound对应的最近Delivered_Date的累积统计（下界）
    lower_merge = pd.merge_asof(
        hpo_df.sort_values('lower_bound'),
        stats_df.sort_values('Delivered_Date'),
        left_on='lower_bound',
        right_on='Delivered_Date',
        by=list(group_cols),
        direction='backward'
    )
    
    # 5. 计算区间均值
    upper_merge = upper_merge.set_index(hpo_df.index)
    lower_merge = lower_merge.set_index(hpo_df.index)
    
    total_sum = upper_merge['cum_sum'] - lower_merge['cum_sum'].fillna(0)
    total_count = upper_merge['cum_count'] - lower_merge['cum_count'].fillna(0)
    
    new_col_name = f'mean_lead_time_{"_".join(group_cols)}'
    hpo_df[new_col_name] = total_sum / total_count.where(total_count > 0, np.nan)
    
    # 清理临时列
    hpo_df.drop('lower_bound', axis=1, inplace=True)

效果说明 #

优化后的代码利用了Pandas的分组排序、累积统计和二分查找/merge_asof的矢量化能力，时间复杂度降至O(n log n)，能够在合理时间内处理25万行的数据量，避免内核卡顿问题。

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