diff --git a/docs/user_guide/encoding/StringSimilarityEncoder.rst b/docs/user_guide/encoding/StringSimilarityEncoder.rst index 3fbfaec27..0397db198 100644 --- a/docs/user_guide/encoding/StringSimilarityEncoder.rst +++ b/docs/user_guide/encoding/StringSimilarityEncoder.rst @@ -299,6 +299,37 @@ Below, we see the resulting dataframe: 393 0.0 0.437500 0.666667 0.666667 +With polars +----------- + +:class:`StringSimilarityEncoder()` works the same way with polars dataframes: + +.. code:: python + + import polars as pl + from feature_engine.encoding import StringSimilarityEncoder + + df = pl.DataFrame({"words": ["dog", "dig", "cat"]}) + + encoder = StringSimilarityEncoder() + dft = encoder.fit_transform(df) + dft + +We see the same similarity values as with the pandas dataframe: + +.. code:: text + + shape: (3, 3) + ┌───────────┬───────────┬───────────┐ + │ words_dog ┆ words_dig ┆ words_cat │ + │ --- ┆ --- ┆ --- │ + │ f64 ┆ f64 ┆ f64 │ + ╞═══════════╪═══════════╪═══════════╡ + │ 1.0 ┆ 0.666667 ┆ 0.0 │ + │ 0.666667 ┆ 1.0 ┆ 0.0 │ + │ 0.0 ┆ 0.0 ┆ 1.0 │ + └───────────┴───────────┴───────────┘ + Additional resources -------------------- diff --git a/feature_engine/encoding/similarity_encoder.py b/feature_engine/encoding/similarity_encoder.py index f15f87003..f56219bd6 100644 --- a/feature_engine/encoding/similarity_encoder.py +++ b/feature_engine/encoding/similarity_encoder.py @@ -1,8 +1,9 @@ from difflib import SequenceMatcher from typing import List, Optional, Union +import narwhals as nw import numpy as np -import pandas as pd +from narwhals.typing import IntoDataFrame, IntoSeries from sklearn.utils.validation import check_is_fitted from feature_engine._docstrings.fit_attributes import ( @@ -183,6 +184,26 @@ class StringSimilarityEncoder(CategoricalMethodsMixin, CategoricalInitMixin): 1 2 0.666667 1.000000 0.444444 0.4 2 3 0.444444 0.444444 1.000000 0.0 3 4 0.000000 0.400000 0.000000 1.0 + + With polars + + >>> import polars as pl + >>> from feature_engine.encoding import StringSimilarityEncoder + >>> X = pl.DataFrame(dict(x1 = [1,2,3,4], x2 = ["dog", "dig", "dagger", "hi"])) + >>> sse = StringSimilarityEncoder() + >>> sse.fit(X) + >>> sse.transform(X) + shape: (4, 5) + ┌─────┬──────────┬──────────┬───────────┬───────┐ + │ x1 ┆ x2_dog ┆ x2_dig ┆ x2_dagger ┆ x2_hi │ + │ --- ┆ --- ┆ --- ┆ --- ┆ --- │ + │ i64 ┆ f64 ┆ f64 ┆ f64 ┆ f64 │ + ╞═════╪══════════╪══════════╪═══════════╪═══════╡ + │ 1 ┆ 1.0 ┆ 0.666667 ┆ 0.444444 ┆ 0.0 │ + │ 2 ┆ 0.666667 ┆ 1.0 ┆ 0.444444 ┆ 0.4 │ + │ 3 ┆ 0.444444 ┆ 0.444444 ┆ 1.0 ┆ 0.0 │ + │ 4 ┆ 0.0 ┆ 0.4 ┆ 0.0 ┆ 1.0 │ + └─────┴──────────┴──────────┴───────────┴───────┘ """ def __init__( @@ -198,7 +219,11 @@ def __init__( raise ValueError( f"top_categories takes only integers. Got {top_categories!r} instead." ) - if missing_values not in ("raise", "impute", "ignore"): + if not isinstance(missing_values, str) or missing_values not in ( + "raise", + "impute", + "ignore", + ): raise ValueError( "missing_values should be one of 'raise', 'impute' or 'ignore'." f" Got {missing_values!r} instead." @@ -217,7 +242,7 @@ def __init__( self.missing_values = missing_values self.keywords = keywords - def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): + def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): """ Learns the unique categories per variable. If top_categories is indicated, it will learn the most popular categories. Alternatively, it learns all @@ -226,15 +251,15 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The training input samples. Can be the entire dataframe, not just the variables to encode. - y: pandas series, default=None + y: Series, default=None Target. It is not needed in this encoder. You can pass y or None. """ - X = check_X(X) + nw_X = check_X(X) variables_ = self._check_or_select_variables(X) if self.keywords and not all( @@ -249,115 +274,102 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): if self.missing_values == "raise": _check_contains_na(X, variables_, error_msg="optional") - self.encoder_dict_ = {} + encoder_dict_ = {} if self.keywords: - self.encoder_dict_.update(self.keywords) + encoder_dict_.update(self.keywords) cols_to_iterate = [x for x in variables_ if x not in self.keywords] else: cols_to_iterate = variables_ - if self.missing_values == "raise": - for var in cols_to_iterate: - self.encoder_dict_[var] = ( - X[var] - .astype(str) - .value_counts() - .head(self.top_categories) - .index.tolist() - ) - elif self.missing_values == "impute": - for var in cols_to_iterate: - series = X[var] - self.encoder_dict_[var] = ( - series.astype(str) - .mask(series.isna(), "") - .value_counts() - .head(self.top_categories) - .index.tolist() - ) - elif self.missing_values == "ignore": - for var in cols_to_iterate: - self.encoder_dict_[var] = ( - X[var] - .dropna() - .astype(str) - .value_counts(dropna=True) - .head(self.top_categories) - .index.tolist() - ) - else: - raise ValueError( - "Unrecognized value for missing_values. It should be 'raise', 'ignore' " - f"or 'impute'. Got {self.missing_values} instead." - ) + # cast(nw.String) keeps nulls as nulls, unlike pandas' astype(str) + for var in cols_to_iterate: + col = nw_X.get_column(var) + if self.missing_values == "impute": + col = col.cast(nw.String).fill_null("") + elif self.missing_values == "ignore": + col = col.drop_nulls().cast(nw.String) + else: + col = col.cast(nw.String) + + # sort=True mirrors pandas' own value_counts() default order + # (descending by count, ties broken by first appearance), so + # encoder_dict_ keeps the same category order as before. + counts = col.value_counts(sort=True) + categories = counts.get_column(counts.columns[0]).to_list() + encoder_dict_[var] = categories[: self.top_categories] # assign underscore parameters at the end in case code above fails self.variables_ = variables_ + self.encoder_dict_ = encoder_dict_ self._get_feature_names_in(X) return self - def transform(self, X: pd.DataFrame) -> pd.DataFrame: + def transform(self, X: IntoDataFrame) -> IntoDataFrame: """ Replaces the categorical variables with the similarity variables. 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. + X_new: dataframe. The transformed dataframe. The shape of the dataframe will be different from the original as it includes the similarity variables in place of the original categorical ones. """ check_is_fitted(self) - X = self._check_transform_input_and_state(X) + nw_X = self._check_transform_input_and_state(X) if self.missing_values == "raise": _check_contains_na(X, self.variables_, error_msg="optional") if len(self.variables_) == 0: - return X + return nw_X.to_native() - new_values = [] + # difflib has no vectorised equivalent, so similarities are computed in numpy + # once per unique value, then mapped back to the rows + new_series = [] for var in self.variables_: + col = nw_X.get_column(var) + categories = self.encoder_dict_[var] + + null_mask = None if self.missing_values == "impute": - series = X[var] - series = series.astype(str).mask(series.isna(), "") + str_col = col.cast(nw.String).fill_null("") else: - series = X[var].astype(str) - - categories = series.unique() - column_encoder_dict = { - x: _gpm_fast_vec(x, self.encoder_dict_[var]) for x in categories - } - # Ensure map result is always an array of the correct size. - # Missing values in categories or unknown categories will map to NaN. - default_nan = np.full(len(self.encoder_dict_[var]), np.nan) - if "nan" not in column_encoder_dict: - column_encoder_dict["nan"] = default_nan - if "" not in column_encoder_dict: - column_encoder_dict[""] = default_nan - - encoded_series = series.map(column_encoder_dict) - - # Robust stacking: replace any float NaNs (from unknown values) with arrays - encoded_list = [ - v if isinstance(v, (list, np.ndarray)) else default_nan - for v in encoded_series - ] - encoded = np.vstack(encoded_list) - if self.missing_values == "ignore": - encoded[X[var].isna(), :] = np.nan - new_values.append(encoded) - - new_features = self._get_new_features_name() - X.loc[:, new_features] = np.hstack(new_values) - - return X.drop(self.variables_, axis=1) + str_col = col.cast(nw.String) + if self.missing_values == "ignore": + null_mask = np.array(col.is_null().to_list()) + + values = np.asarray(str_col.to_list(), dtype=object) + if null_mask is not None: + # placeholder value for null rows: overwritten with NaN + # below, the string itself is never used. + values = np.where(null_mask, "", values) + + unique_vals, inverse = np.unique(values, return_inverse=True) + cats_arr = np.asarray(categories, dtype=object) + sim_matrix = _gpm_fast_vec( + unique_vals.reshape(-1, 1), cats_arr.reshape(1, -1) + ) + encoded = sim_matrix[inverse] + + if null_mask is not None: + encoded[null_mask, :] = np.nan + + for j, category in enumerate(categories): + name = f"{var}_nan" if category == "" else f"{var}_{category}" + new_series.append( + nw.new_series(name, encoded[:, j], backend=nw_X.implementation) + ) + + nw_X = nw_X.with_columns(*new_series).drop(self.variables_) + + return nw_X.to_native() def _get_new_features_name(self) -> List[str]: """Return names of the created features.""" @@ -378,7 +390,7 @@ def _add_new_feature_names(self, feature_names: List[str]) -> List[str]: return feature_names - def inverse_transform(self, X: pd.DataFrame): + def inverse_transform(self, X: IntoDataFrame): """inverse_transform is not implemented for this transformer.""" raise NotImplementedError( "inverse_transform is not implemented for this transformer." diff --git a/tests/test_encoding/test_similarity_encoder.py b/tests/test_encoding/test_similarity_encoder.py index 09c17443b..c11bc1fc2 100644 --- a/tests/test_encoding/test_similarity_encoder.py +++ b/tests/test_encoding/test_similarity_encoder.py @@ -1,3 +1,4 @@ +import re from difflib import SequenceMatcher import numpy as np @@ -6,8 +7,73 @@ from feature_engine.encoding import StringSimilarityEncoder from feature_engine.encoding.similarity_encoder import _gpm_fast +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." +) + + +# init parameters +@pytest.mark.parametrize("top_cat", ["hello", 0.5, [1]]) +def test_error_if_top_categories_not_integer(top_cat): + msg = f"top_categories takes only integers. Got {top_cat!r} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + StringSimilarityEncoder(top_categories=top_cat) + + +@pytest.mark.parametrize( + "missing_values", + ["error", "propagate", "Raise", ["raise"], ("impute",), 1, 0.1, False, None], +) +def test_error_if_missing_values_not_allowed(missing_values): + msg = ( + "missing_values should be one of 'raise', 'impute' or 'ignore'. " + f"Got {missing_values!r} instead." + ) + with pytest.raises(ValueError, match=re.escape(msg)): + StringSimilarityEncoder(missing_values=missing_values) + + +@pytest.mark.parametrize("keywords", ["hello", 0.5, [1]]) +def test_keywords_bad_type(keywords): + msg = f"keywords should be a dictionary or None. Got {keywords!r} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + StringSimilarityEncoder(keywords=keywords) + + +@pytest.mark.parametrize("item", ["hello", 0.5, 1]) +def test_keywords_bad_items(item): + keywords = {"var_A": item} + msg = f"The items in keywords should be lists. Got {keywords.values()!r} instead." + with pytest.raises(ValueError, match=re.escape(msg)): + StringSimilarityEncoder(keywords=keywords) +@pytest.mark.parametrize( + "top_categories, keywords, missing_values, ignore_format", + [ + (None, None, "impute", False), + (2, {"var_A": ["XYZ"]}, "raise", True), + (10, {"var_A": ["X"], "var_B": ["Y", "Z"]}, "ignore", False), + ], +) +def test_init_param_assignment(top_categories, keywords, missing_values, ignore_format): + encoder = StringSimilarityEncoder( + top_categories=top_categories, + keywords=keywords, + missing_values=missing_values, + ignore_format=ignore_format, + ) + assert encoder.top_categories == top_categories + assert encoder.keywords == keywords + assert encoder.missing_values == missing_values + assert encoder.ignore_format is ignore_format + + +# fit and transform @pytest.mark.parametrize( "strings", [("hola", "chau"), ("hi there", "hi here"), (100, 1000)] ) @@ -18,39 +84,10 @@ def test_gpm_fast(strings): ) -def test_encode_top_categories(): - df = pd.DataFrame( - { - "var_A": ["A"] * 5 - + ["B"] * 11 - + ["C"] * 4 - + ["D"] * 9 - + ["E"] * 2 - + ["F"] * 2 - + ["G"] * 7, - "var_B": ["A"] * 11 - + ["B"] * 7 - + ["C"] * 4 - + ["D"] * 9 - + ["E"] * 2 - + ["F"] * 2 - + ["G"] * 5, - "var_C": ["A"] * 4 - + ["B"] * 5 - + ["C"] * 11 - + ["D"] * 9 - + ["E"] * 2 - + ["F"] * 2 - + ["G"] * 7, - } - ) - +def test_encode_top_categories(make_df, data_enc_top): encoder = StringSimilarityEncoder(top_categories=4) - X = encoder.fit_transform(df) + X = encoder.fit_transform(make_df(data_enc_top)) - # test init params - assert encoder.top_categories == 4 - # test fit attr transf = { "var_A_D": 9, "var_A_B": 11, @@ -75,69 +112,38 @@ def test_encode_top_categories(): "var_C": ["C", "D", "G", "B"], } # test transform output - for col in transf.keys(): - assert X[col].sum() == transf[col] - assert "var_B" not in X.columns - assert "var_B_F" not in X.columns - + assert isinstance(X, make_df) + result = frame_to_dict(X) + assert {col: sum(result[col]) for col in transf} == transf + assert "var_B" not in result + assert "var_B_F" not in result -@pytest.mark.parametrize("top_cat", ["hello", 0.5, [1]]) -def test_error_if_top_categories_not_integer(top_cat): - with pytest.raises(ValueError): - StringSimilarityEncoder(top_categories=top_cat) - - -@pytest.mark.parametrize( - "handle_missing", ["error", "propagate", ["raise"], 1, 0.1, False] -) -def test_error_if_handle_missing_invalid(handle_missing): - with pytest.raises(ValueError): - StringSimilarityEncoder(missing_values=handle_missing) - - -@pytest.mark.parametrize("missing_vals", ["other", False, 1]) -def test_error_if_missing_values_not_recognized_in_fit(missing_vals, df_enc): - enc = StringSimilarityEncoder() - enc.missing_values = missing_vals - with pytest.raises(ValueError): - enc.fit(df_enc) - -def test_nan_behaviour_error_fit(df_enc_big_na): +def test_nan_behaviour_error_fit(make_df, data_enc_big_na): encoder = StringSimilarityEncoder(missing_values="raise") - with pytest.raises(ValueError) as record: - encoder.fit(df_enc_big_na) - - msg = ( - "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." - ) - assert str(record.value) == msg + with pytest.raises(ValueError, match=re.escape(MSG_NA)): + encoder.fit(make_df(data_enc_big_na)) +# pandas offers several NA sentinels (np.nan, pd.NA, None); polars only has +# a single null representation, so this stays pandas-only. @pytest.mark.parametrize("nan_value", [np.nan, pd.NA, None]) -def test_nan_behaviour_error_transform(df_enc_big, nan_value): +def test_nan_behaviour_error_transform(nan_value, data_enc_big): + df_enc_big = pd.DataFrame(data_enc_big) encoder = StringSimilarityEncoder(missing_values="raise") encoder.fit(df_enc_big) df_enc_big_na = df_enc_big.copy() df_enc_big_na.loc[0, "var_A"] = nan_value - with pytest.raises(ValueError) as record: + with pytest.raises(ValueError, match=re.escape(MSG_NA)): encoder.transform(df_enc_big_na) - msg = ( - "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." - ) - assert str(record.value) == msg +# pandas-only: several NA sentinels, see above. @pytest.mark.parametrize("nan_value", [np.nan, pd.NA, None]) -def test_nan_behaviour_impute(df_enc_big, nan_value): - - df_enc_big_na = df_enc_big.copy() +def test_nan_behaviour_impute(nan_value, data_enc_big): + df_enc_big_na = pd.DataFrame(data_enc_big) df_enc_big_na.loc[0, "var_A"] = nan_value encoder = StringSimilarityEncoder(missing_values="impute") @@ -151,9 +157,10 @@ def test_nan_behaviour_impute(df_enc_big, nan_value): } +# pandas-only: several NA sentinels, see above. @pytest.mark.parametrize("nan_value", [np.nan, pd.NA, None]) -def test_nan_behaviour_ignore(df_enc_big, nan_value): - df_enc_big_na = df_enc_big.copy() +def test_nan_behaviour_ignore(nan_value, data_enc_big): + df_enc_big_na = pd.DataFrame(data_enc_big) df_enc_big_na.loc[0, "var_A"] = nan_value encoder = StringSimilarityEncoder(missing_values="ignore") @@ -167,18 +174,17 @@ def test_nan_behaviour_ignore(df_enc_big, nan_value): def test_string_dtype_with_pd_na(): - # Test StringDtype with pd.NA to hit "" branch in transform + # pandas nullable "string" dtype is pandas-specific. df = pd.DataFrame({"var_A": ["A", "B", pd.NA]}, dtype="string") encoder = StringSimilarityEncoder(missing_values="impute") X = encoder.fit_transform(df) assert (X.isna().sum() == 0).all(axis=None) - # The categories will include "" or the string version of it assert "" in encoder.encoder_dict_["var_A"] def test_string_dtype_with_literal_nan_strings(): - # Test with literal "nan" and "" strings to hit skips in - # transform (line 339, 341 False) + # literal "nan"/"" strings (not real nulls) must be treated as + # ordinary categories; pandas nullable "string" dtype is pandas-specific. df = pd.DataFrame({"var_A": ["nan", "", "A", "B"]}, dtype="string") encoder = StringSimilarityEncoder(missing_values="impute") X = encoder.fit_transform(df) @@ -187,15 +193,17 @@ def test_string_dtype_with_literal_nan_strings(): assert "" in encoder.encoder_dict_["var_A"] -def test_inverse_transform_error(df_enc_big): +def test_inverse_transform_error(make_df, data_enc_big): encoder = StringSimilarityEncoder() - X = encoder.fit_transform(df_enc_big) - with pytest.raises(NotImplementedError): + X = encoder.fit_transform(make_df(data_enc_big)) + msg = "inverse_transform is not implemented for this transformer." + with pytest.raises(NotImplementedError, match=re.escape(msg)): encoder.inverse_transform(X) -def test_get_feature_names_out(df_enc_big): - input_features = df_enc_big.columns.tolist() +def test_get_feature_names_out(make_df, data_enc_big): + df_enc_big = make_df(data_enc_big) + input_features = list(data_enc_big) tr = StringSimilarityEncoder() tr.fit(df_enc_big) @@ -236,18 +244,20 @@ def test_get_feature_names_out(df_enc_big): assert tr.get_feature_names_out(input_features=None) == out assert tr.get_feature_names_out(input_features=input_features) == out - with pytest.raises(ValueError): + msg = "input_features must be a list or an array. Got var_A instead." + with pytest.raises(ValueError, match=re.escape(msg)): tr.get_feature_names_out("var_A") - with pytest.raises(ValueError): + msg = "input_features is not equal to feature_names_in_" + with pytest.raises(ValueError, match=re.escape(msg)): tr.get_feature_names_out(["var_A", "hola"]) -def test_get_feature_names_out_na(df_enc_big_na): - input_features = df_enc_big_na.columns.tolist() +def test_get_feature_names_out_na(make_df, data_enc_big_na): + input_features = list(data_enc_big_na) tr = StringSimilarityEncoder() - tr.fit(df_enc_big_na) + tr.fit(make_df(data_enc_big_na)) out = [ "var_A_B", @@ -284,58 +294,18 @@ def test_get_feature_names_out_na(df_enc_big_na): assert tr.get_feature_names_out(input_features=input_features) == out -@pytest.mark.parametrize("keywords", ["hello", 0.5, [1]]) -def test_keywords_bad_type(keywords): - with pytest.raises(ValueError): - StringSimilarityEncoder(keywords=keywords) - - -@pytest.mark.parametrize("item", ["hello", 0.5, 1]) -def test_keywords_bad_items(item): - with pytest.raises(ValueError): - StringSimilarityEncoder(keywords={"var_A": item}) - - @pytest.mark.parametrize("key", ["hello", 0.5, 1]) -def test_keywords_bad_keys(df_enc_big, key): +def test_keywords_bad_keys(key, make_df, data_enc_big): encoder = StringSimilarityEncoder(keywords={key: ["A"]}) - with pytest.raises(ValueError): - encoder.fit(df_enc_big) - - -def test_encode_partial_keywords(): - df = pd.DataFrame( - { - "var_A": ["A"] * 5 - + ["B"] * 11 - + ["C"] * 4 - + ["D"] * 9 - + ["E"] * 2 - + ["F"] * 2 - + ["G"] * 7, - "var_B": ["A"] * 11 - + ["B"] * 7 - + ["C"] * 4 - + ["D"] * 9 - + ["E"] * 2 - + ["F"] * 2 - + ["G"] * 5, - "var_C": ["A"] * 4 - + ["B"] * 5 - + ["C"] * 11 - + ["D"] * 9 - + ["E"] * 2 - + ["F"] * 2 - + ["G"] * 7, - } - ) + msg = "There are variables in keywords that are not present in the dataset." + with pytest.raises(ValueError, match=re.escape(msg)): + encoder.fit(make_df(data_enc_big)) + +def test_encode_partial_keywords(make_df, data_enc_top): encoder = StringSimilarityEncoder(top_categories=2, keywords={"var_A": ["XYZ"]}) - X = encoder.fit_transform(df) + X = encoder.fit_transform(make_df(data_enc_top)) - # test init params - assert encoder.top_categories == 2 - # test fit attr transf = { "var_A_XYZ": 0, "var_B_A": 11, @@ -353,43 +323,18 @@ def test_encode_partial_keywords(): "var_C": ["C", "D"], } # test transform output - for col in transf.keys(): - assert X[col].sum() == transf[col] - assert "var_B" not in X.columns - assert "var_B_F" not in X.columns - - -def test_encode_complete_keywords(): - df = pd.DataFrame( - { - "var_A": ["A"] * 5 - + ["B"] * 11 - + ["C"] * 4 - + ["D"] * 9 - + ["E"] * 2 - + ["F"] * 2 - + ["G"] * 7, - "var_B": ["A"] * 11 - + ["B"] * 7 - + ["C"] * 4 - + ["D"] * 9 - + ["E"] * 2 - + ["F"] * 2 - + ["G"] * 5, - "var_C": ["A"] * 4 - + ["B"] * 5 - + ["C"] * 11 - + ["D"] * 9 - + ["E"] * 2 - + ["F"] * 2 - + ["G"] * 7, - } - ) + assert isinstance(X, make_df) + result = frame_to_dict(X) + assert {col: sum(result[col]) for col in transf} == transf + assert "var_B" not in result + assert "var_B_F" not in result + +def test_encode_complete_keywords(make_df, data_enc_top): encoder = StringSimilarityEncoder( keywords={"var_A": ["X"], "var_B": ["Y"], "var_C": ["Z"]} ) - X = encoder.fit_transform(df) + X = encoder.fit_transform(make_df(data_enc_top)) # test fit attr transf = { @@ -407,17 +352,18 @@ def test_encode_complete_keywords(): "var_C": ["Z"], } # test transform output - for col in transf.keys(): - assert X[col].sum() == transf[col] - assert "var_B" not in X.columns - assert "var_B_F" not in X.columns + assert isinstance(X, make_df) + result = frame_to_dict(X) + assert {col: sum(result[col]) for col in transf} == transf + assert "var_B" not in result + assert "var_B_F" not in result -def test_get_feature_names_out_w_keywords(df_enc_big_na): - input_features = df_enc_big_na.columns.tolist() +def test_get_feature_names_out_w_keywords(make_df, data_enc_big_na): + input_features = list(data_enc_big_na) tr = StringSimilarityEncoder(keywords={"var_A": ["XYZ"]}) - tr.fit(df_enc_big_na) + tr.fit(make_df(data_enc_big_na)) out = [ "var_A_XYZ",