Skip to content

Narwhals categorical imputer #1009

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
solegalli wants to merge 2 commits into
base: narwhals-migration
Choose a base branch
from
Open

Narwhals categorical imputer #1009

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
fb23ba6
Migrate BaseImputer to narwhals, add polars support
solegalli Aug 24, 2026
87361c1
Migrate CategoricalImputer to narwhals, add polars support
solegalli Aug 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
86 changes: 85 additions & 1 deletion docs/user_guide/imputation/CategoricalImputer.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -292,7 +292,91 @@ We see that this variable has 3 categories with similar maximum number of observ
0 Ex
1 Fa
2 Gd
Name: PoolQC, dtype: object
Name: PoolQC, dtype: str

With polars
-----------

:class:`CategoricalImputer()` works in the same way with a polars dataframe:

.. code:: python

import polars as pl
from feature_engine.imputation import CategoricalImputer

df = pl.DataFrame({
"City": ["London", "Manchester", None, "Bristol", "London", None],
"Studies": ["Bachelor", None, "Bachelor", "PhD", None, "Masters"],
})

imputer = CategoricalImputer(imputation_method="frequent")
print(imputer.fit_transform(df))

The most frequent category imputation gives the same result as with pandas:

.. code:: text

shape: (6, 2)
┌────────────┬──────────┐
│ City ┆ Studies │
│ --- ┆ --- │
│ str ┆ str │
╞════════════╪══════════╡
│ London ┆ Bachelor │
│ Manchester ┆ Bachelor │
│ London ┆ Bachelor │
│ Bristol ┆ PhD │
│ London ┆ Bachelor │
│ London ┆ Masters │
└────────────┴──────────┘

Imputing with an arbitrary string also works the same way:

.. code:: python

imputer = CategoricalImputer(fill_value="Missing")
print(imputer.fit_transform(df))

.. code:: text

shape: (6, 2)
┌────────────┬──────────┐
│ City ┆ Studies │
│ --- ┆ --- │
│ str ┆ str │
╞════════════╪══════════╡
│ London ┆ Bachelor │
│ Manchester ┆ Missing │
│ Missing ┆ Bachelor │
│ Bristol ┆ PhD │
│ London ┆ Missing │
│ Missing ┆ Masters │
└────────────┴──────────┘

.. note::

polars' ``Categorical`` dtype accepts a brand-new fill value automatically,
unlike pandas' ``category`` dtype, which needs its categories widened first
(:class:`CategoricalImputer()` handles that difference for you on both
backends). polars' ``Enum`` dtype, however, has a *fixed* set of categories
that cannot be widened. If you impute a fixed-category ``Enum`` column with
a `fill_value` that isn't already one of its categories, the transformer
raises a clear error instead of silently writing null:

.. code:: python

enum_dtype = pl.Enum(["London", "Manchester", "Bristol"])
df_enum = df.with_columns(pl.col("City").cast(enum_dtype))

imputer = CategoricalImputer(fill_value="Missing", variables=["City"])
imputer.fit_transform(df_enum)

.. code:: text

ValueError: Cannot fill variable 'City' with 'Missing': it is a polars
Enum with fixed categories ('London', 'Manchester', 'Bristol') that do
not include the fill value. Cast the column to Categorical or String
before imputing.

Considerations
--------------
Expand Down
69 changes: 49 additions & 20 deletions feature_engine/imputation/base_imputer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,18 +1,18 @@
import pandas as pd
import narwhals as nw
import narwhals.dependencies as nwd
from narwhals.typing import IntoDataFrame
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.utils.validation import check_is_fitted

from feature_engine._base_transformers.mixins import GetFeatureNamesOutMixin
from feature_engine.dataframe_checks import _check_X_matches_training_df, check_X
from feature_engine.tags import _return_tags

_PANDAS_LT_3 = int(pd.__version__.split(".")[0]) < 3


class BaseImputer(TransformerMixin, BaseEstimator, GetFeatureNamesOutMixin):
"""shared set-up checks and methods across imputers"""

def _transform(self, X: pd.DataFrame) -> pd.DataFrame:
def _transform(self, X: IntoDataFrame) -> IntoDataFrame:
"""
Common checks before transforming data:

Expand All @@ -23,11 +23,11 @@ def _transform(self, X: pd.DataFrame) -> pd.DataFrame:

Parameters
----------
X: Pandas DataFrame
X: dataframe of shape = [n_samples, n_features]

Returns
-------
X: Pandas DataFrame
X: dataframe.
The same dataframe entered by the user.
"""
# Check method fit has been called
Expand All @@ -40,42 +40,71 @@ def _transform(self, X: pd.DataFrame) -> pd.DataFrame:
_check_X_matches_training_df(X, self.n_features_in_)

# reorder df to match train set
X = X[self.feature_names_in_]
is_pandas = nwd.is_pandas_dataframe(X)
if is_pandas is True:
X = X[self.feature_names_in_]
else:
X = (
nw.from_native(X, eager_only=True)
.select(self.feature_names_in_)
.to_native()
)

return X

def transform(self, X: pd.DataFrame) -> pd.DataFrame:
def transform(self, X: IntoDataFrame) -> IntoDataFrame:
"""
Replace missing data with the learned parameters.

Parameters
----------
X: pandas dataframe of shape = [n_samples, n_features]
X: dataframe of shape = [n_samples, n_features]
The data to be transformed.

Returns
-------
X_new: pandas dataframe of shape = [n_samples, n_features]
X_new: dataframe of shape = [n_samples, n_features]
The dataframe without missing values in the selected variables.
"""

X = self._transform(X)

# Replace missing data with learned parameters. In pandas < 3, fillna
# downcasts object columns and warns; the option applies the pandas 3
# behavior: no downcasting, and infer_objects restores numeric dtypes.
if _PANDAS_LT_3:
with pd.option_context("future.no_silent_downcasting", True):
# Benchmarked: pandas-native fillna is ~1.3-1.6x faster than the
# narwhals-generic fill_null equivalent at the 10k-100k row sizes
# imputers are typically used at (the gap narrows to parity only
# past ~1M rows), so pandas keeps its own fast path here.
is_pandas = nwd.is_pandas_dataframe(X)
if is_pandas is True:
# Namespace of the dataframe already in hand, not a fresh import:
# pandas can only reach this branch already imported by the caller.
pd = nw.from_native(X, eager_only=True).__native_namespace__()
pandas_lt_3 = int(pd.__version__.split(".")[0]) < 3
# In pandas < 3, fillna downcasts object columns and warns; the
# option applies the pandas 3 behavior: no downcasting, and
# infer_objects restores numeric dtypes.
if pandas_lt_3 is True:
with pd.option_context("future.no_silent_downcasting", True):
X = X.fillna(value=self.imputer_dict_)
else:
X = X.fillna(value=self.imputer_dict_)
X = X.infer_objects()
else:
X = X.fillna(value=self.imputer_dict_)
return X.infer_objects()
nw_X = nw.from_native(X, eager_only=True)
nw_X = nw_X.with_columns(
nw.col(var).fill_null(value)
for var, value in self.imputer_dict_.items()
)
X = nw_X.to_native()

return X

def _get_feature_names_in(self, X):
"""Get the names and number of features in the train set (the dataframe
used during fit)."""

self.feature_names_in_ = X.columns.to_list()
is_pandas = nwd.is_pandas_dataframe(X)
if is_pandas is True:
self.feature_names_in_ = list(X.columns)
else:
self.feature_names_in_ = nw.from_native(X, eager_only=True).columns
self.n_features_in_ = X.shape[1]

return self
Expand Down
129 changes: 85 additions & 44 deletions feature_engine/imputation/categorical.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,9 @@

from typing import List, Optional, Union

import pandas as pd
import narwhals as nw
import narwhals.dependencies as nwd
from narwhals.typing import IntoDataFrame, IntoSeries

from feature_engine._check_init_parameters.check_variables import (
_check_variables_input_value,
Expand Down Expand Up @@ -162,16 +164,17 @@ def __init__(
_check_return_empty_is_bool(return_empty)
self.return_empty = return_empty

def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None):
"""
Learn the most frequent category if the imputation method is set to frequent.

Parameters
----------
X: pandas dataframe of shape = [n_samples, n_features]
The training dataset.
X: dataframe of shape = [n_samples, n_features]
The training dataset. Can be a pandas, polars, or any other dataframe
supported by narwhals.

y: pandas Series, default=None
y: Series, default=None
y is not needed in this imputation. You can pass None or y.
"""

Expand All @@ -194,46 +197,52 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
imputer_dict_ = {var: self.fill_value for var in variables_}

elif self.imputation_method == "frequent":
# if imputing only 1 variable:
if len(variables_) == 1:
var = variables_[0]
mode_vals = X[var].mode()
# Benchmarked (10k-100k rows x 1-10 cols): a per-variable mode()
# loop is not slower than pandas' batch X[variables_].mode(), so
# both branches loop the same way - only the mode() call differs.
is_pandas = nwd.is_pandas_dataframe(X)
imputer_dict_ = {}
multi_mode_vars = []
if is_pandas is True:
for var in variables_:
mode_vals = X[var].mode()
if len(mode_vals) > 1:
multi_mode_vars.append(var)
else:
imputer_dict_[var] = mode_vals[0]
else:
nw_X = nw.from_native(X, eager_only=True)
for var in variables_:
# Unlike pandas' mode(dropna=True default), polars' mode()
# does not drop nulls, so a column whose nulls outnumber
# any single category would otherwise make null "the mode".
mode_vals = nw_X[var].drop_nulls().mode(keep="all")
if len(mode_vals) > 1:
multi_mode_vars.append(var)
else:
imputer_dict_[var] = mode_vals[0]

# Some variables may contain more than 1 mode:
if len(mode_vals) > 1:
# Some variables may contain more than 1 mode:
if len(multi_mode_vars) > 0:
varnames_str = ", ".join(str(v) for v in multi_mode_vars)
if len(variables_) == 1:
raise ValueError(
f"The variable {var} contains multiple frequent categories."
f"The variable {varnames_str} contains multiple frequent "
"categories."
)

imputer_dict_ = {var: mode_vals[0]}

# imputing multiple variables:
else:
# Returns a dataframe with 1 row if there is one mode per
# variable, or more rows if there are more modes:
mode_vals = X[variables_].mode()

# Careful: some variables contain multiple modes
if len(mode_vals) > 1:
varnames = mode_vals.dropna(axis=1).columns.to_list()
if len(varnames) > 1:
varnames_str = ", ".join(varnames)
else:
varnames_str = varnames[0]
else:
raise ValueError(
f"The variable(s) {varnames_str} contain(s) multiple frequent "
f"categories."
f"The variable(s) {varnames_str} contain(s) multiple "
"frequent categories."
)

imputer_dict_ = mode_vals.iloc[0].to_dict()

self.variables_ = variables_
self.imputer_dict_ = imputer_dict_
self._get_feature_names_in(X)

return self

def transform(self, X: pd.DataFrame) -> pd.DataFrame:
def transform(self, X: IntoDataFrame) -> IntoDataFrame:
# Frequent category imputation
if self.imputation_method == "frequent":
X = super().transform(X)
Expand All @@ -242,19 +251,51 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame:
else:
X = self._transform(X)

# if variable is of type category, we need to add the new
# category, before filling in the nan
for variable in self.variables_:
if X[variable].dtype.name == "category":
X[variable] = X[variable].cat.add_categories(
self.imputer_dict_[variable]
)

X = X.fillna(self.imputer_dict_)
is_pandas = nwd.is_pandas_dataframe(X)
if is_pandas is True:
# if variable is of type category, we need to add the new
# category, before filling in the nan
for variable in self.variables_:
if X[variable].dtype.name == "category":
X[variable] = X[variable].cat.add_categories(
self.imputer_dict_[variable]
)

X = X.fillna(self.imputer_dict_)
else:
nw_X = nw.from_native(X, eager_only=True)
schema = nw_X.schema
for variable in self.variables_:
dtype = schema[variable]
fill_value = self.imputer_dict_[variable]
# polars' Categorical widens itself on fill_null, but its
# Enum has a fixed category set and silently fills with
# null (no error) if fill_value isn't already a member.
if isinstance(dtype, nw.Enum) and (
fill_value not in dtype.categories
):
raise ValueError(
f"Cannot fill variable '{variable}' with "
f"'{fill_value}': it is a polars Enum with fixed "
f"categories {dtype.categories} that do not include "
"the fill value. Cast the column to Categorical or "
"String before imputing."
)

nw_X = nw_X.with_columns(
nw.col(var).fill_null(value)
for var, value in self.imputer_dict_.items()
)
X = nw_X.to_native()

# add additional step to return variables cast as object
if self.return_object:
X[self.variables_] = X[self.variables_].astype("O")
if self.return_object is True:
is_pandas = nwd.is_pandas_dataframe(X)
if is_pandas is True:
X[self.variables_] = X[self.variables_].astype("O")
# polars/narwhals backends never silently upcast a string-typed
# column back to numeric (unlike pandas' fillna+infer_objects),
# so there is nothing to recast there.

return X

Expand Down
Loading