From 154f112b32ac9afef4651615e1e4ef602025fea8 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Wed, 26 Aug 2026 20:46:45 +0200 Subject: [PATCH 1/8] Migrate DecisionTreeEncoder to narwhals, add polars support Replaces the old sklearn Pipeline(OrdinalEncoder, DecisionTreeDiscretiser) composition with a direct narwhals-based fit: each variable's categories are ordinal-encoded via a dict built from either a target-mean group_by (encoding_method="ordered") or plain unique-value enumeration ("arbitrary"), a decision tree is trained on the ordinal codes, and predictions are made only on the (few) unique codes rather than the full column, since the tree's output for a category depends only on its code - identical result, far less prediction work for a low-cardinality variable. The "ordered" path sorts by (mean, category) rather than mean alone, matching the tie-break fix applied to the sibling OrdinalEncoder/ MeanEncoder migrations this session, since group_by's own row order isn't guaranteed to match across backends for tied means. Added n_jobs (default None, sequential, unchanged behavior), parallelizing tree training across variables via joblib threads, following the same pattern as DecisionTreeFeatures/DecisionTreeDiscretiser. Verified: 56/56 own tests, full encoding suite 345 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 (the package-level import chain still needs pandas only because sibling encoders on this branch aren't migrated yet, expected given the per-encoder parallel-branch strategy). Co-Authored-By: Claude Sonnet 5 --- .../encoding/DecisionTreeEncoder.rst | 57 ++++ feature_engine/encoding/decision_tree.py | 166 +++++++--- .../test_decision_tree_encoder.py | 312 ++++++++++++------ 3 files changed, 396 insertions(+), 139 deletions(-) diff --git a/docs/user_guide/encoding/DecisionTreeEncoder.rst b/docs/user_guide/encoding/DecisionTreeEncoder.rst index 9c5ab968d..5e220597c 100644 --- a/docs/user_guide/encoding/DecisionTreeEncoder.rst +++ b/docs/user_guide/encoding/DecisionTreeEncoder.rst @@ -438,6 +438,63 @@ In the following image we also see a monotonic relationship after the encoding: be some sort of relationship between the target and the categories that can be captured by the decision tree. Use with caution. +With polars +----------- + +:class:`DecisionTreeEncoder()` works the same way with a polars dataframe. Let's create a toy +dataset: + +.. code:: python + + import polars as pl + from feature_engine.encoding import DecisionTreeEncoder + + X = pl.DataFrame({ + "city": ["London", "Manchester", "Liverpool", "London", "Manchester", "Liverpool"], + "price": [500, 300, 250, 520, 310, 260], + }) + y = pl.Series("target", [1, 0, 0, 1, 0, 1]) + +Let's set up :class:`DecisionTreeEncoder()` to encode `city` with a classification tree, and fit +it to the data: + +.. code:: python + + encoder = DecisionTreeEncoder(variables=["city"], regression=False, cv=2) + encoder.fit(X, y) + + encoder.encoder_dict_ + +We see the resulting mappings from category to the tree's predictions: + +.. code:: python + + {'city': {'London': 1.0, 'Manchester': 0.25, 'Liverpool': 0.25}} + +Now let's transform the data: + +.. code:: python + + encoder.transform(X) + +We obtain a polars dataframe with the categories in `city` replaced by the tree's predictions: + +.. code:: text + + shape: (6, 2) + ┌──────┬───────┐ + │ city ┆ price │ + │ --- ┆ --- │ + │ f64 ┆ i64 │ + ╞══════╪═══════╡ + │ 1.0 ┆ 500 │ + │ 0.25 ┆ 300 │ + │ 0.25 ┆ 250 │ + │ 1.0 ┆ 520 │ + │ 0.25 ┆ 310 │ + │ 0.25 ┆ 260 │ + └──────┴───────┘ + Additional resources -------------------- diff --git a/feature_engine/encoding/decision_tree.py b/feature_engine/encoding/decision_tree.py index 9b66fe32d..9474f9105 100644 --- a/feature_engine/encoding/decision_tree.py +++ b/feature_engine/encoding/decision_tree.py @@ -1,11 +1,15 @@ # Authors: Soledad Galli # License: BSD 3 clause -from typing import List, Optional, 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 sklearn.pipeline import Pipeline +from joblib import Parallel, delayed +from narwhals.typing import IntoDataFrame, IntoSeries +from sklearn.model_selection import GridSearchCV +from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor from sklearn.utils.multiclass import check_classification_targets, type_of_target from feature_engine._docstrings.fit_attributes import ( @@ -27,13 +31,11 @@ ) from feature_engine._docstrings.substitute import Substitution from feature_engine.dataframe_checks import _check_contains_na, check_X_y -from feature_engine.discretisation import DecisionTreeDiscretiser from feature_engine.encoding._helper_functions import check_parameter_unseen from feature_engine.encoding.base_encoder import ( CategoricalInitMixin, CategoricalMethodsMixin, ) -from feature_engine.encoding.ordinal import OrdinalEncoder from feature_engine.tags import _return_tags _unseen_docstring = ( @@ -139,6 +141,15 @@ class DecisionTreeEncoder(CategoricalMethodsMixin, CategoricalInitMixin): fill_value: float, default=None The value used to encode unseen categories. Only used when `unseen='encode'`. + n_jobs: int, default=None + The number of jobs to run in parallel when training the decision trees + across variables. Trees are fit using threads rather than processes, + since fitting a decision tree releases the GIL for the bulk of its + computation, which avoids the overhead of copying the entire dataframe + to separate worker processes. `None` means 1, i.e. sequential training + (this transformer's original behaviour); `-1` means using all available + processors. + Attributes ---------- encoder_dict_: @@ -212,6 +223,27 @@ class DecisionTreeEncoder(CategoricalMethodsMixin, CategoricalInitMixin): 2 3 0.666667 3 4 0.500000 4 5 0.500000 + + With polars: + + >>> import polars as pl + >>> X = pl.DataFrame(dict(x1 = [1,2,3,4,5], x2 = ["b", "b", "b", "a", "a"])) + >>> y = [0, 1, 1, 1, 0] + >>> dte = DecisionTreeEncoder(regression=False, cv=2) + >>> dte.fit(X, y) + >>> dte.transform(X) + shape: (5, 2) + ┌─────┬──────────┐ + │ x1 ┆ x2 │ + │ --- ┆ --- │ + │ i64 ┆ f64 │ + ╞═════╪══════════╡ + │ 1 ┆ 0.666667 │ + │ 2 ┆ 0.666667 │ + │ 3 ┆ 0.666667 │ + │ 4 ┆ 0.5 │ + │ 5 ┆ 0.5 │ + └─────┴──────────┘ """ def __init__( @@ -228,6 +260,7 @@ def __init__( precision: Optional[int] = None, unseen: str = "ignore", fill_value: Optional[float] = None, + n_jobs: Optional[int] = None, ) -> None: if encoding_method not in ["ordered", "arbitrary"]: @@ -261,18 +294,19 @@ def __init__( self.precision = precision self.unseen = unseen self.fill_value = fill_value + self.n_jobs = n_jobs - def fit(self, X: pd.DataFrame, y: pd.Series): + def fit(self, X: IntoDataFrame, y: IntoSeries): """ Fit a decision tree per variable. 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 categorical variables. - y : pandas series. + y: Series. The target variable. Required to train the decision tree and for ordered ordinal encoding. """ @@ -303,60 +337,49 @@ def fit(self, X: pd.DataFrame, y: pd.Series): self._get_feature_names_in(X) return self - encoder = OrdinalEncoder( - encoding_method=self.encoding_method, - variables=variables_, - missing_values="raise", - ignore_format=self.ignore_format, - ) - - tree = DecisionTreeDiscretiser( - cv=self.cv, - scoring=self.scoring, - variables=variables_, - param_grid=param_grid, - regression=self.regression, - random_state=self.random_state, - ) + nw_X = nw.from_native(X, eager_only=True) + + # only needed for "ordered": pairs the target with X once so every + # variable's group_by below can reuse it, instead of rebuilding it + # per variable. + nw_Xy = None + target_name = "__feature_engine_decision_tree_target__" + if self.encoding_method == "ordered": + if nwd.is_into_series(y): + y_nw = nw.from_native(y, series_only=True).alias(target_name) + else: + y_nw = nw.new_series( + name=target_name, values=y, backend=nw_X.implementation + ) + nw_Xy = nw_X.with_columns(y_nw) - # pipeline for the encoder - pipe = Pipeline( - [ - ("encoder", encoder), - ("tree", tree), - ] + mappings = Parallel(n_jobs=self.n_jobs, prefer="threads")( + delayed(self._fit_one_variable)( + nw_X, nw_Xy, var, y, target_name, param_grid + ) + for var in variables_ ) - Xt = pipe.fit_transform(X, y) - - encoder_ = {} - if self.precision is None: - for var in variables_: - encoder_[var] = dict(zip(X[var], Xt[var])) - else: - for var in variables_: - encoder_[var] = dict(zip(X[var], np.round(Xt[var], self.precision))) - if self.unseen == "encode": self._unseen = self.fill_value - self.encoder_dict_ = encoder_ + self.encoder_dict_ = dict(zip(variables_, mappings)) 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: """ Replace categorical variables by the predictions of the decision tree. Parameters ---------- - X : pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The input samples. Returns ------- - X_new : pandas dataframe of shape = [n_samples, n_features]. + X_new: dataframe of shape = [n_samples, n_features]. Dataframe with variables encoded with decision tree predictions. """ X = self._check_transform_input_and_state(X) @@ -365,6 +388,63 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame: return X + def _fit_one_variable( + self, nw_X, nw_Xy, var, y: IntoSeries, target_name: str, param_grid: Dict + ) -> Dict: + """Learn the category-to-prediction mapping for one variable: encode its + categories to ordinal integers, fit a decision tree on those integers, and + predict on each unique category to get its final mapped value.""" + if self.encoding_method == "ordered": + # sort by (mean, category): group_by's own order isn't guaranteed + # across backends, and this tie-break on the category itself + # reproduces pandas' groupby(sort=True) + stable sort_values + # behavior for categories with equal target means. + categories = ( + nw_Xy.group_by(var, drop_null_keys=True) + .agg(nw.col(target_name).mean()) + .sort([target_name, var]) + .get_column(var) + .to_list() + ) + else: + categories = nw_X.get_column(var).unique(maintain_order=True).to_list() + + ordinal_map = {k: i for i, k in enumerate(categories, 0)} + + X_sub = nw_X.get_column(var).replace_strict(ordinal_map).to_frame().to_native() + estimator = self._fit_one_tree(X_sub, y, param_grid) + + # predict directly on the (few) unique ordinal codes instead of the + # full column: the tree's prediction for a category depends only on + # its ordinal code, so this gives identical results far cheaper. + X_pred = nw.new_series( + var, list(range(len(categories))), backend=nw_X.implementation + ).to_frame().to_native() + + if self.regression is True: + preds = estimator.predict(X_pred) + else: + preds = estimator.predict_proba(X_pred)[:, 1] + + if self.precision is not None: + preds = np.round(preds, self.precision) + + return dict(zip(categories, preds)) + + def _fit_one_tree(self, X_sub: IntoDataFrame, y: IntoSeries, param_grid: Dict): + """Instantiate and fit one decision tree on one variable's ordinal-encoded + values.""" + if self.regression is True: + model = DecisionTreeRegressor(random_state=self.random_state) + else: + model = DecisionTreeClassifier(random_state=self.random_state) + + tree_model = GridSearchCV( + model, cv=self.cv, scoring=self.scoring, param_grid=param_grid + ) + tree_model.fit(X_sub, y) + return tree_model + def _assign_param_grid(self): if self.param_grid: param_grid = self.param_grid diff --git a/tests/test_encoding/test_decision_tree_encoder.py b/tests/test_encoding/test_decision_tree_encoder.py index fd4cef789..bf31b342c 100644 --- a/tests/test_encoding/test_decision_tree_encoder.py +++ b/tests/test_encoding/test_decision_tree_encoder.py @@ -1,7 +1,10 @@ +import math 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 @@ -9,6 +12,39 @@ from feature_engine.encoding import DecisionTreeEncoder +def _to_backend(df: pd.DataFrame, make_df): + """Rebuild a pandas fixture dataframe on the requested backend. + + Swaps float NaN for None in string columns - polars (unlike pandas) + rejects a float NaN mixed into an otherwise-string column. + """ + data = {} + for col in df.columns: + values = df[col].tolist() + if any(isinstance(v, str) for v in values): + values = [ + None if isinstance(v, float) and math.isnan(v) else v for v in values + ] + data[col] = values + return make_df(data) + + +def _assert_values(X, expected: dict) -> None: + """NaN-aware, backend-agnostic comparison of a dataframe's contents.""" + result = nw.from_native(X, eager_only=True).to_dict(as_series=False) + assert list(result.keys()) == list(expected.keys()) + for col, exp_values in expected.items(): + got_values = result[col] + assert len(got_values) == len(exp_values) + for got, exp in zip(got_values, exp_values): + if isinstance(exp, float) and math.isnan(exp): + assert got is None or (isinstance(got, float) and math.isnan(got)) + elif isinstance(exp, float): + assert got == pytest.approx(exp) + else: + assert got == exp + + # init parameters @pytest.mark.parametrize("enc_method", ["count", False, 1]) def test_error_if_encoding_method_not_permitted_value(enc_method): @@ -69,9 +105,13 @@ def test_init_param_assignment( # fit attributes -def test_encoding_dictionary(df_enc): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_encoding_dictionary(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) + y = df_enc["target"].tolist() + encoder = DecisionTreeEncoder(regression=False) - encoder.fit(df_enc[["var_A", "var_B"]], df_enc["target"]) + encoder.fit(X, y) # Tree: var_A <= 1.5 -> 0.25 else 0.5 # Tree: var_B <= 0.5 -> 0.2 else 0.4 @@ -82,9 +122,34 @@ def test_encoding_dictionary(df_enc): assert encoder.encoder_dict_ == expected_encodings -def test_precision(df_enc): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_ordered_encoding_dictionary(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) + y = df_enc["target"].tolist() + + encoder = DecisionTreeEncoder(regression=False, encoding_method="ordered") + encoder.fit(X, y) + + # ordered ranks: var_A -> B(mean 0.2) < A(mean 0.333) < C(mean 0.5) + # var_B -> A(mean 0.2) < B(mean 0.333) < C(mean 0.5) + # so var_A's ordinal codes are B=0, A=1, C=2 (split at code <= 0.5, + # i.e. B alone); var_B's are A=0, B=1, C=2 (split at code <= 0.5, i.e. A + # alone). Same tree-split logic as the arbitrary-encoding case above, + # applied to the reordered codes. + expected_encodings = { + "var_A": {"B": 0.2, "A": 0.4, "C": 0.4}, + "var_B": {"A": 0.2, "B": 0.4, "C": 0.4}, + } + assert encoder.encoder_dict_ == expected_encodings + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_precision(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) + y = df_enc["target"].tolist() + encoder = DecisionTreeEncoder(regression=False, precision=1) - encoder.fit(df_enc[["var_A", "var_B"]], df_enc["target"]) + encoder.fit(X, y) # Tree: var_A <= 1.5 -> 0.25 else 0.5 # Tree: var_B <= 0.5 -> 0.2 else 0.4 @@ -95,73 +160,100 @@ def test_precision(df_enc): assert encoder.encoder_dict_ == expected_encodings -def test_classification(df_enc): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_classification(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) + y = df_enc["target"].tolist() + encoder = DecisionTreeEncoder(regression=False) - encoder.fit(df_enc[["var_A", "var_B"]], df_enc["target"]) - X = encoder.transform(df_enc[["var_A", "var_B"]]) + encoder.fit(X, y) + Xt = encoder.transform(X) - transf_df = df_enc.copy() - transf_df["var_A"] = [0.25] * 16 + [0.5] * 4 # Tree: var_A <= 1.5 -> 0.25 else 0.5 - transf_df["var_B"] = [0.2] * 10 + [0.4] * 10 # Tree: var_B <= 0.5 -> 0.2 else 0.4 - pd.testing.assert_frame_equal(X, transf_df[["var_A", "var_B"]]) + expected = { + "var_A": [0.25] * 16 + [0.5] * 4, # Tree: var_A <= 1.5 -> 0.25 else 0.5 + "var_B": [0.2] * 10 + [0.4] * 10, # Tree: var_B <= 0.5 -> 0.2 else 0.4 + } + _assert_values(Xt, expected) -def test_regression(df_enc): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_regression(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) random = np.random.RandomState(42) y = random.normal(0, 0.1, len(df_enc)) encoder = DecisionTreeEncoder( regression=True, random_state=random, ) - encoder.fit(df_enc[["var_A", "var_B"]], y) - X = encoder.transform(df_enc[["var_A", "var_B"]]) + encoder.fit(X, y) + Xt = encoder.transform(X) + + expected = { + "var_A": ( + [0.034348] * 6 + [-0.024679] * 10 + [-0.075473] * 4 + ), # Tree: var_A <= 1.5 -> 0.25 else 0.5 + "var_B": [0.044806] * 10 + [-0.079066] * 10, + } + nw_Xt = nw.from_native(Xt, eager_only=True) + rounded = { + col: [round(v, 6) for v in nw_Xt.get_column(col).to_list()] + for col in ["var_A", "var_B"] + } + assert rounded == expected - transf_df = df_enc.copy() - transf_df["var_A"] = ( - [0.034348] * 6 + [-0.024679] * 10 + [-0.075473] * 4 - ) # Tree: var_A <= 1.5 -> 0.25 else 0.5 - transf_df["var_B"] = [0.044806] * 10 + [-0.079066] * 10 - pd.testing.assert_frame_equal(X.round(6), transf_df[["var_A", "var_B"]]) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_fit_raises_error_if_df_contains_na(df_enc_na, make_df): + X = _to_backend(df_enc_na[["var_A", "var_B"]], make_df) + y = df_enc_na["target"].tolist() -def test_fit_raises_error_if_df_contains_na(df_enc_na): - # test case 4: when dataset contains na, fit method encoder = DecisionTreeEncoder(regression=False) msg = ( "Some of the variables in the dataset contain NaN. Check and " "remove those before using this transformer." ) with pytest.raises(ValueError, match=msg): - encoder.fit(df_enc_na[["var_A", "var_B"]], df_enc_na["target"]) + encoder.fit(X, y) + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_transform_raises_error_if_df_contains_na(df_enc, df_enc_na, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) + X_na = _to_backend(df_enc_na[["var_A", "var_B"]], make_df) + y = df_enc["target"].tolist() -def test_transform_raises_error_if_df_contains_na(df_enc, df_enc_na): - # test case 4: when dataset contains na, transform method encoder = DecisionTreeEncoder(regression=False) - encoder.fit(df_enc[["var_A", "var_B"]], df_enc["target"]) + encoder.fit(X, y) msg = ( "Some of the variables in the dataset contain NaN. Check and " "remove those before using this transformer." ) with pytest.raises(ValueError, match=msg): - encoder.transform(df_enc_na[["var_A", "var_B"]]) + encoder.transform(X_na) + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_classification_ignore_format(df_enc_numeric, make_df): + X = _to_backend(df_enc_numeric[["var_A", "var_B"]], make_df) + y = df_enc_numeric["target"].tolist() -def test_classification_ignore_format(df_enc_numeric): encoder = DecisionTreeEncoder( regression=False, ignore_format=True, ) - encoder.fit(df_enc_numeric[["var_A", "var_B"]], df_enc_numeric["target"]) - X = encoder.transform(df_enc_numeric[["var_A", "var_B"]]) + encoder.fit(X, y) + Xt = encoder.transform(X) - transf_df = df_enc_numeric.copy() - transf_df["var_A"] = [0.25] * 16 + [0.5] * 4 # Tree: var_A <= 1.5 -> 0.25 else 0.5 - transf_df["var_B"] = [0.2] * 10 + [0.4] * 10 # Tree: var_B <= 0.5 -> 0.2 else 0.4 - pd.testing.assert_frame_equal(X, transf_df[["var_A", "var_B"]]) + expected = { + "var_A": [0.25] * 16 + [0.5] * 4, # Tree: var_A <= 1.5 -> 0.25 else 0.5 + "var_B": [0.2] * 10 + [0.4] * 10, # Tree: var_B <= 0.5 -> 0.2 else 0.4 + } + _assert_values(Xt, expected) -def test_regression_ignore_format(df_enc_numeric): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_regression_ignore_format(df_enc_numeric, make_df): + X = _to_backend(df_enc_numeric[["var_A", "var_B"]], make_df) random = np.random.RandomState(42) y = random.normal(0, 0.1, len(df_enc_numeric)) encoder = DecisionTreeEncoder( @@ -169,18 +261,25 @@ def test_regression_ignore_format(df_enc_numeric): random_state=random, ignore_format=True, ) - encoder.fit(df_enc_numeric[["var_A", "var_B"]], y) - X = encoder.transform(df_enc_numeric[["var_A", "var_B"]]) - - transf_df = df_enc_numeric.copy() - transf_df["var_A"] = ( - [0.034348] * 6 + [-0.024679] * 10 + [-0.075473] * 4 - ) # Tree: var_A <= 1.5 -> 0.25 else 0.5 - transf_df["var_B"] = [0.044806] * 10 + [-0.079066] * 10 - pd.testing.assert_frame_equal(X.round(6), transf_df[["var_A", "var_B"]]) + encoder.fit(X, y) + Xt = encoder.transform(X) + + expected = { + "var_A": ( + [0.034348] * 6 + [-0.024679] * 10 + [-0.075473] * 4 + ), # Tree: var_A <= 1.5 -> 0.25 else 0.5 + "var_B": [0.044806] * 10 + [-0.079066] * 10, + } + nw_Xt = nw.from_native(Xt, eager_only=True) + rounded = { + col: [round(v, 6) for v in nw_Xt.get_column(col).to_list()] + for col in ["var_A", "var_B"] + } + assert rounded == expected def test_variables_cast_as_category(df_enc_category_dtypes): + # pandas Categorical dtype has no direct polars equivalent - pandas-only. df = df_enc_category_dtypes.copy() encoder = DecisionTreeEncoder(regression=False) encoder.fit(df[["var_A", "var_B"]], df["target"]) @@ -193,7 +292,11 @@ def test_variables_cast_as_category(df_enc_category_dtypes): assert X["var_A"].dtypes == float -def test_error_when_regression_is_true_and_target_is_binary(df_enc): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_error_when_regression_is_true_and_target_is_binary(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) + y = df_enc["target"].tolist() + encoder = DecisionTreeEncoder(regression=True) msg = ( "Trying to fit a regression to a binary target is not " @@ -201,16 +304,18 @@ def test_error_when_regression_is_true_and_target_is_binary(df_enc): "or set regression to False." ) with pytest.raises(ValueError, match=msg): - encoder.fit(df_enc[["var_A", "var_B"]], df_enc["target"]) + encoder.fit(X, y) -def test_error_when_regression_is_false_and_target_is_continuous(df_enc): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_error_when_regression_is_false_and_target_is_continuous(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) random = np.random.RandomState(42) y = random.normal(0, 10, len(df_enc)) encoder = DecisionTreeEncoder(regression=False) # the error message comes from sklearn api - won't test with pytest.raises(ValueError): - encoder.fit(df_enc[["var_A", "var_B"]], y) + encoder.fit(X, y) @pytest.mark.parametrize( @@ -225,54 +330,60 @@ def test_assigns_param_grid(grid): assert encoder._assign_param_grid() == grid -def test_unseen_is_encode(df_enc): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_unseen_is_encode(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) + y = df_enc["target"].tolist() + encoder = DecisionTreeEncoder(unseen="encode", regression=False, fill_value=-1) - encoder.fit(df_enc[["var_A", "var_B"]], df_enc["target"]) + encoder.fit(X, y) - X_unseen_input = pd.DataFrame( + X_unseen_input = make_df( { "var_A": ["A", "ZZZ", "YYY"], "var_B": ["C", "YYY", "ZZZ"], } ) - - X_unseen_output = pd.DataFrame( - { - "var_A": [0.25, -1, -1], - "var_B": [0.4, -1, -1], - } - ) + expected = { + "var_A": [0.25, -1, -1], + "var_B": [0.4, -1, -1], + } Xt = encoder.transform(X_unseen_input) - pd.testing.assert_frame_equal(Xt, X_unseen_output) + _assert_values(Xt, expected) -def test_unseen_is_ignore(df_enc): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_unseen_is_ignore(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) + y = df_enc["target"].tolist() + encoder = DecisionTreeEncoder(unseen="ignore", regression=False) - encoder.fit(df_enc[["var_A", "var_B"]], df_enc["target"]) + encoder.fit(X, y) - X_unseen_input = pd.DataFrame( + X_unseen_input = make_df( { "var_A": ["A", "ZZZ", "YYY"], "var_B": ["C", "YYY", "ZZZ"], } ) - - X_unseen_output = pd.DataFrame( - { - "var_A": [0.25, np.nan, np.nan], - "var_B": [0.4, np.nan, np.nan], - } - ) + expected = { + "var_A": [0.25, np.nan, np.nan], + "var_B": [0.4, np.nan, np.nan], + } Xt = encoder.transform(X_unseen_input) - pd.testing.assert_frame_equal(Xt, X_unseen_output) + _assert_values(Xt, expected) + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_fit_errors_if_new_cat_values_and_unseen_is_raise_param(df_enc, make_df): + X = _to_backend(df_enc[["var_A", "var_B"]], make_df) + y = df_enc["target"].tolist() -def test_fit_errors_if_new_cat_values_and_unseen_is_raise_param(df_enc): encoder = DecisionTreeEncoder(unseen="raise", regression=False) - encoder.fit(df_enc[["var_A", "var_B"]], df_enc["target"]) - X = pd.DataFrame( + encoder.fit(X, y) + X_unseen = make_df( { "var_A": ["A", "ZZZ", "YYY"], "var_B": ["C", "YYY", "ZZZ"], @@ -285,56 +396,65 @@ def test_fit_errors_if_new_cat_values_and_unseen_is_raise_param(df_enc): ) # new categories will raise an error with pytest.raises(ValueError, match=msg): - encoder.transform(X) + encoder.transform(X_unseen) -def test_inverse_transform_when_no_unseen(): - X = pd.DataFrame({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) - y = pd.Series([0, 0, 1, 1, 1, 1, 0]) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_inverse_transform_when_no_unseen(make_df): + X = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) + y = [0, 0, 1, 1, 1, 1, 0] enc = DecisionTreeEncoder(regression=False) enc.fit(X, y) dft = enc.transform(X) - pd.testing.assert_frame_equal(enc.inverse_transform(dft), X) + Xi = enc.inverse_transform(dft) + _assert_values(Xi, {"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) -def test_inverse_transform_when_ignore_unseen(): - X = pd.DataFrame({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) - y = pd.Series([0, 0, 1, 1, 1, 1, 0]) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_inverse_transform_when_ignore_unseen(make_df): + X = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) + y = [0, 0, 1, 1, 1, 1, 0] enc = DecisionTreeEncoder(regression=False, unseen="ignore") enc.fit(X, y) - df1 = pd.DataFrame({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "frog"]}) - df2 = pd.DataFrame({"words": ["dog", "dog", "dog", "cat", "cat", "cat", np.nan]}) + df1 = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "frog"]}) dft = enc.transform(df1) - pd.testing.assert_frame_equal(enc.inverse_transform(dft), df2) + Xi = enc.inverse_transform(dft) + _assert_values( + Xi, {"words": ["dog", "dog", "dog", "cat", "cat", "cat", np.nan]} + ) -def test_inverse_transform_when_encode_unseen(): - X = pd.DataFrame({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) - y = pd.Series([0, 0, 1, 1, 1, 1, 0]) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_inverse_transform_when_encode_unseen(make_df): + X = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) + y = [0, 0, 1, 1, 1, 1, 0] enc = DecisionTreeEncoder(regression=False, unseen="encode", fill_value=1000) enc.fit(X, y) - df1 = pd.DataFrame({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "frog"]}) - df2 = pd.DataFrame({"words": ["dog", "dog", "dog", "cat", "cat", "cat", np.nan]}) + df1 = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "frog"]}) dft = enc.transform(df1) - pd.testing.assert_frame_equal(enc.inverse_transform(dft), df2) + Xi = enc.inverse_transform(dft) + _assert_values( + Xi, {"words": ["dog", "dog", "dog", "cat", "cat", "cat", np.nan]} + ) -def test_inverse_transform_raises_non_fitted_error(): - X = pd.DataFrame({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) - y = pd.Series([0, 0, 1, 1, 1, 1, 0]) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_inverse_transform_raises_non_fitted_error(make_df): + X = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) + y = [0, 0, 1, 1, 1, 1, 0] enc = DecisionTreeEncoder() # Test when fit is not called prior to transform. with pytest.raises(NotFittedError): enc.inverse_transform(X) - X.loc[len(X) - 1] = np.nan + X_na = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", None]}) with pytest.raises(ValueError): - enc.fit(X, y) + enc.fit(X_na, y) # Test when fit is not called prior to transform. with pytest.raises(NotFittedError): - enc.inverse_transform(X) + enc.inverse_transform(X_na) From 592e0a42bfecbb2fc207e753c0ad9d07ab221b94 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Mon, 31 Aug 2026 00:49:22 +0200 Subject: [PATCH 2/8] Adapt DecisionTreeEncoder to narwhals-returning check_X check_X_y now returns a narwhals frame, so bind that 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 redundant nw.from_native(X) in fit(); the parallel _fit_one_variable calls reuse nw_X from check_X_y. In transform(), bind _check_transform_input_and_state to nw_X, keep native X for _check_contains_na, and pass nw_X to _encode (which now expects narwhals). Co-Authored-By: Claude Sonnet 5 --- feature_engine/encoding/decision_tree.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/feature_engine/encoding/decision_tree.py b/feature_engine/encoding/decision_tree.py index 9474f9105..0ff35d8a4 100644 --- a/feature_engine/encoding/decision_tree.py +++ b/feature_engine/encoding/decision_tree.py @@ -310,7 +310,7 @@ def fit(self, X: IntoDataFrame, y: IntoSeries): The target variable. Required to train the decision tree and for ordered ordinal encoding. """ - X, y = check_X_y(X, y) + nw_X, y = check_X_y(X, y) # confirm model type and target variables are compatible. if self.regression is True: @@ -337,8 +337,6 @@ def fit(self, X: IntoDataFrame, y: IntoSeries): self._get_feature_names_in(X) return self - nw_X = nw.from_native(X, eager_only=True) - # only needed for "ordered": pairs the target with X once so every # variable's group_by below can reuse it, instead of rebuilding it # per variable. @@ -382,9 +380,9 @@ def transform(self, X: IntoDataFrame) -> IntoDataFrame: X_new: dataframe of shape = [n_samples, n_features]. Dataframe with variables encoded with decision tree predictions. """ - X = self._check_transform_input_and_state(X) + nw_X = self._check_transform_input_and_state(X) _check_contains_na(X, self.variables_) - X = self._encode(X) + X = self._encode(nw_X) return X From 9d76c81bb90a08ed49a0eef9adeb218c96cb4f15 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Mon, 14 Sep 2026 22:19:45 +0200 Subject: [PATCH 3/8] Use shared backend test fixtures and helpers in DecisionTreeEncoder tests Replace the file-local _to_backend/_assert_values helpers with the shared test structure: make_df and data_enc* fixtures, y built with make_series on the backend under test, isinstance(X, make_df) plus to_dict() checks, and pytest.raises/warns(match=re.escape(msg)). Add a test passing the target as a list and as a numpy array, which take a different code path than a Series. Co-Authored-By: Claude Opus 5 --- .../test_decision_tree_encoder.py | 280 +++++++----------- 1 file changed, 112 insertions(+), 168 deletions(-) diff --git a/tests/test_encoding/test_decision_tree_encoder.py b/tests/test_encoding/test_decision_tree_encoder.py index bf31b342c..5eb59625d 100644 --- a/tests/test_encoding/test_decision_tree_encoder.py +++ b/tests/test_encoding/test_decision_tree_encoder.py @@ -1,48 +1,29 @@ -import math 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.encoding import DecisionTreeEncoder - - -def _to_backend(df: pd.DataFrame, make_df): - """Rebuild a pandas fixture dataframe on the requested backend. - - Swaps float NaN for None in string columns - polars (unlike pandas) - rejects a float NaN mixed into an otherwise-string column. - """ - data = {} - for col in df.columns: - values = df[col].tolist() - if any(isinstance(v, str) for v in values): - values = [ - None if isinstance(v, float) and math.isnan(v) else v for v in values - ] - data[col] = values - return make_df(data) - - -def _assert_values(X, expected: dict) -> None: - """NaN-aware, backend-agnostic comparison of a dataframe's contents.""" - result = nw.from_native(X, eager_only=True).to_dict(as_series=False) - assert list(result.keys()) == list(expected.keys()) - for col, exp_values in expected.items(): - got_values = result[col] - assert len(got_values) == len(exp_values) - for got, exp in zip(got_values, exp_values): - if isinstance(exp, float) and math.isnan(exp): - assert got is None or (isinstance(got, float) and math.isnan(got)) - elif isinstance(exp, float): - assert got == pytest.approx(exp) - else: - assert got == exp +from tests.backend_helpers import make_series, to_dict + +# Tree: var_A <= 1.5 -> 0.25 else 0.5 +# Tree: var_B <= 0.5 -> 0.2 else 0.4 +ENCODED = { + "var_A": [0.25] * 16 + [0.5] * 4, + "var_B": [0.2] * 10 + [0.4] * 10, +} +ENCODED_REGRESSION = { + "var_A": [0.034348] * 6 + [-0.024679] * 10 + [-0.075473] * 4, + "var_B": [0.044806] * 10 + [-0.079066] * 10, +} + + +def _rounded(X, decimals=6): + return { + col: [round(v, decimals) for v in values] for col, values in to_dict(X).items() + } # init parameters @@ -52,7 +33,7 @@ def test_error_if_encoding_method_not_permitted_value(enc_method): "`encoding_method` takes only values 'ordered' and 'arbitrary'." f" Got {enc_method} instead." ) - with pytest.raises(ValueError, match=msg): + with pytest.raises(ValueError, match=re.escape(msg)): DecisionTreeEncoder(encoding_method=enc_method) @@ -60,11 +41,11 @@ def test_error_if_encoding_method_not_permitted_value(enc_method): "unseen", ["string", False, ("raise", "ignore"), ["ignore"], np.nan] ) def test_error_if_unseen_gets_not_permitted_value(unseen): - msg = re.escape( + msg = ( "Parameter `unseen` takes only values ignore, raise, encode. " - rf"Got {unseen} instead." + f"Got {unseen} instead." ) - with pytest.raises(ValueError, match=msg): + with pytest.raises(ValueError, match=re.escape(msg)): DecisionTreeEncoder(unseen=unseen) @@ -73,14 +54,14 @@ def test_error_if_unseen_is_encode_and_fill_value_is_none(): "When `unseen='encode'` you need to pass a number to `fill_value`. " f"Got {None} instead." ) - with pytest.raises(ValueError, match=msg): + with pytest.raises(ValueError, match=re.escape(msg)): DecisionTreeEncoder(unseen="encode", fill_value=None) @pytest.mark.parametrize("precision", ["string", 0.1, -1, np.nan]) def test_error_if_precision_gets_not_permitted_value(precision): msg = "Parameter `precision` takes integers or None. " f"Got {precision} instead." - with pytest.raises(ValueError, match=msg): + with pytest.raises(ValueError, match=re.escape(msg)): DecisionTreeEncoder(precision=precision) @@ -105,10 +86,9 @@ def test_init_param_assignment( # fit attributes -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_encoding_dictionary(df_enc, make_df): - X = _to_backend(df_enc[["var_A", "var_B"]], make_df) - y = df_enc["target"].tolist() +def test_encoding_dictionary(make_df, data_enc): + X = make_df(data_enc)[["var_A", "var_B"]] + y = make_series(make_df, data_enc["target"]) encoder = DecisionTreeEncoder(regression=False) encoder.fit(X, y) @@ -122,10 +102,9 @@ def test_encoding_dictionary(df_enc, make_df): assert encoder.encoder_dict_ == expected_encodings -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_ordered_encoding_dictionary(df_enc, make_df): - X = _to_backend(df_enc[["var_A", "var_B"]], make_df) - y = df_enc["target"].tolist() +def test_ordered_encoding_dictionary(make_df, data_enc): + X = make_df(data_enc)[["var_A", "var_B"]] + y = make_series(make_df, data_enc["target"]) encoder = DecisionTreeEncoder(regression=False, encoding_method="ordered") encoder.fit(X, y) @@ -143,10 +122,9 @@ def test_ordered_encoding_dictionary(df_enc, make_df): assert encoder.encoder_dict_ == expected_encodings -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_precision(df_enc, make_df): - X = _to_backend(df_enc[["var_A", "var_B"]], make_df) - y = df_enc["target"].tolist() +def test_precision(make_df, data_enc): + X = make_df(data_enc)[["var_A", "var_B"]] + y = make_series(make_df, data_enc["target"]) encoder = DecisionTreeEncoder(regression=False, precision=1) encoder.fit(X, y) @@ -160,27 +138,36 @@ def test_precision(df_enc, make_df): assert encoder.encoder_dict_ == expected_encodings -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_classification(df_enc, make_df): - X = _to_backend(df_enc[["var_A", "var_B"]], make_df) - y = df_enc["target"].tolist() +def test_classification(make_df, data_enc): + X = make_df(data_enc)[["var_A", "var_B"]] + y = make_series(make_df, data_enc["target"]) encoder = DecisionTreeEncoder(regression=False) encoder.fit(X, y) Xt = encoder.transform(X) - expected = { - "var_A": [0.25] * 16 + [0.5] * 4, # Tree: var_A <= 1.5 -> 0.25 else 0.5 - "var_B": [0.2] * 10 + [0.4] * 10, # Tree: var_B <= 0.5 -> 0.2 else 0.4 - } - _assert_values(Xt, expected) + assert isinstance(Xt, make_df) + assert to_dict(Xt) == ENCODED + +@pytest.mark.parametrize("to_target", [list, np.array]) +def test_target_as_list_or_array(make_df, data_enc, to_target): + # a list or numpy array target takes a different code path than a Series + X = make_df(data_enc)[["var_A", "var_B"]] + y = to_target(data_enc["target"]) -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_regression(df_enc, make_df): - X = _to_backend(df_enc[["var_A", "var_B"]], make_df) + encoder = DecisionTreeEncoder(regression=False) + encoder.fit(X, y) + Xt = encoder.transform(X) + + assert isinstance(Xt, make_df) + assert to_dict(Xt) == ENCODED + + +def test_regression(make_df, data_enc): + X = make_df(data_enc)[["var_A", "var_B"]] random = np.random.RandomState(42) - y = random.normal(0, 0.1, len(df_enc)) + y = make_series(make_df, random.normal(0, 0.1, len(data_enc["target"]))) encoder = DecisionTreeEncoder( regression=True, random_state=random, @@ -188,39 +175,27 @@ def test_regression(df_enc, make_df): encoder.fit(X, y) Xt = encoder.transform(X) - expected = { - "var_A": ( - [0.034348] * 6 + [-0.024679] * 10 + [-0.075473] * 4 - ), # Tree: var_A <= 1.5 -> 0.25 else 0.5 - "var_B": [0.044806] * 10 + [-0.079066] * 10, - } - nw_Xt = nw.from_native(Xt, eager_only=True) - rounded = { - col: [round(v, 6) for v in nw_Xt.get_column(col).to_list()] - for col in ["var_A", "var_B"] - } - assert rounded == expected + assert isinstance(Xt, make_df) + assert _rounded(Xt) == ENCODED_REGRESSION -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_fit_raises_error_if_df_contains_na(df_enc_na, make_df): - X = _to_backend(df_enc_na[["var_A", "var_B"]], make_df) - y = df_enc_na["target"].tolist() +def test_fit_raises_error_if_df_contains_na(make_df, data_enc_na): + X = make_df(data_enc_na)[["var_A", "var_B"]] + y = make_series(make_df, data_enc_na["target"]) encoder = DecisionTreeEncoder(regression=False) msg = ( "Some of the variables in the dataset contain NaN. Check and " "remove those before using this transformer." ) - with pytest.raises(ValueError, match=msg): + with pytest.raises(ValueError, match=re.escape(msg)): encoder.fit(X, y) -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_transform_raises_error_if_df_contains_na(df_enc, df_enc_na, make_df): - X = _to_backend(df_enc[["var_A", "var_B"]], make_df) - X_na = _to_backend(df_enc_na[["var_A", "var_B"]], make_df) - y = df_enc["target"].tolist() +def test_transform_raises_error_if_df_contains_na(make_df, data_enc, data_enc_na): + X = make_df(data_enc)[["var_A", "var_B"]] + X_na = make_df(data_enc_na)[["var_A", "var_B"]] + y = make_series(make_df, data_enc["target"]) encoder = DecisionTreeEncoder(regression=False) encoder.fit(X, y) @@ -228,14 +203,13 @@ def test_transform_raises_error_if_df_contains_na(df_enc, df_enc_na, make_df): "Some of the variables in the dataset contain NaN. Check and " "remove those before using this transformer." ) - with pytest.raises(ValueError, match=msg): + with pytest.raises(ValueError, match=re.escape(msg)): encoder.transform(X_na) -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_classification_ignore_format(df_enc_numeric, make_df): - X = _to_backend(df_enc_numeric[["var_A", "var_B"]], make_df) - y = df_enc_numeric["target"].tolist() +def test_classification_ignore_format(make_df, data_enc_numeric): + X = make_df(data_enc_numeric)[["var_A", "var_B"]] + y = make_series(make_df, data_enc_numeric["target"]) encoder = DecisionTreeEncoder( regression=False, @@ -244,18 +218,14 @@ def test_classification_ignore_format(df_enc_numeric, make_df): encoder.fit(X, y) Xt = encoder.transform(X) - expected = { - "var_A": [0.25] * 16 + [0.5] * 4, # Tree: var_A <= 1.5 -> 0.25 else 0.5 - "var_B": [0.2] * 10 + [0.4] * 10, # Tree: var_B <= 0.5 -> 0.2 else 0.4 - } - _assert_values(Xt, expected) + assert isinstance(Xt, make_df) + assert to_dict(Xt) == ENCODED -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_regression_ignore_format(df_enc_numeric, make_df): - X = _to_backend(df_enc_numeric[["var_A", "var_B"]], make_df) +def test_regression_ignore_format(make_df, data_enc_numeric): + X = make_df(data_enc_numeric)[["var_A", "var_B"]] random = np.random.RandomState(42) - y = random.normal(0, 0.1, len(df_enc_numeric)) + y = make_series(make_df, random.normal(0, 0.1, len(data_enc_numeric["target"]))) encoder = DecisionTreeEncoder( regression=True, random_state=random, @@ -264,18 +234,8 @@ def test_regression_ignore_format(df_enc_numeric, make_df): encoder.fit(X, y) Xt = encoder.transform(X) - expected = { - "var_A": ( - [0.034348] * 6 + [-0.024679] * 10 + [-0.075473] * 4 - ), # Tree: var_A <= 1.5 -> 0.25 else 0.5 - "var_B": [0.044806] * 10 + [-0.079066] * 10, - } - nw_Xt = nw.from_native(Xt, eager_only=True) - rounded = { - col: [round(v, 6) for v in nw_Xt.get_column(col).to_list()] - for col in ["var_A", "var_B"] - } - assert rounded == expected + assert isinstance(Xt, make_df) + assert _rounded(Xt) == ENCODED_REGRESSION def test_variables_cast_as_category(df_enc_category_dtypes): @@ -292,10 +252,9 @@ def test_variables_cast_as_category(df_enc_category_dtypes): assert X["var_A"].dtypes == float -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_error_when_regression_is_true_and_target_is_binary(df_enc, make_df): - X = _to_backend(df_enc[["var_A", "var_B"]], make_df) - y = df_enc["target"].tolist() +def test_error_when_regression_is_true_and_target_is_binary(make_df, data_enc): + X = make_df(data_enc)[["var_A", "var_B"]] + y = make_series(make_df, data_enc["target"]) encoder = DecisionTreeEncoder(regression=True) msg = ( @@ -303,15 +262,14 @@ def test_error_when_regression_is_true_and_target_is_binary(df_enc, make_df): "allowed by this transformer. Check the target values " "or set regression to False." ) - with pytest.raises(ValueError, match=msg): + with pytest.raises(ValueError, match=re.escape(msg)): encoder.fit(X, y) -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_error_when_regression_is_false_and_target_is_continuous(df_enc, make_df): - X = _to_backend(df_enc[["var_A", "var_B"]], make_df) +def test_error_when_regression_is_false_and_target_is_continuous(make_df, data_enc): + X = make_df(data_enc)[["var_A", "var_B"]] random = np.random.RandomState(42) - y = random.normal(0, 10, len(df_enc)) + y = make_series(make_df, random.normal(0, 10, len(data_enc["target"]))) encoder = DecisionTreeEncoder(regression=False) # the error message comes from sklearn api - won't test with pytest.raises(ValueError): @@ -330,10 +288,9 @@ def test_assigns_param_grid(grid): assert encoder._assign_param_grid() == grid -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_unseen_is_encode(df_enc, make_df): - X = _to_backend(df_enc[["var_A", "var_B"]], make_df) - y = df_enc["target"].tolist() +def test_unseen_is_encode(make_df, data_enc): + X = make_df(data_enc)[["var_A", "var_B"]] + y = make_series(make_df, data_enc["target"]) encoder = DecisionTreeEncoder(unseen="encode", regression=False, fill_value=-1) encoder.fit(X, y) @@ -344,19 +301,15 @@ def test_unseen_is_encode(df_enc, make_df): "var_B": ["C", "YYY", "ZZZ"], } ) - expected = { - "var_A": [0.25, -1, -1], - "var_B": [0.4, -1, -1], - } - Xt = encoder.transform(X_unseen_input) - _assert_values(Xt, expected) + + assert isinstance(Xt, make_df) + assert to_dict(Xt) == {"var_A": [0.25, -1, -1], "var_B": [0.4, -1, -1]} -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_unseen_is_ignore(df_enc, make_df): - X = _to_backend(df_enc[["var_A", "var_B"]], make_df) - y = df_enc["target"].tolist() +def test_unseen_is_ignore(make_df, data_enc): + X = make_df(data_enc)[["var_A", "var_B"]] + y = make_series(make_df, data_enc["target"]) encoder = DecisionTreeEncoder(unseen="ignore", regression=False) encoder.fit(X, y) @@ -367,19 +320,15 @@ def test_unseen_is_ignore(df_enc, make_df): "var_B": ["C", "YYY", "ZZZ"], } ) - expected = { - "var_A": [0.25, np.nan, np.nan], - "var_B": [0.4, np.nan, np.nan], - } - Xt = encoder.transform(X_unseen_input) - _assert_values(Xt, expected) + + assert isinstance(Xt, make_df) + assert to_dict(Xt) == {"var_A": [0.25, None, None], "var_B": [0.4, None, None]} -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_fit_errors_if_new_cat_values_and_unseen_is_raise_param(df_enc, make_df): - X = _to_backend(df_enc[["var_A", "var_B"]], make_df) - y = df_enc["target"].tolist() +def test_fit_errors_if_new_cat_values_and_unseen_is_raise_param(make_df, data_enc): + X = make_df(data_enc)[["var_A", "var_B"]] + y = make_series(make_df, data_enc["target"]) encoder = DecisionTreeEncoder(unseen="raise", regression=False) encoder.fit(X, y) @@ -389,61 +338,56 @@ def test_fit_errors_if_new_cat_values_and_unseen_is_raise_param(df_enc, make_df) "var_B": ["C", "YYY", "ZZZ"], } ) - var_ls = "var_A, var_B" msg = ( "During the encoding, NaN values were introduced in the " - rf"feature\(s\) {var_ls}." + "feature(s) var_A, var_B." ) # new categories will raise an error - with pytest.raises(ValueError, match=msg): + with pytest.raises(ValueError, match=re.escape(msg)): encoder.transform(X_unseen) -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) def test_inverse_transform_when_no_unseen(make_df): - X = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) - y = [0, 0, 1, 1, 1, 1, 0] + words = ["dog", "dog", "dog", "cat", "cat", "cat", "bird"] + X = make_df({"words": words}) + y = make_series(make_df, [0, 0, 1, 1, 1, 1, 0]) enc = DecisionTreeEncoder(regression=False) enc.fit(X, y) dft = enc.transform(X) Xi = enc.inverse_transform(dft) - _assert_values(Xi, {"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) + assert isinstance(Xi, make_df) + assert to_dict(Xi) == {"words": words} -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) def test_inverse_transform_when_ignore_unseen(make_df): X = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) - y = [0, 0, 1, 1, 1, 1, 0] + y = make_series(make_df, [0, 0, 1, 1, 1, 1, 0]) enc = DecisionTreeEncoder(regression=False, unseen="ignore") enc.fit(X, y) df1 = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "frog"]}) dft = enc.transform(df1) Xi = enc.inverse_transform(dft) - _assert_values( - Xi, {"words": ["dog", "dog", "dog", "cat", "cat", "cat", np.nan]} - ) + assert isinstance(Xi, make_df) + assert to_dict(Xi) == {"words": ["dog", "dog", "dog", "cat", "cat", "cat", None]} -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) def test_inverse_transform_when_encode_unseen(make_df): X = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) - y = [0, 0, 1, 1, 1, 1, 0] + y = make_series(make_df, [0, 0, 1, 1, 1, 1, 0]) enc = DecisionTreeEncoder(regression=False, unseen="encode", fill_value=1000) enc.fit(X, y) df1 = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "frog"]}) dft = enc.transform(df1) Xi = enc.inverse_transform(dft) - _assert_values( - Xi, {"words": ["dog", "dog", "dog", "cat", "cat", "cat", np.nan]} - ) + assert isinstance(Xi, make_df) + assert to_dict(Xi) == {"words": ["dog", "dog", "dog", "cat", "cat", "cat", None]} -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) def test_inverse_transform_raises_non_fitted_error(make_df): X = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) - y = [0, 0, 1, 1, 1, 1, 0] + y = make_series(make_df, [0, 0, 1, 1, 1, 1, 0]) enc = DecisionTreeEncoder() # Test when fit is not called prior to transform. From 0f338809bd5e06284167a0dcb85be3e7631e7fe1 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Tue, 15 Sep 2026 12:06:06 +0200 Subject: [PATCH 4/8] Use frame_to_dict after the shared helper rename in #1045 Co-Authored-By: Claude Opus 5 --- .../test_decision_tree_encoder.py | 28 ++++++++++++------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/tests/test_encoding/test_decision_tree_encoder.py b/tests/test_encoding/test_decision_tree_encoder.py index 5eb59625d..fce8098cc 100644 --- a/tests/test_encoding/test_decision_tree_encoder.py +++ b/tests/test_encoding/test_decision_tree_encoder.py @@ -6,7 +6,7 @@ from sklearn.exceptions import NotFittedError from feature_engine.encoding import DecisionTreeEncoder -from tests.backend_helpers import make_series, to_dict +from tests.backend_helpers import make_series, frame_to_dict # Tree: var_A <= 1.5 -> 0.25 else 0.5 # Tree: var_B <= 0.5 -> 0.2 else 0.4 @@ -22,7 +22,8 @@ def _rounded(X, decimals=6): return { - col: [round(v, decimals) for v in values] for col, values in to_dict(X).items() + col: [round(v, decimals) for v in values] + for col, values in frame_to_dict(X).items() } @@ -147,7 +148,7 @@ def test_classification(make_df, data_enc): Xt = encoder.transform(X) assert isinstance(Xt, make_df) - assert to_dict(Xt) == ENCODED + assert frame_to_dict(Xt) == ENCODED @pytest.mark.parametrize("to_target", [list, np.array]) @@ -161,7 +162,7 @@ def test_target_as_list_or_array(make_df, data_enc, to_target): Xt = encoder.transform(X) assert isinstance(Xt, make_df) - assert to_dict(Xt) == ENCODED + assert frame_to_dict(Xt) == ENCODED def test_regression(make_df, data_enc): @@ -219,7 +220,7 @@ def test_classification_ignore_format(make_df, data_enc_numeric): Xt = encoder.transform(X) assert isinstance(Xt, make_df) - assert to_dict(Xt) == ENCODED + assert frame_to_dict(Xt) == ENCODED def test_regression_ignore_format(make_df, data_enc_numeric): @@ -304,7 +305,7 @@ def test_unseen_is_encode(make_df, data_enc): Xt = encoder.transform(X_unseen_input) assert isinstance(Xt, make_df) - assert to_dict(Xt) == {"var_A": [0.25, -1, -1], "var_B": [0.4, -1, -1]} + assert frame_to_dict(Xt) == {"var_A": [0.25, -1, -1], "var_B": [0.4, -1, -1]} def test_unseen_is_ignore(make_df, data_enc): @@ -323,7 +324,10 @@ def test_unseen_is_ignore(make_df, data_enc): Xt = encoder.transform(X_unseen_input) assert isinstance(Xt, make_df) - assert to_dict(Xt) == {"var_A": [0.25, None, None], "var_B": [0.4, None, None]} + assert frame_to_dict(Xt) == { + "var_A": [0.25, None, None], + "var_B": [0.4, None, None], + } def test_fit_errors_if_new_cat_values_and_unseen_is_raise_param(make_df, data_enc): @@ -356,7 +360,7 @@ def test_inverse_transform_when_no_unseen(make_df): dft = enc.transform(X) Xi = enc.inverse_transform(dft) assert isinstance(Xi, make_df) - assert to_dict(Xi) == {"words": words} + assert frame_to_dict(Xi) == {"words": words} def test_inverse_transform_when_ignore_unseen(make_df): @@ -369,7 +373,9 @@ def test_inverse_transform_when_ignore_unseen(make_df): dft = enc.transform(df1) Xi = enc.inverse_transform(dft) assert isinstance(Xi, make_df) - assert to_dict(Xi) == {"words": ["dog", "dog", "dog", "cat", "cat", "cat", None]} + assert frame_to_dict(Xi) == { + "words": ["dog", "dog", "dog", "cat", "cat", "cat", None] + } def test_inverse_transform_when_encode_unseen(make_df): @@ -382,7 +388,9 @@ def test_inverse_transform_when_encode_unseen(make_df): dft = enc.transform(df1) Xi = enc.inverse_transform(dft) assert isinstance(Xi, make_df) - assert to_dict(Xi) == {"words": ["dog", "dog", "dog", "cat", "cat", "cat", None]} + assert frame_to_dict(Xi) == { + "words": ["dog", "dog", "dog", "cat", "cat", "cat", None] + } def test_inverse_transform_raises_non_fitted_error(make_df): From 67ee9f7abcb6d5df3e756185ae8da162d58df397 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Tue, 15 Sep 2026 13:34:24 +0200 Subject: [PATCH 5/8] Use add_target_to_X in DecisionTreeEncoder, check encoding_method type, tidy tests Co-Authored-By: Claude Opus 5 --- feature_engine/encoding/decision_tree.py | 43 ++++----- .../test_decision_tree_encoder.py | 94 +++++++++++++------ 2 files changed, 84 insertions(+), 53 deletions(-) diff --git a/feature_engine/encoding/decision_tree.py b/feature_engine/encoding/decision_tree.py index 0ff35d8a4..e731e5ebb 100644 --- a/feature_engine/encoding/decision_tree.py +++ b/feature_engine/encoding/decision_tree.py @@ -4,7 +4,6 @@ from typing import Dict, List, Optional, Union import narwhals as nw -import narwhals.dependencies as nwd import numpy as np from joblib import Parallel, delayed from narwhals.typing import IntoDataFrame, IntoSeries @@ -31,7 +30,11 @@ ) from feature_engine._docstrings.substitute import Substitution from feature_engine.dataframe_checks import _check_contains_na, check_X_y -from feature_engine.encoding._helper_functions import check_parameter_unseen +from feature_engine.encoding._helper_functions import ( + TARGET_NAME, + add_target_to_X, + check_parameter_unseen, +) from feature_engine.encoding.base_encoder import ( CategoricalInitMixin, CategoricalMethodsMixin, @@ -263,7 +266,10 @@ def __init__( n_jobs: Optional[int] = None, ) -> None: - if encoding_method not in ["ordered", "arbitrary"]: + if not isinstance(encoding_method, str) or encoding_method not in [ + "ordered", + "arbitrary", + ]: raise ValueError( "`encoding_method` takes only values 'ordered' and 'arbitrary'." f" Got {encoding_method} instead." @@ -337,24 +343,11 @@ def fit(self, X: IntoDataFrame, y: IntoSeries): self._get_feature_names_in(X) return self - # only needed for "ordered": pairs the target with X once so every - # variable's group_by below can reuse it, instead of rebuilding it - # per variable. - nw_Xy = None - target_name = "__feature_engine_decision_tree_target__" - if self.encoding_method == "ordered": - if nwd.is_into_series(y): - y_nw = nw.from_native(y, series_only=True).alias(target_name) - else: - y_nw = nw.new_series( - name=target_name, values=y, backend=nw_X.implementation - ) - nw_Xy = nw_X.with_columns(y_nw) + # the target is only needed to order the categories + nw_Xy = add_target_to_X(nw_X, y) if self.encoding_method == "ordered" else None mappings = Parallel(n_jobs=self.n_jobs, prefer="threads")( - delayed(self._fit_one_variable)( - nw_X, nw_Xy, var, y, target_name, param_grid - ) + delayed(self._fit_one_variable)(nw_X, nw_Xy, var, y, param_grid) for var in variables_ ) @@ -387,20 +380,18 @@ def transform(self, X: IntoDataFrame) -> IntoDataFrame: return X def _fit_one_variable( - self, nw_X, nw_Xy, var, y: IntoSeries, target_name: str, param_grid: Dict + self, nw_X, nw_Xy, var, y: IntoSeries, param_grid: Dict ) -> Dict: """Learn the category-to-prediction mapping for one variable: encode its categories to ordinal integers, fit a decision tree on those integers, and predict on each unique category to get its final mapped value.""" if self.encoding_method == "ordered": - # sort by (mean, category): group_by's own order isn't guaranteed - # across backends, and this tie-break on the category itself - # reproduces pandas' groupby(sort=True) + stable sort_values - # behavior for categories with equal target means. + # sort by mean, then category, so ties get the same order + # in every backend categories = ( nw_Xy.group_by(var, drop_null_keys=True) - .agg(nw.col(target_name).mean()) - .sort([target_name, var]) + .agg(nw.col(TARGET_NAME).mean()) + .sort([TARGET_NAME, var]) .get_column(var) .to_list() ) diff --git a/tests/test_encoding/test_decision_tree_encoder.py b/tests/test_encoding/test_decision_tree_encoder.py index fce8098cc..6893ae346 100644 --- a/tests/test_encoding/test_decision_tree_encoder.py +++ b/tests/test_encoding/test_decision_tree_encoder.py @@ -28,7 +28,10 @@ def _rounded(X, decimals=6): # init parameters -@pytest.mark.parametrize("enc_method", ["count", False, 1]) +@pytest.mark.parametrize( + "enc_method", + ["count", "Ordered", "", False, 1, None, ["ordered"], ("arbitrary",)], +) def test_error_if_encoding_method_not_permitted_value(enc_method): msg = ( "`encoding_method` takes only values 'ordered' and 'arbitrary'." @@ -67,23 +70,53 @@ def test_error_if_precision_gets_not_permitted_value(precision): @pytest.mark.parametrize( - "encoding_method,ignore_format,precision,unseen,fill_value", + "params", [ - ("arbitrary", True, 1, "raise", None), - ("ordered", False, 2, "ignore", 1), - ("ordered", False, None, "encode", 0.1), + { + "encoding_method": "arbitrary", + "cv": 3, + "scoring": "neg_mean_squared_error", + "regression": True, + "param_grid": None, + "random_state": None, + "ignore_format": True, + "precision": 1, + "unseen": "raise", + "fill_value": None, + "n_jobs": None, + }, + { + "encoding_method": "ordered", + "cv": 5, + "scoring": "roc_auc", + "regression": False, + "param_grid": {"max_depth": [1, 2]}, + "random_state": 0, + "ignore_format": False, + "precision": None, + "unseen": "encode", + "fill_value": 0.1, + "n_jobs": -1, + }, + { + "encoding_method": "ordered", + "cv": 2, + "scoring": "accuracy", + "regression": False, + "param_grid": {"max_depth": [3]}, + "random_state": 42, + "ignore_format": False, + "precision": 2, + "unseen": "ignore", + "fill_value": 1, + "n_jobs": 2, + }, ], ) -def test_init_param_assignment( - encoding_method, ignore_format, precision, unseen, fill_value -): - DecisionTreeEncoder( - encoding_method=encoding_method, - ignore_format=ignore_format, - precision=precision, - unseen=unseen, - fill_value=fill_value, - ) +def test_init_param_assignment(params): + encoder = DecisionTreeEncoder(**params) + for param, value in params.items(): + assert getattr(encoder, param) == value # fit attributes @@ -110,12 +143,8 @@ def test_ordered_encoding_dictionary(make_df, data_enc): encoder = DecisionTreeEncoder(regression=False, encoding_method="ordered") encoder.fit(X, y) - # ordered ranks: var_A -> B(mean 0.2) < A(mean 0.333) < C(mean 0.5) - # var_B -> A(mean 0.2) < B(mean 0.333) < C(mean 0.5) - # so var_A's ordinal codes are B=0, A=1, C=2 (split at code <= 0.5, - # i.e. B alone); var_B's are A=0, B=1, C=2 (split at code <= 0.5, i.e. A - # alone). Same tree-split logic as the arbitrary-encoding case above, - # applied to the reordered codes. + # codes by target mean: var_A B=0, A=1, C=2 and var_B A=0, B=1, C=2 + # both trees split code 0 from the rest expected_encodings = { "var_A": {"B": 0.2, "A": 0.4, "C": 0.4}, "var_B": {"A": 0.2, "B": 0.4, "C": 0.4}, @@ -272,8 +301,11 @@ def test_error_when_regression_is_false_and_target_is_continuous(make_df, data_e random = np.random.RandomState(42) y = make_series(make_df, random.normal(0, 10, len(data_enc["target"]))) encoder = DecisionTreeEncoder(regression=False) - # the error message comes from sklearn api - won't test - with pytest.raises(ValueError): + msg = ( + "Unknown label type: continuous. Maybe you are trying to fit a classifier, " + "which expects discrete classes on a regression target with continuous values." + ) + with pytest.raises(ValueError, match=re.escape(msg)): encoder.fit(X, y) @@ -396,17 +428,25 @@ def test_inverse_transform_when_encode_unseen(make_df): def test_inverse_transform_raises_non_fitted_error(make_df): X = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", "bird"]}) y = make_series(make_df, [0, 0, 1, 1, 1, 1, 0]) - enc = DecisionTreeEncoder() + enc = DecisionTreeEncoder(regression=False) + msg = ( + "This DecisionTreeEncoder instance is not fitted yet. Call 'fit' with " + "appropriate arguments before using this estimator." + ) + msg_na = ( + "Some of the variables in the dataset contain NaN. Check and " + "remove those before using this transformer." + ) # Test when fit is not called prior to transform. - with pytest.raises(NotFittedError): + with pytest.raises(NotFittedError, match=re.escape(msg)): enc.inverse_transform(X) X_na = make_df({"words": ["dog", "dog", "dog", "cat", "cat", "cat", None]}) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match=re.escape(msg_na)): enc.fit(X_na, y) # Test when fit is not called prior to transform. - with pytest.raises(NotFittedError): + with pytest.raises(NotFittedError, match=re.escape(msg)): enc.inverse_transform(X_na) From ec15072e0799124ce1d197b941d185ac9f79f59c Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Tue, 15 Sep 2026 15:10:26 +0200 Subject: [PATCH 6/8] Shorten n_jobs docstring in DecisionTreeEncoder Co-Authored-By: Claude Opus 5 --- feature_engine/encoding/decision_tree.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/feature_engine/encoding/decision_tree.py b/feature_engine/encoding/decision_tree.py index e731e5ebb..7ce3136a8 100644 --- a/feature_engine/encoding/decision_tree.py +++ b/feature_engine/encoding/decision_tree.py @@ -145,13 +145,9 @@ class DecisionTreeEncoder(CategoricalMethodsMixin, CategoricalInitMixin): The value used to encode unseen categories. Only used when `unseen='encode'`. n_jobs: int, default=None - The number of jobs to run in parallel when training the decision trees - across variables. Trees are fit using threads rather than processes, - since fitting a decision tree releases the GIL for the bulk of its - computation, which avoids the overhead of copying the entire dataframe - to separate worker processes. `None` means 1, i.e. sequential training - (this transformer's original behaviour); `-1` means using all available - processors. + The number of jobs to run in parallel. `fit` is parallelized over the variables, + training one decision tree per variable. `None` means 1 unless in a + `joblib.parallel_backend` context. `-1` means using all processors. Attributes ---------- From 12cb0b9be1eebeb889559e04d16ee10fa2413986 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Fri, 18 Sep 2026 14:35:37 +0200 Subject: [PATCH 7/8] Build DecisionTreeEncoder on OrdinalEncoder and DecisionTreeDiscretiser Co-Authored-By: Claude Opus 5 --- feature_engine/encoding/decision_tree.py | 108 ++++++++--------------- 1 file changed, 37 insertions(+), 71 deletions(-) diff --git a/feature_engine/encoding/decision_tree.py b/feature_engine/encoding/decision_tree.py index 7ce3136a8..ee6dfbbd8 100644 --- a/feature_engine/encoding/decision_tree.py +++ b/feature_engine/encoding/decision_tree.py @@ -1,14 +1,12 @@ # Authors: Soledad Galli # License: BSD 3 clause -from typing import Dict, List, Optional, Union +from typing import List, Optional, Union import narwhals as nw import numpy as np -from joblib import Parallel, delayed from narwhals.typing import IntoDataFrame, IntoSeries -from sklearn.model_selection import GridSearchCV -from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor +from sklearn.pipeline import Pipeline from sklearn.utils.multiclass import check_classification_targets, type_of_target from feature_engine._docstrings.fit_attributes import ( @@ -30,15 +28,13 @@ ) from feature_engine._docstrings.substitute import Substitution from feature_engine.dataframe_checks import _check_contains_na, check_X_y -from feature_engine.encoding._helper_functions import ( - TARGET_NAME, - add_target_to_X, - check_parameter_unseen, -) +from feature_engine.discretisation import DecisionTreeDiscretiser +from feature_engine.encoding._helper_functions import check_parameter_unseen from feature_engine.encoding.base_encoder import ( CategoricalInitMixin, CategoricalMethodsMixin, ) +from feature_engine.encoding.ordinal import OrdinalEncoder from feature_engine.tags import _return_tags _unseen_docstring = ( @@ -339,18 +335,43 @@ def fit(self, X: IntoDataFrame, y: IntoSeries): self._get_feature_names_in(X) return self - # the target is only needed to order the categories - nw_Xy = add_target_to_X(nw_X, y) if self.encoding_method == "ordered" else None - - mappings = Parallel(n_jobs=self.n_jobs, prefer="threads")( - delayed(self._fit_one_variable)(nw_X, nw_Xy, var, y, param_grid) - for var in variables_ + encoder = OrdinalEncoder( + encoding_method=self.encoding_method, + variables=variables_, + missing_values="raise", + ignore_format=self.ignore_format, + ) + tree = DecisionTreeDiscretiser( + variables=variables_, + cv=self.cv, + scoring=self.scoring, + param_grid=param_grid, + regression=self.regression, + random_state=self.random_state, + n_jobs=self.n_jobs, ) + Xt = Pipeline([("encoder", encoder), ("tree", tree)]).fit_transform(X, y) + nw_Xt = nw.from_native(Xt, eager_only=True) + + # map each category to the prediction of its tree + encoder_dict_ = {} + for var in variables_: + pairs = ( + nw_X.get_column(var) + .alias("__category__") + .to_frame() + .with_columns(nw_Xt.get_column(var).alias("__prediction__")) + .unique() + ) + preds = pairs["__prediction__"].to_numpy() + if self.precision is not None: + preds = np.round(preds, self.precision) + encoder_dict_[var] = dict(zip(pairs["__category__"].to_list(), preds)) if self.unseen == "encode": self._unseen = self.fill_value - self.encoder_dict_ = dict(zip(variables_, mappings)) + self.encoder_dict_ = encoder_dict_ self.variables_ = variables_ self._get_feature_names_in(X) return self @@ -375,61 +396,6 @@ def transform(self, X: IntoDataFrame) -> IntoDataFrame: return X - def _fit_one_variable( - self, nw_X, nw_Xy, var, y: IntoSeries, param_grid: Dict - ) -> Dict: - """Learn the category-to-prediction mapping for one variable: encode its - categories to ordinal integers, fit a decision tree on those integers, and - predict on each unique category to get its final mapped value.""" - if self.encoding_method == "ordered": - # sort by mean, then category, so ties get the same order - # in every backend - categories = ( - nw_Xy.group_by(var, drop_null_keys=True) - .agg(nw.col(TARGET_NAME).mean()) - .sort([TARGET_NAME, var]) - .get_column(var) - .to_list() - ) - else: - categories = nw_X.get_column(var).unique(maintain_order=True).to_list() - - ordinal_map = {k: i for i, k in enumerate(categories, 0)} - - X_sub = nw_X.get_column(var).replace_strict(ordinal_map).to_frame().to_native() - estimator = self._fit_one_tree(X_sub, y, param_grid) - - # predict directly on the (few) unique ordinal codes instead of the - # full column: the tree's prediction for a category depends only on - # its ordinal code, so this gives identical results far cheaper. - X_pred = nw.new_series( - var, list(range(len(categories))), backend=nw_X.implementation - ).to_frame().to_native() - - if self.regression is True: - preds = estimator.predict(X_pred) - else: - preds = estimator.predict_proba(X_pred)[:, 1] - - if self.precision is not None: - preds = np.round(preds, self.precision) - - return dict(zip(categories, preds)) - - def _fit_one_tree(self, X_sub: IntoDataFrame, y: IntoSeries, param_grid: Dict): - """Instantiate and fit one decision tree on one variable's ordinal-encoded - values.""" - if self.regression is True: - model = DecisionTreeRegressor(random_state=self.random_state) - else: - model = DecisionTreeClassifier(random_state=self.random_state) - - tree_model = GridSearchCV( - model, cv=self.cv, scoring=self.scoring, param_grid=param_grid - ) - tree_model.fit(X_sub, y) - return tree_model - def _assign_param_grid(self): if self.param_grid: param_grid = self.param_grid From 10e728be4ab30bb5ef18dce4a203ac4bd16867ff Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Fri, 18 Sep 2026 14:36:05 +0200 Subject: [PATCH 8/8] Test DecisionTreeEncoder with integer column names Co-Authored-By: Claude Opus 5 --- tests/test_encoding/test_decision_tree_encoder.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/test_encoding/test_decision_tree_encoder.py b/tests/test_encoding/test_decision_tree_encoder.py index 6893ae346..0cd7d48cc 100644 --- a/tests/test_encoding/test_decision_tree_encoder.py +++ b/tests/test_encoding/test_decision_tree_encoder.py @@ -282,6 +282,20 @@ def test_variables_cast_as_category(df_enc_category_dtypes): assert X["var_A"].dtypes == float +def test_integer_column_names(data_enc): + # integer column names are pandas-only + X = pd.DataFrame({0: data_enc["var_A"], 1: data_enc["var_B"]}) + y = pd.Series(data_enc["target"]) + + encoder = DecisionTreeEncoder(regression=False).fit(X, y) + expected = DecisionTreeEncoder(regression=False).fit(X.rename(columns=str), y) + + assert encoder.encoder_dict_ == { + 0: expected.encoder_dict_["0"], + 1: expected.encoder_dict_["1"], + } + + def test_error_when_regression_is_true_and_target_is_binary(make_df, data_enc): X = make_df(data_enc)[["var_A", "var_B"]] y = make_series(make_df, data_enc["target"])