diff --git a/docs/user_guide/preprocessing/MatchVariables.rst b/docs/user_guide/preprocessing/MatchVariables.rst index 325aaafa0..2203aacc9 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:: text + + 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..8b92e0923 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 @@ -85,10 +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. - Methods ------- fit: @@ -116,12 +117,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 +137,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 +153,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 +188,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 +213,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 +250,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 +285,90 @@ 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)) + 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) + 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..dbd1437d2 100644 --- a/tests/test_preprocessing/test_match_columns.py +++ b/tests/test_preprocessing/test_match_columns.py @@ -1,370 +1,469 @@ +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("