Pandas分组嵌套apply性能优化求助：计算前5行最大2值均值

问题描述 #

我有一个包含category和number列的DataFrame，需要新增avg_of_largest_2_from_prev_5列：按category分组后，计算当前行除外的前5行number值中最大2个值的均值。

原实现及性能瓶颈 #

原代码 #

np.random.seed(123)
n_rows = 10000
data = {'category': np.random.randint(1, 1000, n_rows), 'number': np.random.randint(1, n_rows, n_rows)}
df = pd.DataFrame(data)

%timeit df['avg_of_largest_2_from_prev_5'] = df.groupby('category')['number'].apply(lambda x: x.shift(1).rolling(5, min_periods=0).apply(lambda y: pd.Series(y).nlargest(2).mean()))

df = df[df['category'] == df['category'].values[0]]
df

原性能表现 #

4.55 s ± 34.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
10000行数据、1000个分类时耗时36秒，百万级数据需约8分钟

输出示例 #

category        number   avg_of_largest_2_from_prev_5
0           511         4179    NaN
392         511         2878    4179.0
1292        511         5834    3528.5
1350        511         1054    5006.5
1639        511         8673    5006.5
3145        511         8506    7253.5
4176        511         947     8589.5
4471        511         151     8589.5
4735        511         5326    8589.5
4965        511         4827    8589.5
5046        511         9792    6916.0
5316        511         3772    7559.0
5535        511         1095    7559.0
5722        511         5619    7559.0
5732        511         700     7705.5
6825        511         1156    7705.5
6877        511         7240    4695.5
8100        511         2381    6429.5
8398        511         2376    6429.5

原方案的核心瓶颈在于嵌套lambda调用和rolling.apply的逐行Python循环，每次窗口计算都要转换为Series并调用nlargest，大量Python层操作导致效率极低。

高效优化方案 #

方案1：Numpy向量化实现 #

利用numpy的np.partition快速获取窗口内前两大值，结合rolling.apply(raw=True)直接操作numpy数组，避免不必要的类型转换。

import numpy as np
import pandas as pd

np.random.seed(123)
n_rows = 10000
data = {'category': np.random.randint(1, 1000, n_rows), 'number': np.random.randint(1, n_rows, n_rows)}
df = pd.DataFrame(data)

def compute_top2_mean(x):
    shifted = x.shift(1).values
    
    def window_func(arr):
        # 过滤NaN值
        valid_vals = arr[~np.isnan(arr)]
        val_count = len(valid_vals)
        
        if val_count == 0:
            return np.nan
        elif val_count == 1:
            return valid_vals[0]
        # 取最大的两个值计算均值
        top_two = np.partition(valid_vals, -2)[-2:]
        return top_two.mean()
    
    return pd.Series(shifted).rolling(5, min_periods=0).apply(window_func, raw=True)

# 分组计算并计时
%timeit df['avg_of_largest_2_from_prev_5'] = df.groupby('category')['number'].apply(compute_top2_mean)

# 查看单个分类结果
df_single = df[df['category'] == df['category'].values[0]]
print(df_single)

方案2：Numba编译加速 #

对于超大规模数据，用numba编译窗口计算函数，将Python循环转为机器码执行，进一步提升速度。

import numpy as np
import pandas as pd
from numba import jit

np.random.seed(123)
n_rows = 10000
data = {'category': np.random.randint(1, 1000, n_rows), 'number': np.random.randint(1, n_rows, n_rows)}
df = pd.DataFrame(data)

@jit(nopython=True)
def numba_window_calc(arr):
    max1 = -np.inf
    max2 = -np.inf
    count = 0
    
    for val in arr:
        if np.isnan(val):
            continue
        count += 1
        if val > max1:
            max2 = max1
            max1 = val
        elif val > max2:
            max2 = val
    
    if count == 0:
        return np.nan
    elif count == 1:
        return max1
    else:
        return (max1 + max2) / 2

def compute_top2_mean_numba(x):
    shifted = x.shift(1).values
    return pd.Series(shifted).rolling(5, min_periods=0).apply(numba_window_calc, raw=True)

# 分组计算并计时
%timeit df['avg_of_largest_2_from_prev_5'] = df.groupby('category')['number'].apply(compute_top2_mean_numba)

# 查看单个分类结果
df_single = df[df['category'] == df['category'].values[0]]
print(df_single)

性能对比 #

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