使用Pandas为大型数据集分组计算新增特征的高效方法

优化实现代码 #

import pandas as pd

# 提前把timedelta类型的天数转为整数，避免重复计算属性
joinedData['DaysSincePurchase'] = joinedData['DaysSincePurchase'].dt.days

# 直接用pandas内置的分组聚合接口，全程矢量化运算无Python循环开销
customer_agg = joinedData.groupby('CustomerID', as_index=False).agg(
    TotalDemand = ('Amount', 'sum'),
    DaysSinceLastPurchase = ('DaysSincePurchase', 'min'),
    nrDepartments = ('Department', 'nunique')
)

# 替换首列为1开始的连续整数
customer_agg.insert(0, 'CustomerID', range(1, len(customer_agg) + 1))

# 得到的customer_agg就是你需要的最终数据集

核心优化说明 #

你原来的代码耗时极长的核心原因是用Python层的循环逐次调用get_group筛选用户数据，每一次循环都要遍历全量数据集匹配用户ID，1800万行规模下时间开销会指数级上升。
而pandas内置的agg聚合是底层C语言实现的矢量化操作，没有Python循环的额外开销，1800万行数据的聚合通常仅需几秒到几十秒即可完成。
另外你原来的循环写法存在逻辑漏洞：如果你的CustomerID不是从1开始的连续整数，会直接报错或者计算结果完全错位，优化后的写法先按真实CustomerID聚合再统一生成连续序号，逻辑更严谨。

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