[Python] How to preserve data types when working with pandas and sklearn transformers?

While working with a large sklearn Pipeline (fit using a DataFrame) I ran into an error that lead back to a wrong data type of my input. The problem occurred on an a single observation coming from an API that is supposed to interface the model in production. Missing information in a single line makes pandas (obviously) incapable of inferring the correct dtype but I thought that my fit transformers will handle the conversions. Apparently, I am mistaken.

import pandas as pd
from sklearn.impute import SimpleImputer

X_tr = pd.DataFrame({"A": [5, 8, None, 4], "B": [3, None, 9.9, 12]})
print(X_tr.dtypes)

#>> A float64
#>> B float64

x = pd.DataFrame({"A": [10.1], "B": [None]})
print(x.dtypes)

#>> A float64
#>> B object


The above shows clearly that pandas infers the float64 types for column A and B in the training dataset, however (again obviously) for the single observation it doesn't know the dtype for the column B so it assigns object. No issue so far. But let's imagine a SimpleImputer somewhere within a Pipeline to replace the missing values:

imputer = SimpleImputer(
fill_value=0, strategy="constant", missing_values=pd.NA
).set_output(transform="pandas")

X_tr_im = imputer.fit_transform(X_tr) # training
print(X_tr_im.dtypes)

#>> A float64
#>> B float64

x_im = imputer.transform(x)
print(x_im.dtypes)

#>> A object
#>> B object


The imputer does replace the None values with zeros in all cases, however, two things happened that I did not expect:

• Column B was NOT converted to the dtype that it was fit on
• Column A was converted to the unwanted dtype of object

This creates two unwanted non-numeric data types that lead to errors further down the pipeline. Even if it is not the task of transformers to preserve dtypes, in my case it would still be very helpful.

Am I doing something fundamentally wrong? Are there any solutions available?

Continue reading...