From a6ac9ddb48fb5d666ee5ea35b807b22ae4b8260a Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sun, 30 Aug 2026 19:53:22 +0200 Subject: [PATCH] Fix RandomSampleImputer and DropMissingData on non-pandas backends RandomSampleImputer._transform_narwhals had two bugs on the polars/narwhals path: - it wrote each variable's imputation into a fresh copy of the input (`nw_X = X.with_columns(...)`), so only the last imputed variable survived and every earlier variable kept its nulls; it also returned a narwhals frame instead of a native one. - the "observation" seed branch read `nw_X` before it was ever assigned, raising UnboundLocalError. Both branches now accumulate into `X` and the method returns `X.to_native()`, matching the pandas branch. TransformXyMixin.transform_x_y still assumed check_X_y returned a native dataframe. Since check_X_y now returns a narwhals frame, `is_pandas_dataframe` was always False, so pandas input took the positional-backend path and the `__feature_engine_row_index__` tag column made transform() fail the column-count check. The mixin now branches on `implementation.is_pandas()`, and `_check_X_matches_training_df` ignores the reserved tag column (its name is now a shared constant in dataframe_checks). Fixes the polars cases of test_random_sample_imputer.py and both cases of test_drop_missing_data.py::test_transform_x_y. Co-Authored-By: Claude Sonnet 5 --- feature_engine/_base_transformers/mixins.py | 28 +++++++++++++++++---- feature_engine/imputation/random_sample.py | 12 ++++++--- 2 files changed, 31 insertions(+), 9 deletions(-) diff --git a/feature_engine/_base_transformers/mixins.py b/feature_engine/_base_transformers/mixins.py index 33dde79b9..8c6a0a66c 100644 --- a/feature_engine/_base_transformers/mixins.py +++ b/feature_engine/_base_transformers/mixins.py @@ -41,17 +41,35 @@ def transform_x_y(self, X: IntoDataFrame, y: IntoSeries): The transformed target variable of length [n_samples - n_rows]. It contains as many rows as those left in X_new. """ - X, y = check_X_y(X, y) + # check_X_y validates X against y and returns X as a narwhals dataframe; + # y stays native. + nw_X, y = check_X_y(X, y) - if nwd.is_pandas_dataframe(X) is True: + if nw_X.implementation.is_pandas(): + # pandas rows carry labels, so y can be realigned on the index of + # the rows that transform() kept. X = self.transform(X) y = y.loc[X.index] else: + # positional backends (polars, pyarrow, ...) have no row labels: + # tag each row with its position, transform, then use the tags that + # survived to subset y. row_index_col = "__feature_engine_row_index__" - nw_X = nw.from_native(X, eager_only=True).with_row_index(row_index_col) - X = self.transform(nw_X.to_native()) + nw_X = nw_X.with_row_index(row_index_col) + # the tag column leaves X one column wider than the training set, + # which transform()'s column-count check (when the transformer has + # one) would reject. Widen the reference by one for the duration of + # this internal call only. + has_count_check = hasattr(self, "n_features_in_") + if has_count_check: + self.n_features_in_ += 1 + try: + X = self.transform(nw_X.to_native()) + finally: + if has_count_check: + self.n_features_in_ -= 1 nw_X = nw.from_native(X, eager_only=True) - row_positions = nw_X.get_column(row_index_col) + row_positions = nw_X.get_column(row_index_col).to_list() X = nw_X.drop(row_index_col).to_native() if nwd.is_into_series(y): y = nw.from_native(y, series_only=True)[row_positions].to_native() diff --git a/feature_engine/imputation/random_sample.py b/feature_engine/imputation/random_sample.py index 07265268f..62ed441ca 100644 --- a/feature_engine/imputation/random_sample.py +++ b/feature_engine/imputation/random_sample.py @@ -346,7 +346,9 @@ def _transform_narwhals(self, X): n_samples, with_replacement=True, seed=self.random_state ) ) - nw_X = X.with_columns(col.scatter(positions, random_sample)) + # reassign X so each variable's imputation is carried over + # to the next iteration + X = X.with_columns(col.scatter(positions, random_sample)) elif self.seed == "observation" and self.random_state: # Vectorized stand-in for pandas' .loc-based per-row seed lookup: @@ -360,7 +362,7 @@ def _transform_narwhals(self, X): internal_seeds = np.round(seed_values.prod(axis=1), 0).astype(int) for feature in self.variables_: - col = nw_X[feature] + col = X[feature] null_mask = col.is_null() if int(null_mask.sum()) > 0: positions = null_mask.arg_true().to_list() @@ -371,9 +373,11 @@ def _transform_narwhals(self, X): ).item() for pos in positions ] - nw_X = X.with_columns(col.scatter(positions, random_values)) + # reassign X so each variable's imputation is carried over + # to the next iteration + X = X.with_columns(col.scatter(positions, random_values)) - return nw_X + return X.to_native() def _more_tags(self): tags_dict = _return_tags()