Groupby滚动均值设置min_periods=1且不忽略NaN的实现问题

我需要在分组计算滚动均值时，只要滚动窗口内存在一个NaN，计算出的均值就为NaN，但目前找到的资料都是讲如何忽略NaN的。

输入数据 #

grouping    value_to_avg
0   1           1.0
1   1           2.0
2   1           3.0
3   1           NaN
4   1           4.0
5   2           5.0
6   2           NaN
7   2           6.0
8   2           7.0
9   2           8.0

生成输入数据的代码：

import numpy as np
import pandas as pd

data = {'grouping': [1,1,1,1,1,2,2,2,2,2], 'value_to_avg': [1,2,3,np.nan,4,5,np.nan,6,7,8]}  
db = pd.DataFrame(data)

尝试的代码 #

db['rolling_mean_actual'] = db.groupby('grouping')['value_to_avg'].transform(lambda s: s.rolling(window=3, center=True, min_periods=1).mean(skipna=False))

实际与预期输出对比 #

可以看到，设置mean(skipna=False)并未达到预期效果——当窗口内非NaN值数量满足min_periods时，仍会计算均值。

解决方案 #

使用rolling.apply()自定义计算逻辑，先检查窗口内是否存在NaN，存在则返回NaN，否则计算均值：

def rolling_mean_with_nan(window):
    if np.isnan(window).any():
        return np.nan
    return window.mean()

db['rolling_mean_expected'] = db.groupby('grouping')['value_to_avg'].transform(
    lambda s: s.rolling(window=3, center=True, min_periods=1).apply(rolling_mean_with_nan, raw=True)
)

说明 #

• rolling.apply()支持自定义窗口计算逻辑，完全适配需求；
• raw=True让函数直接处理numpy数组，提升计算性能；
• 函数内通过np.isnan(window).any()判断窗口是否含NaN，是则返回NaN，否则计算均值。

运行后得到的rolling_mean_expected列将完全符合预期。

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