使用Pandas与Sklearn转换器时如何保留数据类型？

你遇到的问题是Pandas和Scikit-learn的行为特性导致的：

• 当单条观测的某列全为pd.NA时，Pandas无法推断出数值类型，会默认设为object；
• Scikit-learn的SimpleImputer仅负责填充缺失值，不会自动将列转换回训练时的数值类型，甚至会因为输入中存在object列，导致原本是数值类型的列在转换后也变成object，进而引发后续Pipeline的类型错误。

可行解决方案 #

1. 生产环境输入时强制指定列类型 #

训练阶段先保存各列的类型，生产环境接收数据后，直接用该类型创建DataFrame：

# 训练时保存列类型（可序列化存储，比如用joblib）
train_dtypes = X_tr.dtypes

# 生产环境构造输入DataFrame时指定类型
x = pd.DataFrame({"A": [10.1], "B": [None]}, dtype=train_dtypes)
print(x.dtypes)
#>> A    float64
#>> B    float64

# 再用imputer转换
x_im = imputer.transform(x)
print(x_im.dtypes)
#>> A    float64
#>> B    float64

2. 在Pipeline中加入类型强制转换步骤 #

自定义一个转换器，在Imputer之前将列转换为训练时的类型，确保输入到Imputer的数据类型与训练集一致：

from sklearn.preprocessing import FunctionTransformer
from sklearn.pipeline import Pipeline

# 定义类型转换函数
def enforce_dtypes(X, dtypes):
    return X.astype(dtypes)

# 训练时保存列类型
train_dtypes = X_tr.dtypes

# 创建类型转换器
type_enforcer = FunctionTransformer(enforce_dtypes, kw_args={"dtypes": train_dtypes})

# 构建Pipeline
pipeline = Pipeline([
    ("type_enforcer", type_enforcer),
    ("imputer", SimpleImputer(
        fill_value=0, strategy="constant", missing_values=pd.NA
    ).set_output(transform="pandas"))
])

# 训练Pipeline
pipeline.fit(X_tr)

# 生产环境测试
x = pd.DataFrame({"A": [10.1], "B": [None]})
x_im = pipeline.transform(x)
print(x_im.dtypes)
#>> A    float64
#>> B    float64

3. 用ColumnTransformer针对数值列统一处理 #

如果你的数据包含多种类型（数值、类别），可以用ColumnTransformer单独对数值列做类型转换+缺失值填充，确保类型安全：

from sklearn.compose import ColumnTransformer

# 定义数值列处理流程：先转float再填充
numeric_transformer = Pipeline([
    ("cast_to_float", FunctionTransformer(lambda x: x.astype(float))),
    ("imputer", SimpleImputer(fill_value=0, strategy="constant", missing_values=pd.NA))
])

# 构建预处理器
preprocessor = ColumnTransformer(
    transformers=[
        ("numeric", numeric_transformer, ["A", "B"])
    ]).set_output(transform="pandas")

# 训练
preprocessor.fit(X_tr)

# 测试
x_im = preprocessor.transform(x)
print(x_im.dtypes)
#>> A    float64
#>> B    float64

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