From fb23ba6fdfc6186a7b05ddd38e3de7ab41146935 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Mon, 24 Aug 2026 21:53:38 +0200 Subject: [PATCH 1/2] Migrate BaseImputer to narwhals, add polars support Shared base for the imputation module: _transform() (fit-state checks + column reorder) and transform() (fillna via imputer_dict_) are now dataframe-agnostic, with _get_feature_names_in() reading columns through narwhals on non-pandas input. Benchmarked the fillna step (select + fill from a per-column value dict) at 10k/100k/1M rows x 1/2/10 columns: pandas-native fillna runs ~1.3-1.6x faster than the narwhals-generic fill_null equivalent at the 10k-100k row sizes imputers are normally used at (the gap narrows to ~1.0x only past ~1M rows) - a real, not minimal, loss, so pandas keeps its own fast path (is_pandas = nwd.is_pandas_dataframe(X); if is_pandas is True: ... else narwhals fill_null per column). Also benchmarked a numpy rewrite (to_numpy + np.where per column, mirroring RelativeFeatures) but it did not beat pandas-native and was consistently slower than narwhals fill_null on polars, so it wasn't adopted here - unlike RelativeFeatures' arithmetic, a plain value fill is already close to a no-op for both pandas and narwhals/polars, leaving no room for a numpy win. The pandas<3 fillna-downcasting workaround (option_context + infer_objects) is preserved on the pandas branch but no longer imports pandas at module level - the module is fetched via nw.from_native(X).__native_namespace__() only once X is already confirmed to be a pandas dataframe, so no import is attempted on a polars-only install. Verified: tests/test_imputation full suite unchanged (95 passed, 7 pre-existing failures in test_check_estimator_imputers.py - sklearn's check_estimator feeds raw numpy arrays, which check_X() has always rejected per the narwhals migration's dataframe-only contract, predates this change). flake8 and mypy clean on the file. Module imports with pandas import blocked. sphinx -W build clean (only the pre-existing unrelated linkcode_resolve warning). Co-Authored-By: Claude Sonnet 5 --- feature_engine/imputation/base_imputer.py | 69 ++++++++++++++++------- 1 file changed, 49 insertions(+), 20 deletions(-) diff --git a/feature_engine/imputation/base_imputer.py b/feature_engine/imputation/base_imputer.py index f9c3a2fea..50bd93f86 100644 --- a/feature_engine/imputation/base_imputer.py +++ b/feature_engine/imputation/base_imputer.py @@ -1,4 +1,6 @@ -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 @@ -6,13 +8,11 @@ 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: @@ -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 @@ -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 From 2bc5c86f63b8b8baba3a9ac89ccf00ac315d956d Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Tue, 25 Aug 2026 01:46:27 +0200 Subject: [PATCH 2/2] Migrate MissingIndicator/AddMissingIndicator to narwhals, add polars support Removed the module-level `import pandas as pd`; X/y type hints now use narwhals' IntoDataFrame/IntoSeries. This file overrides transform() rather than extending BaseImputer's, so both the fit() null-count filter and the transform() indicator-column step needed their own narwhals path. Benchmarked both operations at 10k/50k/100k rows x 1/2/10 columns (varying how many columns need indicators), plus a mixed string+numeric-dtype dataset matching MissingIndicator's real "all variable types" usage: - fit()'s `[var for var in variables_ if X[var].isnull().sum() > 0]` loop is ~2-5x faster on pandas than a narwhals-generic `null_count()` call (e.g. 100k rows x 10 cols: 0.41ms loop vs 0.78ms narwhals-on-pandas). A vectorized `X[variables_].isnull().sum()` alternative didn't beat the loop either. narwhals-on-polars was consistently fastest of all (its own native path), so the split is pandas-loop vs narwhals-generic (used for polars/other backends), matching BaseImputer's is_pandas branch pattern. - transform()'s `X[vars].isna().astype("int8").add_suffix("_na")` + `pd.concat` is ~2-5x faster on pandas than narwhals' with_columns equivalent (100k rows x 10 cols: 0.28ms concat vs 1.27ms narwhals-on- pandas), and also beats `assign()`-per-column (0.91ms) and `join()` (0.44ms) alternatives - concat already batches all new columns in one op. So transform() keeps the same pandas fast path, split from a narwhals with_columns path for other backends. Both losses are >1.7x, past the "keep pandas fast path" threshold, so merging into one narwhals-generic path (as BaseImputer's docstring discusses for its own fillna step) was not justified here either. Numpy: converting columns via `.to_numpy()` + `pd.isna()` (the only numpy op that works across MissingIndicator's mixed string/numeric columns, since np.isnan raises on object arrays) was consistently ~1.7-2x slower than pandas-native isnull()/isna() for both fit and transform on mixed dtypes - the extra .to_numpy() copy plus pd.isna() dispatch outweighs any gain, same conclusion as BaseImputer's fillna numpy experiment. Tests: converted tests/test_imputation/test_missing_indicator.py from the pandas-only `df_na` fixture to a plain DATA dict parametrized over `make_df` in [pd.DataFrame, pl.DataFrame], asserting identical variables_ selection and identical `_na` column values on both backends for the same input (one cross-backend PerformanceWarning regression test stays pandas-only, since it targets the pandas fast path specifically). Docs: docs/user_guide/imputation/MissingIndicator.rst has no inline printed output to go stale (it references a screenshot image instead of doctest-style text) - verified its house_prices code example's logic against the migrated transformer with a synthetic stand-in dataset (no network access in this environment) and it behaves identically. Added a verified "With polars" example to the class docstring. Verified: tests/test_imputation/test_missing_indicator.py 29 passed. tests/test_imputation full suite: 107 passed / 7 pre-existing failures in test_check_estimator_imputers.py (confirmed identical failures against a baseline run of origin/narwhals-imputation-base: 95 passed / same 7 failures - sklearn's check_estimator feeds raw numpy arrays, which check_X() has always rejected per the narwhals migration's dataframe-only contract; predates this change). flake8 and mypy clean. Module imports with pandas import blocked. sphinx -W build clean (only the pre-existing unrelated linkcode_resolve warning). Co-Authored-By: Claude Sonnet 5 --- .../imputation/missing_indicator.py | 87 ++++++++++++--- .../test_imputation/test_missing_indicator.py | 105 +++++++++++++----- 2 files changed, 152 insertions(+), 40 deletions(-) diff --git a/feature_engine/imputation/missing_indicator.py b/feature_engine/imputation/missing_indicator.py index 012cf7b23..57b49bce2 100644 --- a/feature_engine/imputation/missing_indicator.py +++ b/feature_engine/imputation/missing_indicator.py @@ -3,7 +3,10 @@ from typing import List, Optional, Union import warnings -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, @@ -105,6 +108,30 @@ class MissingIndicator(BaseImputer): 2 1.0 b 0 0 3 0.0 NaN 0 1 4 NaN a 1 0 + + With polars: + + >>> import polars as pl + >>> from feature_engine.imputation import MissingIndicator + >>> X = pl.DataFrame(dict( + ... x1 = [None, 1, 1, 0, None], + ... x2 = ["a", None, "b", None, "a"], + ... )) + >>> ami = MissingIndicator() + >>> ami.fit(X) + >>> ami.transform(X) + shape: (5, 4) + ┌──────┬──────┬───────┬───────┐ + │ x1 ┆ x2 ┆ x1_na ┆ x2_na │ + │ --- ┆ --- ┆ --- ┆ --- │ + │ i64 ┆ str ┆ i8 ┆ i8 │ + ╞══════╪══════╪═══════╪═══════╡ + │ null ┆ a ┆ 1 ┆ 0 │ + │ 1 ┆ null ┆ 0 ┆ 1 │ + │ 1 ┆ b ┆ 0 ┆ 0 │ + │ 0 ┆ null ┆ 0 ┆ 1 │ + │ null ┆ a ┆ 1 ┆ 0 │ + └──────┴──────┴───────┴───────┘ """ def __init__( @@ -123,16 +150,16 @@ 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 variables for which the missing indicators will be created. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The training dataset. - y: pandas Series, default=None + y: Series, default=None y is not needed in this imputation. You can pass None or y. """ @@ -146,38 +173,68 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): variables_ = check_all_variables(X, self.variables) if self.missing_only is True: - variables_ = [var for var in variables_ if X[var].isnull().sum() > 0] + # Benchmarked: a per-column isnull().sum() loop is ~2-5x faster + # than narwhals' single null_count() call on pandas input (the + # loop calls straight into pandas' C implementation with no + # narwhals overhead), so pandas keeps its own fast path here. + is_pandas = nwd.is_pandas_dataframe(X) + if is_pandas is True: + variables_ = [ + var for var in variables_ if X[var].isnull().sum() > 0 + ] + else: + nw_X = nw.from_native(X, eager_only=True) + null_counts = nw_X.select(variables_).null_count().row(0) + variables_ = [ + var + for var, count in zip(variables_, null_counts) + if count > 0 + ] self.variables_ = variables_ self._get_feature_names_in(X) return self - def transform(self, X: pd.DataFrame) -> pd.DataFrame: + def transform(self, X: IntoDataFrame) -> IntoDataFrame: """ Add the binary missing indicators. Parameters ---------- - X : pandas dataframe of shape = [n_samples, n_features] + X : dataframe of shape = [n_samples, n_features] The dataframe 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 containing the additional binary variables. """ X = self._transform(X) - X_indicators = ( - X[self.variables_] - .isna() - .astype("int8") - .add_suffix("_na") - ) - X = pd.concat([X, X_indicators], axis=1) + + # Benchmarked: building a separate indicator frame and concatenating + # it (pandas-native) is ~2-5x faster than narwhals' with_columns + # equivalent on pandas input, so pandas keeps its own fast path here. + is_pandas = nwd.is_pandas_dataframe(X) + if is_pandas is True: + pd = nw.from_native(X, eager_only=True).__native_namespace__() + X_indicators = ( + X[self.variables_] + .isna() + .astype("int8") + .add_suffix("_na") + ) + X = pd.concat([X, X_indicators], axis=1) + else: + nw_X = nw.from_native(X, eager_only=True) + nw_X = nw_X.with_columns( + nw.col(var).is_null().cast(nw.Int8).alias(f"{var}_na") + for var in self.variables_ + ) + X = nw_X.to_native() return X diff --git a/tests/test_imputation/test_missing_indicator.py b/tests/test_imputation/test_missing_indicator.py index 386d3b61e..b7fdaaca2 100644 --- a/tests/test_imputation/test_missing_indicator.py +++ b/tests/test_imputation/test_missing_indicator.py @@ -1,24 +1,67 @@ +import datetime import warnings +import narwhals as nw import numpy as np import pandas as pd +import polars as pl import pytest from sklearn.pipeline import Pipeline from feature_engine.imputation import MissingIndicator, AddMissingIndicator - +DATA = { + "Name": ["tom", "nick", "krish", None, "peter", None, "fred", "sam"], + "City": [ + "London", + "Manchester", + None, + None, + "London", + "London", + "Bristol", + "Manchester", + ], + "Studies": [ + "Bachelor", + "Bachelor", + None, + None, + "Bachelor", + "PhD", + "None", + "Masters", + ], + "Age": [20, 21, 19, None, 23, 40, 41, 37], + "Marks": [0.9, 0.8, 0.7, None, 0.3, None, 0.8, 0.6], + "dob": [ + datetime.datetime(2020, 2, 24) + datetime.timedelta(minutes=i) + for i in range(8) + ], +} + + +def _cols(X): + return list(nw.from_native(X, eager_only=True).columns) + + +def _col_sum(X, col): + return sum(nw.from_native(X, eager_only=True).get_column(col).to_list()) + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize( "indicator_cls", [MissingIndicator, AddMissingIndicator], ) def test_detect_variables_with_missing_data_when_variables_is_none( - df_na, indicator_cls + make_df, indicator_cls ): + X = make_df(DATA) # test case 1: automatically detect variables with missing data imputer = indicator_cls(missing_only=True, variables=None) - X_transformed = imputer.fit_transform(df_na) + X_transformed = imputer.fit_transform(X) # init params assert imputer.missing_only is True @@ -30,20 +73,22 @@ def test_detect_variables_with_missing_data_when_variables_is_none( # transform outputs assert X_transformed.shape == (8, 11) - assert "Name_na" in X_transformed.columns - assert X_transformed["Name_na"].sum() == 2 + assert "Name_na" in _cols(X_transformed) + assert _col_sum(X_transformed, "Name_na") == 2 +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize( "indicator_cls", [MissingIndicator, AddMissingIndicator], ) def test_add_indicators_to_all_variables_when_variables_is_none( - df_na, indicator_cls + make_df, indicator_cls ): + X = make_df(DATA) imputer = indicator_cls(missing_only=False, variables=None) - X_transformed = imputer.fit_transform(df_na) + X_transformed = imputer.fit_transform(X) assert imputer.variables_ == [ "Name", @@ -54,45 +99,49 @@ def test_add_indicators_to_all_variables_when_variables_is_none( "dob", ] assert X_transformed.shape == (8, 12) - assert "dob_na" in X_transformed.columns - assert X_transformed["dob_na"].sum() == 0 + assert "dob_na" in _cols(X_transformed) + assert _col_sum(X_transformed, "dob_na") == 0 +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize( "indicator_cls", [MissingIndicator, AddMissingIndicator], ) -def test_add_indicators_to_one_variable(df_na, indicator_cls): +def test_add_indicators_to_one_variable(make_df, indicator_cls): + X = make_df(DATA) imputer = indicator_cls(variables="Name") - X_transformed = imputer.fit_transform(df_na) + X_transformed = imputer.fit_transform(X) assert imputer.variables_ == ["Name"] assert X_transformed.shape == (8, 7) - assert "Name_na" in X_transformed.columns - assert X_transformed["Name_na"].sum() == 2 + assert "Name_na" in _cols(X_transformed) + assert _col_sum(X_transformed, "Name_na") == 2 +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize( "indicator_cls", [MissingIndicator, AddMissingIndicator], ) def test_detect_variables_with_missing_data_in_variables_entered_by_user( - df_na, indicator_cls + make_df, indicator_cls ): + X = make_df(DATA) imputer = indicator_cls( missing_only=True, variables=["City", "Studies", "Age", "dob"], ) - X_transformed = imputer.fit_transform(df_na) + X_transformed = imputer.fit_transform(X) assert imputer.variables == ["City", "Studies", "Age", "dob"] assert imputer.variables_ == ["City", "Studies", "Age"] assert X_transformed.shape == (8, 9) - assert "City_na" in X_transformed.columns - assert "dob_na" not in X_transformed.columns - assert X_transformed["City_na"].sum() == 2 + assert "City_na" in _cols(X_transformed) + assert "dob_na" not in _cols(X_transformed) + assert _col_sum(X_transformed, "City_na") == 2 @pytest.mark.parametrize( @@ -104,15 +153,17 @@ def test_error_when_missing_only_not_bool(indicator_cls): indicator_cls(missing_only="missing_only") +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize( "indicator_cls", [MissingIndicator, AddMissingIndicator], ) -def test_get_feature_names_out(df_na, indicator_cls): - original_features = df_na.columns.to_list() +def test_get_feature_names_out(make_df, indicator_cls): + X = make_df(DATA) + original_features = _cols(X) tr = indicator_cls(missing_only=False) - tr.fit(df_na) + tr.fit(X) out = [f + "_na" for f in original_features] feat_out = original_features + out @@ -121,7 +172,7 @@ def test_get_feature_names_out(df_na, indicator_cls): assert tr.get_feature_names_out(input_features=original_features) == feat_out tr = indicator_cls(missing_only=True) - tr.fit(df_na) + tr.fit(X) out = [f + "_na" for f in original_features[0:-1]] feat_out = original_features + out @@ -136,18 +187,20 @@ def test_get_feature_names_out(df_na, indicator_cls): tr.get_feature_names_out(["Name", "hola"]) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize( "indicator_cls", [MissingIndicator, AddMissingIndicator], ) -def test_get_feature_names_out_from_pipeline(df_na, indicator_cls): - original_features = df_na.columns.to_list() +def test_get_feature_names_out_from_pipeline(make_df, indicator_cls): + X = make_df(DATA) + original_features = _cols(X) tr = Pipeline( [("transformer", indicator_cls(missing_only=False))] ) - tr.fit(df_na) + tr.fit(X) out = [f + "_na" for f in original_features] feat_out = original_features + out @@ -161,6 +214,8 @@ def test_get_feature_names_out_from_pipeline(df_na, indicator_cls): [MissingIndicator, AddMissingIndicator], ) def test_no_performance_warning_with_many_variables(indicator_cls): + # pandas-only: exercises the pandas fast path's PerformanceWarning + # behaviour specifically, not a cross-backend value comparison. n_cols = 101 df = pd.DataFrame(