Pandas Apply函数性能优化：RCA计算提速方案咨询

问题场景 #

现有如下结构的DataFrame：

需求是：

1. 创建rca列：当(empPeople/区域总员工数)/(行业总员工数/总员工数) >1时取值1，否则为0
2. 将DataFrame转换为以region为索引、activity为列、rca为值的透视表

原实现使用apply方法耗时长达6047秒，需要更高效的方案。

原实现代码 #

def rca_emp(activity:str, region:str , emp:float):
    top = emp / df[df['region'] == region].empPeople.sum()
    bottom = df[df['activity'] == activity].empPeople.sum() / df.empPeople.sum()
    rca = top/bottom
    if rca > 1: 
        return 1
    else:
        return 0

# 生成rca列
df['rca'] = df.apply(lambda x : rca_emp(activity=x['activity'] , region=x['region'] , emp=x['empPeople']) , axis=1)
# 创建透视表
df.pivot(index='region', columns='activity', values='rca')

性能瓶颈分析 #

原代码慢的核心原因：apply是逐行执行Python逻辑，每一行都要重新筛选region和activity并计算求和，相当于对全表做了N次（N为行数）扫描，数据量越大，耗时呈指数级增长。

优化方案：用矢量化操作替代逐行循环 #

利用Pandas的groupby.transform实现批量计算，所有操作都是底层C级别的矢量化运算，性能提升几个数量级。

完整优化代码 #

# 1. 预计算每个region的总员工数，映射到每一行
df['region_total'] = df.groupby('region')['empPeople'].transform('sum')
# 2. 预计算每个activity的总员工数，映射到每一行
df['activity_total'] = df.groupby('activity')['empPeople'].transform('sum')
# 3. 计算全局总员工数
total_emp = df['empPeople'].sum()

# 4. 矢量化计算RCA并生成rca列
df['rca'] = ((df['empPeople'] / df['region_total']) / (df['activity_total'] / total_emp) > 1).astype(int)

# 5. 创建透视表
pivot_df = df.pivot(index='region', columns='activity', values='rca')

优化点说明 #

• groupby.transform会把聚合结果广播到原DataFrame的每一行，仅需两次groupby聚合，而非N次逐行计算
• 所有比较和类型转换都是矢量化操作，完全避免Python层面的循环开销
• 计算逻辑与原需求完全一致，但性能提升巨大，百万级数据也能在几秒内完成

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