From 084354aa1849f0b8be920ebcc8e3c5eac668033e Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sat, 19 Sep 2026 11:16:48 +0200 Subject: [PATCH 1/3] Migrate MatchVariables to narwhals, add polars support fit() and transform() take any narwhals-supported dataframe and return the same backend. pandas adds, drops and reorders the columns with a single reindex (faster than the previous drop/setitem/select and than narwhals at 10k-500k rows); other backends use with_columns + select. With match_dtypes, pandas keeps its dtypes and astype, since narwhals dtypes don't hold the categories of pandas categoricals. Other backends store narwhals dtypes and cast with narwhals, turning values outside Enum categories into nulls and parsing strings into dates, as pandas does. With polars, np.nan fill values add null Float64 columns and integer fill values add Int64 columns. The verbose messages list the variables in a fixed order: training order for added variables, input order for dropped ones. Init error messages now end with "Got {param} instead.". Rewrite the tests to the make_df conventions and add a polars example to the docstring and the user guide, refreshing its outputs. Co-Authored-By: Claude Opus 5 --- .../preprocessing/MatchVariables.rst | 104 ++- feature_engine/preprocessing/match_columns.py | 226 ++++-- .../test_preprocessing/test_match_columns.py | 677 ++++++++++-------- 3 files changed, 607 insertions(+), 400 deletions(-) diff --git a/docs/user_guide/preprocessing/MatchVariables.rst b/docs/user_guide/preprocessing/MatchVariables.rst index 325aaafa0..8d083c1f5 100644 --- a/docs/user_guide/preprocessing/MatchVariables.rst +++ b/docs/user_guide/preprocessing/MatchVariables.rst @@ -86,11 +86,11 @@ We see that `sex` and `age` are no longer in the dataframe: .. code:: python pclass survived sibsp parch fare cabin embarked - 1000 3 1 0 0 7.7500 n Q - 1001 3 1 2 0 23.2500 n Q - 1002 3 1 2 0 23.2500 n Q - 1003 3 1 2 0 23.2500 n Q - 1004 3 1 0 0 7.7875 n Q + 1000 3 1 0 0 7.7500 NaN Q + 1001 3 1 2 0 23.2500 NaN Q + 1002 3 1 2 0 23.2500 NaN Q + 1003 3 1 2 0 23.2500 NaN Q + 1004 3 1 0 0 7.7875 NaN Q If we transform the dataframe with the dropped columns using :class:`MatchVariables()`, we see that the new dataframe contains all the variables, and those that were missing @@ -107,13 +107,13 @@ Indeed, `sex` and `age` are back, filled with missing values: .. code:: python - The following variables are added to the DataFrame: ['age', 'sex'] + The following variables are added to the DataFrame: ['sex', 'age'] pclass survived sex age sibsp parch fare cabin embarked - 1000 3 1 NaN NaN 0 0 7.7500 n Q - 1001 3 1 NaN NaN 2 0 23.2500 n Q - 1002 3 1 NaN NaN 2 0 23.2500 n Q - 1003 3 1 NaN NaN 2 0 23.2500 n Q - 1004 3 1 NaN NaN 0 0 7.7875 n Q + 1000 3 1 NaN NaN 0 0 7.7500 NaN Q + 1001 3 1 NaN NaN 2 0 23.2500 NaN Q + 1002 3 1 NaN NaN 2 0 23.2500 NaN Q + 1003 3 1 NaN NaN 2 0 23.2500 NaN Q + 1004 3 1 NaN NaN 0 0 7.7875 NaN Q Note how the missing columns were added back to the transformed test set, with missing values, in the position (i.e., order) in which they were in the train set. @@ -133,11 +133,11 @@ We now have 2 extra columns, `var_a` and `var_b`, that were not present in the t .. code:: python pclass survived sibsp parch fare cabin embarked var_a var_b - 1000 3 1 0 0 7.7500 n Q 0 0 - 1001 3 1 2 0 23.2500 n Q 0 0 - 1002 3 1 2 0 23.2500 n Q 0 0 - 1003 3 1 2 0 23.2500 n Q 0 0 - 1004 3 1 0 0 7.7875 n Q 0 0 + 1000 3 1 0 0 7.7500 NaN Q 0 0 + 1001 3 1 2 0 23.2500 NaN Q 0 0 + 1002 3 1 2 0 23.2500 NaN Q 0 0 + 1003 3 1 2 0 23.2500 NaN Q 0 0 + 1004 3 1 0 0 7.7875 NaN Q 0 0 And now, we transform the data with :class:`MatchVariables()`: @@ -152,14 +152,14 @@ the additional columns from the resulting dataset: .. code:: python - The following variables are added to the DataFrame: ['age', 'sex'] - The following variables are dropped from the DataFrame: ['var_b', 'var_a'] + The following variables are added to the DataFrame: ['sex', 'age'] + The following variables are dropped from the DataFrame: ['var_a', 'var_b'] pclass survived sex age sibsp parch fare cabin embarked - 1000 3 1 NaN NaN 0 0 7.7500 n Q - 1001 3 1 NaN NaN 2 0 23.2500 n Q - 1002 3 1 NaN NaN 2 0 23.2500 n Q - 1003 3 1 NaN NaN 2 0 23.2500 n Q - 1004 3 1 NaN NaN 0 0 7.7875 n Q + 1000 3 1 NaN NaN 0 0 7.7500 NaN Q + 1001 3 1 NaN NaN 2 0 23.2500 NaN Q + 1002 3 1 NaN NaN 2 0 23.2500 NaN Q + 1003 3 1 NaN NaN 2 0 23.2500 NaN Q + 1004 3 1 NaN NaN 0 0 7.7875 NaN Q However, if we look closely, the dtypes for the `sex` variable do not match. This could cause issues if other transformations depend upon having the correct dtypes. This is the @@ -173,7 +173,7 @@ Which is: .. code:: python - dtype('O') + And this is the dtype in the transformed test set: @@ -201,8 +201,8 @@ We see in the messages that the `sex` dtype was changed to match that of the tra .. code:: python The following variables are added to the DataFrame: ['sex', 'age'] - The following variables are dropped from the DataFrame: ['var_b', 'var_a'] - The sex dtype is changing from float64 to object + The following variables are dropped from the DataFrame: ['var_a', 'var_b'] + The sex dtype is changing from float64 to str Now the dtype matches: @@ -214,11 +214,61 @@ Which is: .. code:: python - dtype('O') + By default, :class:`MatchVariables()` will print out messages indicating which variables were added, removed and altered. We can switch off the messages through the parameter `verbose`. +Working with polars +^^^^^^^^^^^^^^^^^^^ + +:class:`MatchVariables()` also works with polars dataframes, and returns a polars +dataframe. With polars, the variables added with the default `fill_value` contain +nulls, which is how polars represents missing data: + +.. code:: python + + import polars as pl + from feature_engine.preprocessing import MatchVariables + + train = pl.DataFrame({ + "pclass": [1, 1, 3, 2], + "sex": ["female", "male", "male", "female"], + "age": [29.0, 2.0, 30.0, 25.0], + "fare": [211.34, 151.55, 7.75, 26.0], + }) + + test = pl.DataFrame({ + "fare": [7.75, 23.25, 7.78], + "pclass": [3, 3, 3], + "var_a": [0, 0, 0], + }) + + match_cols = MatchVariables(match_dtypes=True) + match_cols.fit(train) + + test_t = match_cols.transform(test) + print(test_t) + +The transformer added `sex` and `age`, removed `var_a`, sorted the variables as in the +train set and cast `sex` to the string dtype that it had in the train set: + +.. code:: python + + The following variables are added to the DataFrame: ['sex', 'age'] + The following variables are dropped from the DataFrame: ['var_a'] + The sex dtype is changing from Float64 to String + shape: (3, 4) + ┌────────┬──────┬──────┬───────┐ + │ pclass ┆ sex ┆ age ┆ fare │ + │ --- ┆ --- ┆ --- ┆ --- │ + │ i64 ┆ str ┆ f64 ┆ f64 │ + ╞════════╪══════╪══════╪═══════╡ + │ 3 ┆ null ┆ null ┆ 7.75 │ + │ 3 ┆ null ┆ null ┆ 23.25 │ + │ 3 ┆ null ┆ null ┆ 7.78 │ + └────────┴──────┴──────┴───────┘ + When to use the transformer ^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/feature_engine/preprocessing/match_columns.py b/feature_engine/preprocessing/match_columns.py index 4a1598d49..38f9c9471 100644 --- a/feature_engine/preprocessing/match_columns.py +++ b/feature_engine/preprocessing/match_columns.py @@ -1,11 +1,16 @@ -from typing import Dict, List, Union +from typing import Dict, 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 sklearn.base import BaseEstimator, TransformerMixin from sklearn.utils.validation import check_is_fitted from feature_engine._base_transformers.mixins import GetFeatureNamesOutMixin +from feature_engine._check_init_parameters.check_init_input_params import ( + _check_param_missing_values, +) from feature_engine.dataframe_checks import _check_contains_na, check_X from feature_engine.tags import _return_tags @@ -47,11 +52,10 @@ class MatchVariables(TransformerMixin, BaseEstimator, GetFeatureNamesOutMixin): df_transformed - Name City Age Marks - 0 tom np.nan 20 0.9 - 1 sam np.nan 22 0.7 - 2 nick np.nan 23 0.6 - + Name City Age Marks + 0 tom NaN 20 0.9 + 1 sam NaN 22 0.7 + 2 nick NaN 23 0.6 The order of the variables in the transformed dataset is also adjusted to match that observed in the train set. @@ -62,6 +66,7 @@ class MatchVariables(TransformerMixin, BaseEstimator, GetFeatureNamesOutMixin): ---------- fill_value: integer, float or string. Default=np.nan The values for the variables that will be added to the transformed dataset. + With polars dataframes, np.nan adds the variables as nulls. missing_values: string, default='raise' Indicates if missing values should be ignored or raised. If 'raise' the @@ -87,7 +92,9 @@ class MatchVariables(TransformerMixin, BaseEstimator, GetFeatureNamesOutMixin): dtype_dict_: If `match_dtypes` is set to `True`, then this attribute will exist, and it will - contain a dictionary of variables and their corresponding dtypes. + contain a dictionary of variables and their corresponding dtypes. With pandas, + these are pandas dtypes. With other dataframe libraries, like polars, these + are narwhals dtypes, which carry the same names as the polars dtypes. Methods ------- @@ -116,12 +123,12 @@ class MatchVariables(TransformerMixin, BaseEstimator, GetFeatureNamesOutMixin): >>> from feature_engine.preprocessing import MatchVariables >>> X_train = pd.DataFrame(dict(x1 = ["a","b","c"], x2 = [4,5,6])) >>> X_test = pd.DataFrame(dict(x1 = ["c","b","a","d"], - >>> x2 = [5,6,4,7], - >>> x3 = [1,1,1,1])) + ... x2 = [5,6,4,7], + ... x3 = [1,1,1,1])) >>> mv = MatchVariables(missing_values="ignore") >>> mv.fit(X_train) >>> mv.transform(X_train) - x1 x2 + x1 x2 0 a 4 1 b 5 2 c 6 @@ -136,7 +143,7 @@ class MatchVariables(TransformerMixin, BaseEstimator, GetFeatureNamesOutMixin): >>> import pandas as pd >>> from feature_engine.preprocessing import MatchVariables >>> X_train = pd.DataFrame(dict(x1 = ["a","b","c"], - >>> x2 = [4,5,6], x3 = [1,1,1])) + ... x2 = [4,5,6], x3 = [1,1,1])) >>> X_test = pd.DataFrame(dict(x1 = ["c","b","a","d"], x2 = [5,6,4,7])) >>> mv = MatchVariables(missing_values="ignore") >>> mv.fit(X_train) @@ -152,6 +159,32 @@ class MatchVariables(TransformerMixin, BaseEstimator, GetFeatureNamesOutMixin): 1 b 6 NaN 2 a 4 NaN 3 d 7 NaN + + With polars: + + >>> import polars as pl + >>> from feature_engine.preprocessing import MatchVariables + >>> X_train = pl.DataFrame(dict(x1 = ["a","b","c"], + ... x2 = [4,5,6], x3 = [1,1,1])) + >>> X_test = pl.DataFrame(dict(x2 = [5,6,4,7], + ... x1 = ["c","b","a","d"], + ... x4 = [0,0,0,0])) + >>> mv = MatchVariables(missing_values="ignore") + >>> mv.fit(X_train) + >>> mv.transform(X_test) + The following variables are added to the DataFrame: ['x3'] + The following variables are dropped from the DataFrame: ['x4'] + shape: (4, 3) + ┌─────┬─────┬──────┐ + │ x1 ┆ x2 ┆ x3 │ + │ --- ┆ --- ┆ --- │ + │ str ┆ i64 ┆ f64 │ + ╞═════╪═════╪══════╡ + │ c ┆ 5 ┆ null │ + │ b ┆ 6 ┆ null │ + │ a ┆ 4 ┆ null │ + │ d ┆ 7 ┆ null │ + └─────┴─────┴──────┘ """ def __init__( @@ -161,28 +194,24 @@ def __init__( match_dtypes: bool = False, verbose: bool = True, ): - if missing_values not in ["raise", "ignore"]: - raise ValueError( - "missing_values takes only values 'raise' or 'ignore'." - f"Got '{missing_values} instead." - ) + _check_param_missing_values(missing_values) if not isinstance(match_dtypes, bool): raise ValueError( "match_dtypes takes only booleans True and False. " - f"Got '{match_dtypes} instead." + f"Got {match_dtypes} instead." ) if not isinstance(verbose, bool): raise ValueError( - f"verbose takes only booleans True and False. Got '{verbose} instead." + f"verbose takes only booleans True and False. Got {verbose} instead." ) # note: np.nan is an instance of float!!! if not isinstance(fill_value, (str, int, float)): raise ValueError( - "fill_value takes integers, floats or strings." - f"Got '{fill_value} instead." + "fill_value takes integers, floats or strings. " + f"Got {fill_value} instead." ) self.fill_value = fill_value @@ -190,35 +219,36 @@ def __init__( self.match_dtypes = match_dtypes self.verbose = verbose - def fit(self, X: pd.DataFrame, y: pd.Series = None): + def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): """Learns and stores the names of the variables in the training dataset. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The input dataframe. y: None y is not needed for this transformer. You can pass y or None. """ - X = check_X(X) + nw_X = check_X(X) if self.missing_values == "raise": - # check if dataset contains na - _check_contains_na(X, X.columns) + _check_contains_na(X, nw_X.columns) - # save input features - self.feature_names_in_: List[Union[str, int]] = X.columns.tolist() + self.feature_names_in_: List[Union[str, int]] = nw_X.columns + self.n_features_in_ = nw_X.shape[1] - self.n_features_in_ = X.shape[1] - - if self.match_dtypes: - self.dtype_dict_: Dict = X.dtypes.to_dict() + if self.match_dtypes is True: + # narwhals dtypes don't carry the categories of pandas categoricals. + if nwd.is_pandas_dataframe(X) is True: + self.dtype_dict_: Dict = X.dtypes.to_dict() + else: + self.dtype_dict_ = dict(nw_X.schema) return self - def transform(self, X: pd.DataFrame) -> pd.DataFrame: + def transform(self, X: IntoDataFrame) -> IntoDataFrame: """ Drops variables that were not seen in the train set and adds variables that were in the train set but not in the data to transform. In other words, it @@ -226,29 +256,30 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame: Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The data to transform. Returns ------- - X_new: Pandas dataframe, shape = [n_samples, n_features] - The dataframe with variables that match those observed in the train set. + X_new: dataframe of shape = [n_samples, n_features] + The dataframe with variables that match those observed in the train set. """ check_is_fitted(self) - X = check_X(X) + nw_X = check_X(X) + columns = set(nw_X.columns) if self.missing_values == "raise": - # Some variables from the train set may not be present in the test set - # and vice versa. We'll check for nan only in the variables seen during - # training. - vars = [var for var in self.feature_names_in_ if var in X.columns] - _check_contains_na(X, vars) + # Variables from the train set may be missing from X. + _check_contains_na( + X, [var for var in self.feature_names_in_ if var in columns] + ) - _columns_to_drop = list(set(X.columns) - set(self.feature_names_in_)) - _columns_to_add = list(set(self.feature_names_in_) - set(X.columns)) + train_columns = set(self.feature_names_in_) + _columns_to_add = [var for var in self.feature_names_in_ if var not in columns] + _columns_to_drop = [var for var in nw_X.columns if var not in train_columns] - if self.verbose: + if self.verbose is True: if len(_columns_to_add) > 0: print( "The following variables are added to the DataFrame: " @@ -260,36 +291,81 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame: f"{_columns_to_drop}" ) - X = X.drop(_columns_to_drop, axis=1) - - # Add missing columns first and then reorder to avoid - # Pandas 3 StringDtype reindex issue (before we used reindex) - X[_columns_to_add] = self.fill_value - X = X[self.feature_names_in_] - - if self.match_dtypes: - _current_dtypes = X.dtypes.to_dict() - _columns_to_update = { - column: new_dtype - for column, new_dtype in self.dtype_dict_.items() - if new_dtype != _current_dtypes[column] - } - - for column, new_dtype in _columns_to_update.items(): - if self.verbose: - print( - f"The {column} dtype is changing from ", - f"{_current_dtypes[column]} to {new_dtype}", - ) - - # Handle pandas 4 future warning - if isinstance(new_dtype, pd.CategoricalDtype): - cats = new_dtype.categories - X[column] = X[column].where(X[column].isin(cats)) - - X = X.astype(_columns_to_update) - - return X + if nwd.is_pandas_dataframe(X) is True: + # pandas is faster than narwhals. + X = X.reindex(columns=self.feature_names_in_, fill_value=self.fill_value) + if self.match_dtypes is True: + X = self._match_dtypes_pandas(X) + return X + + if len(_columns_to_add) > 0: + fill_value = self._fill_value_expression() + nw_X = nw_X.with_columns(fill_value.alias(var) for var in _columns_to_add) + nw_X = nw_X.select(self.feature_names_in_) + + if self.match_dtypes is True: + nw_X = self._match_dtypes_narwhals(nw_X) + + return nw_X.to_native() + + def _fill_value_expression(self) -> nw.Expr: + if isinstance(self.fill_value, float) and np.isnan(self.fill_value): + # polars treats NaN as a value, not as missing data. + return nw.lit(None, dtype=nw.Float64()) + if isinstance(self.fill_value, int) and not isinstance(self.fill_value, bool): + # polars would store integers as Int32, pandas uses int64. + return nw.lit(self.fill_value, dtype=nw.Int64()) + return nw.lit(self.fill_value) + + def _dtypes_to_update(self, current_dtypes: Dict) -> Dict: + dtypes_to_update = { + column: new_dtype + for column, new_dtype in self.dtype_dict_.items() + if new_dtype != current_dtypes[column] + } + if self.verbose is True: + for column, new_dtype in dtypes_to_update.items(): + print( + f"The {column} dtype is changing from ", + f"{current_dtypes[column]} to {new_dtype}", + ) + return dtypes_to_update + + def _match_dtypes_pandas(self, X): + dtypes_to_update = self._dtypes_to_update(X.dtypes.to_dict()) + + for column, new_dtype in dtypes_to_update.items(): + # Handle pandas 4 future warning + if new_dtype.name == "category": + X[column] = X[column].where(X[column].isin(new_dtype.categories)) + + return X.astype(dtypes_to_update) + + def _match_dtypes_narwhals(self, nw_X: nw.DataFrame) -> nw.DataFrame: + current_dtypes = nw_X.schema + dtypes_to_update = self._dtypes_to_update(current_dtypes) + if len(dtypes_to_update) == 0: + return nw_X + + expressions = [] + for column, new_dtype in dtypes_to_update.items(): + expression = nw.col(column) + if isinstance(new_dtype, nw.Enum): + # polars raises on values outside the categories, pandas sets NaN. + # Strings, because is_in on an Enum rejects values it doesn't have. + expression = expression.cast(nw.String()) + expression = nw.when(expression.is_in(new_dtype.categories)).then( + expression + ) + elif current_dtypes[column] == nw.String: + # polars does not parse strings when casting them to dates. + if new_dtype == nw.Datetime: + expression = expression.str.to_datetime() + elif new_dtype == nw.Date: + expression = expression.str.to_date() + expressions.append(expression.cast(new_dtype)) + + return nw_X.with_columns(expressions) # for the check_estimator tests def _more_tags(self): diff --git a/tests/test_preprocessing/test_match_columns.py b/tests/test_preprocessing/test_match_columns.py index 6726b33f9..840c31ca1 100644 --- a/tests/test_preprocessing/test_match_columns.py +++ b/tests/test_preprocessing/test_match_columns.py @@ -1,370 +1,451 @@ +import datetime +import re + +import narwhals as nw import numpy as np import pandas as pd +import polars as pl import pytest from sklearn.exceptions import NotFittedError from feature_engine.preprocessing import MatchVariables +from tests.backend_helpers import frame_to_dict, null_count + +DOB = [datetime.datetime(2020, 2, 24, 0, minute) for minute in range(4)] + +DATA_TRAIN = { + "Name": ["tom", "nick", "krish", "jack"], + "City": ["London", "Manchester", "Liverpool", "Bristol"], + "Age": [20, 21, 19, 18], + "Marks": [0.9, 0.8, 0.7, 0.6], + "dob": DOB, +} + +# lacks City and Age, has two extra variables and a different column order +DATA_TEST = { + "extra_1": ["a", "b", "c", "d"], + "Marks": [0.5, 0.4, 0.3, 0.2], + "Name": ["sam", "fred", "peter", "bob"], + "dob": DOB, + "extra_2": [1, 2, 3, 4], +} + +DATA_TRAIN_NA = { + "Name": ["tom", None, "krish", "jack"], + "City": ["London", "Manchester", None, "Bristol"], + "Age": [20, 21, 19, 18], + "Marks": [0.9, None, 0.7, 0.6], +} + +MSG_NA = ( + "Some of the variables in the dataset contain NaN. Check and " + "remove those before using this transformer." +) -_params_fill_value = [ - (1, [1, 1, 1, 1], [1, 1, 1, 1]), - (0.1, [0.1, 0.1, 0.1, 0.1], [0.1, 0.1, 0.1, 0.1]), - ("none", ["none", "none", "none", "none"], ["none", "none", "none", "none"]), - (np.nan, [np.nan, np.nan, np.nan, np.nan], [np.nan, np.nan, np.nan, np.nan]), -] -_params_allowed = [ - ([0, 1], "ignore", True, True), - ("nan", "hola", True, True), - ("nan", "ignore", True, "hallo"), - ("nan", "ignore", "hallo", True), -] +# init parameters +@pytest.mark.parametrize("fill_value", [[0, 1], None, {"a": 1}, (1,)]) +def test_error_if_fill_value_not_allowed(fill_value): + msg = f"fill_value takes integers, floats or strings. Got {fill_value} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + MatchVariables(fill_value=fill_value) -@pytest.mark.parametrize( - "fill_value, expected_studies, expected_age", _params_fill_value -) -def test_drop_and_add_columns( - fill_value, expected_studies, expected_age, df_vartypes, df_na -): - train = df_na.copy() - test = df_vartypes.copy() - test = test.drop("Age", axis=1) # to add more than one column - - # adding columns to test if they are removed - for new_col in ["test1", "test2"]: - test.loc[:, new_col] = new_col - - match_columns = MatchVariables( - fill_value=fill_value, - missing_values="ignore", +@pytest.mark.parametrize("missing_values", ["hola", 1, None, ["raise"]]) +def test_error_if_missing_values_not_allowed(missing_values): + msg = ( + "missing_values takes only values 'raise' or 'ignore'. " + f"Got {missing_values} instead." ) - match_columns.fit(train) + with pytest.raises(ValueError, match=re.escape(msg)): + MatchVariables(missing_values=missing_values) - transformed_df = match_columns.transform(test) - expected_result = pd.DataFrame( - { - "Name": ["tom", "nick", "krish", "jack"], - "City": ["London", "Manchester", "Liverpool", "Bristol"], - "Studies": expected_studies, - "Age": expected_age, - "Marks": [0.9, 0.8, 0.7, 0.6], - "dob": pd.date_range("2020-02-24", periods=4, freq="min"), - } +@pytest.mark.parametrize("match_dtypes", ["hallo", 1, None, [True]]) +def test_error_if_match_dtypes_not_bool(match_dtypes): + msg = ( + "match_dtypes takes only booleans True and False. " + f"Got {match_dtypes} instead." ) + with pytest.raises(ValueError, match=re.escape(msg)): + MatchVariables(match_dtypes=match_dtypes) - # test init params - if fill_value is np.nan: - assert match_columns.fill_value is np.nan - else: - assert match_columns.fill_value == fill_value - assert match_columns.verbose is True - assert match_columns.missing_values == "ignore" - assert match_columns.match_dtypes is False - # test fit attrs - assert list(match_columns.feature_names_in_) == list(train.columns) - assert match_columns.n_features_in_ == 6 - # test transform output - pd.testing.assert_frame_equal(expected_result, transformed_df) + +@pytest.mark.parametrize("verbose", ["hallo", 1, None, [True]]) +def test_error_if_verbose_not_bool(verbose): + msg = f"verbose takes only booleans True and False. Got {verbose} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + MatchVariables(verbose=verbose) @pytest.mark.parametrize( - "fill_value, expected_studies, expected_age", _params_fill_value + "fill_value, missing_values, match_dtypes, verbose", + [ + (np.nan, "raise", False, True), + (1, "ignore", True, False), + (0.1, "raise", True, True), + ("none", "ignore", False, False), + ], ) -def test_columns_addition_when_more_columns_in_train_than_test( - fill_value, expected_studies, expected_age, df_vartypes, df_na -): - train = df_na.copy() - test = df_vartypes.copy() - test = test.drop("Age", axis=1) # to add more than one column - - match_columns = MatchVariables( +def test_init_param_assignment(fill_value, missing_values, match_dtypes, verbose): + transformer = MatchVariables( fill_value=fill_value, - missing_values="ignore", + missing_values=missing_values, + match_dtypes=match_dtypes, + verbose=verbose, ) - match_columns.fit(train) + assert transformer.fill_value is fill_value + assert transformer.missing_values == missing_values + assert transformer.match_dtypes is match_dtypes + assert transformer.verbose is verbose - transformed_df = match_columns.transform(test) - expected_result = pd.DataFrame( - { - "Name": ["tom", "nick", "krish", "jack"], - "City": ["London", "Manchester", "Liverpool", "Bristol"], - "Studies": expected_studies, - "Age": expected_age, - "Marks": [0.9, 0.8, 0.7, 0.6], - "dob": pd.date_range("2020-02-24", periods=4, freq="min"), - } - ) +# fit and transform +def test_fit_attributes(make_df): + transformer = MatchVariables().fit(make_df(DATA_TRAIN)) + assert transformer.feature_names_in_ == ["Name", "City", "Age", "Marks", "dob"] + assert transformer.n_features_in_ == 5 + assert not hasattr(transformer, "dtype_dict_") - # test init params - if fill_value is np.nan: - assert match_columns.fill_value is np.nan - else: - assert match_columns.fill_value == fill_value - assert match_columns.verbose is True - assert match_columns.missing_values == "ignore" - assert match_columns.match_dtypes is False - # test fit attrs - assert list(match_columns.feature_names_in_) == list(train.columns) - assert match_columns.n_features_in_ == 6 - # test transform output - pd.testing.assert_frame_equal(expected_result, transformed_df) - - -def test_drop_columns_when_more_columns_in_test_than_train(df_vartypes, df_na): - train = df_vartypes.copy() - train = train.drop("City", axis=1) # to remove more than one column - test = df_na.copy() - - match_columns = MatchVariables(missing_values="ignore") - match_columns.fit(train) - - transformed_df = match_columns.transform(test) - - expected_result = test.drop(columns=["Studies", "City"]) - - # test init params - assert match_columns.fill_value is np.nan - assert match_columns.verbose is True - assert match_columns.missing_values == "ignore" - assert match_columns.match_dtypes is False - # test fit attrs - assert list(match_columns.feature_names_in_) == list(train.columns) - assert match_columns.n_features_in_ == 4 - # test transform output - pd.testing.assert_frame_equal(expected_result, transformed_df) - - -def test_match_dtypes_string_to_numbers(df_vartypes): - train = df_vartypes.copy().select_dtypes("number") - test = train.copy().astype("string") - - match_columns = MatchVariables(match_dtypes=True) - match_columns.fit(train) - - transformed_df = match_columns.transform(test) - - # test init params - assert match_columns.match_dtypes is True - # test fit attrs - assert match_columns.dtype_dict_ == { - "Age": np.dtype("int64"), - "Marks": np.dtype("float64"), + +@pytest.mark.parametrize( + "fill_value, expected", + [(np.nan, None), (1, 1), (0.1, 0.1), ("none", "none")], +) +def test_add_drop_and_reorder_variables(make_df, fill_value, expected): + transformer = MatchVariables(fill_value=fill_value, verbose=False) + transformer.fit(make_df(DATA_TRAIN)) + Xt = transformer.transform(make_df(DATA_TEST)) + + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "Name": ["sam", "fred", "peter", "bob"], + "City": [expected] * 4, + "Age": [expected] * 4, + "Marks": [0.5, 0.4, 0.3, 0.2], + "dob": DOB, } + assert transformer.get_feature_names_out() == [ + "Name", + "City", + "Age", + "Marks", + "dob", + ] - # test transform output - pd.testing.assert_series_equal(train.dtypes, transformed_df.dtypes) - pd.testing.assert_frame_equal(transformed_df, train) +@pytest.mark.parametrize( + "fill_value, expected_dtype", + [(np.nan, nw.Float64), (1, nw.Int64), (0.1, nw.Float64), ("none", nw.String)], +) +def test_dtype_of_added_variables(make_df, fill_value, expected_dtype): + transformer = MatchVariables(fill_value=fill_value, verbose=False) + transformer.fit(make_df(DATA_TRAIN)) + Xt = transformer.transform(make_df(DATA_TEST)) + + schema = nw.from_native(Xt).schema + assert schema["City"] == expected_dtype + assert schema["Age"] == expected_dtype + + +def test_nan_fill_value_adds_missing_data(make_df): + # polars treats NaN as a value, so the added variables must hold nulls. + transformer = MatchVariables(verbose=False).fit(make_df(DATA_TRAIN)) + Xt = transformer.transform(make_df(DATA_TEST)) + assert null_count(Xt, "City") == 4 + assert null_count(Xt, "Age") == 4 + + +def test_only_reorder_variables(make_df): + X = make_df({"Age": [1, 2], "Name": ["a", "b"], "Marks": [0.1, 0.2]}) + train = make_df({"Name": ["c"], "Marks": [0.3], "Age": [3]}) + transformer = MatchVariables().fit(train) + Xt = transformer.transform(X) + + assert isinstance(Xt, make_df) + assert list(Xt.columns) == ["Name", "Marks", "Age"] + assert frame_to_dict(Xt) == { + "Name": ["a", "b"], + "Marks": [0.1, 0.2], + "Age": [1, 2], + } -def test_match_dtypes_numbers_to_string(df_vartypes): - train = df_vartypes.copy().select_dtypes("number").astype("string") - test = df_vartypes.copy().select_dtypes("number") - match_columns = MatchVariables(match_dtypes=True) - match_columns.fit(train) +def test_no_variable_in_common(make_df): + transformer = MatchVariables(verbose=False).fit(make_df({"a": [1], "b": ["x"]})) + Xt = transformer.transform(make_df({"c": [1, 2]})) - transformed_df = match_columns.transform(test) + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == {"a": [None, None], "b": [None, None]} - # test init params - assert match_columns.match_dtypes is True - # test fit attrs - assert isinstance(match_columns.dtype_dict_, dict) - # test transform output - pd.testing.assert_series_equal(train.dtypes, transformed_df.dtypes) - pd.testing.assert_frame_equal(transformed_df, train) +def test_transform_does_not_modify_input(make_df): + X = make_df(DATA_TEST) + transformer = MatchVariables(fill_value=0, verbose=False) + transformer.fit(make_df(DATA_TRAIN)) + Xt = transformer.transform(X) -def test_match_dtypes_string_to_datetime(df_vartypes): - train = df_vartypes.copy().loc[:, ["dob"]] - test = train.copy().astype("string") + assert list(X.columns) == ["extra_1", "Marks", "Name", "dob", "extra_2"] + assert frame_to_dict(X) == DATA_TEST + assert frame_to_dict(Xt)["City"] == [0, 0, 0, 0] - match_columns = MatchVariables(match_dtypes=True, verbose=False) - match_columns.fit(train) - transformed_df = match_columns.transform(test) +def test_verbose_print_out(capsys, make_df): + transformer = MatchVariables(verbose=True).fit(make_df(DATA_TRAIN)) + transformer.transform(make_df(DATA_TEST)) - # test init params - assert match_columns.match_dtypes is True - assert match_columns.verbose is False - # test fit attrs - # TODO: Remove pandas < 3 support when dropping older pandas versions - if pd.__version__ >= "3": - assert match_columns.dtype_dict_ == {"dob": np.dtype(" Date: Sat, 19 Sep 2026 11:36:14 +0200 Subject: [PATCH 2/3] Keep the dtypes learned by MatchVariables private and allow missing data in integer and boolean variables Co-Authored-By: Claude Opus 5 --- feature_engine/preprocessing/match_columns.py | 21 +++++++++------- .../test_preprocessing/test_match_columns.py | 24 ++++++++++++++++--- 2 files changed, 33 insertions(+), 12 deletions(-) diff --git a/feature_engine/preprocessing/match_columns.py b/feature_engine/preprocessing/match_columns.py index 38f9c9471..8b92e0923 100644 --- a/feature_engine/preprocessing/match_columns.py +++ b/feature_engine/preprocessing/match_columns.py @@ -90,12 +90,6 @@ class MatchVariables(TransformerMixin, BaseEstimator, GetFeatureNamesOutMixin): n_features_in_: The number of features in the train set used in fit. - dtype_dict_: - If `match_dtypes` is set to `True`, then this attribute will exist, and it will - contain a dictionary of variables and their corresponding dtypes. With pandas, - these are pandas dtypes. With other dataframe libraries, like polars, these - are narwhals dtypes, which carry the same names as the polars dtypes. - Methods ------- fit: @@ -242,9 +236,9 @@ def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): if self.match_dtypes is True: # narwhals dtypes don't carry the categories of pandas categoricals. if nwd.is_pandas_dataframe(X) is True: - self.dtype_dict_: Dict = X.dtypes.to_dict() + self._dtype_dict: Dict = X.dtypes.to_dict() else: - self.dtype_dict_ = dict(nw_X.schema) + self._dtype_dict = dict(nw_X.schema) return self @@ -320,7 +314,7 @@ def _fill_value_expression(self) -> nw.Expr: def _dtypes_to_update(self, current_dtypes: Dict) -> Dict: dtypes_to_update = { column: new_dtype - for column, new_dtype in self.dtype_dict_.items() + for column, new_dtype in self._dtype_dict.items() if new_dtype != current_dtypes[column] } if self.verbose is True: @@ -338,9 +332,18 @@ def _match_dtypes_pandas(self, X): # Handle pandas 4 future warning if new_dtype.name == "category": X[column] = X[column].where(X[column].isin(new_dtype.categories)) + elif new_dtype.kind in "iub" and X[column].hasnans is True: + # numpy integers can't hold NaN and booleans turn it into True. + dtypes_to_update[column] = self._nullable_dtype(new_dtype) return X.astype(dtypes_to_update) + def _nullable_dtype(self, dtype) -> str: + if dtype.kind == "b": + return "boolean" + prefix = "UInt" if dtype.kind == "u" else "Int" + return f"{prefix}{dtype.itemsize * 8}" + def _match_dtypes_narwhals(self, nw_X: nw.DataFrame) -> nw.DataFrame: current_dtypes = nw_X.schema dtypes_to_update = self._dtypes_to_update(current_dtypes) diff --git a/tests/test_preprocessing/test_match_columns.py b/tests/test_preprocessing/test_match_columns.py index 840c31ca1..dbd1437d2 100644 --- a/tests/test_preprocessing/test_match_columns.py +++ b/tests/test_preprocessing/test_match_columns.py @@ -105,7 +105,7 @@ def test_fit_attributes(make_df): transformer = MatchVariables().fit(make_df(DATA_TRAIN)) assert transformer.feature_names_in_ == ["Name", "City", "Age", "Marks", "dob"] assert transformer.n_features_in_ == 5 - assert not hasattr(transformer, "dtype_dict_") + assert not hasattr(transformer, "_dtype_dict") @pytest.mark.parametrize( @@ -345,6 +345,24 @@ def test_match_dtypes_of_added_variables(make_df): assert nw.from_native(Xt).schema["City"] == nw.String +def test_match_dtypes_of_added_integer_and_boolean_variables(make_df): + # the added variables are missing, so they need types that allow missing data + train = {"int": [1, 2, 3], "bool": [True, False, True], "float": [0.1, 0.2, 0.3]} + transformer = MatchVariables(match_dtypes=True, verbose=False) + transformer.fit(make_df(train)) + Xt = transformer.transform(make_df({"float": [0.4, 0.5]})) + + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "int": [None, None], + "bool": [None, None], + "float": [0.4, 0.5], + } + schema = nw.from_native(Xt).schema + assert schema["int"] == nw.Int64 + assert schema["bool"] == nw.Boolean + + @pytest.mark.parametrize("test_is_categorical", [False, True]) def test_match_dtypes_categories(make_df, test_is_categorical): # values not in the categories seen in fit become missing data. @@ -390,7 +408,7 @@ def test_pandas_dtype_dict_and_categories(): transformer = MatchVariables(match_dtypes=True, verbose=False).fit(train) Xt = transformer.transform(test) - assert transformer.dtype_dict_ == { + assert transformer._dtype_dict == { "Name": pd.CategoricalDtype( categories=["jack", "krish", "nick", "tom"], ordered=False ), @@ -405,7 +423,7 @@ def test_pandas_dtype_dict_and_categories(): def test_polars_dtype_dict(): train = pl.DataFrame({"a": [1], "b": ["x"], "c": [DOB[0]]}) transformer = MatchVariables(match_dtypes=True).fit(train) - assert transformer.dtype_dict_ == { + assert transformer._dtype_dict == { "a": nw.Int64(), "b": nw.String(), "c": nw.Datetime("us"), From 92e61c846b66be513f50b0c455d2d7952ce18424 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sat, 19 Sep 2026 12:50:05 +0200 Subject: [PATCH 3/3] Mark polars output blocks in the user guide as text Co-Authored-By: Claude Opus 5 --- docs/user_guide/preprocessing/MatchVariables.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/user_guide/preprocessing/MatchVariables.rst b/docs/user_guide/preprocessing/MatchVariables.rst index 8d083c1f5..2203aacc9 100644 --- a/docs/user_guide/preprocessing/MatchVariables.rst +++ b/docs/user_guide/preprocessing/MatchVariables.rst @@ -253,7 +253,7 @@ nulls, which is how polars represents missing data: The transformer added `sex` and `age`, removed `var_a`, sorted the variables as in the train set and cast `sex` to the string dtype that it had in the train set: -.. code:: python +.. code:: text The following variables are added to the DataFrame: ['sex', 'age'] The following variables are dropped from the DataFrame: ['var_a']