如何加速Pandas分组循环？千万级数据处理遇性能瓶颈

哇，1000万行数据、循环900万次跑10小时，这效率确实让人头疼！我帮你分析下原代码的核心问题，再给出几个能大幅提速的优化方案。

原代码的核心瓶颈 #

1. 无差别循环所有分组：大部分id只有1条访问记录，原代码却对每个组都执行了一套复杂操作，完全是做无用功。
2. 循环中反复concat：每次循环都把结果拼到d里，这会不断创建新的DataFrame——相当于每次都要复制一遍已有的数据，越到后面内存和时间开销越爆炸。
3. 组内操作冗余：用列表推导式生成c再concat，多次重置和设置索引，做了很多不必要的中间计算。

针对性优化方案 #

第一步：拆分数据，只处理需要的分组 #

先把只有1条记录的id单独拎出来直接处理，只对有多个访问记录的id执行复杂操作，瞬间减少90%以上的循环次数！

第二步：用向量化操作替代低效循环 #

用numpy或pandas的内置向量方法代替Python级别的循环，同时避免反复拼接数据。

完整优化代码 #

import pandas as pd
import numpy as np

# 读取数据（保持原读取逻辑）
data = pd.read_csv('https://raw.githubusercontent.com/skiler07/data/master/so_data.csv', index_col=0).set_index(['id','date'])

# 1. 拆分单条记录和多条记录的分组
grouped = data.groupby(level='id')

# 处理单条记录的id：直接添加visits=0，无需复杂操作
single_ids = grouped.filter(lambda x: len(x) == 1).reset_index()
single_ids['visits'] = 0
single_result = single_ids.set_index(['id', 'visits', 'date'])

# 2. 高效处理多条记录的id
def process_multi_group(group):
    n_rows = len(group)
    # 用numpy生成visits列：每个起始位置i对应n_rows-i条记录
    visits = np.repeat(np.arange(n_rows), np.arange(n_rows, 0, -1))
    # 生成重复的行索引，避免多次concat子DataFrame
    repeated_idx = group.index.repeat(np.arange(n_rows, 0, -1))
    # 构造结果并设置索引
    result = group.loc[repeated_idx.get_level_values('date')].reset_index()
    result['visits'] = visits
    return result.set_index(['id', 'visits', 'date'])

# 批量处理多记录分组，最后一次性拼接
multi_groups = grouped.filter(lambda x: len(x) > 1)
multi_result = pd.concat([process_multi_group(g) for _, g in multi_groups.groupby(level='id')])

# 3. 合并最终结果
final_d = pd.concat([single_result, multi_result])

为什么这能大幅提速？ #

• 减少循环次数：只处理有多个访问记录的id，大部分id直接跳过循环，节省了90%以上的迭代时间。
• 避免反复拼接：先收集所有多组的结果，最后一次concat，而不是每次循环都复制整个结果集。
• 向量化操作：用numpy的repeat替代Python循环生成visits列，速度快了几个数量级。
• 简化索引操作：减少了中间不必要的索引重置和设置步骤，降低计算开销。

额外优化建议 #

• 数据类型优化：把id设为int32，date转为datetime64类型，减少内存占用，进一步提升运算速度。
• 分块处理：如果内存压力大，可以把多记录的id分成若干批次处理，避免一次性加载过多数据。

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