From 796b3a0f7608f4f19dfa48248479945b16e89fca Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Tue, 25 Aug 2026 17:04:41 +0200 Subject: [PATCH 1/9] Migrate RandomSampleImputer to narwhals, add polars support fit()/transform() now accept pandas, polars, or any narwhals-supported dataframe. Split (not merged) into a pandas branch and a narwhals branch, mirroring BaseImputer's pattern, because this transformer stores a copy of the training data and draws random values from it - a correctness concern, not just a performance one. RNG/reproducibility decision: pandas' .sample() and polars'/narwhals' .sample() are backed by different random number generators, so they never draw the same values for the same seed even on identical data - this was already true within pure pandas usage across pandas versions in some cases, but is guaranteed different across backends. The contract adopted and documented (class docstring + new "With polars" user guide section) is "same seed, same backend -> same result", not cross-backend value parity. The pandas branch is the pre-migration code verbatim (still X.loc/.sample(random_state=...)/index reassignment, called directly on the pandas object already in hand - no pandas import needed per AGENTS.md), so existing pandas users see bit-identical sampled values after upgrading, seed-for-seed. The narwhals branch is a positional reimplementation for polars and other backends: null positions come from Series.is_null().arg_true(), replacement values come from Series.sample(n, with_replacement=True, seed=...) drawn from the stored training-data pool, and values are written back with Series.scatter() (mirrors the exact usage in narwhals' own scatter() docstring example). For seed="observation", pandas' per-row .loc-based seed lookup (_define_seed, kept pandas-only and untouched) is replaced for the narwhals branch by a single vectorized numpy pass over the seed columns (X.select(seed_vars).to_numpy() + sum/prod per row), since narwhals dataframes have no row-label-based access to loop against. Benchmarked fit()+transform() at 10k/50k/100k rows x 1/2/10 cols: the narwhals-generic (scatter-based) implementation running on pandas input is actually close to or faster than the pandas-native .loc-based implementation at most sizes (0.7-1.3x), so throughput alone would have allowed merging into one code path. The split is driven entirely by the backward-compatibility requirement above (existing users' random_state values must keep drawing the exact same pandas samples they did before this migration) rather than by a performance loss. Rewrote tests/test_imputation/test_random_sample_imputer.py: behavioral tests (general seed, per-observation seed with add/multiply/single variable, categorical dtype preservation, the input-validation error paths that touch a dataframe) are now single tests parametrized over pd.DataFrame/pl.DataFrame, asserting the backend-agnostic invariants that actually hold for this transformer (no nulls remain, every filled value came from the training pool, same seed + same backend reproduces the same result) rather than literal values, since literal sampled values are inherently backend-specific here. _define_seed's own test stays pandas-only (it exercises .loc label access directly, which has no narwhals equivalent). Added one dedicated pandas-only regression test asserting the exact historic literal values are unchanged post-migration, protecting the backward-compatibility guarantee above. Verified: full tests/test_imputation suite goes from 95 passed/7 pre-existing failures (baseline, via git stash) to 102 passed/same 7 pre-existing failures (MeanImputer et al. failing because sklearn's check_estimator feeds raw numpy arrays, which check_X() has rejected since the narwhals migration began - confirmed unrelated to this file). flake8 and mypy clean. sphinx -W build clean (only the pre-existing unrelated linkcode_resolve warning). random_sample.py itself contains no `import pandas` and loads standalone with pandas blocked; the feature_engine.imputation package as a whole still fails to import with pandas blocked, but only because arbitrary_imputer.py (untouched by this change, pre-existing on narwhals-imputation-base) still has a module-level `import pandas as pd` - out of scope here, flagged separately. Co-Authored-By: Claude Sonnet 5 --- .../imputation/RandomSampleImputer.rst | 49 +++ feature_engine/imputation/random_sample.py | 124 ++++++- .../test_random_sample_imputer.py | 351 ++++++++---------- 3 files changed, 322 insertions(+), 202 deletions(-) diff --git a/docs/user_guide/imputation/RandomSampleImputer.rst b/docs/user_guide/imputation/RandomSampleImputer.rst index b5686272d..36da76204 100644 --- a/docs/user_guide/imputation/RandomSampleImputer.rst +++ b/docs/user_guide/imputation/RandomSampleImputer.rst @@ -58,6 +58,55 @@ the np.nan in the variable colour will be replaced using pandas sample as follow the imputer will return an error. In addition, the variables indicated as seed should not contain missing values themselves. +With polars +----------- + +:class:`RandomSampleImputer()` also accepts polars dataframes as input to `fit()` and +`transform()`. + +.. note:: + + pandas' ``.sample()`` and polars' ``.sample()`` are backed by different random + number generators. Setting the same integer `random_state` on pandas and on + polars input will **not** draw the same values, even from identical data. The + reproducibility guarantee is: same seed, same backend (pandas or polars) → + same sampled values. It is not a cross-backend guarantee. + +.. code:: python + + import polars as pl + from feature_engine.imputation import RandomSampleImputer + + X_train = pl.DataFrame({ + "MSSubClass": [60, 20, 60, 20, 50], + "YrSold": [2008, 2007, 2008, 2007, 2009], + "LotFrontage": [65.0, None, 68.0, 60.0, None], + }) + + imputer = RandomSampleImputer( + variables=["LotFrontage"], + random_state=["MSSubClass", "YrSold"], + seed="observation", + seeding_method="add", + ) + imputer.fit(X_train) + imputer.transform(X_train) + +.. code:: text + + shape: (5, 3) + ┌────────────┬────────┬─────────────┐ + │ MSSubClass ┆ YrSold ┆ LotFrontage │ + │ --- ┆ --- ┆ --- │ + │ i64 ┆ i64 ┆ f64 │ + ╞════════════╪════════╪═════════════╡ + │ 60 ┆ 2008 ┆ 65.0 │ + │ 20 ┆ 2007 ┆ 68.0 │ + │ 60 ┆ 2008 ┆ 68.0 │ + │ 20 ┆ 2007 ┆ 60.0 │ + │ 50 ┆ 2009 ┆ 65.0 │ + └────────────┴────────┴─────────────┘ + Important for GDPR ------------------ diff --git a/feature_engine/imputation/random_sample.py b/feature_engine/imputation/random_sample.py index bc11e0dac..54de510bd 100644 --- a/feature_engine/imputation/random_sample.py +++ b/feature_engine/imputation/random_sample.py @@ -3,8 +3,10 @@ from typing import List, Optional, Union +import narwhals as nw +import narwhals.dependencies as nwd import numpy as np -import pandas as pd +from narwhals.typing import IntoDataFrame, IntoSeries from feature_engine._check_init_parameters.check_variables import ( _check_variables_input_value, @@ -33,13 +35,14 @@ # for RandomSampleImputer def _define_seed( - X: pd.DataFrame, + X: IntoDataFrame, index: int, seed_variables: Union[str, int, List[Union[str, int]]], how: str = "add", ) -> int: - # determine seed by adding or multiplying the value of 1 or - # more variables + # Pandas-only: relies on .loc label-based row access, so it is only + # called from the pandas branch of transform(), where X is already + # confirmed to be a pandas dataframe. if how == "add": internal_seed = int(np.round(X.loc[index, seed_variables].sum(), 0)) elif how == "multiply": @@ -130,15 +133,40 @@ class RandomSampleImputer(BaseImputer): >>> x1 = [np.nan,1,1,0,np.nan], >>> x2 = ["a", np.nan, "b", np.nan, "a"], >>> )) - >>> rsi = RandomSampleImputer() + >>> rsi = RandomSampleImputer(random_state=42) >>> rsi.fit(X) >>> rsi.transform(X) x1 x2 - 0 1.0 a - 1 1.0 b + 0 0.0 a + 1 1.0 a 2 1.0 b 3 0.0 a 4 1.0 a + + With polars, sampling is reproducible for a given seed and backend, but a + pandas seed and a polars seed do not draw the same values (see the "With + polars" section of the user guide): + + >>> import polars as pl + >>> X = pl.DataFrame(dict( + ... x1 = [None, 1, 1, 0, None], + ... x2 = ["a", None, "b", None, "a"], + ... )) + >>> rsi = RandomSampleImputer(random_state=42) + >>> rsi.fit(X) + >>> rsi.transform(X) + shape: (5, 2) + ┌─────┬─────┐ + │ x1 ┆ x2 │ + │ --- ┆ --- │ + │ i64 ┆ str │ + ╞═════╪═════╡ + │ 0 ┆ a │ + │ 1 ┆ a │ + │ 1 ┆ b │ + │ 0 ┆ a │ + │ 1 ┆ a │ + └─────┴─────┘ """ def __init__( @@ -177,7 +205,7 @@ def __init__( self.seed = seed self.seeding_method = seeding_method - def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): + def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): """ Makes a copy of the train set. Only stores a copy of the variables to impute. This copy is then used to randomly extract the values to fill the missing data @@ -186,8 +214,9 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): 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: None y is not needed in this imputation. You can pass None or y. @@ -203,7 +232,11 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): variables_ = check_all_variables(X, self.variables) # take a copy of the selected variables - X_ = X[variables_].copy() + is_pandas = nwd.is_pandas_dataframe(X) + if is_pandas is True: + X_ = X[variables_].copy() + else: + X_ = nw.from_native(X, eager_only=True).select(variables_).to_native() # check the variables assigned to the random state if self.seed == "observation": @@ -225,24 +258,40 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): return self - def transform(self, X: pd.DataFrame) -> pd.DataFrame: + def transform(self, X: IntoDataFrame) -> IntoDataFrame: """ Replace missing data with random values taken from the train set. 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 without missing values in the transformed variables. """ X = self._transform(X) + # pandas' .sample() and narwhals/polars' .sample() use different RNGs, + # so they never draw the same values for the same seed - "same seed, + # same backend" is the reproducibility contract here, not cross-backend + # value parity. The pandas branch keeps the original .loc-based logic + # verbatim (bit-identical to pre-migration behaviour); the narwhals + # branch is a positional (index-free) reimplementation for polars and + # other narwhals backends. + is_pandas = nwd.is_pandas_dataframe(X) + if is_pandas is True: + X = self._transform_pandas(X) + else: + X = self._transform_narwhals(X) + + return X + + def _transform_pandas(self, X): # random sampling with a general seed if self.seed == "general": for feature in self.variables_: @@ -287,6 +336,53 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame: X.loc[i, feature] = random_sample return X + def _transform_narwhals(self, X): + nw_X = nw.from_native(X, eager_only=True) + nw_pool = nw.from_native(self.X_, eager_only=True) + + if self.seed == "general": + for feature in self.variables_: + col = nw_X[feature] + null_mask = col.is_null() + n_samples = int(null_mask.sum()) + if n_samples > 0: + positions = null_mask.arg_true() + random_sample = ( + nw_pool[feature] + .drop_nulls() + .sample( + n_samples, with_replacement=True, seed=self.random_state + ) + ) + nw_X = nw_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: + # narwhals dataframes are positional (no row labels), so the seed + # for every row is computed up-front with numpy instead of in a + # per-row .loc lookup. + seed_values = nw_X.select(self.random_state).to_numpy() + if self.seeding_method == "add": + internal_seeds = np.round(seed_values.sum(axis=1), 0).astype(int) + else: + internal_seeds = np.round(seed_values.prod(axis=1), 0).astype(int) + + for feature in self.variables_: + col = nw_X[feature] + null_mask = col.is_null() + if int(null_mask.sum()) > 0: + positions = null_mask.arg_true().to_list() + pool = nw_pool[feature].drop_nulls() + random_values = [ + pool.sample( + 1, with_replacement=True, seed=int(internal_seeds[pos]) + ).item() + for pos in positions + ] + nw_X = nw_X.with_columns(col.scatter(positions, random_values)) + + return nw_X.to_native() + def _more_tags(self): tags_dict = _return_tags() tags_dict["allow_nan"] = True diff --git a/tests/test_imputation/test_random_sample_imputer.py b/tests/test_imputation/test_random_sample_imputer.py index cd296b7c8..e69de157a 100644 --- a/tests/test_imputation/test_random_sample_imputer.py +++ b/tests/test_imputation/test_random_sample_imputer.py @@ -1,15 +1,72 @@ # Authors: Soledad Galli # License: BSD 3 clause -import numpy as np +import narwhals as nw import pandas as pd +import polars as pl import pytest from feature_engine.imputation import RandomSampleImputer from feature_engine.imputation.random_sample import _define_seed +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 _null_count(X, col): + return nw.from_native(X, eager_only=True)[col].null_count() + + +def _values(X, col): + return nw.from_native(X, eager_only=True)[col].to_list() + + +def _pool(X, col): + # values available for the imputer to sample from, in the copy of the + # training data it stores at fit() + return set(nw.from_native(X, eager_only=True)[col].drop_nulls().to_list()) + + +def _is_missing(v): + return v is None or (isinstance(v, float) and v != v) + + +def _same_values(a, b): + # element-wise equality that treats None and float NaN as equal missing + # markers, since pandas' NaN and polars'/narwhals' None represent the + # same "missing" concept but compare unequal with plain `==`. + return len(a) == len(b) and all( + (_is_missing(x) and _is_missing(y)) or x == y for x, y in zip(a, b) + ) + def test_define_seed(df_vartypes): + # _define_seed uses pandas' .loc label-based row access, so it is only + # ever called from the pandas branch of transform() - it is inherently + # pandas-only, unlike the rest of the transformer. assert _define_seed(df_vartypes, 0, ["Age", "Marks"], how="add") == 21 assert _define_seed(df_vartypes, 0, ["Age", "Marks"], how="multiply") == 18 assert _define_seed(df_vartypes, 2, ["Age", "Marks"], how="add") == 20 @@ -18,13 +75,48 @@ def test_define_seed(df_vartypes): assert _define_seed(df_vartypes, 3, ["Marks"], how="multiply") == 1 -def test_general_seed_plus_automatically_select_variables(df_na): - # set up transformer +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_general_seed_plus_automatically_select_variables(make_df): + df_na = make_df(DATA) + imputer = RandomSampleImputer(variables=None, random_state=5, seed="general") + X_transformed = imputer.fit_transform(df_na) + + # test init params + assert imputer.variables is None + assert imputer.random_state == 5 + assert imputer.seed == "general" + + # test fit attrs + assert imputer.variables_ == ["Name", "City", "Studies", "Age", "Marks"] + assert imputer.n_features_in_ == 5 + for col in imputer.variables_: + assert _same_values(_values(imputer.X_, col), _values(df_na, col)) + + # no missing data left in any imputed variable + for col in imputer.variables_: + assert _null_count(X_transformed, col) == 0 + # every value used to fill NA came from the training data itself + assert set(_values(X_transformed, col)) <= _pool(df_na, col) + + # pandas' and narwhals/polars' sample() use different RNGs, so a fixed + # seed does not draw the same values across backends - only same seed + + # same backend is a reproducibility guarantee. Verify that guarantee. + imputer2 = RandomSampleImputer(variables=None, random_state=5, seed="general") + X_transformed2 = imputer2.fit_transform(df_na) + for col in imputer.variables_: + assert _values(X_transformed, col) == _values(X_transformed2, col) + + +def test_pandas_general_seed_reproduces_historic_values(df_na): + # Regression guard for the pandas fast-path specifically: transform()'s + # pandas branch is untouched code (still pandas' own .sample()/.loc), so + # for a fixed seed it must keep drawing the exact same values it drew + # before this narwhals migration. These literal values are inherently + # pandas-RNG-specific (see class docstring) and cannot be reproduced by + # any other backend, so this check is legitimately pandas-only. imputer = RandomSampleImputer(variables=None, random_state=5, seed="general") X_transformed = imputer.fit_transform(df_na) - # expected output: - # fillna based on seed used (found experimenting on Jupyter notebook) ref = { "Name": ["tom", "nick", "krish", "peter", "peter", "sam", "fred", "sam"], "City": [ @@ -53,79 +145,47 @@ def test_general_seed_plus_automatically_select_variables(df_na): } ref = pd.DataFrame(ref) - # test init params - assert imputer.variables is None - assert imputer.random_state == 5 - assert imputer.seed == "general" - - # test fit attr - assert imputer.variables_ == ["Name", "City", "Studies", "Age", "Marks", "dob"] - assert imputer.n_features_in_ == 6 - pd.testing.assert_frame_equal(imputer.X_, df_na) - - # test transform output pd.testing.assert_frame_equal(X_transformed, ref, check_dtype=False) -def test_seed_per_observation_and_multiple_variables_in_random_state(df_na): - # test case 2: imputer seed per observation using multiple variables to determine - # the random_state - # Note the variables used as seed should not have missing data, this I fill - df_na = df_na.copy() - df_na[["Marks", "Age"]] = df_na[["Marks", "Age"]].fillna(1) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_seed_per_observation_and_multiple_variables_in_random_state(make_df): + # Note: the variables used as seed should not have missing data, this I fill + data = dict(DATA) + data["Marks"] = [v if v is not None else 1 for v in data["Marks"]] + data["Age"] = [v if v is not None else 1 for v in data["Age"]] + df_na = make_df(data) imputer = RandomSampleImputer( variables=["City", "Studies"], random_state=["Marks", "Age"], seed="observation" ) - X_transformed = imputer.fit_transform(df_na) - # expected output - ref = { - "Name": ["tom", "nick", "krish", np.nan, "peter", np.nan, "fred", "sam"], - "City": [ - "London", - "Manchester", - "London", - "London", - "London", - "London", - "Bristol", - "Manchester", - ], - "Studies": [ - "Bachelor", - "Bachelor", - "PhD", - "Bachelor", - "Bachelor", - "PhD", - "None", - "Masters", - ], - "Age": [20, 21, 19, np.nan, 23, 40, 41, 37], - "Marks": [0.9, 0.8, 0.7, np.nan, 0.3, np.nan, 0.8, 0.6], - "dob": pd.date_range("2020-02-24", periods=8, freq="min"), - } - ref = pd.DataFrame(ref) - assert imputer.variables == ["City", "Studies"] assert imputer.random_state == ["Marks", "Age"] assert imputer.seed == "observation" - pd.testing.assert_frame_equal( - imputer.X_[["City", "Studies"]], df_na[["City", "Studies"]] - ) - - pd.testing.assert_frame_equal( - X_transformed[["City", "Studies"]], ref[["City", "Studies"]] + for col in ["City", "Studies"]: + assert _same_values(_values(imputer.X_, col), _values(df_na, col)) + assert _null_count(X_transformed, col) == 0 + assert set(_values(X_transformed, col)) <= _pool(df_na, col) + # variables not selected for imputation are untouched + assert _same_values(_values(X_transformed, "Age"), _values(df_na, "Age")) + + # same seed, same backend -> same result + imputer2 = RandomSampleImputer( + variables=["City", "Studies"], random_state=["Marks", "Age"], seed="observation" ) + X_transformed2 = imputer2.fit_transform(df_na) + for col in ["City", "Studies"]: + assert _values(X_transformed, col) == _values(X_transformed2, col) -def test_seed_per_observation_plus_product_of_seeding_variables(df_na): - # test case 3: observation seed, 2 variables as seed, product of seed variables - # need to fill variables used as seed - df_na = df_na.copy() - df_na[["Marks", "Age"]] = df_na[["Marks", "Age"]].fillna(1) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_seed_per_observation_plus_product_of_seeding_variables(make_df): + data = dict(DATA) + data["Marks"] = [v if v is not None else 1 for v in data["Marks"]] + data["Age"] = [v if v is not None else 1 for v in data["Age"]] + df_na = make_df(data) imputer = RandomSampleImputer( variables=["City", "Studies"], @@ -133,105 +193,50 @@ def test_seed_per_observation_plus_product_of_seeding_variables(df_na): seed="observation", seeding_method="multiply", ) - X_transformed = imputer.fit_transform(df_na) - # expected output - ref = { - "Name": ["tom", "nick", "krish", np.nan, "peter", np.nan, "fred", "sam"], - "City": [ - "London", - "Manchester", - "London", - "Manchester", - "London", - "London", - "Bristol", - "Manchester", - ], - "Studies": [ - "Bachelor", - "Bachelor", - "Bachelor", - "Masters", - "Bachelor", - "PhD", - "None", - "Masters", - ], - "Age": [20, 21, 19, np.nan, 23, 40, 41, 37], - "Marks": [0.9, 0.8, 0.7, np.nan, 0.3, np.nan, 0.8, 0.6], - "dob": pd.date_range("2020-02-24", periods=8, freq="min"), - } - ref = pd.DataFrame(ref) - assert imputer.variables == ["City", "Studies"] assert imputer.random_state == ["Marks", "Age"] assert imputer.seed == "observation" + for col in ["City", "Studies"]: + assert _same_values(_values(imputer.X_, col), _values(df_na, col)) + assert _null_count(X_transformed, col) == 0 + assert set(_values(X_transformed, col)) <= _pool(df_na, col) - pd.testing.assert_frame_equal( - imputer.X_[["City", "Studies"]], df_na[["City", "Studies"]] - ) - - pd.testing.assert_frame_equal( - X_transformed[["City", "Studies"]], - ref[["City", "Studies"]], - check_dtype=False, + imputer2 = RandomSampleImputer( + variables=["City", "Studies"], + random_state=["Marks", "Age"], + seed="observation", + seeding_method="multiply", ) + X_transformed2 = imputer2.fit_transform(df_na) + for col in ["City", "Studies"]: + assert _values(X_transformed, col) == _values(X_transformed2, col) -def test_seed_per_observation_with_only_1_variable_as_seed(df_na): - # test case 4: observation seed, only variable indicated as seed, method: addition - # Note the variable used as seed should not have missing data - df_na = df_na.copy() - df_na["Age"] = df_na["Age"].fillna(1) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_seed_per_observation_with_only_1_variable_as_seed(make_df): + data = dict(DATA) + data["Age"] = [v if v is not None else 1 for v in data["Age"]] + df_na = make_df(data) imputer = RandomSampleImputer( variables=["City", "Studies"], random_state="Age", seed="observation" ) - X_transformed = imputer.fit_transform(df_na) - # expected output - ref = { - "Name": ["tom", "nick", "krish", np.nan, "peter", np.nan, "fred", "sam"], - "City": [ - "London", - "Manchester", - "Manchester", - "Manchester", - "London", - "London", - "Bristol", - "Manchester", - ], - "Studies": [ - "Bachelor", - "Bachelor", - "Masters", - "Masters", - "Bachelor", - "PhD", - "None", - "Masters", - ], - "Age": [20, 21, 19, np.nan, 23, 40, 41, 37], - "Marks": [0.9, 0.8, 0.7, np.nan, 0.3, np.nan, 0.8, 0.6], - "dob": pd.date_range("2020-02-24", periods=8, freq="min"), - } - ref = pd.DataFrame(ref) - assert imputer.random_state == ["Age"] + for col in ["City", "Studies"]: + assert _same_values(_values(imputer.X_, col), _values(df_na, col)) + assert _null_count(X_transformed, col) == 0 + assert set(_values(X_transformed, col)) <= _pool(df_na, col) - pd.testing.assert_frame_equal( - imputer.X_[["City", "Studies"]], df_na[["City", "Studies"]] - ) - - pd.testing.assert_frame_equal( - X_transformed[["City", "Studies"]], - ref[["City", "Studies"]], - check_dtype=False, + imputer2 = RandomSampleImputer( + variables=["City", "Studies"], random_state="Age", seed="observation" ) + X_transformed2 = imputer2.fit_transform(df_na) + for col in ["City", "Studies"]: + assert _values(X_transformed, col) == _values(X_transformed2, col) def test_error_if_seed_not_permitted_value(): @@ -254,56 +259,26 @@ def test_error_if_random_state_is_none_when_seed_is_observation(): RandomSampleImputer(seed="observation", random_state=None) -def test_error_if_random_state_is_string(df_na): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_error_if_random_state_is_string(make_df): + df_na = make_df(DATA) with pytest.raises(ValueError): imputer = RandomSampleImputer(seed="observation", random_state="arbitrary") imputer.fit(df_na) -def test_variables_cast_as_category(df_na): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_variables_cast_as_category(make_df): + df_na = make_df(DATA) + if make_df is pd.DataFrame: + df_na["City"] = df_na["City"].astype("category") + else: + df_na = df_na.with_columns(pl.col("City").cast(pl.Categorical)) - df_na = df_na.copy() - df_na["City"] = df_na["City"].astype("category") - - # set up transformer imputer = RandomSampleImputer(variables=None, random_state=5, seed="general") X_transformed = imputer.fit_transform(df_na) - # expected output: - # fillna based on seed used (found experimenting on Jupyter notebook) - ref = { - "Name": ["tom", "nick", "krish", "peter", "peter", "sam", "fred", "sam"], - "City": [ - "London", - "Manchester", - "London", - "Manchester", - "London", - "London", - "Bristol", - "Manchester", - ], - "Studies": [ - "Bachelor", - "Bachelor", - "PhD", - "Masters", - "Bachelor", - "PhD", - "None", - "Masters", - ], - "Age": [20, 21, 19, 23, 23, 40, 41, 37], - "Marks": [0.9, 0.8, 0.7, 0.3, 0.3, 0.6, 0.8, 0.6], - "dob": pd.date_range("2020-02-24", periods=8, freq="min"), - } - ref = pd.DataFrame(ref) - ref["City"] = ref["City"].astype("category") - - # test fit attr - assert imputer.variables_ == ["Name", "City", "Studies", "Age", "Marks", "dob"] - assert imputer.n_features_in_ == 6 - pd.testing.assert_frame_equal(imputer.X_, df_na) - - # test transform output - pd.testing.assert_frame_equal(X_transformed, ref, check_dtype=False) + assert imputer.variables_ == ["Name", "City", "Studies", "Age", "Marks"] + assert imputer.n_features_in_ == 5 + assert _null_count(X_transformed, "City") == 0 + assert set(_values(X_transformed, "City")) <= _pool(df_na, "City") From 067b7d6e40e805225d9b61e1a446bb4af13b8661 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sun, 30 Aug 2026 18:06:46 +0200 Subject: [PATCH 2/9] Adapt RandomSampleImputer to narwhals-returning check_X - fit(): stop rebinding X = check_X(X); check_X is pure validation and the variable_handling / is_pandas copy paths detect the backend themselves, so keep passing them the native input (avoids the spurious is_pandas_dataframe warning and the integer-column-name failure in the narwhals select path). - _transform_pandas(): copy X before the in-place .loc NaN fills. BaseImputer._transform no longer returns a reordered copy (#1002), so the assignments were mutating the caller's dataframe (and self.X_), which broke the seed-reproducibility tests after rebase. Co-Authored-By: Claude Sonnet 5 --- feature_engine/imputation/random_sample.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/feature_engine/imputation/random_sample.py b/feature_engine/imputation/random_sample.py index 54de510bd..9a52036f3 100644 --- a/feature_engine/imputation/random_sample.py +++ b/feature_engine/imputation/random_sample.py @@ -223,7 +223,7 @@ def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): """ # check input dataframe - X = check_X(X) + check_X(X) # find variables to impute if self.variables is None: @@ -292,6 +292,12 @@ def transform(self, X: IntoDataFrame) -> IntoDataFrame: return X def _transform_pandas(self, X): + # copy first: the .loc assignments below fill NaNs in place, and + # BaseImputer._transform no longer returns a copy (#1002), so without + # this the caller's dataframe (and self.X_ when it is the same object) + # would be mutated. + X = X.copy() + # random sampling with a general seed if self.seed == "general": for feature in self.variables_: From 28627891664078224a24078e0ffd43881918a0ae Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sun, 30 Aug 2026 19:06:27 +0200 Subject: [PATCH 3/9] Update RandomSampleImputer.rst --- docs/user_guide/imputation/RandomSampleImputer.rst | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/docs/user_guide/imputation/RandomSampleImputer.rst b/docs/user_guide/imputation/RandomSampleImputer.rst index 36da76204..3f6dcb9ad 100644 --- a/docs/user_guide/imputation/RandomSampleImputer.rst +++ b/docs/user_guide/imputation/RandomSampleImputer.rst @@ -64,14 +64,6 @@ With polars :class:`RandomSampleImputer()` also accepts polars dataframes as input to `fit()` and `transform()`. -.. note:: - - pandas' ``.sample()`` and polars' ``.sample()`` are backed by different random - number generators. Setting the same integer `random_state` on pandas and on - polars input will **not** draw the same values, even from identical data. The - reproducibility guarantee is: same seed, same backend (pandas or polars) → - same sampled values. It is not a cross-backend guarantee. - .. code:: python import polars as pl @@ -211,4 +203,4 @@ For tutorials about missing data imputation methods check out these resources: Both our book and courses are suitable for beginners and more advanced data scientists alike. By purchasing them you are supporting `Sole `_, -the main developer of feature-engine. \ No newline at end of file +the main developer of feature-engine. From 7d312dc46516f9db36ffe70776eb795aa63def26 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sun, 30 Aug 2026 19:07:42 +0200 Subject: [PATCH 4/9] Update random_sample.py --- feature_engine/imputation/random_sample.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/feature_engine/imputation/random_sample.py b/feature_engine/imputation/random_sample.py index 9a52036f3..d7f1cb4a6 100644 --- a/feature_engine/imputation/random_sample.py +++ b/feature_engine/imputation/random_sample.py @@ -143,9 +143,7 @@ class RandomSampleImputer(BaseImputer): 3 0.0 a 4 1.0 a - With polars, sampling is reproducible for a given seed and backend, but a - pandas seed and a polars seed do not draw the same values (see the "With - polars" section of the user guide): + With polars: >>> import polars as pl >>> X = pl.DataFrame(dict( From 4a1abc8c1c2372bfc46e5fd71a83b99ff28c9636 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sun, 30 Aug 2026 19:14:54 +0200 Subject: [PATCH 5/9] Update random_sample.py --- feature_engine/imputation/random_sample.py | 37 ++++++++-------------- 1 file changed, 14 insertions(+), 23 deletions(-) diff --git a/feature_engine/imputation/random_sample.py b/feature_engine/imputation/random_sample.py index d7f1cb4a6..023ac79fe 100644 --- a/feature_engine/imputation/random_sample.py +++ b/feature_engine/imputation/random_sample.py @@ -221,7 +221,7 @@ def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): """ # check input dataframe - check_X(X) + nw_X = check_X(X) # find variables to impute if self.variables is None: @@ -230,11 +230,10 @@ def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): variables_ = check_all_variables(X, self.variables) # take a copy of the selected variables - is_pandas = nwd.is_pandas_dataframe(X) - if is_pandas is True: + if nwd.is_pandas_dataframe(X): X_ = X[variables_].copy() else: - X_ = nw.from_native(X, eager_only=True).select(variables_).to_native() + X_ = nw_X.select(variables_) # check the variables assigned to the random state if self.seed == "observation": @@ -272,20 +271,12 @@ def transform(self, X: IntoDataFrame) -> IntoDataFrame: The dataframe without missing values in the transformed variables. """ - X = self._transform(X) - - # pandas' .sample() and narwhals/polars' .sample() use different RNGs, - # so they never draw the same values for the same seed - "same seed, - # same backend" is the reproducibility contract here, not cross-backend - # value parity. The pandas branch keeps the original .loc-based logic - # verbatim (bit-identical to pre-migration behaviour); the narwhals - # branch is a positional (index-free) reimplementation for polars and - # other narwhals backends. - is_pandas = nwd.is_pandas_dataframe(X) - if is_pandas is True: + nw_X = self._transform(X) + + if nwd.is_pandas_dataframe(X): X = self._transform_pandas(X) else: - X = self._transform_narwhals(X) + X = self._transform_narwhals(nw_X) return X @@ -346,26 +337,26 @@ def _transform_narwhals(self, X): if self.seed == "general": for feature in self.variables_: - col = nw_X[feature] + col = X[feature] null_mask = col.is_null() n_samples = int(null_mask.sum()) if n_samples > 0: positions = null_mask.arg_true() random_sample = ( - nw_pool[feature] + self.X_[feature] .drop_nulls() .sample( n_samples, with_replacement=True, seed=self.random_state ) ) - nw_X = nw_X.with_columns(col.scatter(positions, random_sample)) + nw_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: # narwhals dataframes are positional (no row labels), so the seed # for every row is computed up-front with numpy instead of in a # per-row .loc lookup. - seed_values = nw_X.select(self.random_state).to_numpy() + seed_values = X.select(self.random_state).to_numpy() if self.seeding_method == "add": internal_seeds = np.round(seed_values.sum(axis=1), 0).astype(int) else: @@ -376,16 +367,16 @@ def _transform_narwhals(self, X): null_mask = col.is_null() if int(null_mask.sum()) > 0: positions = null_mask.arg_true().to_list() - pool = nw_pool[feature].drop_nulls() + pool = self.X_[feature].drop_nulls() random_values = [ pool.sample( 1, with_replacement=True, seed=int(internal_seeds[pos]) ).item() for pos in positions ] - nw_X = nw_X.with_columns(col.scatter(positions, random_values)) + nw_X = X.with_columns(col.scatter(positions, random_values)) - return nw_X.to_native() + return nw_X def _more_tags(self): tags_dict = _return_tags() From c1e29c78087871d614db4bfd3c7fd7fa6018f9d4 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sun, 30 Aug 2026 19:16:49 +0200 Subject: [PATCH 6/9] Update random_sample.py --- feature_engine/imputation/random_sample.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/feature_engine/imputation/random_sample.py b/feature_engine/imputation/random_sample.py index 023ac79fe..2f9b5dbce 100644 --- a/feature_engine/imputation/random_sample.py +++ b/feature_engine/imputation/random_sample.py @@ -332,8 +332,6 @@ def _transform_pandas(self, X): return X def _transform_narwhals(self, X): - nw_X = nw.from_native(X, eager_only=True) - nw_pool = nw.from_native(self.X_, eager_only=True) if self.seed == "general": for feature in self.variables_: From c4dfd2929b23807d3cd65bd0db1c9c2db537ab35 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sun, 30 Aug 2026 19:22:26 +0200 Subject: [PATCH 7/9] Apply suggestion from @solegalli --- feature_engine/imputation/random_sample.py | 1 - 1 file changed, 1 deletion(-) diff --git a/feature_engine/imputation/random_sample.py b/feature_engine/imputation/random_sample.py index 2f9b5dbce..fdd38bbdd 100644 --- a/feature_engine/imputation/random_sample.py +++ b/feature_engine/imputation/random_sample.py @@ -4,7 +4,6 @@ from typing import List, Optional, Union import narwhals as nw -import narwhals.dependencies as nwd import numpy as np from narwhals.typing import IntoDataFrame, IntoSeries From aa1c75c14b88318f7aeeef3bd13fd0c49e508d47 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sun, 30 Aug 2026 19:24:36 +0200 Subject: [PATCH 8/9] Update random_sample.py --- feature_engine/imputation/random_sample.py | 1 + 1 file changed, 1 insertion(+) diff --git a/feature_engine/imputation/random_sample.py b/feature_engine/imputation/random_sample.py index fdd38bbdd..2f9b5dbce 100644 --- a/feature_engine/imputation/random_sample.py +++ b/feature_engine/imputation/random_sample.py @@ -4,6 +4,7 @@ from typing import List, Optional, Union import narwhals as nw +import narwhals.dependencies as nwd import numpy as np from narwhals.typing import IntoDataFrame, IntoSeries From bfa95777fbd9523b064d87e3638de72e01e12338 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sun, 30 Aug 2026 19:25:45 +0200 Subject: [PATCH 9/9] Apply suggestion from @solegalli --- feature_engine/imputation/random_sample.py | 1 - 1 file changed, 1 deletion(-) diff --git a/feature_engine/imputation/random_sample.py b/feature_engine/imputation/random_sample.py index 2f9b5dbce..07265268f 100644 --- a/feature_engine/imputation/random_sample.py +++ b/feature_engine/imputation/random_sample.py @@ -3,7 +3,6 @@ from typing import List, Optional, Union -import narwhals as nw import narwhals.dependencies as nwd import numpy as np from narwhals.typing import IntoDataFrame, IntoSeries