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 9faba307e05983d262f91bfa03bd88da6b0ac742 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Tue, 25 Aug 2026 17:05:46 +0200 Subject: [PATCH 2/2] Migrate EndTailImputer to narwhals, add polars support fit() now computes the Gaussian/IQR/max end-of-distribution values via a single narwhals aggregation (nw_X.select(...) of per-variable mean/std/ quantile/max expressions), instead of pandas-only .mean()/.std()/.quantile(). transform() already worked cross-backend via the already-migrated BaseImputer. Merge vs split: benchmarked pandas-native vs narwhals-generic (on both pandas and polars) at 10k/50k/100k rows x 1/2/10 columns, with NaNs present (this is an imputer, so skip-NaN semantics matter - mean/std/quantile must skip missing values like pandas' default skipna=True). Results: - gaussian: narwhals-on-pandas is 0.93-1.5x pandas-native's time (parity to a mild loss, narrowing towards 1.0x as rows scale up), and 3-10x *faster* than pandas-native when run on polars. - iqr: narwhals-on-pandas is consistently *faster* than pandas-native (~1.3-2x), on both backends. Nowhere near the "real loss" (1.7x+) split threshold, so one code path (no is_pandas branching) serves both backends - unlike BaseImputer's fillna, which stayed split because it *was* consistently 1.3-1.6x slower via narwhals on pandas. Also benchmarked a numpy rewrite (nanmean/nanstd/nanpercentile per column, mirroring RelativeFeatures' numpy-acceleration pattern) and rejected it: numpy's nan-aware reductions are slow (isnan-mask overhead), and at 10 columns narwhals-on-polars beat numpy-on-polars by ~10x (0.65ms vs 7.2ms at 100k rows x 10 cols) since polars aggregates columns natively/in parallel instead of looping in Python. RelativeFeatures' numpy win doesn't transfer here because that transformer's arithmetic has no NaN-skipping requirement, so plain (non-nan-aware) numpy ops sufficed there. Tests rewritten to one parametrized test per behavior over `@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame])`, replacing the pandas-only test_end_tail_imputer.py. Test data uses `None` for missing values instead of `np.nan`: polars treats a literal np.nan as a real float (not a null), so it would NOT be skipped by mean/std/quantile the way pandas skips NaN by default - `None` becomes a null on both backends and is skipped consistently. Docs: verified the existing house_prices example still runs and produces matching output; added a "With polars" section to both the class docstring and docs/user_guide/imputation/EndTailImputer.rst. No bugs found in the pre-migration code. The 7 pre-existing test_check_estimator_from_sklearn failures in this test module (numpy array input now rejected by check_X, e.g. for MeanImputer) predate this change and are unrelated to EndTailImputer. Co-Authored-By: Claude Sonnet 5 --- docs/user_guide/imputation/EndTailImputer.rst | 47 ++++++ feature_engine/imputation/end_tail.py | 77 +++++---- .../test_imputation/test_end_tail_imputer.py | 149 ++++++++++++------ 3 files changed, 196 insertions(+), 77 deletions(-) diff --git a/docs/user_guide/imputation/EndTailImputer.rst b/docs/user_guide/imputation/EndTailImputer.rst index e908cf518..6725fbfff 100644 --- a/docs/user_guide/imputation/EndTailImputer.rst +++ b/docs/user_guide/imputation/EndTailImputer.rst @@ -119,6 +119,53 @@ imputation (in red the imputed variable): The second peak corresponds to the missing data, which were replaced with a value at that side of the distribution. +With polars +----------- + +:class:`EndTailImputer()` also works with polars dataframes: + +.. code:: python + + import polars as pl + from feature_engine.imputation import EndTailImputer + + X = pl.DataFrame({ + "LotFrontage": [65.0, 80.0, None, 60.0, 84.0, None, 75.0], + "MasVnrArea": [196.0, None, 162.0, 0.0, 350.0, None, 0.0], + }) + + # set up the imputer + tail_imputer = EndTailImputer( + imputation_method='gaussian', + tail='right', + fold=3, + variables=['LotFrontage', 'MasVnrArea'], + ) + + # fit the imputer + tail_imputer.fit(X) + + # transform the data + X_t = tail_imputer.transform(X) + X_t + +.. code:: text + + shape: (7, 2) + ┌─────────────┬────────────┐ + │ LotFrontage ┆ MasVnrArea │ + │ --- ┆ --- │ + │ f64 ┆ f64 │ + ╞═════════════╪════════════╡ + │ 65.0 ┆ 196.0 │ + │ 80.0 ┆ 583.800407 │ + │ 103.053925 ┆ 162.0 │ + │ 60.0 ┆ 0.0 │ + │ 84.0 ┆ 350.0 │ + │ 103.053925 ┆ 583.800407 │ + │ 75.0 ┆ 0.0 │ + └─────────────┴────────────┘ + Additional resources -------------------- diff --git a/feature_engine/imputation/end_tail.py b/feature_engine/imputation/end_tail.py index e52500056..11f92478a 100644 --- a/feature_engine/imputation/end_tail.py +++ b/feature_engine/imputation/end_tail.py @@ -3,7 +3,8 @@ from typing import List, Optional, Union -import pandas as pd +import narwhals as nw +from narwhals.typing import IntoDataFrame, IntoSeries from feature_engine._check_init_parameters.check_variables import ( _check_variables_input_value, @@ -140,6 +141,27 @@ class EndTailImputer(BaseImputer): 2 0.500000 3 0.000000 4 1.199359 + + With polars: + + >>> import polars as pl + >>> from feature_engine.imputation import EndTailImputer + >>> X = pl.DataFrame({"x1": [None, 0.5, 0.5, 0.0, None]}) + >>> eti = EndTailImputer(imputation_method='gaussian', tail='right', fold=3) + >>> eti.fit(X) + >>> eti.transform(X) + shape: (5, 1) + ┌──────────┐ + │ x1 │ + │ --- │ + │ f64 │ + ╞══════════╡ + │ 1.199359 │ + │ 0.5 │ + │ 0.5 │ + │ 0.0 │ + │ 1.199359 │ + └──────────┘ """ def __init__( @@ -170,13 +192,13 @@ 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 values at the end of the variable distribution. 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 @@ -191,33 +213,34 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): else: variables_ = check_numerical_variables(X, self.variables) - # estimate imputation values - if self.imputation_method == "max": - imputer_dict_ = (X[variables_].max() * self.fold).to_dict() - - elif self.imputation_method == "gaussian": - if self.tail == "right": - imputer_dict_ = ( - X[variables_].mean() + self.fold * X[variables_].std() - ).to_dict() - elif self.tail == "left": - imputer_dict_ = ( - X[variables_].mean() - self.fold * X[variables_].std() - ).to_dict() - - elif self.imputation_method == "iqr": - IQR = X[variables_].quantile(0.75) - X[variables_].quantile(0.25) - if self.tail == "right": - imputer_dict_ = ( - X[variables_].quantile(0.75) + (IQR * self.fold) - ).to_dict() - elif self.tail == "left": - imputer_dict_ = ( - X[variables_].quantile(0.25) - (IQR * self.fold) - ).to_dict() + # Narwhals aggregation matches/beats pandas-native on pandas and is + # 3-10x faster on polars (benchmarked), so one path serves both backends. + nw_X = nw.from_native(X, eager_only=True) + exprs = [self._end_value_expr(v) for v in variables_] + agg = nw_X.select(*exprs) + imputer_dict_ = {k: v[0] for k, v in agg.to_dict(as_series=False).items()} self.variables_ = variables_ self.imputer_dict_ = imputer_dict_ self._get_feature_names_in(X) return self + + def _end_value_expr(self, variable: Union[str, int]) -> nw.Expr: + """Build the narwhals expression that computes the end-of-distribution + replacement value for one variable, per `imputation_method` and `tail`.""" + col = nw.col(variable) + + if self.imputation_method == "max": + return (col.max() * self.fold).alias(variable) + + if self.imputation_method == "gaussian": + if self.tail == "right": + return (col.mean() + self.fold * col.std()).alias(variable) + return (col.mean() - self.fold * col.std()).alias(variable) + + # imputation_method == "iqr" + iqr = col.quantile(0.75, "linear") - col.quantile(0.25, "linear") + if self.tail == "right": + return (col.quantile(0.75, "linear") + self.fold * iqr).alias(variable) + return (col.quantile(0.25, "linear") - self.fold * iqr).alias(variable) diff --git a/tests/test_imputation/test_end_tail_imputer.py b/tests/test_imputation/test_end_tail_imputer.py index 88998d658..36a1db459 100644 --- a/tests/test_imputation/test_end_tail_imputer.py +++ b/tests/test_imputation/test_end_tail_imputer.py @@ -1,21 +1,69 @@ +import narwhals as nw import numpy as np import pandas as pd +import polars as pl import pytest from feature_engine.imputation import EndTailImputer - -def test_automatically_find_variables_and_gaussian_imputation_on_right_tail(df_na): - # set up transformer +# Missing values are written as `None`, not `np.nan`: polars treats np.nan as +# a real float value (not a null), so mean/std/quantile would NOT skip it, +# unlike pandas' NaN-as-missing default. `None` becomes a null on both +# backends and is skipped by both, keeping the two code paths comparable. +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 _none_to_nan(values): + # Missing values print as None for polars, NaN for pandas float columns + # - both mean "missing" here, so normalize both sides before comparing. + return [np.nan if v is None else v for v in values] + + +def assert_df_equal(X, expected: dict, abs_tol: float = 1e-5) -> None: + result = nw.from_native(X, eager_only=True).to_dict(as_series=False) + assert list(result.keys()) == list(expected.keys()) + for col, values in expected.items(): + assert _none_to_nan(result[col]) == pytest.approx( + _none_to_nan(values), abs=abs_tol, nan_ok=True + ) + + +def _missing_count(X, columns) -> int: + nw_X = nw.from_native(X, eager_only=True) + return sum(int(nw_X.get_column(c).is_null().sum()) for c in columns) + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_automatically_find_variables_and_gaussian_imputation_on_right_tail(make_df): + df = make_df(DATA) imputer = EndTailImputer( imputation_method="gaussian", tail="right", fold=3, variables=None ) - X_transformed = imputer.fit_transform(df_na) - - # set up expected output - X_reference = df_na.copy() - X_reference["Age"] = X_reference["Age"].fillna(58.94908118478389) - X_reference["Marks"] = X_reference["Marks"].fillna(1.3244261503263175) + X_transformed = imputer.fit_transform(df) # test init params assert imputer.imputation_method == "gaussian" @@ -24,64 +72,65 @@ def test_automatically_find_variables_and_gaussian_imputation_on_right_tail(df_n assert imputer.variables is None # test fit attr assert imputer.variables_ == ["Age", "Marks"] - assert imputer.n_features_in_ == 6 - imputer.imputer_dict_ = { - key: round(value, 3) for (key, value) in imputer.imputer_dict_.items() - } - assert imputer.imputer_dict_ == { - "Age": 58.949, - "Marks": 1.324, - } + assert imputer.n_features_in_ == 5 + rounded = {k: round(v, 3) for k, v in imputer.imputer_dict_.items()} + assert rounded == {"Age": 58.949, "Marks": 1.324} + # transform output: indicated vars ==> no NA, not indicated vars with NA - assert X_transformed[["Age", "Marks"]].isnull().sum().sum() == 0 - assert X_transformed[["City", "Name"]].isnull().sum().sum() > 0 - pd.testing.assert_frame_equal(X_transformed, X_reference) + assert _missing_count(X_transformed, ["Age", "Marks"]) == 0 + assert _missing_count(X_transformed, ["City", "Name"]) > 0 + + expected = dict(DATA) + expected["Age"] = [20, 21, 19, 58.94908118478389, 23, 40, 41, 37] + expected["Marks"] = [ + 0.9, 0.8, 0.7, 1.3244261503263175, 0.3, 1.3244261503263175, 0.8, 0.6, + ] + assert_df_equal(X_transformed, expected) -def test_user_enters_variables_and_iqr_imputation_on_right_tail(df_na): - # set up transformer +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_user_enters_variables_and_iqr_imputation_on_right_tail(make_df): + df = make_df(DATA) imputer = EndTailImputer( imputation_method="iqr", tail="right", fold=1.5, variables=["Age", "Marks"] ) - X_transformed = imputer.fit_transform(df_na) + X_transformed = imputer.fit_transform(df) - # set up expected result - X_reference = df_na.copy() - X_reference["Age"] = X_reference["Age"].fillna(65.5) - X_reference["Marks"] = X_reference["Marks"].fillna(1.0625) - - # test fit and transform attr and output assert imputer.imputer_dict_ == {"Age": 65.5, "Marks": 1.0625} - assert X_transformed[["Age", "Marks"]].isnull().sum().sum() == 0 - pd.testing.assert_frame_equal(X_transformed, X_reference) + assert _missing_count(X_transformed, ["Age", "Marks"]) == 0 + + expected = dict(DATA) + expected["Age"] = [20, 21, 19, 65.5, 23, 40, 41, 37] + expected["Marks"] = [0.9, 0.8, 0.7, 1.0625, 0.3, 1.0625, 0.8, 0.6] + assert_df_equal(X_transformed, expected) -def test_user_enters_variables_and_max_value_imputation(df_na): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_user_enters_variables_and_max_value_imputation(make_df): + df = make_df(DATA) imputer = EndTailImputer( imputation_method="max", tail="right", fold=2, variables=["Age", "Marks"] ) - imputer.fit(df_na) + imputer.fit(df) assert imputer.imputer_dict_ == {"Age": 82.0, "Marks": 1.8} -def test_automatically_select_variables_and_gaussian_imputation_on_left_tail(df_na): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_automatically_select_variables_and_gaussian_imputation_on_left_tail(make_df): + df = make_df(DATA) imputer = EndTailImputer(imputation_method="gaussian", tail="left", fold=3) - imputer.fit(df_na) - imputer.imputer_dict_ = { - key: round(value, 3) for (key, value) in imputer.imputer_dict_.items() - } - assert imputer.imputer_dict_ == { - "Age": -1.521, - "Marks": 0.042, - } - - -def test_user_enters_variables_and_iqr_imputation_on_left_tail(df_na): - # test case 5: IQR + left tail + imputer.fit(df) + rounded = {k: round(v, 3) for k, v in imputer.imputer_dict_.items()} + assert rounded == {"Age": -1.521, "Marks": 0.042} + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_user_enters_variables_and_iqr_imputation_on_left_tail(make_df): + df = make_df(DATA) imputer = EndTailImputer( imputation_method="iqr", tail="left", fold=1.5, variables=["Age", "Marks"] ) - imputer.fit(df_na) + imputer.fit(df) assert imputer.imputer_dict_["Age"] == -6.5 assert np.round(imputer.imputer_dict_["Marks"], 3) == np.round( 0.36249999999999993, 3 @@ -89,15 +138,15 @@ def test_user_enters_variables_and_iqr_imputation_on_left_tail(df_na): def test_error_when_imputation_method_is_not_permitted(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="imputation_method takes only values"): EndTailImputer(imputation_method="arbitrary") def test_error_when_tail_is_string(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="tail takes only values"): EndTailImputer(tail="arbitrary") def test_error_when_fold_is_1(): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="fold takes only positive numbers"): EndTailImputer(fold=-1)