基于类别分组的SimpleImputer自定义填充器实现问题排查

报错原因分析 #

你遇到的ValueError: Found array with 0 sample(s)核心问题是：transform阶段出现了fit阶段未见过的分类分组，导致该分组下没有可处理的样本，而SimpleImputer要求至少1个样本才能执行填充逻辑。

自定义转换器在fit时仅保存了训练数据中存在的分组对应的Imputer实例，但transform时遇到新分组（或训练时未出现的分组值）时，会尝试对空样本集执行填充操作，直接触发报错。

修复方案 #

通过以下两点修改解决问题：

1. 在fit阶段额外训练一个全局默认Imputer，用于处理transform时出现的新分组；
2. 在transform阶段判断当前分组是否存在于已训练的Imputer字典中，不存在则使用全局默认Imputer填充。

修复后的完整代码 #

import numpy as np
import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.impute import SimpleImputer

class GroupedSimpleImputer(BaseEstimator, TransformerMixin):
    def __init__(self, group_col, strategy='mean'):
        self.group_col = group_col
        self.strategy = strategy
        self.imputers_ = {}  # 存储各分组的Imputer实例
        self.default_imputer_ = None  # 全局默认Imputer

    def fit(self, X, y=None):
        # 确保输入为DataFrame格式
        if not isinstance(X, pd.DataFrame):
            X = pd.DataFrame(X)
        
        # 训练全局默认Imputer（使用所有数据的统计量）
        self.default_imputer_ = SimpleImputer(strategy=self.strategy)
        self.default_imputer_.fit(X.drop(self.group_col, axis=1))
        
        # 按分组训练专属Imputer
        groups = X.groupby(self.group_col)
        for group_name, group_data in groups:
            imputer = SimpleImputer(strategy=self.strategy)
            imputer.fit(group_data.drop(self.group_col, axis=1))
            self.imputers_[group_name] = imputer
        
        return self

    def transform(self, X, y=None):
        if not isinstance(X, pd.DataFrame):
            X = pd.DataFrame(X)
        
        # 保存原索引，保证输出顺序与输入一致
        original_index = X.index
        # 拆分分组列和待填充列
        group_col_series = X[self.group_col]
        data_cols = X.drop(self.group_col, axis=1)
        
        filled_dfs = []
        # 遍历每个唯一分组
        for group_name in group_col_series.unique():
            mask = group_col_series == group_name
            current_data = data_cols[mask]
            
            # 选择对应Imputer填充
            if group_name in self.imputers_:
                filled_data = self.imputers_[group_name].transform(current_data)
            else:
                filled_data = self.default_imputer_.transform(current_data)
            
            # 组装回DataFrame并恢复分组列与索引
            filled_df = pd.DataFrame(filled_data, columns=data_cols.columns, index=X[mask].index)
            filled_df[self.group_col] = group_name
            filled_dfs.append(filled_df)
        
        # 合并数据并按原索引排序
        result = pd.concat(filled_dfs).reindex(original_index)
        return result

关键修复点说明 #

• 全局默认Imputer：避免因新分组导致的空样本集问题，保证所有分组都能正常填充；
• 原索引恢复：严格保持输出数据与输入数据的行顺序一致，符合sklearn转换器的规范；
• 明确列拆分：将分组列与待填充列分离，避免分组列被传入Imputer引发不必要的错误。

测试示例 #

# 训练数据
train_data = pd.DataFrame({
    'category': ['A', 'A', 'B', 'B', 'C'],
    'col1': [1, np.nan, 3, np.nan, 5],
    'col2': [np.nan, 2, np.nan, 4, np.nan]
})

# 初始化并训练转换器
imputer = GroupedSimpleImputer(group_col='category', strategy='mean')
imputer.fit(train_data)

# 测试数据（包含训练时未出现的分组'D'和缺失值）
test_data = pd.DataFrame({
    'category': ['A', 'B', 'D', np.nan],
    'col1': [np.nan, np.nan, 7, np.nan],
    'col2': [np.nan, np.nan, np.nan, 8]
})

# 执行填充
filled_test = imputer.transform(test_data)
print(filled_test)

内容的提问来源于stack exchange，提问作者Irving Pérez