基于非NaN条件汇总Pandas DataFrame的高效实现方法问询

问题背景 #

现有一个以因子为列、样本为行的DataFrame，其中大量因子值为NaN，另有一列非空的value列。需要对每个因子完成三项汇总：计算因子的均值、统计非NaN行数、计算该因子非NaN行的value总和。原实现中循环遍历每个因子计算value总和的方式效率极低，在25000个因子+750个样本的场景下耗时超2秒，且需处理数千个同类DataFrame。

示例数据：

# 源DataFrame
         p0        p1      value
0  0.588817  0.681094  80.087629
1       NaN  0.984875  52.507150
2  0.722749       NaN  84.599839
3  0.786565  0.917544   2.230228
4       NaN  0.506236  15.199917

# 预期输出
        mean  cnt       value
p0  0.699377    3  166.917697
p1  0.772437    4  150.024925

原低效代码：

df2 = pd.DataFrame(index=factors)
df2['mean'] = df.mean()
df2['cnt'] = df.count()
valuelist = [np.where(pd.isnull(df[f]), 0, df.value).sum() for f in factors]  # 循环遍历，效率瓶颈
df2['value'] = valuelist

优化方案 #

方案1：向量化矩阵乘法（最优性能） #

利用numpy的向量化运算，一次性计算所有因子的value总和，彻底避免Python循环：

import pandas as pd
import numpy as np

# 生成测试数据（同原逻辑）
numFactors = 25000
samples = 750
factors = [f'p{x}' for x in range(numFactors)]
df = pd.DataFrame(np.random.rand(samples, numFactors), columns=factors)
df = df.where(df > 0.5)
df['value'] = np.random.rand(samples, 1) * 100

# 高效计算流程
mask = df[factors].notna()  # 获取因子列的非NaN布尔掩码矩阵
value_sums = mask.T @ df['value']  # 矩阵乘法，一次性得到所有因子的value总和

# 构建结果DataFrame
df2 = pd.DataFrame({
    'mean': df[factors].mean(),
    'cnt': df[factors].count(),
    'value': value_sums
})

性能优势：将Python循环替换为numpy底层优化的C级运算，25000因子+750样本的场景下耗时可降至几十毫秒，效率提升数十倍。

方案2：宽表转长表+分组汇总 #

通过melt将宽表转为长表，过滤NaN后分组统计，代码逻辑更直观：

# 转长表，保留value关联关系
melted = df.melt(id_vars='value', value_vars=factors, var_name='factor', value_name='factor_val')
# 过滤NaN并一次性完成三项汇总
df2 = melted.dropna(subset=['factor_val']).groupby('factor').agg(
    mean=('factor_val', 'mean'),
    cnt=('factor_val', 'count'),
    value=('value', 'sum')
)

适用场景：如果需要叠加其他复杂分组逻辑，该方法可读性更强，但性能略逊于矩阵乘法方案。

核心优化逻辑 #

原循环方案的瓶颈在于Python层面的逐列遍历，而向量化操作将计算逻辑交给numpy的并行化C实现，充分利用CPU计算能力，大幅降低计算开销。

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