0
$\begingroup$

I am having issues with scikit-learn converting dataframes to numpy arrays.

For instance, the following code

from sklearn.impute import SimpleImputer
import pandas as pd

df = pd.DataFrame(dict(
    x=[1, 2, np.nan],
    y=[2, np.nan, 0]
))

SimpleImputer().fit_transform(df)

Returns

array([[1. , 2. ],
       [2. , 1. ],
       [1.5, 0. ]])

Is there a way to use an imputer that returns a pandas dataframe instead of a numpy array? Is there a scikit-learn implementation for that? I am aware of sklearn-pandas, but the interface is kind of different.

$\endgroup$
1
2
$\begingroup$

Might be late but for anyone with the same question the answers (as almost everything with Scikit-learn) is the usage of Pipelines

from sklearn.impute import SimpleImputer
from sklearn.preprocessing import FunctionTransformer
from sklearn.pipeline import Pipeline
import pandas as pd

df = pd.DataFrame(dict(
    x=[1, 2, np.nan],
    y=[2, np.nan, 0]
))

imputer = Pipeline([("imputer", SimpleImputer()),
                    ("pandarizer",FunctionTransformer(lambda x: pd.DataFrame(x, columns = ["x", "y"])))])

imputer.fit_transform(df)
$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.