Python如何结合groupby使用sklearn的KNNImputer填充缺失值

方案选型原则 #

• 禁止全量数据直接拟合插补器：会打破name分组的逻辑边界，把不同主体的样本混在一起计算，插补结果不符合业务要求，且全量KNN计算复杂度高，大数据量下容易触发内存溢出
• 分组独立计算：每个分组单独拟合插补模型，既保证插补逻辑符合分组要求，又天然支持并行拆分，适配大规模数据处理
• 按需选择插补器：固定属性字段（如示例中的age、job）无需用复杂的KNN，用组内非空值填充即可；时序波动字段（如示例中的weight）可根据数据量选择KNN、中位数、线性插值等策略

基础示例数据构造 #

import pandas as pd
import numpy as np
from sklearn.impute import KNNImputer

# 构造带缺失值的测试数据集
data = {'name':  ['Alex', 'Ben', 'Marry','Alex', 'Ben', 'Marry']*2,
        'job': ['teacher', 'doctor', 'engineer']*4,
        'age': [27, 32, 78,27, 32, 78] + [np.nan]*3,
        'weight': [160, 209, 130,164, 206, 132] + [np.nan]*3,
        'date': ['6-12-2022']*3 + ['6-13-2022']*3 + ['6-14-2022']*3
        }
df = pd.DataFrame(data)

中小规模数据实现（代码最简） #

通过pandas原生groupby.apply实现，逻辑清晰，适合十万级以内数据量：

def knn_group_impute(group_df):
    target_cols = ['age', 'weight']
    # 初始化KNN插补器，邻居数根据组内有效样本量调整
    imputer = KNNImputer(n_neighbors=min(2, group_df[target_cols].dropna().shape[0]))
    group_df[target_cols] = imputer.fit_transform(group_df[target_cols])
    return group_df

df_imputed = df.groupby('name', group_keys=False).apply(knn_group_impute)

示例数据插补结果：Alex的age填充为27，weight填充为162；Ben的age填充为32，weight填充为207.5；Marry的age填充为78，weight填充为131，完全符合同一个体的属性逻辑。

大规模数据优化实现（百万级以上） #

针对百万级以上数据，做两点性能优化：一是跳过无缺失值的分组减少无效计算，二是启用多进程并行充分利用CPU资源：

# 导入并行计算工具，需先安装：pip install pandarallel
from pandarallel import pandarallel
# 初始化并行环境，默认调用所有CPU核心
pandarallel.initialize(progress_bar=False)

def optimized_impute(group_df):
    target_cols = ['age', 'weight']
    # 当前组无缺失值直接返回，跳过计算
    if not group_df[target_cols].isna().any().any():
        return group_df
    valid_count = group_df[target_cols].dropna().shape[0]
    # 组内有效样本不足2个时，自动降级为中位数填充，避免KNN报错
    if valid_count < 2:
        group_df[target_cols] = group_df[target_cols].fillna(group_df[target_cols].median())
        return group_df
    # 正常执行KNN插补
    imputer = KNNImputer(n_neighbors=min(5, valid_count))
    group_df[target_cols] = imputer.fit_transform(group_df[target_cols])
    return group_df

df_imputed_large = df.groupby('name', group_keys=False).parallel_apply(optimized_impute)

轻量替代方案（性能最优） #

如果不需要KNN的近邻计算逻辑，可选择更轻量的填充策略，处理速度比KNN快10倍以上：

• 固定属性列（age、job）：同一个体取值固定，直接取组内第一个非空值填充
• 数值波动列（weight）：取组内中位数/均值、或时间序列线性插值填充

# 按规则批量填充
fill_config = {
    'age': 'first',
    'job': 'first',
    'weight': 'median'
}
for col, agg_method in fill_config.items():
    df[col] = df.groupby('name')[col].transform(lambda x: x.fillna(x.agg(agg_method)))

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