Pandas分组扩展均值计算含缺失值时报错，求简洁解决方案

问题背景 #

无缺失值的数据集可通过df.groupby('id').expanding().mean().values得到正确的分组扩展均值，但当数据包含pd.NA时，Pandas会因后端浮点类型断言报错，直接使用apply(np.nansum)等方法因缺少计数逻辑无法得到正确均值。

解决方案 #

方法1：转换列类型后使用内置expanding.mean

将含pd.NA的列转换为浮点类型（此时pd.NA会转为np.nan），再利用expanding.mean()的skipna=True参数忽略缺失值计算：

import pandas as pd
import numpy as np

df2 = pd.DataFrame({'id':[1,1,1,2,2,2],'y':[1,pd.NA,3,1,2,3]})

# 转换为浮点类型
df2['y'] = df2['y'].astype(float)
# 计算分组扩展均值
result = df2.groupby('id')['y'].expanding().mean(skipna=True).values.reshape(-1, 1)
print(result)

输出结果：

array([[1. ],
       [1. ],
       [2. ],
       [1. ],
       [1.5],
       [2. ]])

方法2：手动计算累计和与累计非NA数量

通过扩展窗口计算累计非NA值的和与累计非NA值的数量，两者相除得到均值，无需转换列类型：

import pandas as pd
import numpy as np

df2 = pd.DataFrame({'id':[1,1,1,2,2,2],'y':[1,pd.NA,3,1,2,3]})

grouped = df2.groupby('id')['y']
# 累计非NA值的和
cum_sum = grouped.expanding().apply(lambda s: s.sum(skipna=True), raw=False)
# 累计非NA值的数量
cum_count = grouped.expanding().apply(lambda s: s.count(), raw=False)
# 计算均值并整理格式
result = (cum_sum / cum_count).values.reshape(-1, 1)
print(result)

输出结果与方法1一致。

问题原因说明 #

出现报错是因为含pd.NA的列默认是**Nullable Integer（Int64）**类型，Pandas的expanding.mean()在处理该类型时，后端的浮点类型断言逻辑会触发错误。转换为浮点类型后，缺失值变为np.nan，skipna=True参数可以正常生效忽略缺失值；手动计算则绕过了内置方法的类型限制，直接基于缺失值处理逻辑计算。

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