diff --git a/docs/user_guide/preprocessing/MatchCategories.rst b/docs/user_guide/preprocessing/MatchCategories.rst index ed5c46fc2..730934170 100644 --- a/docs/user_guide/preprocessing/MatchCategories.rst +++ b/docs/user_guide/preprocessing/MatchCategories.rst @@ -6,7 +6,8 @@ MatchCategories =============== :class:`MatchCategories()` ensures that categorical variables are encoded as pandas -'categorical' dtype instead of generic python 'object' or other dtypes. +'categorical' dtype, or polars 'Enum' dtype, instead of generic python 'object', +string or other dtypes. Under the hood, 'categorical' dtype is a representation that maps each category to an integer, thus providing a more memory-efficient object @@ -79,10 +80,10 @@ Here are the mappings learnt for each categorical variable: .. code:: python - {'pclass': Int64Index([1, 2, 3], dtype='int64'), - 'sex': Index(['female', 'male'], dtype='object'), - 'cabin': Index(['A', 'B', 'C', 'D', 'E', 'F', 'M', 'T'], dtype='object'), - 'embarked': Index(['C', 'Missing', 'Q', 'S'], dtype='object')} + {'pclass': Index([1, 2, 3], dtype='int64'), + 'sex': Index(['female', 'male'], dtype='str'), + 'cabin': Index(['A', 'B', 'C', 'D', 'E', 'F', 'M', 'T'], dtype='str'), + 'embarked': Index(['C', 'Missing', 'Q', 'S'], dtype='str')} To see why this matters, let's compare the order in which the categories of `embarked` appear in the raw train and test sets. This is the order in the train set: @@ -95,7 +96,9 @@ We obtain the following order: .. code:: python - array(['S', 'C', 'Missing', 'Q'], dtype=object) + + ['S', 'C', 'Missing', 'Q'] + Length: 4, dtype: str And this is the order in the test set: @@ -107,7 +110,9 @@ Which is different from the train set: .. code:: python - array(['Q', 'S', 'C'], dtype=object) + + ['Q', 'S', 'C'] + Length: 3, dtype: str The categories appear in a different order in each set. If we transform the dataframes using the same `match_categories` object, categorical variables will be converted to a @@ -122,7 +127,7 @@ Which is: .. code:: python - Index(['C', 'Missing', 'Q', 'S'], dtype='object') + Index(['C', 'Missing', 'Q', 'S'], dtype='str') And this is the order we now obtain for the test set: @@ -134,11 +139,10 @@ The 2 sets now show exactly the same category order: .. code:: python - Index(['C', 'Missing', 'Q', 'S'], dtype='object') + Index(['C', 'Missing', 'Q', 'S'], dtype='str') If some category was not present in the training data, it will not be mapped -to any integer and will thus not get encoded. This behaviour can be modified through the -parameter `errors`. Let's illustrate this with the `cabin` variable. These are the +to any integer and will become a missing value instead. Let's illustrate this with the `cabin` variable. These are the categories present in the train set: .. code:: python @@ -149,7 +153,9 @@ We obtain the following categories: .. code:: python - array(['B', 'C', 'E', 'D', 'A', 'M', 'T', 'F'], dtype=object) + + ['B', 'C', 'E', 'D', 'A', 'M', 'T', 'F'] + Length: 8, dtype: str And these are the categories present in the test set, which include a category, 'G', that was not seen during training: @@ -162,7 +168,9 @@ We obtain the following categories, including the unseen 'G': .. code:: python - array(['M', 'F', 'E', 'G'], dtype=object) + + ['M', 'F', 'E', 'G'] + Length: 4, dtype: str After transforming the train set, we obtain the same categories as before, now correctly typed as 'category' dtype: @@ -176,7 +184,7 @@ Which are: .. code:: python ['B', 'C', 'E', 'D', 'A', 'M', 'T', 'F'] - Categories (8, object): ['A', 'B', 'C', 'D', 'E', 'F', 'M', 'T'] + Categories (8, str): ['A', 'B', 'C', 'D', 'E', 'F', 'M', 'T'] But when we transform the test set, the unseen category 'G' is not mapped to any integer, and becomes a missing value instead: @@ -190,7 +198,130 @@ We see that 'G' has been replaced by a missing value: .. code:: python ['M', 'F', 'E', NaN] - Categories (8, object): ['A', 'B', 'C', 'D', 'E', 'F', 'M', 'T'] + Categories (8, str): ['A', 'B', 'C', 'D', 'E', 'F', 'M', 'T'] + +Because we set `missing_values="ignore"`, :class:`MatchCategories()` warns us that +missing values were introduced: + +.. code:: python + + UserWarning: During the encoding, NaN values were introduced in the feature(s) cabin. + +With the default `missing_values="raise"`, :class:`MatchCategories()` raises an error +instead, both when the data contains missing values and when unseen categories would +introduce them. + +With polars +^^^^^^^^^^^ + +:class:`MatchCategories()` also works with polars dataframes. In polars, the variables +are cast to the 'Enum' dtype, which, like pandas 'categorical', holds a fixed list of +categories. Let's create a toy train set and test set: + +.. code:: python + + import polars as pl + from feature_engine.preprocessing import MatchCategories + + train = pl.DataFrame({ + "city": ["London", "Paris", "Madrid", "Paris"], + "rooms": [2, 3, 1, 3], + }) + test = pl.DataFrame({ + "city": ["Madrid", "Rome", "London", "Paris"], + "rooms": [1, 2, 4, 3], + }) + +We fit :class:`MatchCategories()` to the train set: + +.. code:: python + + match_categories = MatchCategories(missing_values="ignore") + match_categories.fit(train) + + match_categories.category_dict_ + +With polars, the categories are stored in lists: + +.. code:: python + + {'city': ['London', 'Madrid', 'Paris']} + +Now we transform the test set: + +.. code:: python + + test_t = match_categories.transform(test) + test_t + +The variable `city` is now an 'Enum', and the unseen category 'Rome' became a missing +value: + +.. code:: text + + shape: (4, 2) + ┌────────┬───────┐ + │ city ┆ rooms │ + │ --- ┆ --- │ + │ enum ┆ i64 │ + ╞════════╪═══════╡ + │ Madrid ┆ 1 │ + │ null ┆ 2 │ + │ London ┆ 4 │ + │ Paris ┆ 3 │ + └────────┴───────┘ + +We can check the categories in the schema: + +.. code:: python + + test_t.schema + +The categories are the ones learned from the train set: + +.. code:: python + + Schema({'city': Enum(categories=['London', 'Madrid', 'Paris']), 'rooms': Int64}) + +The polars 'Enum' dtype only takes strings. Hence, if we cast numerical variables by +setting `ignore_format=True`, their values become strings. The categories are sorted +in numerical order: + +.. code:: python + + match_categories = MatchCategories( + variables=["rooms"], ignore_format=True, missing_values="ignore" + ) + match_categories.fit(train) + + match_categories.category_dict_ + +We see the categories of `rooms` as strings: + +.. code:: python + + {'rooms': ['1', '2', '3']} + +And these are the values after the transformation, where the unseen value 4 became a +missing value: + +.. code:: python + + match_categories.transform(test) + +.. code:: text + + shape: (4, 2) + ┌────────┬───────┐ + │ city ┆ rooms │ + │ --- ┆ --- │ + │ str ┆ enum │ + ╞════════╪═══════╡ + │ Madrid ┆ 1 │ + │ Rome ┆ 2 │ + │ London ┆ null │ + │ Paris ┆ 3 │ + └────────┴───────┘ When to use the transformer diff --git a/feature_engine/preprocessing/match_categories.py b/feature_engine/preprocessing/match_categories.py index 9d66c3d6c..981f1cc0b 100644 --- a/feature_engine/preprocessing/match_categories.py +++ b/feature_engine/preprocessing/match_categories.py @@ -1,7 +1,9 @@ import warnings from typing import List, Optional, Union -import pandas as pd +import narwhals as nw +import narwhals.dependencies as nwd +from narwhals.typing import IntoDataFrame, IntoSeries from feature_engine._base_transformers.mixins import GetFeatureNamesOutMixin from feature_engine._check_init_parameters.check_init_input_params import ( @@ -19,7 +21,7 @@ ) from feature_engine._docstrings.init_parameters.encoders import _ignore_format_docstring from feature_engine._docstrings.substitute import Substitution -from feature_engine.dataframe_checks import _check_contains_na, check_X +from feature_engine.dataframe_checks import check_X from feature_engine.encoding.base_encoder import ( CategoricalInitMixinNA, CategoricalMethodsMixin, @@ -40,7 +42,8 @@ class MatchCategories( ): """ MatchCategories() ensures that categorical variables are encoded as pandas - `'categorical'` dtype, instead of generic python `'object'` or other dtypes. + `'categorical'` dtype, or polars `'Enum'` dtype, instead of generic python + `'object'`, string or other dtypes. Under the hood, `'categorical'` dtype is a representation that maps each category to an integer, thus providing a more memory-efficient object @@ -51,7 +54,11 @@ class MatchCategories( category, and can thus be used to ensure that the correct encoding gets applied when passing categorical data to modelling packages that support this dtype, or to prevent unseen categories from reaching a further transformer - or estimator in a pipeline, for example. + or estimator in a pipeline, for example. Categories not seen during fit become + missing values. + + The polars `'Enum'` dtype only takes strings, so with polars, numerical + variables cast with `ignore_format=True` become strings. More details in the :ref:`User Guide `. @@ -68,7 +75,8 @@ class MatchCategories( Attributes ---------- category_dict_: - Dictionary with the category encodings assigned to each variable. + Dictionary with the categories learned for each variable. With polars, the + categories are stored as lists of strings. {variables_} @@ -94,7 +102,7 @@ class MatchCategories( Set the parameters of this estimator. transform: - Enforce the type of categorical variables as dtype `categorical`. + Cast the categorical variables to a categorical dtype. Examples -------- @@ -131,84 +139,103 @@ def __init__( super().__init__(variables, missing_values, ignore_format) self.return_empty = return_empty - def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): + def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): """ - Learn the encodings or levels to use for representing categorical variables. + Learn the categories of each categorical variable. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The training dataset. Can be the entire dataframe, not just the variables to be transformed. - y: pandas Series, default = None - y is not needed in this encoder. You can pass y or None. + y: Series, default = None + y is not needed in this transformer. You can pass y or None. """ - X = check_X(X) + nw_X = check_X(X) variables_ = self._check_or_select_variables(X) - - if self.missing_values == "raise": - _check_contains_na(X, variables_, error_msg="optional") - - self.category_dict_ = dict() - for var in variables_: - self.category_dict_[var] = pd.Categorical(X[var]).categories + self._check_na(X, variables_) + + if nwd.is_pandas_dataframe(X) is True: + # pandas is faster than narwhals. + self.category_dict_ = { + var: X[var].astype("category").cat.categories for var in variables_ + } + else: + self.category_dict_ = { + var: self._find_categories(nw_X.get_column(var)) for var in variables_ + } self.variables_ = variables_ self._get_feature_names_in(X) return self - def transform(self, X: pd.DataFrame) -> pd.DataFrame: + def transform(self, X: IntoDataFrame) -> IntoDataFrame: """ - Encode categorical variables as pandas categorical dtype. + Cast the categorical variables to a categorical dtype with the categories + learned during fit. Categories not seen during fit become missing values. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features]. - The dataset to encode. + X: dataframe of shape = [n_samples, n_features]. + The dataset to transform. Returns ------- - X_new: pandas dataframe of shape = [n_samples, n_features]. - The dataframe with the variables encoded as pandas categorical dtype. + X_new: dataframe of shape = [n_samples, n_features]. + The dataframe with the variables cast to pandas `category` or polars + `Enum` dtype. """ - X = self._check_transform_input_and_state(X) - - if self.missing_values == "raise": - _check_contains_na(X, self.variables_, error_msg="optional") - - for feature, levels in self.category_dict_.items(): - X[feature] = pd.Categorical( - X[feature].where(X[feature].isin(levels)), - categories=levels + nw_X = self._check_transform_input_and_state(X) + self._check_na(X, self.variables_) + + if nwd.is_pandas_dataframe(X) is True: + # pandas is faster than narwhals. + X = X.copy() + categorical = nw.get_native_namespace(nw_X).Categorical + for feature, levels in self.category_dict_.items(): + # get_indexer returns -1 for unseen categories, which from_codes + # turns into NaN. + X[feature] = categorical.from_codes( + levels.get_indexer(X[feature]), categories=levels + ) + nw_X = nw.from_native(X, eager_only=True) + else: + nw_X = nw_X.with_columns( + *[ + nw.when(nw.col(feature).cast(nw.String).is_in(levels)) + .then(nw.col(feature).cast(nw.String)) + .cast(nw.Enum(levels)) + for feature, levels in self.category_dict_.items() + ] ) - self._check_nas_in_result(X) - return X + self._check_nas_in_result(nw_X) + return nw_X.to_native() - def _check_nas_in_result(self, X: pd.DataFrame): - # check if NaN values were introduced by the encoding - if X[self.category_dict_.keys()].isnull().sum().sum() > 0: + def _check_nas_in_result(self, nw_X: nw.DataFrame): + nan_columns = [ + str(feature) + for feature in self.category_dict_ + if nw_X.get_column(feature).null_count() > 0 + ] - # obtain the name(s) of the columns that have null values - nan_columns = ( - X[self.category_dict_.keys()] - .columns[X[self.category_dict_.keys()].isnull().any()] - .tolist() + if len(nan_columns) > 0: + msg = ( + "During the encoding, NaN values were introduced in the feature(s) " + f"{', '.join(nan_columns)}." ) - - if len(nan_columns) > 1: - nan_columns_str = ", ".join(nan_columns) - else: - nan_columns_str = nan_columns[0] - if self.missing_values == "ignore": - warnings.warn( - "During the encoding, NaN values were introduced in the feature(s) " - f"{nan_columns_str}." - ) + warnings.warn(msg) elif self.missing_values == "raise": - raise ValueError( - "During the encoding, NaN values were introduced in the feature(s) " - f"{nan_columns_str}." - ) + raise ValueError(msg) + + def _find_categories(self, series: nw.Series) -> List[str]: + if series.dtype == nw.Enum: + return list(series.dtype.categories) + if series.dtype.is_float() is True: + # NaN is a missing value in pandas, but a regular value in polars. + series = series.filter(~series.is_nan()) + # polars Enum only takes strings, so categories are sorted in the original + # dtype, to keep numbers in numeric order, and then cast to string. + return series.drop_nulls().unique().sort().cast(nw.String).to_list() diff --git a/tests/test_preprocessing/test_match_categories.py b/tests/test_preprocessing/test_match_categories.py index dff612c3c..bb6a2d3e0 100644 --- a/tests/test_preprocessing/test_match_categories.py +++ b/tests/test_preprocessing/test_match_categories.py @@ -1,67 +1,310 @@ -import warnings +import re -import numpy as np +import narwhals as nw import pandas as pd +import polars as pl import pytest +from sklearn.exceptions import NotFittedError from feature_engine.preprocessing import MatchCategories +from tests.backend_helpers import frame_to_dict +MSG_NA = ( + "Some of the variables in the dataset contain NaN. Check and " + "remove those before using this transformer or set the parameter " + "`missing_values='ignore'` when initialising this transformer." +) -def test_category_encoder_outputs_correct_dtype(): - df_str = pd.DataFrame({"col1": ["a", "b", "c"]}) - res_str = MatchCategories().fit(df_str).transform(df_str) - assert res_str.dtypes["col1"] == "category" +MSG_NA_INTRODUCED = ( + "During the encoding, NaN values were introduced in the feature(s) {}." +) - df_float = pd.DataFrame({"col1": [1.0, 2.0, 3.0]}) - tr = MatchCategories(variables=["col1"], ignore_format=True) - res_float = tr.fit(df_float).transform(df_float) - assert res_float.dtypes["col1"] == "category" +TRAIN = {"x1": ["b", "a", "c", "a"], "x2": [4, 5, 6, 7], "x3": ["z", "y", "z", "y"]} +TEST = {"x1": ["c", "d", "a", "b"], "x2": [5, 6, 4, 7], "x3": ["y", "w", "z", "y"]} - df_obj = pd.DataFrame({"col1": ["a", None, -1.0]}) - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - res_obj = MatchCategories(missing_values="ignore").fit(df_obj).transform(df_obj) - assert res_obj.dtypes["col1"] == "category" - df_categ = pd.DataFrame({"col1": pd.Categorical(pd.Series(["a", "b", "c"]))}) - res_categ = MatchCategories().fit(df_categ).transform(df_categ) - assert res_categ.dtypes["col1"] == "category" +def dtype_categories(X, variable): + if isinstance(X, pd.DataFrame): + return list(X[variable].cat.categories) + return list(X.schema[variable].categories) -def test_category_encoder_handles_missing(): - df_no_nas = pd.DataFrame({"col1": ["a", "b", "c"]}) - df_nas = pd.DataFrame({"col1": ["a", "b", None]}) - df_new = pd.DataFrame({"col1": ["a", "b", "d"]}) +# init parameters +@pytest.mark.parametrize( + "missing_values", ["other", "Raise", "", 1, 0.5, True, 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." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + MatchCategories(missing_values=missing_values) + + +@pytest.mark.parametrize("ignore_format", ["True", 1, 0, None, [True]]) +def test_error_if_ignore_format_not_bool(ignore_format): + msg = ( + "ignore_format takes only booleans True and False. " + f"Got {ignore_format} instead." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + MatchCategories(ignore_format=ignore_format) + + +@pytest.mark.parametrize("return_empty", ["True", 1, 0, None, [True]]) +def test_error_if_return_empty_not_bool(return_empty): + msg = ( + "return_empty takes only boolean values True and False. " + f"Got {return_empty} instead." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + MatchCategories(return_empty=return_empty) + + +@pytest.mark.parametrize( + "missing_values, ignore_format", [("raise", False), ("ignore", True)] +) +def test_init_param_assignment(missing_values, ignore_format): + transformer = MatchCategories( + missing_values=missing_values, ignore_format=ignore_format + ) + assert transformer.missing_values == missing_values + assert transformer.ignore_format is ignore_format + + +# fit and transform +def test_learns_categories_and_casts_to_categorical(make_df): + transformer = MatchCategories() + transformer.fit(make_df(TRAIN)) + Xt = transformer.transform(make_df(TRAIN)) + + assert transformer.variables_ == ["x1", "x3"] + assert {k: list(v) for k, v in transformer.category_dict_.items()} == { + "x1": ["a", "b", "c"], + "x3": ["y", "z"], + } + assert transformer.n_features_in_ == 3 + assert transformer.feature_names_in_ == ["x1", "x2", "x3"] + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == TRAIN + assert dtype_categories(Xt, "x1") == ["a", "b", "c"] + assert dtype_categories(Xt, "x3") == ["y", "z"] + + +def test_categories_are_the_same_in_train_and_test(make_df): + train = make_df({"x1": ["b", "a", "c"]}) + test = make_df({"x1": ["c", "b", "c"]}) + transformer = MatchCategories().fit(train) + + assert dtype_categories(transformer.transform(train), "x1") == ["a", "b", "c"] + assert dtype_categories(transformer.transform(test), "x1") == ["a", "b", "c"] - # check that it fails for missing values when using 'raise' - tr = MatchCategories(missing_values="raise") - with pytest.raises(ValueError): - tr.fit(df_nas) - tr.fit(df_no_nas) - with pytest.raises(ValueError): - tr.transform(df_nas) +def test_unseen_categories_become_nan_and_warn(make_df): + transformer = MatchCategories(missing_values="ignore").fit(make_df(TRAIN)) - # check that it doens't fail for missing values when using 'ignore' - tr = MatchCategories(missing_values="ignore").fit(df_nas) - with pytest.warns(UserWarning): - tr.transform(df_nas) + with pytest.warns(UserWarning, match=re.escape(MSG_NA_INTRODUCED.format("x1, x3"))): + Xt = transformer.transform(make_df(TEST)) - # check that it doesn't fail at transforming new values when using 'ignore' - tr = MatchCategories(missing_values="ignore").fit(df_no_nas) - with pytest.warns(UserWarning): - tr.transform(df_new) + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "x1": ["c", None, "a", "b"], + "x2": [5, 6, 4, 7], + "x3": ["y", None, "z", "y"], + } + assert dtype_categories(Xt, "x1") == ["a", "b", "c"] -def test_category_outputs_correct_results(): - df = pd.DataFrame({"col1": ["a", "b", "c"], "col2": [1.0, 2.0, 3.0]}) - res = MatchCategories(variables=["col1", "col2"], ignore_format=True).fit_transform( - df +def test_error_if_unseen_categories_when_missing_values_raise(make_df): + transformer = MatchCategories().fit(make_df(TRAIN)) + with pytest.raises(ValueError, match=re.escape(MSG_NA_INTRODUCED.format("x1, x3"))): + transformer.transform(make_df(TEST)) + + +def test_error_if_nan_in_fit_when_missing_values_raise(make_df): + X = make_df({"x1": ["a", None, "b"]}) + with pytest.raises(ValueError, match=re.escape(MSG_NA)): + MatchCategories().fit(X) + + +def test_error_if_nan_in_transform_when_missing_values_raise(make_df): + transformer = MatchCategories().fit(make_df({"x1": ["a", "b", "b"]})) + with pytest.raises(ValueError, match=re.escape(MSG_NA)): + transformer.transform(make_df({"x1": ["a", None, "b"]})) + + +def test_nan_is_not_a_category_when_missing_values_ignore(make_df): + X = make_df({"x1": ["b", None, "a", "b"], "x2": [1.0, None, 2.0, 3.0]}) + transformer = MatchCategories(missing_values="ignore").fit(X) + + with pytest.warns(UserWarning, match=re.escape(MSG_NA_INTRODUCED.format("x1"))): + Xt = transformer.transform(X) + + assert {k: list(v) for k, v in transformer.category_dict_.items()} == { + "x1": ["a", "b"] + } + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "x1": ["b", None, "a", "b"], + "x2": [1.0, None, 2.0, 3.0], + } + + +@pytest.mark.parametrize("variables", ["x3", ["x3"]]) +def test_transforms_only_selected_variables(make_df, variables): + transformer = MatchCategories(variables=variables, missing_values="ignore") + transformer.fit(make_df(TRAIN)) + + with pytest.warns(UserWarning, match=re.escape(MSG_NA_INTRODUCED.format("x3"))): + Xt = transformer.transform(make_df(TEST)) + + assert transformer.variables_ == ["x3"] + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == { + "x1": ["c", "d", "a", "b"], + "x2": [5, 6, 4, 7], + "x3": ["y", None, "z", "y"], + } + assert nw.from_native(Xt).schema["x1"] == nw.String + + +def test_keeps_categories_of_categorical_input(make_df): + # the categories come from the dtype, so unused ones and their order are kept. + X = ( + nw.from_native(make_df({"x1": ["b", "a", "b"]})) + .with_columns(nw.col("x1").cast(nw.Enum(["z", "b", "a"]))) + .to_native() ) - pd.testing.assert_frame_equal(df, res, check_dtype=False, check_categorical=False) + transformer = MatchCategories().fit(X) + Xt = transformer.transform(make_df({"x1": ["a", "b", "a"]})) + + assert list(transformer.category_dict_["x1"]) == ["z", "b", "a"] + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == {"x1": ["a", "b", "a"]} + assert dtype_categories(Xt, "x1") == ["z", "b", "a"] + + +def test_ignore_format_casts_numerical_variables(make_df): + # polars categorical dtypes only take strings, so numbers become strings. + X = make_df({"x1": [3, 1, 2, 10], "x2": [1.5, float("nan"), 2.5, 10.0]}) + X_test = make_df({"x1": [1, 5, 10, 3], "x2": [2.5, 1.5, 7.0, 10.0]}) + transformer = MatchCategories(ignore_format=True, missing_values="ignore") + transformer.fit(X) + + with pytest.warns(UserWarning, match=re.escape(MSG_NA_INTRODUCED.format("x1, x2"))): + Xt = transformer.transform(X_test) + + expected_categories = { + pd.DataFrame: {"x1": [1, 2, 3, 10], "x2": [1.5, 2.5, 10.0]}, + pl.DataFrame: {"x1": ["1", "2", "3", "10"], "x2": ["1.5", "2.5", "10.0"]}, + } + expected_values = { + pd.DataFrame: {"x1": [1.0, None, 10.0, 3.0], "x2": [2.5, 1.5, None, 10.0]}, + pl.DataFrame: { + "x1": ["1", None, "10", "3"], + "x2": ["2.5", "1.5", None, "10.0"], + }, + } + assert { + k: list(v) for k, v in transformer.category_dict_.items() + } == expected_categories[make_df] + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == expected_values[make_df] + assert dtype_categories(Xt, "x1") == expected_categories[make_df]["x1"] + + +def test_return_empty_when_no_categorical_variables(make_df): + X = make_df({"x1": [1, 2, 3], "x2": [1.0, 2.0, 3.0]}) + transformer = MatchCategories(return_empty=True) + + with pytest.warns( + UserWarning, + match=re.escape( + "No categorical variables found in this dataframe. " + "Returning an empty list." + ), + ): + transformer.fit(X) + Xt = transformer.transform(X) + + assert transformer.variables_ == [] + assert transformer.category_dict_ == {} + assert isinstance(Xt, make_df) + assert frame_to_dict(Xt) == {"x1": [1, 2, 3], "x2": [1.0, 2.0, 3.0]} + + +def test_error_if_no_categorical_variables(make_df): + msg = ( + "No categorical variables found in this dataframe. Check variable " + "dtypes or set return_empty to True to return an empty list instead." + ) + with pytest.raises(TypeError, match=re.escape(msg)): + MatchCategories().fit(make_df({"x1": [1, 2, 3]})) + + +def test_does_not_modify_input(make_df): + X = make_df(TEST) + transformer = MatchCategories(missing_values="ignore").fit(make_df(TRAIN)) + + with pytest.warns(UserWarning, match=re.escape(MSG_NA_INTRODUCED.format("x1, x3"))): + transformer.transform(X) + + assert frame_to_dict(X) == TEST + assert nw.from_native(X).schema["x1"] == nw.String + + +def test_error_if_transform_before_fit(make_df): + msg = ( + "This MatchCategories instance is not fitted yet. Call 'fit' with " + "appropriate arguments before using this estimator." + ) + with pytest.raises(NotFittedError, match=re.escape(msg)): + MatchCategories().transform(make_df(TRAIN)) + + +def test_output_dtype_is_pandas_category(): + Xt = MatchCategories(missing_values="ignore").fit(pd.DataFrame(TRAIN)) + + with pytest.warns(UserWarning, match=re.escape(MSG_NA_INTRODUCED.format("x1, x3"))): + Xt = Xt.transform(pd.DataFrame(TEST)) + + expected = pd.DataFrame( + { + "x1": pd.Categorical(["c", None, "a", "b"], categories=["a", "b", "c"]), + "x2": [5, 6, 4, 7], + "x3": pd.Categorical(["y", None, "z", "y"], categories=["y", "z"]), + } + ) + pd.testing.assert_frame_equal(Xt, expected) + + +def test_output_dtype_is_polars_enum(): + Xt = MatchCategories().fit_transform(pl.DataFrame(TRAIN)) + assert Xt.schema == pl.Schema( + {"x1": pl.Enum(["a", "b", "c"]), "x2": pl.Int64, "x3": pl.Enum(["y", "z"])} + ) + + +def test_integer_column_names(): + X = pd.DataFrame({0: ["a", "b", "c"], 1: ["x", "y", "x"], "n": [1, 2, 3]}) + X_test = pd.DataFrame({0: ["a", "q", "c"], 1: ["x", "y", "w"], "n": [1, 2, 3]}) + transformer = MatchCategories(missing_values="ignore").fit(X) + + with pytest.warns(UserWarning, match=re.escape(MSG_NA_INTRODUCED.format("0, 1"))): + Xt = transformer.transform(X_test) + + expected = pd.DataFrame( + { + 0: pd.Categorical(["a", None, "c"], categories=["a", "b", "c"]), + 1: pd.Categorical(["x", "y", None], categories=["x", "y"]), + "n": [1, 2, 3], + } + ) + pd.testing.assert_frame_equal(Xt, expected) + - df = pd.DataFrame({"col1": ["a", "b", "d"], "col2": [1.0, 2.0, np.nan]}) - res = MatchCategories( - variables=["col1", "col2"], ignore_format=True, missing_values="ignore" - ).fit_transform(df) - pd.testing.assert_frame_equal(df, res, check_dtype=False, check_categorical=False) +def test_keeps_pandas_index(): + X = pd.DataFrame(TRAIN, index=[10, 20, 30, 40]) + Xt = MatchCategories().fit_transform(X) + pd.testing.assert_index_equal(Xt.index, X.index)