如何加速pd.apply()函数？百万级userid场景性能优化

问题背景 #

示例数据 #

import numpy as np
import pandas as pd

users = 5
size = users*6

df = pd.DataFrame(
    {'userid': np.random.choice(np.arange(0, users), size), 
    'a_time': np.random.normal(loc = 1.5, scale = 0.5, size = size), 
    'b_time': np.random.normal(loc = 1.5, scale = 0.5, size = size), 
    }
)
df['focus'] = np.where(df.userid % 2 == 0, 'a', 'b')

test_dat = df[['userid', 'focus', 'a_time', 'b_time']].sort_values('userid').copy(deep = True).reset_index(drop = True)

核心需求 #

针对每个userid，根据其focus列的取值情况，统计对应条件的次数：

• 若该用户的focus全为a，统计b_time > a_time的次数
• 若该用户的focus全为b，统计a_time > b_time的次数
• 若该用户的focus全为both，统计次数设为0
• 若该用户的focus存在多种取值（混合），则标记为None
  同时要统计每个用户对应的记录行数（命名为days）。

现有问题 #

用户原本通过groupby.apply调用自定义函数实现，但当唯一userid数量超过10万时，运行速度极慢：

def some_func(x):
    if (x.focus == 'a').all():
        a = x.a_time
        b = x.b_time
        x['changes'] = (b > a).sum()
        x['days'] = len(a)
    elif (x.focus == 'b').all():
        a = x.a_time
        b = x.b_time
        x['changes'] = (a > b).sum()
        x['days'] = len(a)
    elif (x.focus == 'both').all():
        x['changes'] = 0
        x['days'] = len(a)
    else: 
        x['changes'] = None
        x['days'] = None
    
    return x

test_dat.groupby(['userid', 'focus']).apply(some_func).reset_index(name = 'n_changes')

优化方案 #

groupby.apply本质是Python层面的逐组循环，在十万级分组场景下效率极低。改用矢量化操作（基于pandas底层C实现）可以大幅提升速度，具体步骤如下：

1. 标记每个用户的focus类型 #

先统计每个userid下focus的唯一取值，确定其属于哪种类型：

# 统计每个userid的focus唯一值集合
focus_groups = test_dat.groupby('userid')['focus'].agg(lambda x: set(x.unique()))

# 映射为明确的类型标签
focus_type = focus_groups.map(
    lambda s: 'a_only' if s == {'a'} 
              else 'b_only' if s == {'b'} 
              else 'both_only' if s == {'both'} 
              else 'mixed'
).reset_index(name='focus_type')

2. 批量计算基础统计量 #

提前用矢量化方式计算所有用户的行数、两种条件的计数：

# 计算每个userid的记录行数
user_days = test_dat.groupby('userid').size().reset_index(name='days')

# 先创建全局的条件列（完全矢量化）
test_dat['b_gt_a'] = test_dat['b_time'] > test_dat['a_time']
test_dat['a_gt_b'] = test_dat['a_time'] > test_dat['b_time']

# 分组求和得到每个用户的条件计数
cnt_b_gt_a = test_dat.groupby('userid')['b_gt_a'].sum().reset_index(name='cnt_b_gt_a')
cnt_a_gt_b = test_dat.groupby('userid')['a_gt_b'].sum().reset_index(name='cnt_a_gt_b')

3. 合并结果并生成最终数据 #

将所有统计结果合并，根据focus_type映射对应的changes值：

# 合并所有统计数据
result = focus_type.merge(user_days, on='userid')\
                   .merge(cnt_b_gt_a, on='userid')\
                   .merge(cnt_a_gt_b, on='userid')

# 根据focus_type生成最终的changes列
result['changes'] = np.select(
    [
        result['focus_type'] == 'a_only',
        result['focus_type'] == 'b_only',
        result['focus_type'] == 'both_only',
        result['focus_type'] == 'mixed'
    ],
    [
        result['cnt_b_gt_a'],
        result['cnt_a_gt_b'],
        0,
        None
    ]
)

# 保留需要的列
final_result = result[['userid', 'focus_type', 'changes', 'days']]

效果说明 #

这种方案把原来的逐组Python循环，替换为批量预处理+矢量化计算，在十万级userid的场景下，速度能提升几十到上百倍，完全规避了apply带来的性能瓶颈。

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