Keep column names after using scikit-learn ColumnTransfer on pandas DataFrame

Viewed 485

I am trying to write my preprocessing with SKLearn Pipelines and ColumnTransformer. However, the fact that the transformers return an array (and not a dataframe) is throwing me off a bit. I would like to be able to use the column names on the processed df also. Imagine the following data and pipeline:

import pandas as pd
import numpy as np
from sklearn.compose import ColumnTransformer
from sklearn.preprocessing import OneHotEncoder
from sklearn.preprocessing import StandardScaler, MinMaxScaler
from sklearn.pipeline import Pipeline
from sklearn.impute import SimpleImputer

df = pd.DataFrame(np.random.randn(6, 4), columns=list("ABCD"))
df["E"] = pd.Categorical(["test", "train", "test", "train", "test", "train"])
df["F"] = "foo"

num_columns = ['A', 'B', 'C']
num_transformer = Pipeline(
    steps=[
        ('imputer', SimpleImputer(strategy='median')),
        ('scaler', StandardScaler()),
    ]
)
cat_columns = ['E', 'F']
cat_transformer = Pipeline(
    steps = [
        ('imputer', SimpleImputer(strategy='most_frequent')),
        ('onhotencoder', OneHotEncoder())
    ]
)
preprocessor = ColumnTransformer(
    transformers=[
        ('num', num_transformer, num_columns),
        ('cat', cat_transformer, cat_columns)
    ]
)

Now I can get the transformed array

preprocessed_array = preprocessor.fit_transform(df)

and without the 'onhotencoder' step, I can even do something like this:

preprocessed_df = pd.DataFrame(preprocessed_array, columns=df.columns)

However, this is obviously not possible because the 'onehotencoder' step added columns to the array.

There is also the preprocessor.get_feature_names() method that should in principle work, but it doesn't because some of the transformers like the SimpleImputer do not have this method.

Does anyone know of a solution for getting a dataframe back from transforming another df?

0 Answers
Related