From ae0a9c2696b07dbfcfc4dac4cb6a3427d0699e34 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Tue, 25 Aug 2026 01:42:56 +0200 Subject: [PATCH] Migrate CategoricalMethodsMixin (encoding base) to narwhals, add polars support Shared base for all 8 encoders. _get_feature_names_in() and _check_transform_input_and_state() follow the same is_pandas-gated column-reorder pattern as BaseImputer/DecisionTreeFeatures. _check_or_select_variables() needed no change: the variable_handling helpers it calls are already fully narwhals-generic. The hot path is _encode()/inverse_transform(), a per-column dict-based map applied on every transform() call across every encoder. Benchmarked pandas-native .map(dict) vs narwhals Series.replace_strict(dict, default=...) at 10k/50k/100k rows x 1/2/10 columns x 5/50 categories (warmed up first to remove first-call JIT/import overhead): narwhals-on-pandas lands at ~1.06x-1.2x of pandas-native at realistic sizes (50k-100k rows), i.e. minimal loss - merged into a single narwhals path per the established decision rule, no pandas fast-path split. narwhals-on- polars is consistently ~4-5x faster than pandas-native at 100k rows. replace_strict() also *simplifies* the old logic: pandas' plain .map() leaves category-dtype columns as category dtype after mapping, which the old code corrected with a manual "cast to int if all-int else float" step. Verified narwhals' replace_strict resolves straight to a plain numeric dtype on both a pandas category column and a polars Categorical column, so that dtype fixup is dead code once replace_strict replaces .map() - dropped it entirely rather than porting it. Used Series.get_column().replace_strict() (not nw.col(), which only accepts string names) throughout, same as DecisionTreeFeatures' precedent for pandas integer column names - nw.col(feature) blew up on int-named columns (caught by the existing test_column_names_are_numbers test, which polars can't cover since it has no integer-column-name concept). _check_nan_values_after_transformation() rewritten off pandas' .isnull().sum().sum()/.columns[...] chain onto per-column Series.null_count(), for the same int-column-name reason. Verified: tests/test_encoding full suite unchanged (17 pre-existing failures - numpy-array-input rejection per the narwhals check_X() contract, plus 3 MeanEncoder inverse_transform failures caused by a pre-existing bug in mean_encoding.py's still-unmigrated fit() passing a numpy y into y.groupby(); reproduced identically against the unmodified base_encoder.py to confirm neither predates nor is introduced by this change - 326 passed both before and after, same failing test IDs). flake8 and mypy clean on the file. Module imports with pandas blocked (loaded standalone, since sibling encoder files in this package are not yet migrated and still import pandas at their own module level). sphinx -W build clean (only the pre-existing unrelated linkcode_resolve warning). Manually verified CountEncoder end-to-end on polars input (fit still pandas-only until its own migration, transform/inverse_transform now backend-agnostic via this mixin) produces identical values to the pandas path, including a pre-existing quirk where count-encoding inverse_transform is ambiguous for categories that share a count (confirmed identical, not a regression, on the old code too). _helper_functions.py checked: pure-python parameter validation, no dataframe interaction, no pandas import - left untouched. Co-Authored-By: Claude Sonnet 5 --- feature_engine/encoding/base_encoder.py | 115 ++++++++++++++---------- 1 file changed, 69 insertions(+), 46 deletions(-) diff --git a/feature_engine/encoding/base_encoder.py b/feature_engine/encoding/base_encoder.py index 53eca3095..e2c4395f1 100644 --- a/feature_engine/encoding/base_encoder.py +++ b/feature_engine/encoding/base_encoder.py @@ -1,7 +1,9 @@ import warnings from typing import List, Union -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 @@ -121,11 +123,11 @@ class CategoricalMethodsMixin(TransformerMixin, BaseEstimator, GetFeatureNamesOu - GetFeatureNamesOutMixin brings method get_feature_names_out(). """ - def _check_na(self, X: pd.DataFrame, variables): + def _check_na(self, X: IntoDataFrame, variables): if self.missing_values == "raise": _check_contains_na(X, variables, error_msg="optional") - def _check_or_select_variables(self, X: pd.DataFrame): + def _check_or_select_variables(self, X: IntoDataFrame): """ Finds categorical variables, or alternatively checks that the variables entered by the user are of type object (categorical). @@ -133,7 +135,7 @@ def _check_or_select_variables(self, X: pd.DataFrame): Parameters ---------- - X: Pandas DataFrame + X: dataframe Raises ------ @@ -159,37 +161,41 @@ def _check_or_select_variables(self, X: pd.DataFrame): return variables_ - def _get_feature_names_in(self, X: pd.DataFrame): + def _get_feature_names_in(self, X: IntoDataFrame): """ Returns attributes `featrure_names_in_` and `n_feature_names_in_`, which are standard for all transformers in the library. """ # save input features - self.feature_names_in_ = X.columns.tolist() + 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 # save train set shape self.n_features_in_ = X.shape[1] - def _check_transform_input_and_state(self, X: pd.DataFrame) -> pd.DataFrame: + def _check_transform_input_and_state(self, X: IntoDataFrame) -> IntoDataFrame: """ Checks that the input is a dataframe and of the same size than the one used in the fit method. Checks absence of NA. Parameters ---------- - X: Pandas DataFrame + X: dataframe Raises ------ TypeError - If the input is not a Pandas DataFrame + If the input is not a dataframe ValueError - If the variable(s) contain null values. - If the df has different number of features than the df used in fit() Returns ------- - X: Pandas DataFrame + X: dataframe The same dataframe entered by the user. """ @@ -203,21 +209,29 @@ def _check_transform_input_and_state(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 categories with the learned parameters. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features]. + X: dataframe of shape = [n_samples, n_features]. The dataset to transform. Returns ------- - X_new: pandas dataframe of shape = [n_samples, n_features]. + X_new: dataframe of shape = [n_samples, n_features]. The dataframe containing the categories replaced by numbers. """ @@ -231,22 +245,25 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame: return X - def _encode(self, X: pd.DataFrame) -> pd.DataFrame: - # replace categories by the learned parameters - for feature in self.encoder_dict_.keys(): - X[feature] = X[feature].map(self.encoder_dict_[feature]) - - # if original variables are cast as categorical, they will remain - # categorical after the encoding, and this is probably not desired - if X[feature].dtype.name == "category": - if all(isinstance(x, int) for x in X[feature]): - X[feature] = X[feature].astype("int") - else: - X[feature] = X[feature].astype("float") - - if self.unseen == "encode": - X[self.variables_] = X[self.variables_].fillna(self._unseen) - else: + def _encode(self, X: IntoDataFrame) -> IntoDataFrame: + # replace categories by the learned parameters. + # narwhals' replace_strict() lets one expression both map known + # categories and fill unseen/missing ones via `default`, so the + # pandas-only category-dtype fixup this used to need (map() leaves + # category dtype behind) is no longer necessary: replace_strict + # already resolves to a plain numeric dtype on both backends. + # get_column()/Series.replace_strict() (rather than nw.col(), which + # only accepts string names) is what lets this handle pandas + # integer column names too, same as DecisionTreeFeatures. + default = self._unseen if self.unseen == "encode" else None + nw_X = nw.from_native(X, eager_only=True) + new_series = [ + nw_X.get_column(feature).replace_strict(mapping, default=default) + for feature, mapping in self.encoder_dict_.items() + ] + X = nw_X.with_columns(*new_series).to_native() + + if self.unseen != "encode": # check if nan values were introduced by the transformation self._check_nan_values_after_transformation(X) @@ -255,19 +272,19 @@ def _encode(self, X: pd.DataFrame) -> pd.DataFrame: def _check_nan_values_after_transformation(self, X): # check if NaN values were introduced by the encoding - if X[self.variables_].isnull().sum().sum() > 0: + nw_X = nw.from_native(X, eager_only=True) + nan_columns = [ + feature + for feature in self.encoder_dict_.keys() + if nw_X.get_column(feature).null_count() > 0 + ] - # obtain the name(s) of the columns have null values - nan_columns = ( - X[self.encoder_dict_.keys()] - .columns[X[self.encoder_dict_.keys()].isnull().any()] - .tolist() - ) + if len(nan_columns) > 0: if len(nan_columns) > 1: - nan_columns_str = ", ".join(nan_columns) + nan_columns_str = ", ".join(str(col) for col in nan_columns) else: - nan_columns_str = nan_columns[0] + nan_columns_str = str(nan_columns[0]) if self.unseen == "ignore": warnings.warn( @@ -280,27 +297,33 @@ def _check_nan_values_after_transformation(self, X): f"{nan_columns_str}." ) - def inverse_transform(self, X: pd.DataFrame) -> pd.DataFrame: + def inverse_transform(self, X: IntoDataFrame) -> IntoDataFrame: """Convert the encoded variable back to the original values. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features]. + X: dataframe of shape = [n_samples, n_features]. The transformed dataframe. Returns ------- - X_tr: pandas dataframe of shape = [n_samples, n_features]. + X_tr: dataframe of shape = [n_samples, n_features]. The un-transformed dataframe, with the categorical variables containing the original values. """ X = self._check_transform_input_and_state(X) - # replace encoded categories by the original values - for feature in self.encoder_dict_.keys(): - inv_map = {v: k for k, v in self.encoder_dict_[feature].items()} - X[feature] = X[feature].map(inv_map) + # replace encoded categories by the original values. get_column() + # rather than nw.col() again, to support pandas integer column names. + nw_X = nw.from_native(X, eager_only=True) + new_series = [ + nw_X.get_column(feature).replace_strict( + {v: k for k, v in mapping.items()}, default=None + ) + for feature, mapping in self.encoder_dict_.items() + ] + X = nw_X.with_columns(*new_series).to_native() return X