sklearn ColumnTransformer/Pipeline能否直接返回DataFrame？

你想要让预处理环节直接返回带正确标识的DataFrame的思路完全合理，完全不需要必须把Pipeline尾部挂模型才能正常使用，之前手动转DataFrame耗时久的问题，本质是操作方式不对，不是这个需求本身存在问题。

你之前的手动转换方案为什么慢 #

• 直接拿preprocessor.fit_transform()返回的结果硬转成空列名的DataFrame，本质是没有匹配对应的列名和行索引。如果你的OneHotEncoder默认返回稀疏矩阵，pandas会先把稀疏结构全量展开成稠密数组，遇到高基数字类别的场景，内存占用会暴涨数十倍，转换自然会慢到无法接受。
• 就算最后转成了DataFrame，因为没有对应列名，后续做特征工程、特征筛选的时候，还要靠列的整数下标反向猜每一列对应哪个原始特征，反而会增加额外的工作量。

最优实现方案（和原生Pipeline完全兼容） #

从sklearn 1.2版本开始，框架本身就原生支持所有Transformer直接输出pandas DataFrame，完全不需要手动做格式转换，速度和原来返回numpy数组几乎没有差异：

1. 先开启全局的pandas输出配置：

   from sklearn import set_config
   # 强制所有预处理步骤输出保留列名的pandas DataFrame
   set_config(transform_output="pandas")
   

2. 调整原有的pipeline参数，避免稀疏结构阻塞pandas输出流程：

   # 数值特征管线保持原有逻辑即可
   numerical_transformer = make_pipeline(
           KNNImputer(n_neighbors=3), 
           StandardScaler()
       )
   
   # 分类特征管线补充关闭稀疏输出的参数
   categorical_transformer = make_pipeline(
       SimpleImputer(strategy='most_frequent'),
       OneHotEncoder(handle_unknown='ignore', sparse_output=False)
   )
   
   # 初始化ColumnTransformer时补充参数，关闭内部稀疏矩阵拼接
   preprocessor = ColumnTransformer(
       transformers=[
           ('numerical', numerical_transformer, numerical_cols),
           ('categorical', categorical_transformer, categorical_cols)
       ],
       sparse_threshold=0
   )
   

3. 这时候直接执行preprocessor.fit_transform(train_df2)，返回的结果就是完整的DataFrame：
   • 数值列会自动保留原来的列名
   • OneHot编码生成的衍生列会自动标注来源，比如原来的城市列生成的编码列会自动命名为categorical__city_北京、categorical__city_上海，不需要手动对应
   • 行索引会和输入的训练数据完全对齐，完全不影响后续做特征衍生、特征筛选的操作

低版本sklearn的兼容方案 #

如果使用的sklearn版本低于1.2，没有全局pandas输出的配置，可以通过框架自带的特征名获取接口，快速生成DataFrame，不会有转换慢的问题：

# 先完成拟合并拿到转换结果
processed_arr = preprocessor.fit_transform(train_df2)
# 直接调用接口拿到和输出列顺序完全匹配的特征名列表
matched_cols = preprocessor.get_feature_names_out()
# 直接指定列名和原数据的行索引生成DataFrame
processed_df = pd.DataFrame(processed_arr, columns=matched_cols, index=train_df2.index)

这种方式因为直接给pandas指定了明确的列和行索引，不需要pandas做额外的类型推断和索引对齐，转换速度和直接使用numpy数组几乎没有差异，完全不会出现之前空转DataFrame的卡顿问题。

补充说明：这种把预处理步骤单独拿出来跑、拿到结构化的预处理结果再做后续特征工程的流程，本身就是特征工程阶段的标准操作，和后续把整个预处理管线挂到模型前面组成完整Pipeline的流程完全不冲突——等做完特征筛选，直接把预处理步骤+新增的特征工程步骤+模型拼到同一个Pipeline里，就可以直接做交叉验证、训练模型，完全不需要推倒重来。

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