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 7b8f6d528ff975a1b372b0901eb985381540fe7c Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Tue, 25 Aug 2026 17:03:09 +0200 Subject: [PATCH 2/2] Migrate MeanImputer/MeanMedianImputer to narwhals, add polars support Fit's mean()/median() computation is split by backend and, on the pandas branch, additionally rewritten to use NumPy directly. Benchmarked (10k-100k rows x 1-10 cols): narwhals-on-pandas vs pandas-native .mean()/.median() showed the same real, not minimal, loss (1.0-3.0x) already documented for BaseImputer's fillna and CategoricalImputer's mode(), so pandas keeps its own fast path. Going further, benchmarked a bulk NumPy nanmean/nanmedian pass (to_numpy() + axis=0 reduction, mirroring MathFeatures' reducer pattern) against pandas-native .mean()/.median() and found NumPy consistently as fast or faster (ratios 0.5-1.05x) - a real win, so the pandas branch now uses NumPy instead of pandas' own methods. For polars, the equivalent NumPy round-trip was benchmarked too and lost to narwhals' native per-column mean()/median() expressions (1.8-3.5x slower for mean; mixed but trending slower for median at scale), so the polars/narwhals branch computes stats with a single narwhals select() of one expression per variable instead - benchmarked against a per-column loop and against select()+to_native().to_dicts() and found select()+rows(named=True) is equal-or-faster and backend-agnostic (no reliance on a polars-only to_dicts() method). All-NaN/all-null columns produce matching values on both backends (verified directly): NumPy's nanmean/nanmedian warn on all-NaN slices where pandas' methods don't, so those warnings are suppressed the same way MathFeatures does. Nullable extension dtypes that would produce object arrays fall back to pandas' native .mean()/.median(), same guard as MathFeatures' dtype.kind check. Found and fixed a real crash: narwhals' select() with zero expressions collapses row count to 0 too, so stats.rows(named=True)[0] would IndexError when return_empty=True yields no numerical variables on polars input. Added an explicit empty-variables guard that skips the backend branch entirely instead of relying on backend-specific zero-column behaviour. Rewrote tests as one parametrized test per behaviour over pd.DataFrame/pl.DataFrame (a self-contained DATA dict replacing the pandas-only df_na fixture, matching the CategoricalImputer migration's pattern), keeping the MeanImputer/MeanMedianImputer deprecation-warning parametrization on top. Verified: tests/test_imputation full suite - 99 passed (up from 95 pre-migration, same tests plus new polars parametrizations), same 7 pre-existing failures in test_check_estimator_imputers.py (sklearn's check_estimator feeds raw numpy arrays, rejected by check_X's dataframe-only contract from the base migration - confirmed identical root cause against the pre-migration baseline via git stash). flake8 and mypy clean. mean_median.py's actual import chain (base_imputer, dataframe_checks, variable_handling) verified pandas-free with pandas blocked, using direct module loading to bypass the sibling not-yet-migrated imputers in imputation/__init__.py. sphinx -W build clean (only the pre-existing unrelated linkcode_resolve warning). Every doc example (docstring pandas/polars examples and the new "With polars" section in MeanImputer.rst) re-run against live output. Co-Authored-By: Claude Sonnet 5 --- docs/user_guide/imputation/MeanImputer.rst | 49 ++++++++++ feature_engine/imputation/mean_median.py | 83 ++++++++++++++-- .../test_mean_median_imputer.py | 96 +++++++++++++------ 3 files changed, 190 insertions(+), 38 deletions(-) diff --git a/docs/user_guide/imputation/MeanImputer.rst b/docs/user_guide/imputation/MeanImputer.rst index b0d0df877..66d8a0873 100644 --- a/docs/user_guide/imputation/MeanImputer.rst +++ b/docs/user_guide/imputation/MeanImputer.rst @@ -310,6 +310,55 @@ center of the distribution: Because of the increase in the number of observations at the center, the variance of the variable decreases, and the kurtosis coefficient increases. +With polars +----------- + +:class:`MeanImputer()` works in the same way with a polars dataframe: + +.. code:: python + + import polars as pl + from feature_engine.imputation import MeanImputer + + df = pl.DataFrame({ + "Age": [20, 21, 19, None, 23, 40, 41, 37], + "Marks": [0.9, 0.8, 0.7, None, 0.3, None, 0.8, 0.6], + }) + + transformer = MeanImputer(imputation_method="mean") + transformer.fit(df) + + print(transformer.imputer_dict_) + +The learned mean values match those found with pandas: + +.. code:: text + + {'Age': 28.714285714285715, 'Marks': 0.6833333333333332} + +.. code:: python + + print(transformer.transform(df)) + +.. code:: text + + shape: (8, 2) + ┌───────────┬──────────┐ + │ Age ┆ Marks │ + │ --- ┆ --- │ + │ f64 ┆ f64 │ + ╞═══════════╪══════════╡ + │ 20.0 ┆ 0.9 │ + │ 21.0 ┆ 0.8 │ + │ 19.0 ┆ 0.7 │ + │ 28.714286 ┆ 0.683333 │ + │ 23.0 ┆ 0.3 │ + │ 40.0 ┆ 0.683333 │ + │ 41.0 ┆ 0.8 │ + │ 37.0 ┆ 0.6 │ + └───────────┴──────────┘ + + Additional resources -------------------- diff --git a/feature_engine/imputation/mean_median.py b/feature_engine/imputation/mean_median.py index 049b768e3..fe8479d93 100644 --- a/feature_engine/imputation/mean_median.py +++ b/feature_engine/imputation/mean_median.py @@ -4,7 +4,10 @@ import warnings from typing import List, Optional, Union -import pandas as pd +import narwhals as nw +import narwhals.dependencies as nwd +import numpy as np +from narwhals.typing import IntoDataFrame, IntoSeries from feature_engine._check_init_parameters.check_variables import ( _check_variables_input_value, @@ -102,6 +105,30 @@ class MeanImputer(BaseImputer): 2 1.0 b 3 0.0 NaN 4 1.0 a + + With polars: + + >>> import polars as pl + >>> from feature_engine.imputation import MeanImputer + >>> X = pl.DataFrame(dict( + >>> x1 = [None, 1, 1, 0, None], + >>> x2 = ["a", None, "b", None, "a"], + >>> )) + >>> mmi = MeanImputer(imputation_method='median') + >>> mmi.fit(X) + >>> mmi.transform(X) + shape: (5, 2) + ┌─────┬──────┐ + │ x1 ┆ x2 │ + │ --- ┆ --- │ + │ f64 ┆ str │ + ╞═════╪══════╡ + │ 1.0 ┆ a │ + │ 1.0 ┆ null │ + │ 1.0 ┆ b │ + │ 0.0 ┆ null │ + │ 1.0 ┆ a │ + └─────┴──────┘ """ def __init__( @@ -120,16 +147,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 mean or median values. 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 or None, default=None + y: Series or None, default=None y is not needed in this imputation. You can pass None or y. """ @@ -143,11 +171,46 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): variables_ = check_numerical_variables(X, self.variables) # find imputation parameters: mean or median - if self.imputation_method == "mean": - imputer_dict_ = X[variables_].mean().to_dict() - - elif self.imputation_method == "median": - imputer_dict_ = X[variables_].median().to_dict() + if len(variables_) == 0: + # narwhals' select() with no expressions collapses rows too, so + # skip the backend branches entirely rather than special-case that. + imputer_dict_ = {} + else: + # Benchmarked (10k-100k rows x 1-10 cols): pandas' bulk .mean()/ + # .median() is consistently slower than a single NumPy + # nanmean/nanmedian pass over the same values (0.5-1.05x, mostly + # a real win), so the pandas branch takes that route. Polars' + # native aggregation already beats a NumPy round-trip (1.8-3.5x + # for mean, competitive-to-faster for median), so it keeps using + # narwhals expressions directly instead. + is_pandas = nwd.is_pandas_dataframe(X) + if is_pandas is True: + values = X[variables_].to_numpy() + reducer = ( + np.nanmean if self.imputation_method == "mean" else np.nanmedian + ) + # Nullable extension dtypes can produce object arrays; keep + # those on the pandas-native fallback path below. + if values.dtype.kind in "biuf": + # pandas' mean()/median() do not warn for all-missing + # columns; NumPy's equivalents do, so silence only those. + with warnings.catch_warnings(): + warnings.simplefilter("ignore", RuntimeWarning) + result = reducer(values, axis=0) + imputer_dict_ = dict(zip(variables_, result)) + elif self.imputation_method == "mean": + imputer_dict_ = X[variables_].mean().to_dict() + else: + imputer_dict_ = X[variables_].median().to_dict() + else: + nw_X = nw.from_native(X, eager_only=True) + stats = nw_X.select( + *[ + getattr(nw.col(var), self.imputation_method)() + for var in variables_ + ] + ) + imputer_dict_ = stats.rows(named=True)[0] self.variables_ = variables_ self.imputer_dict_ = imputer_dict_ diff --git a/tests/test_imputation/test_mean_median_imputer.py b/tests/test_imputation/test_mean_median_imputer.py index c3603ecf7..6f4782a96 100644 --- a/tests/test_imputation/test_mean_median_imputer.py +++ b/tests/test_imputation/test_mean_median_imputer.py @@ -1,6 +1,8 @@ import re +import narwhals as nw import pandas as pd +import polars as pl import pytest from feature_engine.imputation import MeanImputer, MeanMedianImputer @@ -11,6 +13,43 @@ "use MeanImputer instead." ) +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], +} + + +def _cols(X, columns): + # to_dict(as_series=False) is a convenient, backend-agnostic way to read + # values back out for comparison, regardless of pandas vs polars. + result = nw.from_native(X, eager_only=True).to_dict(as_series=False) + return {c: result[c] for c in columns} + + +def _null_count(X, col): + return nw.from_native(X, eager_only=True)[col].null_count() + @pytest.fixture( params=[MeanImputer, MeanMedianImputer], @@ -32,59 +71,60 @@ def test_mean_median_imputer_raises_future_warning(): MeanMedianImputer() -def test_mean_imputation_and_automatically_select_variables(df_na, imputer_class): - # set up transformer +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_mean_imputation_and_automatically_select_variables(make_df, imputer_class): + df_na = make_df(DATA) imputer = make_imputer(imputer_class, imputation_method="mean", variables=None) X_transformed = imputer.fit_transform(df_na) - # set up reference result - X_reference = df_na.copy() - X_reference["Age"] = X_reference["Age"].fillna(28.714285714285715) - X_reference["Marks"] = X_reference["Marks"].fillna(0.6833333333333332) - # test init params assert imputer.imputation_method == "mean" assert imputer.variables is None # test fit attributes assert imputer.variables_ == ["Age", "Marks"] - imputer.imputer_dict_ = { + rounded_dict = { key: round(value, 3) for (key, value) in imputer.imputer_dict_.items() } - assert imputer.imputer_dict_ == { - "Age": 28.714, - "Marks": 0.683, - } - assert imputer.n_features_in_ == 6 + assert rounded_dict == {"Age": 28.714, "Marks": 0.683} + assert imputer.n_features_in_ == 5 # test transform output: # selected variables should have no NA # not selected variables should still have NA - assert X_transformed[["Age", "Marks"]].isnull().sum().sum() == 0 - assert X_transformed[["Name", "City"]].isnull().sum().sum() > 0 - pd.testing.assert_frame_equal(X_transformed, X_reference) - - -def test_median_imputation_when_user_enters_single_variables(df_na, imputer_class): - # set up trasnformer - imputer = make_imputer(imputer_class, imputation_method="median", variables=["Age"]) + assert _null_count(X_transformed, "Age") == 0 + assert _null_count(X_transformed, "Marks") == 0 + assert _null_count(X_transformed, "Name") > 0 + assert _null_count(X_transformed, "City") > 0 + result = _cols(X_transformed, ["Age", "Marks"]) + assert result["Age"] == pytest.approx( + [20, 21, 19, 28.714285714285715, 23, 40, 41, 37] + ) + assert result["Marks"] == pytest.approx( + [0.9, 0.8, 0.7, 0.6833333333333332, 0.3, 0.6833333333333332, 0.8, 0.6] + ) + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_median_imputation_when_user_enters_single_variables(make_df, imputer_class): + df_na = make_df(DATA) + imputer = make_imputer( + imputer_class, imputation_method="median", variables=["Age"] + ) X_transformed = imputer.fit_transform(df_na) - # set up reference output - X_reference = df_na.copy() - X_reference["Age"] = X_reference["Age"].fillna(23.0) - # test init params assert imputer.imputation_method == "median" assert imputer.variables == ["Age"] # test fit attributes - assert imputer.n_features_in_ == 6 + assert imputer.n_features_in_ == 5 assert imputer.imputer_dict_ == {"Age": 23.0} # test transform output - assert X_transformed["Age"].isnull().sum() == 0 - pd.testing.assert_frame_equal(X_transformed, X_reference) + assert _null_count(X_transformed, "Age") == 0 + result = _cols(X_transformed, ["Age"]) + assert result["Age"] == [20, 21, 19, 23.0, 23, 40, 41, 37] def test_error_with_wrong_imputation_method(imputer_class):