基于Pandas与Scikit-Learn的混合数据类型缺失值动态插补

下面是一个无需强制统一缺失值编码、自动适配列类型的预处理实现，能同时处理np.nan、None、pd.NaN，且插补后严格保留原列的数据类型：

import pandas as pd
import numpy as np
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.pipeline import Pipeline

class SmartImputer(BaseEstimator, TransformerMixin):
    def fit(self, X, y=None):
        # 为每列计算插补值，同时记录原列数据类型
        self.fill_values_ = {}
        self.dtypes_ = X.dtypes.to_dict()
        
        for col in X.columns:
            col_data = X[col]
            # 用pd.isna统一识别所有缺失值类型
            mask = pd.isna(col_data)
            
            if pd.api.types.is_numeric_dtype(col_data):
                # 数值列：中位数插补，整数列强制保留整数类型
                fill_val = col_data.median()
                if pd.api.types.is_integer_dtype(col_data):
                    fill_val = int(fill_val)
            else:
                # 非数值列：众数插补，无明确众数时取第一个非缺失值兜底
                fill_val = col_data.mode().iloc[0] if not col_data.mode().empty else col_data.dropna().iloc[0]
            
            self.fill_values_[col] = fill_val
        return self
    
    def transform(self, X):
        X_transformed = X.copy()
        for col in X_transformed.columns:
            # 定位所有缺失值并插补
            mask = pd.isna(X_transformed[col])
            X_transformed.loc[mask, col] = self.fill_values_[col]
            # 强制恢复原列数据类型
            X_transformed[col] = X_transformed[col].astype(self.dtypes_[col])
        return X_transformed

# 示例测试
if __name__ == "__main__":
    # 构造包含多种缺失值的测试数据集
    test_df = pd.DataFrame({
        'int_col': [1, 2, None, 4, 5],
        'float_col': [1.1, np.nan, 3.3, 4.4, pd.NaN],
        'str_col': ['a', 'b', None, 'b', 'c'],
        'cat_col': pd.Categorical(['x', 'y', pd.NaN, 'y', 'x'])
    })
    
    # 构建并运行插补管道
    imputer_pipeline = Pipeline([('smart_imputer', SmartImputer())])
    imputed_df = imputer_pipeline.fit_transform(test_df)
    
    print("原数据类型：")
    print(test_df.dtypes)
    print("\n插补后结果：")
    print(imputed_df)
    print("\n插补后数据类型（与原类型完全一致）：")
    print(imputed_df.dtypes)

核心设计细节： #

• 统一缺失值识别：用pd.isna()替代单独判断np.nan/None/pd.NaN，该函数原生支持识别所有Pandas兼容的缺失值编码，无需手动转换。
• 类型严格保留：
  • 数值列计算中位数后，若原列为整数类型，强制将中位数转为整数再插补，避免插补后列类型从int变为float。
  • 所有列插补完成后，强制转回原列的数据类型，彻底规避缺失值编码转换导致的类型变更问题。
• 鲁棒的众数处理：针对非数值列，优先取出现次数最多的值；若所有值出现频率相同（无明确众数），则取第一个非缺失值作为兜底方案。
• Sklearn生态兼容：自定义Transformer继承了Sklearn的基础类，可无缝集成到Sklearn Pipeline中，后续可串联标准化、编码等其他预处理步骤。

为什么不直接用SimpleImputer？ #

原生SimpleImputer默认仅识别np.nan，即便设置missing_values=pd.NA兼容所有缺失值，仍会将含None的整数列自动转为浮点型。上述自定义实现通过手动记录并恢复原类型，彻底解决了这个痛点。

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