Pandas性能优化：迭代处理超200万条数据过慢的解决方法

数据集 #

表格形式 #

pandas构造代码 #

import pandas as pd

x = pd.DataFrame({
    'id': ['001', '002', '003', '004', '005', '006', '007', '008'],
    'open_month': ['2023-01-31', '2023-01-31', '2023-01-31', '2023-02-28', '2023-02-28', '2023-02-28', '2023-03-31', '2023-03-31'],
    'end_month': ['2023-02-28', '2023-03-31', '2023-04-30', '2023-02-28', '2023-03-31', '2023-04-30', '2023-03-31', '2023-04-30'],
    'quantity': [1, 5, 4, 2, 3, 6, 7, 9]
})

需求说明 #

需要生成一张汇总表：

• 行：开户后的月数（从0开始）
• 列：开户月份
• 值：对应开户月份的记录中，结束时间小于等于「开户月份+N个月」的quantity总和

预期结果 #

原实现代码（性能瓶颈） #

table = pd.DataFrame()
for i in range(len(x['open_month'].unique())):
    for month in x['open_month'].unique():
        date = month + pd.offsets.MonthEnd(i)
        table.at[i, month] = x.query('open_month == @month and end_month <= @date')['quantity'].sum()

原代码在200万行数据上运行极慢，需要性能优化方案。

优化方案 #

核心思路 #

避免嵌套循环和逐次query（每次query都会遍历全表，时间复杂度O(n²)），改用向量化日期计算+分组累积求和，将时间复杂度降到O(n)级别。

步骤1：转换日期类型并计算月差 #

先将日期列转为datetime类型，再计算每条记录的结束时间与开户时间的月差（即该记录会在开户后第几个月被计入汇总）：

# 转换为日期类型
x['open_month'] = pd.to_datetime(x['open_month'])
x['end_month'] = pd.to_datetime(x['end_month'])

# 计算结束时间相对于开户时间的月数差（得到该记录对应的目标行索引）
x['month_diff'] = (x['end_month'].dt.year - x['open_month'].dt.year) * 12 + (x['end_month'].dt.month - x['open_month'].dt.month)

步骤2：分组透视并计算累积求和 #

按开户月份和月差分组求和，再透视成目标格式，最后计算累积求和并补全所有需要的行：

# 分组求和：按开户月份和月差汇总quantity
grouped = x.groupby(['open_month', 'month_diff'])['quantity'].sum().unstack(fill_value=0)

# 生成所有需要的月差行（从0到最大月差）
max_month_diff = grouped.columns.max()
all_month_diffs = pd.Index(range(max_month_diff + 1), name='month_diff')
grouped = grouped.reindex(all_month_diffs, fill_value=0)

# 计算累积求和（第N行需要包含所有<=N月差的记录总和）
result = grouped.cumsum(axis=0)

# 调整索引和列名格式，匹配预期结果
result.index.name = '开户后月数'
result.columns = result.columns.strftime('%Y-%m-%d')

验证结果 #

执行上述代码后，result将与预期结果完全一致，且处理200万行数据的时间会从分钟级压缩到秒级。

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