Pandas按列分组计算统计值新增列时返回NaN该如何解决

问题原因 #

你代码运行后得到全NaN值的核心原因是索引不对齐：
你调用df.groupby('timestamp').apply(calc_some_stat)得到的是一个以分组键timestamp的取值（10、20）为索引的Series，结构如下：

timestamp
10    37
20    44
dtype: int64

而你的原始DataFrame的索引是0~5的连续整数，和分组计算结果的索引完全不匹配。pandas在给列赋值时会自动按索引对齐，找不到对应索引的位置就会填充NaN，因此最终stat列全为空值。

解决方案 #

有两种常用的修复方式：

方案1：用transform代替apply（更推荐） #

transform会将分组计算的结果广播回对应分组的每一行，返回结果的长度、索引和原始DataFrame完全一致，直接赋值即可：

import pandas as pd
import numpy as np

df = pd.DataFrame({'timestamp' : [10,10,10,20,20,20], 'idx': [1,2,3,1,2,3], 'v1' : [1,2,4,5,1,9], 'v2' : [1,2,8,5,1,2]})

def calc_some_stat(d):
    return np.sum(d.v1 * d.v2)

df['stat'] = df.groupby('timestamp').transform(calc_some_stat)

如果要简化代码，也可以不用单独定义函数：

df['stat'] = df.eval('v1 * v2').groupby(df['timestamp']).transform('sum')

方案2：用map匹配分组计算结果 #

如果你一定要用apply做分组计算，可以先把分组统计结果存下来，再用timestamp列匹配统计值：

calc_stat = df.groupby('timestamp').apply(calc_some_stat)
df['stat'] = df['timestamp'].map(calc_stat)

两种方法都可以得到你需要的结果。

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