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 4498c69ad727ede9679ddd060acd21f394bf1b55 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Tue, 25 Aug 2026 01:41:59 +0200 Subject: [PATCH 2/2] Migrate ArbitraryImputer to narwhals, add polars support fit() never touches dataframe values - it only calls the already- narwhals-migrated check_X/check_numerical_variables/find_numerical_variables and builds imputer_dict_ via a plain dict comprehension over column names - so the only change needed was dropping the module-level `import pandas as pd` and swapping the X/y type hints for narwhals' IntoDataFrame/IntoSeries. transform() is fully inherited from the already-migrated BaseImputer. Benchmarked fit()+transform() (via fit_transform) at 10k/50k/100k rows x 1/2/10 cols, pandas vs polars, and old code vs migrated code on pandas input: fit() takes ~0.06-0.13ms regardless of row count, column count, or backend, both before and after the edit (within noise of each other) - confirming fit() truly does no per-row work. No backend split was needed or added; a single narwhals-agnostic path was kept (it already was one). Numpy: not applicable - fit() has no numeric computation over data at all, only dict/list building over variable names, so there is nothing for numpy to accelerate. While touching fit(), changed `if self.imputer_dict:` to `if self.imputer_dict is not None:` per AGENTS.md's ban on truthy container checks; this also fixes a latent edge case where imputer_dict={} was silently treated as "not provided" and fell through to the variables/arbitrary_number branch. Confirmed pre-existing on origin/narwhals-imputation-base (unrelated to this migration, no test previously covered it). Rewrote tests/test_imputation/test_arbitrary_imputer.py to the cross-backend parametrized style (@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame])) in place, replacing the pandas-only df_na fixture and pd.testing.assert_frame_equal/.isnull() assertions with a plain DATA dict and narwhals-based null/value assertions. The deprecation-warning test for ArbitraryNumberImputer and the arbitrary_number-type-validation test stayed single-backend since they never touch a dataframe. Added a "With polars" section to both the class docstring and docs/user_guide/imputation/ArbitraryImputer.rst, output verified by actually running the transformer. No staleness found in the existing rst (it builds its example from fetch_openml, no literal printed dataframe values to go stale). Verified: tests/test_imputation full suite 98 passed / 7 pre-existing unrelated failures in test_check_estimator_imputers.py (same 7 as on origin/narwhals-imputation-base's baseline of 95 passed - the 3 extra passes here are the new cross-backend parametrization, no regressions). 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/ArbitraryImputer.rst | 39 +++++++ .../imputation/arbitrary_imputer.py | 33 ++++-- .../test_imputation/test_arbitrary_imputer.py | 100 +++++++++++------- 3 files changed, 125 insertions(+), 47 deletions(-) diff --git a/docs/user_guide/imputation/ArbitraryImputer.rst b/docs/user_guide/imputation/ArbitraryImputer.rst index 19a10f9f1..4a0084af8 100644 --- a/docs/user_guide/imputation/ArbitraryImputer.rst +++ b/docs/user_guide/imputation/ArbitraryImputer.rst @@ -134,6 +134,45 @@ imputation (in red the imputed variable): .. image:: ../../images/arbitraryvalueimputation.png +With polars +----------- + +:class:`ArbitraryImputer()` works in the same way with a polars dataframe: + +.. code:: python + + import polars as pl + from feature_engine.imputation import ArbitraryImputer + + df = pl.DataFrame({ + "LotFrontage": [65.0, None, 68.0, None, 84.0], + "MasVnrArea": [196.0, None, 162.0, None, 350.0], + }) + + transformer = ArbitraryImputer( + arbitrary_number=-999, + variables=["LotFrontage", "MasVnrArea"], + ) + + print(transformer.fit_transform(df)) + +The resulting values match those found with pandas: + +.. code:: text + + shape: (5, 2) + ┌─────────────┬────────────┐ + │ LotFrontage ┆ MasVnrArea │ + │ --- ┆ --- │ + │ f64 ┆ f64 │ + ╞═════════════╪════════════╡ + │ 65.0 ┆ 196.0 │ + │ -999.0 ┆ -999.0 │ + │ 68.0 ┆ 162.0 │ + │ -999.0 ┆ -999.0 │ + │ 84.0 ┆ 350.0 │ + └─────────────┴────────────┘ + Additional resources -------------------- diff --git a/feature_engine/imputation/arbitrary_imputer.py b/feature_engine/imputation/arbitrary_imputer.py index 333a81918..c0350c1c3 100644 --- a/feature_engine/imputation/arbitrary_imputer.py +++ b/feature_engine/imputation/arbitrary_imputer.py @@ -1,11 +1,10 @@ # Authors: Soledad Galli # License: BSD 3 clause +import warnings from typing import List, Optional, Union -import pandas as pd - -import warnings +from narwhals.typing import IntoDataFrame, IntoSeries from feature_engine._check_init_parameters.check_input_dictionary import ( _check_numerical_dict, @@ -120,6 +119,28 @@ class ArbitraryImputer(BaseImputer): 2 1.0 b 3 0.0 NaN 4 -999.0 a + + With polars: + + >>> import polars as pl + >>> from feature_engine.imputation import ArbitraryImputer + >>> X = pl.DataFrame({"x1": [None, 1, 1, 0, None], + >>> "x2": ["a", None, "b", None, "a"]}) + >>> ai = ArbitraryImputer(arbitrary_number=-999) + >>> ai.fit(X) + >>> ai.transform(X) + shape: (5, 2) + ┌──────┬──────┐ + │ x1 ┆ x2 │ + │ --- ┆ --- │ + │ i64 ┆ str │ + ╞══════╪══════╡ + │ -999 ┆ a │ + │ 1 ┆ null │ + │ 1 ┆ b │ + │ 0 ┆ null │ + │ -999 ┆ a │ + └──────┴──────┘ """ def __init__( @@ -144,13 +165,13 @@ def __init__( self.imputer_dict = imputer_dict - def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): + def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): """ This method does not learn any parameter. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The training dataset. y: None @@ -162,7 +183,7 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): # find or check for numerical variables # create the imputer dictionary - if self.imputer_dict: + if self.imputer_dict is not None: variables_ = check_numerical_variables( X, list(self.imputer_dict.keys()) ) diff --git a/tests/test_imputation/test_arbitrary_imputer.py b/tests/test_imputation/test_arbitrary_imputer.py index a8f2ce1c4..9766204ec 100644 --- a/tests/test_imputation/test_arbitrary_imputer.py +++ b/tests/test_imputation/test_arbitrary_imputer.py @@ -1,18 +1,36 @@ -import pytest +import narwhals as nw import pandas as pd +import polars as pl +import pytest from feature_engine.imputation import ArbitraryImputer, ArbitraryNumberImputer - -def test_impute_with_99_and_automatically_select_variables(df_na): - # set up the transformer +DATA = { + "Name": ["tom", "nick", "krish", None, "peter", None, "fred", "sam"], + "City": [ + "London", + "Manchester", + None, + None, + "London", + "London", + "Bristol", + "Manchester", + ], + "Age": [20.0, 21.0, 19.0, None, 23.0, 40.0, 41.0, 37.0], + "Marks": [0.9, 0.8, 0.7, None, 0.3, None, 0.8, 0.6], +} + + +def _null_count(X, col) -> int: + return nw.from_native(X, eager_only=True)[col].is_null().sum() + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_impute_with_99_and_automatically_select_variables(make_df): + X = make_df(DATA) imputer = ArbitraryImputer(arbitrary_number=99, variables=None) - X_transformed = imputer.fit_transform(df_na) - - # set up output reference - X_reference = df_na.copy() - X_reference["Age"] = X_reference["Age"].fillna(99) - X_reference["Marks"] = X_reference["Marks"].fillna(99) + X_transformed = imputer.fit_transform(X) # test init params assert imputer.arbitrary_number == 99 @@ -20,25 +38,25 @@ def test_impute_with_99_and_automatically_select_variables(df_na): # test fit attributes assert imputer.variables_ == ["Age", "Marks"] - assert imputer.n_features_in_ == 6 + assert imputer.n_features_in_ == 4 assert imputer.imputer_dict_ == {"Age": 99, "Marks": 99} - # test transform output - # selected variables should not contain NA - # non selected variables should still contain 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) + # selected variables should not contain NA, non-selected should still + 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 = nw.from_native(X_transformed, eager_only=True).to_dict(as_series=False) + assert result["Age"] == [20.0, 21.0, 19.0, 99.0, 23.0, 40.0, 41.0, 37.0] + assert result["Marks"] == [0.9, 0.8, 0.7, 99.0, 0.3, 99.0, 0.8, 0.6] -def test_impute_with_1_and_single_variable_entered_by_user(df_na): - # set up transformer - imputer = ArbitraryImputer(arbitrary_number=-1, variables=["Age"]) - X_transformed = imputer.fit_transform(df_na) - # set up output reference - X_reference = df_na.copy() - X_reference["Age"] = X_reference["Age"].fillna(-1) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_impute_with_1_and_single_variable_entered_by_user(make_df): + X = make_df(DATA) + imputer = ArbitraryImputer(arbitrary_number=-1, variables=["Age"]) + X_transformed = imputer.fit_transform(X) # test init params assert imputer.arbitrary_number == -1 @@ -46,12 +64,12 @@ def test_impute_with_1_and_single_variable_entered_by_user(df_na): # test fit attributes assert imputer.variables_ == ["Age"] - assert imputer.n_features_in_ == 6 + assert imputer.n_features_in_ == 4 assert imputer.imputer_dict_ == {"Age": -1} - # 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 = nw.from_native(X_transformed, eager_only=True).to_dict(as_series=False) + assert result["Age"] == [20.0, 21.0, 19.0, -1.0, 23.0, 40.0, 41.0, 37.0] def test_error_when_arbitrary_number_is_string(): @@ -59,24 +77,24 @@ def test_error_when_arbitrary_number_is_string(): ArbitraryImputer(arbitrary_number="arbitrary") -def test_dictionary_of_imputation_values(df_na): - # set up transformer +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_dictionary_of_imputation_values(make_df): + X = make_df(DATA) imputer = ArbitraryImputer(imputer_dict={"Age": -42, "Marks": -999}) - X_transformed = imputer.fit_transform(df_na) - - # set up expected output - X_reference = df_na.copy() - X_reference["Age"] = X_reference["Age"].fillna(-42) - X_reference["Marks"] = X_reference["Marks"].fillna(-999) + X_transformed = imputer.fit_transform(X) # test fit params - assert imputer.n_features_in_ == 6 + assert imputer.n_features_in_ == 4 assert imputer.imputer_dict_ == {"Age": -42, "Marks": -999} - # test transform params - 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) + 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 = nw.from_native(X_transformed, eager_only=True).to_dict(as_series=False) + assert result["Age"] == [20.0, 21.0, 19.0, -42.0, 23.0, 40.0, 41.0, 37.0] + assert result["Marks"] == [0.9, 0.8, 0.7, -999.0, 0.3, -999.0, 0.8, 0.6] def test_imputer_error_when_dictionary_value_is_string():