分组场景下fillna()的高效实现：哪种方法更优？

问题场景 #

我需要对DataFrame按group_by1和group_by2分组后，填充val列的NaN值。目前找到了两种实现方式，但遇到了问题：循环法速度更快，但处理后仍有残留NaN；想明确哪种方法更合适，以及有没有更高效的方案。

先给出测试用的示例DataFrame：

import numpy as np
import pandas as pd

N = 20_000 
df = pd.DataFrame({
    'val': np.random.choice([1, np.nan], p=[0.9, 0.1], size=N), 
    'group_by1': np.random.randint(1, 10, size=N), 
    'group_by2': np.random.randint(1, 10, size=N),
})

现有两种实现 #

1. Transform一行式 #

用groupby+transform，直接用分组均值填充NaN：

df["val"] = df.groupby(["group_by1", "group_by2"]).transform(lambda x: x.fillna(x.mean()))

2. 嵌套循环+拼接 #

先按group_by1分组，再嵌套按group_by2子分组，用分组中位数填充后拼接结果：

l = []
for f, df_f in df.groupby('group_by1'):
    for h, df_h in df_f.groupby('group_by2'):
        df_copy = df_h.copy()
        df_copy['val'].fillna(df_copy['val'].median(), inplace=True)
        l.append(df_copy)

df_output = pd.concat(l)

遇到的问题 #

测试下来，嵌套循环的速度比transform一行式快不少，但处理完的结果里还是有NaN值。

原因分析与优化方案 #

为什么循环后还有NaN？ #

问题出在：如果某个(group_by1, group_by2)的分组里，val列全是NaN，那计算出的中位数（或均值）也是NaN，填充自然没用。

优化后的实现 #

方案1：给全NaN分组加兜底值

不管用哪种方法，都可以给填充值加个兜底，比如用全局的均值/中位数，或者固定值：

• 优化transform版本（均值填充，全NaN分组用全局均值兜底）：

global_mean = df['val'].mean()
df["val"] = df.groupby(["group_by1", "group_by2"])['val'].transform(
    lambda x: x.fillna(x.mean() if not x.isna().all() else global_mean)
)

• 优化循环版本（中位数填充，全NaN分组用全局中位数兜底）：

global_median = df['val'].median()
l = []
for f, df_f in df.groupby('group_by1'):
    for h, df_h in df_f.groupby('group_by2'):
        df_copy = df_h.copy()
        group_median = df_copy['val'].median()
        # 分组中位数为NaN时用全局值兜底
        fill_val = group_median if not pd.isna(group_median) else global_median
        df_copy['val'].fillna(fill_val, inplace=True)
        l.append(df_copy)
df_output = pd.concat(l)

方案2：更高效的transform写法（避免lambda）

不用lambda，直接用transform("mean")配合fillna，性能比lambda版本更好：

# 基础版：分组均值填充
df["val"] = df["val"].fillna(df.groupby(["group_by1", "group_by2"])["val"].transform("mean"))

# 带兜底的版本：全NaN分组用全局均值
df["val"] = df["val"].fillna(
    df.groupby(["group_by1", "group_by2"])["val"].transform("mean")
).fillna(df["val"].mean())

方法选择建议 #

• 如果分组数量很少（比如示例里只有9×9=81个分组），循环法确实更快，但要记得处理全NaN分组的情况。
• 大多数场景下，推荐用优化后的transform写法：代码简洁易维护，性能稳定；当分组数量变多时，transform的效率优势会比循环法更明显。

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