From ec32c15a6bd6c6be22406627e49df6af98d8eb66 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Wed, 26 Aug 2026 20:51:21 +0200 Subject: [PATCH 1/2] Migrate OneHotEncoder to narwhals, add polars support Uses narwhals' to_dummies() for the actual expansion rather than a manual numpy/dict loop, since it's a real vectorized one-hot op on both backends. Handles two edge cases to_dummies() doesn't cover directly: a fixed-length prefix placeholder ("__ohe_tmp__") swapped back out by slicing rather than by using the real column name, since to_dummies() only prefixes with the Series name when it's truthy - a falsy real name (e.g. an int column literally named 0) would otherwise silently drop the prefix; and learned categories absent from (or present-but-unlearned in) a given transform batch, filled with an explicit all-0 column so unseen categories are encoded as 0 across the board, matching the pre-narwhals behavior exactly. fit()'s value_counts()/unique() calls and transform()'s reassembly are a single unified narwhals path - no pandas/polars split needed, verified directly on both backends (identical dummy columns/values for identical input). Rewrote tests/test_encoding/test_onehot_encoder.py to the single cross-backend-parametrized-test convention: local dict fixtures (dropping the pandas-only global df_enc_big/df_enc_numeric/df_enc_binary fixtures) parametrized over make_df in [pd.DataFrame, pl.DataFrame], with narwhals- based column/sum assertions replacing pd.testing.assert_frame_equal. test_variables_cast_as_category stays pandas-only (pandas category dtype has no polars equivalent under test there). Verified: 43/43 own tests, full encoding suite 340 passed/17 pre-existing failures (matches the narwhals-encoding-base baseline exactly), flake8 and mypy clean, sphinx -W build clean (only the pre-existing unrelated linkcode_resolve warning), no pandas import in this file itself. Co-Authored-By: Claude Sonnet 5 --- docs/user_guide/encoding/OneHotEncoder.rst | 33 ++ feature_engine/encoding/one_hot.py | 81 +++-- tests/test_encoding/test_onehot_encoder.py | 380 ++++++++++----------- 3 files changed, 261 insertions(+), 233 deletions(-) diff --git a/docs/user_guide/encoding/OneHotEncoder.rst b/docs/user_guide/encoding/OneHotEncoder.rst index 02901671e..94cc1a284 100644 --- a/docs/user_guide/encoding/OneHotEncoder.rst +++ b/docs/user_guide/encoding/OneHotEncoder.rst @@ -521,6 +521,39 @@ We see the names of the columns below: 'embarked_S', 'embarked_C'] +With polars +----------- + +:class:`OneHotEncoder()` works the same way with a polars dataframe: + +.. code:: python + + import polars as pl + from feature_engine.encoding import OneHotEncoder + + X = pl.DataFrame({"x1": ["b", "b", "b", "a", "a"], "x2": [1, 2, 3, 4, 5]}) + + ohe = OneHotEncoder(variables=["x1"]) + ohe.fit(X) + + print(ohe.transform(X)) + +.. code:: text + + shape: (5, 3) + ┌─────┬──────┬──────┐ + │ x2 ┆ x1_b ┆ x1_a │ + │ --- ┆ --- ┆ --- │ + │ i64 ┆ i8 ┆ i8 │ + ╞═════╪══════╪══════╡ + │ 1 ┆ 1 ┆ 0 │ + │ 2 ┆ 1 ┆ 0 │ + │ 3 ┆ 1 ┆ 0 │ + │ 4 ┆ 0 ┆ 1 │ + │ 5 ┆ 0 ┆ 1 │ + └─────┴──────┴──────┘ + + Considerations -------------- diff --git a/feature_engine/encoding/one_hot.py b/feature_engine/encoding/one_hot.py index 9d028475b..1212a25b1 100644 --- a/feature_engine/encoding/one_hot.py +++ b/feature_engine/encoding/one_hot.py @@ -3,8 +3,8 @@ from typing import List, Optional, Union -import numpy as np -import pandas as pd +import narwhals as nw +from narwhals.typing import IntoDataFrame from feature_engine._docstrings.fit_attributes import ( _feature_names_in_docstring, @@ -196,7 +196,7 @@ def __init__( self.drop_last = drop_last self.drop_last_binary = drop_last_binary - def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): + def fit(self, X: IntoDataFrame, y: Optional[IntoDataFrame] = None): """ Learns the unique categories per variable. If top_categories is indicated, it will learn the most popular categories. Alternatively, it learns all @@ -205,7 +205,7 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: pandas or polars dataframe of shape = [n_samples, n_features] The training input samples. Can be the entire dataframe, not just selected variables. @@ -218,56 +218,56 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): variables_ = self._check_or_select_variables(X) _check_contains_na(X, variables_) + nw_X = nw.from_native(X, eager_only=True) self.encoder_dict_ = {} for var in variables_: + col = nw_X.get_column(var) # make dummies only for the most popular categories if self.top_categories: - self.encoder_dict_[var] = [ - x - for x in X[var] - .value_counts() - .sort_values(ascending=False) - .head(self.top_categories) - .index - ] + top = col.value_counts(sort=True, name="count").head( + self.top_categories + ) + self.encoder_dict_[var] = top.get_column(var).to_list() else: - category_ls = list(X[var].unique()) + category_ls = col.unique(maintain_order=True).to_list() # return k-1 dummies - if self.drop_last: + if self.drop_last is True: self.encoder_dict_[var] = category_ls[:-1] # return k dummies else: self.encoder_dict_[var] = category_ls - self.variables_binary_ = [var for var in variables_ if X[var].nunique() == 2] + self.variables_binary_ = [ + var for var in variables_ if nw_X.get_column(var).n_unique() == 2 + ] # automatically encode binary variables as 1 dummy - if self.drop_last_binary: + if self.drop_last_binary is True: for var in self.variables_binary_: - category = X[var].unique()[0] + category = nw_X.get_column(var).unique(maintain_order=True)[0] self.encoder_dict_[var] = [category] 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: """ Replaces the categorical variables by the binary variables. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: pandas or polars dataframe of shape = [n_samples, n_features] The data to transform. Returns ------- - X_new: pandas dataframe. + X_new: pandas or polars dataframe. The transformed dataframe. The shape of the dataframe will be different from the original as it includes the dummy variables in place of the original categorical ones. @@ -278,20 +278,43 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame: # check if dataset contains na _check_contains_na(X, self.variables_) + nw_X = nw.from_native(X, eager_only=True) + dummy_frames = [] + # a placeholder Series name, swapped back out below by a fixed-length + # prefix slice (never by parsing the category suffix): to_dummies() + # only prefixes with the Series name when it's truthy, so a falsy + # real name (e.g. an int column literally named 0) would otherwise + # silently drop the prefix and make every dummy column look "missing". + tmp_name = "__ohe_tmp__" for feature in self.variables_: - for category in self.encoder_dict_[feature]: - dummy_df = pd.DataFrame( - {f"{feature}_{category}": np.where(X[feature] == category, 1, 0)}, - index=X.index, + desired = [ + f"{feature}_{category}" for category in self.encoder_dict_[feature] + ] + dummies = ( + nw_X.get_column(feature).alias(tmp_name).to_dummies(separator="_") + ) + dummies = dummies.rename( + {c: f"{feature}{c[len(tmp_name):]}" for c in dummies.columns} + ) + # categories learned in fit() but absent from, or unseen categories + # present in, this particular transform batch: to_dummies() only + # creates columns for values it actually finds, so any learned + # category missing here is filled with an all-0 column, and + # selecting just `desired` drops any column for a category that + # wasn't learned (unseen categories are encoded as 0 across the + # board, matching the pre-narwhals behaviour). + missing = [c for c in desired if c not in dummies.columns] + if len(missing) > 0: + dummies = dummies.with_columns( + **{c: nw.lit(0, dtype=nw.Int8) for c in missing} ) - X = pd.concat([X, dummy_df], axis=1) + dummy_frames.append(dummies.select(desired)) - # drop the original non-encoded variables. - X.drop(labels=self.variables_, axis=1, inplace=True) + nw_X = nw.concat([nw_X.drop(*self.variables_), *dummy_frames], how="horizontal") - return X + return nw_X.to_native() - 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_onehot_encoder.py b/tests/test_encoding/test_onehot_encoder.py index aca3448be..726843e0a 100644 --- a/tests/test_encoding/test_onehot_encoder.py +++ b/tests/test_encoding/test_onehot_encoder.py @@ -1,33 +1,73 @@ +import narwhals as nw import pandas as pd +import polars as pl import pytest from sklearn.pipeline import Pipeline from feature_engine.encoding import OneHotEncoder +DF_ENC_BIG = { + "var_A": ["A"] * 6 + ["B"] * 10 + ["C"] * 4 + ["D"] * 10 + ["E"] * 2 + + ["F"] * 2 + ["G"] * 6, + "var_B": ["A"] * 10 + ["B"] * 6 + ["C"] * 4 + ["D"] * 10 + ["E"] * 2 + + ["F"] * 2 + ["G"] * 6, + "var_C": ["A"] * 4 + ["B"] * 6 + ["C"] * 10 + ["D"] * 10 + ["E"] * 2 + + ["F"] * 2 + ["G"] * 6, +} +DF_ENC_NUMERIC = { + "var_A": [1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3], + "var_B": [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3], +} + +DF_ENC_BINARY = { + "var_A": ["A"] * 6 + ["B"] * 10 + ["C"] * 4, + "var_B": ["A"] * 10 + ["B"] * 6 + ["C"] * 4, + "var_C": ["AHA"] * 12 + ["UHU"] * 8, + "var_D": ["OHO"] * 5 + ["EHE"] * 15, + "var_num": [1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 1, 0, 0], +} + + +def _columns(X): + return list(nw.from_native(X, eager_only=True).columns) + + +def _colsum(X, col): + return sum(nw.from_native(X, eager_only=True).get_column(col).to_list()) + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize("index_", [[1, 2, 3], [3, 2, 1], [4, 9, 2]]) -def test_concat_with_non_ordered_index(index_): - df = pd.DataFrame({"varA": ["a", "b", "c"], "varB": ["d", "d", "a"]}, index=index_) +def test_concat_with_non_ordered_index(make_df, index_): + data = {"varA": ["a", "b", "c"], "varB": ["d", "d", "a"]} + if make_df is pd.DataFrame: + df = make_df(data, index=index_) + else: + df = make_df(data) encoder = OneHotEncoder() dft = encoder.fit_transform(df) - df_expected = pd.DataFrame( - { - "varA_a": [1, 0, 0], - "varA_b": [0, 1, 0], - "varA_c": [0, 0, 1], - "varB_d": [1, 1, 0], - "varB_a": [0, 0, 1], - }, - index=index_, - ) - pd.testing.assert_frame_equal(dft, df_expected, check_dtype=False) + + expected = { + "varA_a": [1, 0, 0], + "varA_b": [0, 1, 0], + "varA_c": [0, 0, 1], + "varB_d": [1, 1, 0], + "varB_a": [0, 0, 1], + } + result = nw.from_native(dft, eager_only=True).to_dict(as_series=False) + assert list(result.keys()) == list(expected.keys()) + for col, values in expected.items(): + assert list(result[col]) == values -def test_encode_categories_in_k_binary_plus_select_vars_automatically(df_enc_big): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_encode_categories_in_k_binary_plus_select_vars_automatically(make_df): # test case 1: encode all categories into k binary variables, select variables # automatically + df = make_df(DF_ENC_BIG) encoder = OneHotEncoder(top_categories=None, variables=None, drop_last=False) - X = encoder.fit_transform(df_enc_big) + X = encoder.fit_transform(df) # test init params assert encoder.top_categories is None @@ -35,26 +75,10 @@ def test_encode_categories_in_k_binary_plus_select_vars_automatically(df_enc_big assert encoder.drop_last is False # test fit attr transf = { - "var_A_A": 6, - "var_A_B": 10, - "var_A_C": 4, - "var_A_D": 10, - "var_A_E": 2, - "var_A_F": 2, - "var_A_G": 6, - "var_B_A": 10, - "var_B_B": 6, - "var_B_C": 4, - "var_B_D": 10, - "var_B_E": 2, - "var_B_F": 2, - "var_B_G": 6, - "var_C_A": 4, - "var_C_B": 6, - "var_C_C": 10, - "var_C_D": 10, - "var_C_E": 2, - "var_C_F": 2, + "var_A_A": 6, "var_A_B": 10, "var_A_C": 4, "var_A_D": 10, "var_A_E": 2, + "var_A_F": 2, "var_A_G": 6, "var_B_A": 10, "var_B_B": 6, "var_B_C": 4, + "var_B_D": 10, "var_B_E": 2, "var_B_F": 2, "var_B_G": 6, "var_C_A": 4, + "var_C_B": 6, "var_C_C": 10, "var_C_D": 10, "var_C_E": 2, "var_C_F": 2, "var_C_G": 6, } @@ -67,17 +91,20 @@ def test_encode_categories_in_k_binary_plus_select_vars_automatically(df_enc_big "var_C": ["A", "B", "C", "D", "E", "F", "G"], } # test transform output - assert X.sum().to_dict() == transf - assert "var_A" not in X.columns + for col, expected_sum in transf.items(): + assert _colsum(X, col) == expected_sum + assert "var_A" not in _columns(X) -def test_encode_categories_in_k_minus_1_binary_plus_list_of_variables(df_enc_big): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_encode_categories_in_k_minus_1_binary_plus_list_of_variables(make_df): # test case 2: encode all categories into k-1 binary variables, # pass list of variables + df = make_df(DF_ENC_BIG) encoder = OneHotEncoder( top_categories=None, variables=["var_A", "var_B"], drop_last=True ) - X = encoder.fit_transform(df_enc_big) + X = encoder.fit_transform(df) # test init params assert encoder.top_categories is None @@ -85,18 +112,9 @@ def test_encode_categories_in_k_minus_1_binary_plus_list_of_variables(df_enc_big assert encoder.drop_last is True # test fit attr transf = { - "var_A_A": 6, - "var_A_B": 10, - "var_A_C": 4, - "var_A_D": 10, - "var_A_E": 2, - "var_A_F": 2, - "var_B_A": 10, - "var_B_B": 6, - "var_B_C": 4, - "var_B_D": 10, - "var_B_E": 2, - "var_B_F": 2, + "var_A_A": 6, "var_A_B": 10, "var_A_C": 4, "var_A_D": 10, "var_A_E": 2, + "var_A_F": 2, "var_B_A": 10, "var_B_B": 6, "var_B_C": 4, "var_B_D": 10, + "var_B_E": 2, "var_B_F": 2, } assert encoder.variables_ == ["var_A", "var_B"] @@ -107,41 +125,26 @@ def test_encode_categories_in_k_minus_1_binary_plus_list_of_variables(df_enc_big "var_B": ["A", "B", "C", "D", "E", "F"], } # test transform output + columns = _columns(X) for col in transf.keys(): - assert X[col].sum() == transf[col] - assert "var_B" not in X.columns - assert "var_B_G" not in X.columns - assert "var_C" in X.columns + assert _colsum(X, col) == transf[col] + assert "var_B" not in columns + assert "var_B_G" not in columns + assert "var_C" in columns -def test_encode_top_categories(): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_encode_top_categories(make_df): # test case 3: encode only the most popular 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, - } - ) + data = { + "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, + } + df = make_df(data) encoder = OneHotEncoder(top_categories=4, variables=None, drop_last=False) X = encoder.fit_transform(df) @@ -150,18 +153,9 @@ def test_encode_top_categories(): assert encoder.top_categories == 4 # test fit attr transf = { - "var_A_D": 9, - "var_A_B": 11, - "var_A_A": 5, - "var_A_G": 7, - "var_B_A": 11, - "var_B_D": 9, - "var_B_G": 5, - "var_B_B": 7, - "var_C_D": 9, - "var_C_C": 11, - "var_C_G": 7, - "var_C_B": 5, + "var_A_D": 9, "var_A_B": 11, "var_A_A": 5, "var_A_G": 7, + "var_B_A": 11, "var_B_D": 9, "var_B_G": 5, "var_B_B": 7, + "var_C_D": 9, "var_C_C": 11, "var_C_G": 7, "var_C_B": 5, } # test fit attr @@ -174,10 +168,11 @@ def test_encode_top_categories(): "var_C": ["C", "D", "G", "B"], } # test transform output + columns = _columns(X) 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 _colsum(X, col) == transf[col] + assert "var_B" not in columns + assert "var_B_F" not in columns # init params @@ -199,28 +194,33 @@ def test_raises_error_when_not_allowed_smoothing_param_in_init(drop_binary): OneHotEncoder(drop_last_binary=drop_binary) -def test_raises_error_if_df_contains_na(df_enc_big, df_enc_big_na): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_raises_error_if_df_contains_na(make_df): # test case 4: when dataset contains na, fit method + data_na = dict(DF_ENC_BIG) + data_na["var_A"] = [None] + list(DF_ENC_BIG["var_A"][1:]) + df_na = make_df(data_na) + df = make_df(DF_ENC_BIG) + msg = ( "Some of the variables in the dataset contain NaN. Check and " "remove those before using this transformer." ) encoder = OneHotEncoder() - with pytest.raises(ValueError) as record: - encoder.fit(df_enc_big_na) - - assert str(record.value) == msg + with pytest.raises(ValueError, match=msg): + encoder.fit(df_na) # test case 4: when dataset contains na, transform method encoder = OneHotEncoder() - encoder.fit(df_enc_big) - with pytest.raises(ValueError): - encoder.transform(df_enc_big_na) - assert str(record.value) == msg + encoder.fit(df) + with pytest.raises(ValueError, match=msg): + encoder.transform(df_na) -def test_encode_numerical_variables(df_enc_numeric): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_encode_numerical_variables(make_df): + df = make_df(DF_ENC_NUMERIC) encoder = OneHotEncoder( top_categories=None, variables=None, @@ -228,7 +228,7 @@ def test_encode_numerical_variables(df_enc_numeric): ignore_format=True, ) - X = encoder.fit_transform(df_enc_numeric[["var_A", "var_B"]]) + X = encoder.fit_transform(df) # test fit attr transf = { @@ -240,31 +240,30 @@ def test_encode_numerical_variables(df_enc_numeric): "var_B_3": [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1], } - transf = pd.DataFrame(transf).astype("int32") - X = pd.DataFrame(X).astype("int32") - assert encoder.variables_ == ["var_A", "var_B"] assert encoder.variables_binary_ == [] assert encoder.n_features_in_ == 2 assert encoder.encoder_dict_ == {"var_A": [1, 2, 3], "var_B": [1, 2, 3]} # test transform output - pd.testing.assert_frame_equal(X, transf) + result = nw.from_native(X, eager_only=True).to_dict(as_series=False) + for col, values in transf.items(): + assert list(result[col]) == values + +def test_variables_cast_as_category(): + # pandas-specific: category dtype has no polars equivalent behavior + # under test here (encoding categorical-dtype columns). + df = pd.DataFrame(DF_ENC_NUMERIC) + df[["var_A", "var_B"]] = df[["var_A", "var_B"]].astype("category") -def test_variables_cast_as_category(df_enc_numeric): encoder = OneHotEncoder( top_categories=None, variables=None, drop_last=False, ignore_format=True, ) + X = encoder.fit_transform(df) - df = df_enc_numeric.copy() - df[["var_A", "var_B"]] = df[["var_A", "var_B"]].astype("category") - - X = encoder.fit_transform(df[["var_A", "var_B"]]) - - # test fit attr transf = { "var_A_1": [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], "var_A_2": [0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0], @@ -284,26 +283,13 @@ def test_variables_cast_as_category(df_enc_numeric): pd.testing.assert_frame_equal(X, transf) -@pytest.fixture(scope="module") -def df_enc_binary(): - df = { - "var_A": ["A"] * 6 + ["B"] * 10 + ["C"] * 4, - "var_B": ["A"] * 10 + ["B"] * 6 + ["C"] * 4, - "var_C": ["AHA"] * 12 + ["UHU"] * 8, - "var_D": ["OHO"] * 5 + ["EHE"] * 15, - "var_num": [1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 1, 1, 0, 0], - } - df = pd.DataFrame(df) - - return df - - -def test_encode_into_k_dummy_plus_drop_binary(df_enc_binary): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_encode_into_k_dummy_plus_drop_binary(make_df): + df = make_df(DF_ENC_BINARY) encoder = OneHotEncoder( top_categories=None, variables=None, drop_last=False, drop_last_binary=True ) - X = encoder.fit_transform(df_enc_binary) - X = X.astype("int32") + X = encoder.fit_transform(df) # test fit attr transf = { @@ -317,7 +303,6 @@ def test_encode_into_k_dummy_plus_drop_binary(df_enc_binary): "var_C_AHA": [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0], "var_D_OHO": [1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], } - transf = pd.DataFrame(transf).astype("int32") assert encoder.variables_ == ["var_A", "var_B", "var_C", "var_D"] assert encoder.variables_binary_ == ["var_C", "var_D"] @@ -329,16 +314,20 @@ def test_encode_into_k_dummy_plus_drop_binary(df_enc_binary): "var_D": ["OHO"], } # test transform output - pd.testing.assert_frame_equal(X, transf) - assert "var_C_B" not in X.columns + result = nw.from_native(X, eager_only=True).to_dict(as_series=False) + assert list(result.keys()) == list(transf.keys()) + for col, values in transf.items(): + assert list(result[col]) == values + assert "var_C_B" not in result.keys() -def test_encode_into_kminus1_dummyy_plus_drop_binary(df_enc_binary): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_encode_into_kminus1_dummyy_plus_drop_binary(make_df): + df = make_df(DF_ENC_BINARY) encoder = OneHotEncoder( top_categories=None, variables=None, drop_last=True, drop_last_binary=True ) - X = encoder.fit_transform(df_enc_binary) - X = X.astype("int32") + X = encoder.fit_transform(df) # test fit attr transf = { @@ -350,7 +339,6 @@ def test_encode_into_kminus1_dummyy_plus_drop_binary(df_enc_binary): "var_C_AHA": [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0], "var_D_OHO": [1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], } - transf = pd.DataFrame(transf).astype("int32") assert encoder.variables_ == ["var_A", "var_B", "var_C", "var_D"] assert encoder.variables_binary_ == ["var_C", "var_D"] @@ -362,17 +350,21 @@ def test_encode_into_kminus1_dummyy_plus_drop_binary(df_enc_binary): "var_D": ["OHO"], } # test transform output - pd.testing.assert_frame_equal(X, transf) - assert "var_C_B" not in X.columns + result = nw.from_native(X, eager_only=True).to_dict(as_series=False) + assert list(result.keys()) == list(transf.keys()) + for col, values in transf.items(): + assert list(result[col]) == values + assert "var_C_B" not in result.keys() -def test_encode_into_top_categories_plus_drop_binary(df_enc_binary): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_encode_into_top_categories_plus_drop_binary(make_df): + df = make_df(DF_ENC_BINARY) # top_categories = 1 encoder = OneHotEncoder( top_categories=1, variables=None, drop_last=False, drop_last_binary=True ) - X = encoder.fit_transform(df_enc_binary) - X = X.astype("int32") + X = encoder.fit_transform(df) # test fit attr transf = { @@ -382,7 +374,6 @@ def test_encode_into_top_categories_plus_drop_binary(df_enc_binary): "var_C_AHA": [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0], "var_D_OHO": [1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], } - transf = pd.DataFrame(transf).astype("int32") assert encoder.variables_ == ["var_A", "var_B", "var_C", "var_D"] assert encoder.variables_binary_ == ["var_C", "var_D"] @@ -394,15 +385,17 @@ def test_encode_into_top_categories_plus_drop_binary(df_enc_binary): "var_D": ["OHO"], } # test transform output - pd.testing.assert_frame_equal(X, transf) - assert "var_C_B" not in X.columns + result = nw.from_native(X, eager_only=True).to_dict(as_series=False) + assert list(result.keys()) == list(transf.keys()) + for col, values in transf.items(): + assert list(result[col]) == values + assert "var_C_B" not in result.keys() # top_categories = 2 encoder = OneHotEncoder( top_categories=2, variables=None, drop_last=False, drop_last_binary=True ) - X = encoder.fit_transform(df_enc_binary) - X = X.astype("int32") + X = encoder.fit_transform(df) # test fit attr transf = { @@ -414,7 +407,6 @@ def test_encode_into_top_categories_plus_drop_binary(df_enc_binary): "var_C_AHA": [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0], "var_D_OHO": [1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], } - transf = pd.DataFrame(transf).astype("int32") assert encoder.variables_ == ["var_A", "var_B", "var_C", "var_D"] assert encoder.variables_binary_ == ["var_C", "var_D"] @@ -426,28 +418,25 @@ def test_encode_into_top_categories_plus_drop_binary(df_enc_binary): "var_D": ["OHO"], } # test transform output - pd.testing.assert_frame_equal(X, transf) - assert "var_C_B" not in X.columns + result = nw.from_native(X, eager_only=True).to_dict(as_series=False) + assert list(result.keys()) == list(transf.keys()) + for col, values in transf.items(): + assert list(result[col]) == values + assert "var_C_B" not in result.keys() -def test_get_feature_names_out(df_enc_binary): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_get_feature_names_out(make_df): + df = make_df(DF_ENC_BINARY) original_features = ["var_num"] - input_features = df_enc_binary.columns + input_features = list(DF_ENC_BINARY.keys()) tr = OneHotEncoder() - tr.fit(df_enc_binary) + tr.fit(df) out = [ - "var_A_A", - "var_A_B", - "var_A_C", - "var_B_A", - "var_B_B", - "var_B_C", - "var_C_AHA", - "var_C_UHU", - "var_D_OHO", - "var_D_EHE", + "var_A_A", "var_A_B", "var_A_C", "var_B_A", "var_B_B", "var_B_C", + "var_C_AHA", "var_C_UHU", "var_D_OHO", "var_D_EHE", ] feat_out = original_features + out @@ -456,33 +445,20 @@ def test_get_feature_names_out(df_enc_binary): assert tr.get_feature_names_out(input_features=input_features) == feat_out tr = OneHotEncoder(drop_last=True) - tr.fit(df_enc_binary) + tr.fit(df) - out = [ - "var_A_A", - "var_A_B", - "var_B_A", - "var_B_B", - "var_C_AHA", - "var_D_OHO", - ] + out = ["var_A_A", "var_A_B", "var_B_A", "var_B_B", "var_C_AHA", "var_D_OHO"] feat_out = original_features + out assert tr.get_feature_names_out(input_features=None) == feat_out assert tr.get_feature_names_out(input_features=input_features) == feat_out tr = OneHotEncoder(drop_last_binary=True) - tr.fit(df_enc_binary) + tr.fit(df) out = [ - "var_A_A", - "var_A_B", - "var_A_C", - "var_B_A", - "var_B_B", - "var_B_C", - "var_C_AHA", - "var_D_OHO", + "var_A_A", "var_A_B", "var_A_C", "var_B_A", "var_B_B", "var_B_C", + "var_C_AHA", "var_D_OHO", ] feat_out = original_features + out @@ -490,7 +466,7 @@ def test_get_feature_names_out(df_enc_binary): assert tr.get_feature_names_out(input_features=input_features) == feat_out tr = OneHotEncoder(top_categories=1) - tr.fit(df_enc_binary) + tr.fit(df) out = ["var_A_B", "var_B_A", "var_C_AHA", "var_D_EHE"] feat_out = original_features + out @@ -505,24 +481,18 @@ def test_get_feature_names_out(df_enc_binary): tr.get_feature_names_out(["var_A", "hola"]) -def test_get_feature_names_out_from_pipeline(df_enc_binary): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_get_feature_names_out_from_pipeline(make_df): + df = make_df(DF_ENC_BINARY) original_features = ["var_num"] - input_features = df_enc_binary.columns + input_features = list(DF_ENC_BINARY.keys()) tr = Pipeline([("transformer", OneHotEncoder())]) - tr.fit(df_enc_binary) + tr.fit(df) out = [ - "var_A_A", - "var_A_B", - "var_A_C", - "var_B_A", - "var_B_B", - "var_B_C", - "var_C_AHA", - "var_C_UHU", - "var_D_OHO", - "var_D_EHE", + "var_A_A", "var_A_B", "var_A_C", "var_B_A", "var_B_B", "var_B_C", + "var_C_AHA", "var_C_UHU", "var_D_OHO", "var_D_EHE", ] feat_out = original_features + out @@ -530,7 +500,9 @@ def test_get_feature_names_out_from_pipeline(df_enc_binary): assert tr.get_feature_names_out(input_features=input_features) == feat_out -def test_inverse_transform_raises_not_implemented_error(df_enc_binary): - enc = OneHotEncoder().fit(df_enc_binary) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_inverse_transform_raises_not_implemented_error(make_df): + df = make_df(DF_ENC_BINARY) + enc = OneHotEncoder().fit(df) with pytest.raises(NotImplementedError): - enc.inverse_transform(df_enc_binary) + enc.inverse_transform(df) From 54d6c81fa74d19d688b6458e979d4b3e934a9b18 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Mon, 31 Aug 2026 00:10:07 +0200 Subject: [PATCH 2/2] Adapt OneHotEncoder to narwhals-returning check_X Bind check_X / _check_transform_input_and_state results to nw_X and keep the original native X for _check_or_select_variables, _check_contains_na and _get_feature_names_in (those helpers still expect native input, matching the CategoricalImputer migration on narwhals-migration). Drop the now-redundant nw.from_native(X) round-trips in fit() and transform(); they reuse the narwhals frame returned by check_X / _check_transform_input_and_state. Co-Authored-By: Claude Sonnet 5 --- feature_engine/encoding/one_hot.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/feature_engine/encoding/one_hot.py b/feature_engine/encoding/one_hot.py index 1212a25b1..945f70248 100644 --- a/feature_engine/encoding/one_hot.py +++ b/feature_engine/encoding/one_hot.py @@ -214,11 +214,10 @@ def fit(self, X: IntoDataFrame, y: Optional[IntoDataFrame] = None): None. """ - X = check_X(X) + nw_X = check_X(X) variables_ = self._check_or_select_variables(X) _check_contains_na(X, variables_) - nw_X = nw.from_native(X, eager_only=True) self.encoder_dict_ = {} for var in variables_: @@ -273,12 +272,11 @@ def transform(self, X: IntoDataFrame) -> IntoDataFrame: original categorical ones. """ - X = self._check_transform_input_and_state(X) + nw_X = self._check_transform_input_and_state(X) # check if dataset contains na _check_contains_na(X, self.variables_) - nw_X = nw.from_native(X, eager_only=True) dummy_frames = [] # a placeholder Series name, swapped back out below by a fixed-length # prefix slice (never by parsing the category suffix): to_dummies()