From 7f9bb2b52f9a51414f687564668b3f14ddcc5856 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Mon, 24 Aug 2026 20:01:22 +0200 Subject: [PATCH 1/7] Migrate DecisionTreeFeatures to narwhals, add polars support Follows the same pandas-native / narwhals-generic split established for GeoDistanceFeatures (this transformer also reimplements fit()/transform() directly, not via BaseCreation): .columns extraction, column reorder, prediction-column assignment, and drop_original all split by backend, consistent with every other operation in this module that's been benchmarked as a real (not minimal) loss when routed through narwhals on pandas. Confirmed empirically before designing: sklearn's DecisionTreeRegressor/ Classifier and GridSearchCV accept a polars DataFrame directly for both fit() and predict()/predict_proba(), so the actual tree training/inference calls are unchanged - only the surrounding column selection, extraction, and reassembly needed migrating. Fixed a pre-existing bug found while rewriting the exact code path it lived in: single-feature combos with an integer column name (e.g. DecisionTreeFeatures(features_to_combine=1) on a dataframe with columns 0, 1, ...) crashed, since the original `isinstance(features, str)` check missed the int case and fell through to plain X[features] indexing, which returns a 1D Series rather than the 2D input sklearn requires. Widened to isinstance(features, (str, int)); verified the same single-feature narwhals path (get_column().to_frame()) already handles both cleanly. Regression, binary classification, and multiclass classification paths all verified to produce identical predictions between pandas and polars input. return_empty=True + polars remains untestable here too (same nw.col([]) bug in dataframe_checks.py found during CyclicalFeatures, still tabled) - this is the second transformer it blocks. docs/user_guide/creation/DecisionTreeFeatures.rst is large (511 lines) and built around actual cross-validated tree fitting on the real California housing dataset across many sections - re-verified the cheap, deterministic parts (the raw data table) but did not re-run every tree-fitting example given the cost of repeated grid-search CV fits; unlike the other three creation-module docs this pass touched, the rest of this file's numbers are unverified. Added a self-contained "With polars" section using simple synthetic data instead, fully verified. --- .../creation/DecisionTreeFeatures.rst | 47 ++++++ .../creation/decision_tree_features.py | 148 ++++++++++++------ .../test_decision_tree_features.py | 98 ++++++++++++ 3 files changed, 244 insertions(+), 49 deletions(-) diff --git a/docs/user_guide/creation/DecisionTreeFeatures.rst b/docs/user_guide/creation/DecisionTreeFeatures.rst index 56c6a4056..a7da7115b 100644 --- a/docs/user_guide/creation/DecisionTreeFeatures.rst +++ b/docs/user_guide/creation/DecisionTreeFeatures.rst @@ -485,6 +485,53 @@ are not there: 2670 1.843904 15709 1.843904 +With polars +----------- + +:class:`DecisionTreeFeatures()` works in the same way with a polars dataframe: + +.. code:: python + + import polars as pl + from feature_engine.creation import DecisionTreeFeatures + + X = pl.DataFrame({ + "Age": [20, 44, 19, 33, 51, 40, 41, 37, 30, 54], + "Height": [164, 150, 178, 158, 188, 190, 168, 174, 176, 171], + }) + y = [4.1, 5.8, 3.9, 6.2, 4.3, 4.5, 7.2, 4.4, 4.1, 6.7] + + dtf = DecisionTreeFeatures(features_to_combine=2, drop_original=True) + dtf.fit(X, y) + + print(dtf.transform(X)) + +The resulting values match those found with pandas: + +.. code:: text + + shape: (10, 3) + ┌───────────┬──────────────┬─────────────────────────┐ + │ tree(Age) ┆ tree(Height) ┆ tree(['Age', 'Height']) │ + │ --- ┆ --- ┆ --- │ + │ f64 ┆ f64 ┆ f64 │ + ╞═══════════╪══════════════╪═════════════════════════╡ + │ 4.533333 ┆ 5.366667 ┆ 4.1 │ + │ 6.0 ┆ 5.366667 ┆ 6.475 │ + │ 4.533333 ┆ 4.133333 ┆ 4.0 │ + │ 4.533333 ┆ 5.366667 ┆ 6.475 │ + │ 6.0 ┆ 4.4 ┆ 4.4 │ + │ 4.533333 ┆ 4.4 ┆ 4.4 │ + │ 6.0 ┆ 6.95 ┆ 6.475 │ + │ 4.533333 ┆ 4.133333 ┆ 4.4 │ + │ 4.533333 ┆ 4.133333 ┆ 4.0 │ + │ 6.0 ┆ 6.95 ┆ 6.475 │ + └───────────┴──────────────┴─────────────────────────┘ + +`get_feature_names_out()`, classification, and every other parameter shown +above with pandas work identically with polars. + + Creating features for classification ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/feature_engine/creation/decision_tree_features.py b/feature_engine/creation/decision_tree_features.py index aa76d4bbd..f837482c0 100644 --- a/feature_engine/creation/decision_tree_features.py +++ b/feature_engine/creation/decision_tree_features.py @@ -1,8 +1,10 @@ import itertools from typing import Any, Dict, Iterable, List, Optional, Union +import narwhals as nw +import narwhals.dependencies as nwd import numpy as np -import pandas as pd +from narwhals.typing import IntoDataFrame, IntoSeries from sklearn.base import BaseEstimator, TransformerMixin from sklearn.model_selection import GridSearchCV from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor @@ -210,6 +212,36 @@ class DecisionTreeFeatures(TransformerMixin, BaseEstimator, GetFeatureNamesOutMi 7 4.24 8 4.24 9 6.00 + + With polars: + + >>> import polars as pl + >>> from feature_engine.creation import DecisionTreeFeatures + >>> X = pl.DataFrame({ + ... "Age": [20, 44, 19, 33, 51, 40, 41, 37, 30, 54], + ... "Height": [164, 150, 178, 158, 188, 190, 168, 174, 176, 171], + ... }) + >>> y = [4.1, 5.8, 3.9, 6.2, 4.3, 4.5, 7.2, 4.4, 4.1, 6.7] + >>> dtf = DecisionTreeFeatures(features_to_combine=1) + >>> dtf.fit(X, y) + >>> dtf.transform(X) + shape: (10, 4) + ┌─────┬────────┬───────────┬──────────────┐ + │ Age ┆ Height ┆ tree(Age) ┆ tree(Height) │ + │ --- ┆ --- ┆ --- ┆ --- │ + │ i64 ┆ i64 ┆ f64 ┆ f64 │ + ╞═════╪════════╪═══════════╪══════════════╡ + │ 20 ┆ 164 ┆ 4.533333 ┆ 5.366667 │ + │ 44 ┆ 150 ┆ 6.0 ┆ 5.366667 │ + │ 19 ┆ 178 ┆ 4.533333 ┆ 4.133333 │ + │ 33 ┆ 158 ┆ 4.533333 ┆ 5.366667 │ + │ 51 ┆ 188 ┆ 6.0 ┆ 4.4 │ + │ 40 ┆ 190 ┆ 4.533333 ┆ 4.4 │ + │ 41 ┆ 168 ┆ 6.0 ┆ 6.95 │ + │ 37 ┆ 174 ┆ 4.533333 ┆ 4.133333 │ + │ 30 ┆ 176 ┆ 4.533333 ┆ 4.133333 │ + │ 54 ┆ 171 ┆ 6.0 ┆ 6.95 │ + └─────┴────────┴───────────┴──────────────┘ """ def __init__( @@ -254,18 +286,18 @@ def __init__( self.missing_values = missing_values self.drop_original = drop_original - def fit(self, X: pd.DataFrame, y: pd.Series): + def fit(self, X: IntoDataFrame, y: IntoSeries): """ Fits decision trees based on the input variable combinations with cross-validation and grid-search for hyperparameters. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The training input samples. Can be the entire dataframe, not just the variables to transform. - y: pandas Series or np.array = [n_samples,] + y: Series or np.array = [n_samples,] The target variable that is used to train the decision tree. """ # confirm model type and target variables are compatible. @@ -302,39 +334,48 @@ def fit(self, X: pd.DataFrame, y: pd.Series): how_to_combine=self.features_to_combine, variables=variables_ ) + is_pandas = nwd.is_pandas_dataframe(X) is True + nw_X = nw.from_native(X, eager_only=True) + estimators_ = [] for features in input_features: estimator = self._make_decision_tree(param_grid=param_grid) # single feature models - if isinstance(features, str): - estimator.fit(X[features].to_frame(), y) + if isinstance(features, (str, int)): + X_sub = nw_X.get_column(features).to_frame().to_native() # multi feature models + elif is_pandas is True: + X_sub = X[features] else: - estimator.fit(X[features], y) + X_sub = nw_X.select(features).to_native() + estimator.fit(X_sub, y) estimators_.append(estimator) self.variables_ = variables_ self.input_features_ = input_features self.estimators_ = estimators_ - self.feature_names_in_ = X.columns.tolist() + if is_pandas is True: + self.feature_names_in_ = list(X.columns) + else: + self.feature_names_in_ = nw_X.columns self.n_features_in_ = X.shape[1] return self - def transform(self, X: pd.DataFrame) -> pd.DataFrame: + def transform(self, X: IntoDataFrame) -> IntoDataFrame: """ Create and add new variables. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The data to be transformed. Returns ------- - X_new: pandas dataframe. + X_new: dataframe. Either the original dataframe plus the new features or a dataframe of only the new features. """ @@ -351,50 +392,59 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame: _check_contains_na(X, self.variables_) _check_contains_inf(X, self.variables_) + is_pandas = nwd.is_pandas_dataframe(X) is True + # reorder variables to match train set - X = X[self.feature_names_in_] + if is_pandas is True: + X = X[self.feature_names_in_] + else: + X = ( + nw.from_native(X, eager_only=True) + .select(self.feature_names_in_) + .to_native() + ) + nw_X = nw.from_native(X, eager_only=True) - # create new features and add them to the original dataframe - # if regression or multiclass, we return the output of predict() - if self.regression is True: - for features, estimator in zip(self.input_features_, self.estimators_): - if isinstance(features, str): - preds = estimator.predict(X[features].to_frame()) - if self.precision is not None: - preds = np.round(preds, self.precision) - X.loc[:, f"tree({features})"] = preds - else: - preds = estimator.predict(X[features]) - if self.precision is not None: - preds = np.round(preds, self.precision) - X.loc[:, f"tree({features})"] = preds + def get_x_sub(features): + if isinstance(features, (str, int)): + return nw_X.get_column(features).to_frame().to_native() + if is_pandas is True: + return X[features] + return nw_X.select(features).to_native() + new_series = [] + # if regression or multiclass, we return the output of predict(); # if binary classification, we return the probability - elif self._is_binary == "binary": - for features, estimator in zip(self.input_features_, self.estimators_): - if isinstance(features, str): - preds = estimator.predict_proba(X[features].to_frame()) - if self.precision is not None: - preds = np.round(preds, self.precision) - X.loc[:, f"tree({features})"] = preds[:, 1] - else: - preds = estimator.predict_proba(X[features]) - if self.precision is not None: - preds = np.round(preds, self.precision) - X.loc[:, f"tree({features})"] = preds[:, 1] + for features, estimator in zip(self.input_features_, self.estimators_): + X_sub = get_x_sub(features) + col_name = f"tree({features})" + + if self.regression is True: + preds = estimator.predict(X_sub) + if self.precision is not None: + preds = np.round(preds, self.precision) + elif self._is_binary == "binary": + preds = estimator.predict_proba(X_sub)[:, 1] + if self.precision is not None: + preds = np.round(preds, self.precision) + else: + preds = estimator.predict(X_sub) - # if multiclass, we return the output of predict() - else: - for features, estimator in zip(self.input_features_, self.estimators_): - if isinstance(features, str): - preds = estimator.predict(X[features].to_frame()) - X.loc[:, f"tree({features})"] = preds - else: - preds = estimator.predict(X[features]) - X.loc[:, f"tree({features})"] = preds + if is_pandas is True: + X[col_name] = preds + else: + new_series.append( + nw.new_series(col_name, preds, backend=nw_X.implementation) + ) - if self.drop_original: - X.drop(columns=self.variables_, inplace=True) + if is_pandas is True: + if self.drop_original is True: + X = X.drop(columns=self.variables_) + else: + nw_X = nw.from_native(X, eager_only=True).with_columns(*new_series) + if self.drop_original is True: + nw_X = nw_X.drop(self.variables_) + X = nw_X.to_native() return X diff --git a/tests/test_creation/test_decision_tree_features.py b/tests/test_creation/test_decision_tree_features.py index 4e8a93e8c..80f4b27a9 100644 --- a/tests/test_creation/test_decision_tree_features.py +++ b/tests/test_creation/test_decision_tree_features.py @@ -1,5 +1,7 @@ +import narwhals as nw import numpy as np import pandas as pd +import polars as pl import pytest from sklearn.model_selection import GridSearchCV from sklearn.pipeline import Pipeline @@ -591,3 +593,99 @@ def test_check_return_empty(): # RelativeFeatures and CyclicalFeatures into it), so return_empty is # tested directly here instead. check_return_empty(DecisionTreeFeatures(regression=False)) + + +# ============================ polars support ============================ + +NUMERIC_DATA = { + "Age": [20, 44, 19, 33, 51, 40, 41, 37, 30, 54], + "Height": [164, 150, 178, 158, 188, 190, 168, 174, 176, 171], +} +REGRESSION_Y = [4.1, 5.8, 3.9, 6.2, 4.3, 4.5, 7.2, 4.4, 4.1, 6.7] +BINARY_Y = [0, 1, 0, 1, 1, 0, 1, 0, 0, 1] +MULTICLASS_Y = [0, 1, 2, 0, 1, 2, 0, 1, 2, 0] + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_regression_both_backends_match(make_df): + df = make_df(NUMERIC_DATA) + transformer = DecisionTreeFeatures(features_to_combine=2, random_state=0) + transformer.fit(df, REGRESSION_Y) + Xt = transformer.transform(df) + + result = nw.from_native(Xt, eager_only=True).to_dict(as_series=False) + assert result["tree(['Age', 'Height'])"] == pytest.approx( + [4.1, 6.475, 4.0, 6.475, 4.4, 4.4, 6.475, 4.4, 4.0, 6.475] + ) + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_binary_classification_both_backends_match(make_df): + df = make_df(NUMERIC_DATA) + transformer = DecisionTreeFeatures(regression=False, random_state=0) + transformer.fit(df, BINARY_Y) + Xt = transformer.transform(df) + + pd_df = pd.DataFrame(NUMERIC_DATA) + pd_transformer = DecisionTreeFeatures(regression=False, random_state=0) + pd_transformer.fit(pd_df, BINARY_Y) + expected = pd_transformer.transform(pd_df)["tree(['Age', 'Height'])"].tolist() + + result = nw.from_native(Xt, eager_only=True).get_column( + "tree(['Age', 'Height'])" + ) + assert result.to_list() == pytest.approx(expected) + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_multiclass_classification_both_backends_match(make_df): + df = make_df(NUMERIC_DATA) + transformer = DecisionTreeFeatures(regression=False, random_state=0) + transformer.fit(df, MULTICLASS_Y) + Xt = transformer.transform(df) + + pd_df = pd.DataFrame(NUMERIC_DATA) + pd_transformer = DecisionTreeFeatures(regression=False, random_state=0) + pd_transformer.fit(pd_df, MULTICLASS_Y) + expected = pd_transformer.transform(pd_df)["tree(['Age', 'Height'])"].tolist() + + result = nw.from_native(Xt, eager_only=True).get_column( + "tree(['Age', 'Height'])" + ) + assert result.to_list() == pytest.approx(expected) + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_drop_original_both_backends(make_df): + df = make_df(NUMERIC_DATA) + transformer = DecisionTreeFeatures(features_to_combine=1, drop_original=True) + transformer.fit(df, REGRESSION_Y) + Xt = transformer.transform(df) + assert list(nw.from_native(Xt, eager_only=True).columns) == [ + "tree(Age)", + "tree(Height)", + ] + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_get_feature_names_out_both_backends(make_df): + df = make_df(NUMERIC_DATA) + transformer = DecisionTreeFeatures(features_to_combine=1) + transformer.fit(df, REGRESSION_Y) + Xt = transformer.transform(df) + feat_out = list(nw.from_native(Xt, eager_only=True).columns) + assert transformer.get_feature_names_out(input_features=None) == feat_out + + +def test_single_int_named_feature_combo_both_backends(): + # regression test: a single-variable combo with an integer column name + # used to crash (isinstance(features, str) missed the int case), since + # X[features] for a bare int returns a 1D Series, not the 2D input + # sklearn requires - fixed to check isinstance(features, (str, int)). + df = pd.DataFrame({0: [1.0, 2, 3, 4, 5, 6, 7, 8], 1: [2.0, 3, 4, 5, 6, 7, 8, 9]}) + y = [1.0, 2, 3, 4, 5, 6, 7, 8] + transformer = DecisionTreeFeatures(features_to_combine=1, random_state=0) + transformer.fit(df, y) + Xt = transformer.transform(df) + assert "tree(0)" in Xt.columns + assert "tree(1)" in Xt.columns From ba0d0bfb220545eae95defcf447a467f34afa8a9 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Mon, 24 Aug 2026 20:42:58 +0200 Subject: [PATCH 2/7] Apply suggestion from @solegalli --- feature_engine/creation/decision_tree_features.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feature_engine/creation/decision_tree_features.py b/feature_engine/creation/decision_tree_features.py index f837482c0..ccc15d214 100644 --- a/feature_engine/creation/decision_tree_features.py +++ b/feature_engine/creation/decision_tree_features.py @@ -392,7 +392,7 @@ def transform(self, X: IntoDataFrame) -> IntoDataFrame: _check_contains_na(X, self.variables_) _check_contains_inf(X, self.variables_) - is_pandas = nwd.is_pandas_dataframe(X) is True + is_pandas = nwd.is_pandas_dataframe(X) # reorder variables to match train set if is_pandas is True: From e063d92a957b40311e8a3f4ff1dccad9c9921a0d Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Mon, 24 Aug 2026 20:43:11 +0200 Subject: [PATCH 3/7] Apply suggestion from @solegalli --- feature_engine/creation/decision_tree_features.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feature_engine/creation/decision_tree_features.py b/feature_engine/creation/decision_tree_features.py index ccc15d214..129eee5b4 100644 --- a/feature_engine/creation/decision_tree_features.py +++ b/feature_engine/creation/decision_tree_features.py @@ -334,7 +334,7 @@ def fit(self, X: IntoDataFrame, y: IntoSeries): how_to_combine=self.features_to_combine, variables=variables_ ) - is_pandas = nwd.is_pandas_dataframe(X) is True + is_pandas = nwd.is_pandas_dataframe(X) nw_X = nw.from_native(X, eager_only=True) estimators_ = [] From f0b71d45669a1181cd31624c9e21310d91c9ffe1 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Mon, 24 Aug 2026 20:54:53 +0200 Subject: [PATCH 4/7] docs: clarify is True/is False and cross-backend test conventions in AGENTS.md Two rules made explicit based on recent work: the is True/is False comparison is for flow control only, not variable assignment (per Sole's own simplification of is_pandas = nwd.is_pandas_dataframe(X) is True to just nwd.is_pandas_dataframe(X) in decision_tree_features.py); and dataframe-agnostic transformers get one parametrized test per behavior covering both pandas and polars, never separate per-backend tests. Co-Authored-By: Claude Sonnet 5 --- AGENTS.md | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/AGENTS.md b/AGENTS.md index b5c3378cc..dd3524880 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -33,6 +33,13 @@ object that's already an instance of that module's class. - Check container emptiness with `len(x) == 0`, never `if not x:`. - `isinstance(...)` checks and `in`/`not in` membership tests are already explicit — leave them as-is, this rule isn't about those. +- The explicit `is True`/`is False` comparison is for flow control + (`if`/`while` conditions) only — don't tack it onto a variable + assignment. When a function already returns a strict bool (e.g. + `nwd.is_pandas_dataframe(X)`), assign it directly: + `is_pandas = nwd.is_pandas_dataframe(X)`, not + `is_pandas = nwd.is_pandas_dataframe(X) is True`. The `if`/`while` site + that later consumes `is_pandas` still spells out `if is_pandas is True:`. ## Comments @@ -75,6 +82,16 @@ and easy to miss without an actual comparison. - `pytest.raises(ExceptionType, match=msg)`, never `with pytest.raises() as record: ... assert str(record.value) == msg`. +- Dataframe-agnostic means one test, both backends: parametrize each + behavior over `@pytest.mark.parametrize("make_df", [pd.DataFrame, + pl.DataFrame])` and assert the same input produces the same output + values on both. Never write a separate pandas-only test and a + separate polars-only test for the same behavior — that duplicates + the test and hides the point of being dataframe-agnostic, which is + that the same input gives the same output regardless of backend. + Keep a test single-backend only when the behavior itself is + backend-specific (e.g. integer column names, which polars doesn't + support; pandas nullable extension dtypes). ## API changes From 35b34c388106fe24a083e8a4d42caed721c80c71 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Mon, 24 Aug 2026 20:55:03 +0200 Subject: [PATCH 5/7] feat: add n_jobs for parallel tree training, merge tests to single cross-backend suite Adds an n_jobs parameter to DecisionTreeFeatures that parallelizes tree training across feature combinations via joblib, using threads rather than processes since fitting a decision tree releases the GIL for the bulk of its computation - threads avoid the overhead of copying the whole dataframe to worker processes. Defaults to None (sequential), preserving current behavior. Benchmarked on the committed transformer (5000 rows, 10 vars, features_to_combine=3, 8-point param_grid, 175 trees): 12.17s sequential vs 5.15s at n_jobs=-1, ~2.4x. On small workloads (a handful of feature combinations, the shape of the existing unit tests) parallelizing is a net loss - thread-dispatch overhead outweighs the gain - which is why the default stays sequential. Parallelizing transform()'s predict loop the same way was also benchmarked and found to have no benefit (predict is too cheap per call), so only fit()'s tree training is parallelized. Correctness verified: identical trees/predictions regardless of n_jobs. Also rewrites test_decision_tree_features.py to the single cross-backend-parametrized-test convention used elsewhere in this migration: one test per behavior over make_df=[pd.DataFrame, pl.DataFrame], deleting the separately-added polars-only section that duplicated coverage already present once the original tests are parametrized. Adds n_jobs correctness coverage (parallel vs sequential training gives identical output, both backends). Co-Authored-By: Claude Sonnet 5 --- .../creation/DecisionTreeFeatures.rst | 40 ++ .../creation/decision_tree_features.py | 31 +- .../test_decision_tree_features.py | 621 ++++++------------ 3 files changed, 256 insertions(+), 436 deletions(-) diff --git a/docs/user_guide/creation/DecisionTreeFeatures.rst b/docs/user_guide/creation/DecisionTreeFeatures.rst index a7da7115b..ba59800d2 100644 --- a/docs/user_guide/creation/DecisionTreeFeatures.rst +++ b/docs/user_guide/creation/DecisionTreeFeatures.rst @@ -545,6 +545,46 @@ identical. We just need to set the parameter `regression` to False. classification, on the other hand, the features will contain the prediction of the class. +Training trees in parallel +~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Each tree is trained on its own feature combination independently of the others, so +when there are many combinations (a large number of variables and/or a high +`features_to_combine`) or a large `param_grid` to search, training can be +parallelized across combinations with the `n_jobs` parameter: + +.. code:: python + + import pandas as pd + from feature_engine.creation import DecisionTreeFeatures + + X = pd.DataFrame({ + "Age": [20, 44, 19, 33, 51, 40, 41, 37, 30, 54], + "Height": [164, 150, 178, 158, 188, 190, 168, 174, 176, 171], + "Marks": [1.0, 0.8, 0.6, 0.1, 0.3, 0.4, 0.8, 0.6, 0.5, 0.2], + }) + y = [4.1, 5.8, 3.9, 6.2, 4.3, 4.5, 7.2, 4.4, 4.1, 6.7] + + dtf = DecisionTreeFeatures(features_to_combine=3, n_jobs=2, random_state=0) + dtf.fit(X, y) + + print(dtf.transform(X).columns.tolist()) + +.. code:: text + + ['Age', 'Height', 'Marks', 'tree(Age)', 'tree(Height)', 'tree(Marks)', + "tree(['Age', 'Height'])", "tree(['Age', 'Marks'])", + "tree(['Height', 'Marks'])", "tree(['Age', 'Height', 'Marks'])"] + +`n_jobs` defaults to `None`, which trains the trees sequentially, matching this +transformer's original behaviour. Setting it trains multiple trees at the same +time using threads, which only pays off once there are enough feature +combinations or a large enough `param_grid` to outweigh the overhead of +dispatching work to threads — with just a handful of combinations, sequential +training is faster. The resulting trees and predictions are identical +regardless of `n_jobs`; only training speed changes. + + Additional resources -------------------- diff --git a/feature_engine/creation/decision_tree_features.py b/feature_engine/creation/decision_tree_features.py index 129eee5b4..0f9648c1f 100644 --- a/feature_engine/creation/decision_tree_features.py +++ b/feature_engine/creation/decision_tree_features.py @@ -4,6 +4,7 @@ import narwhals as nw import narwhals.dependencies as nwd import numpy as np +from joblib import Parallel, delayed from narwhals.typing import IntoDataFrame, IntoSeries from sklearn.base import BaseEstimator, TransformerMixin from sklearn.model_selection import GridSearchCV @@ -133,6 +134,18 @@ class DecisionTreeFeatures(TransformerMixin, BaseEstimator, GetFeatureNamesOutMi DecisionTreeClassifier(). For reproducibility it is recommended to set the random_state to an integer. + n_jobs: int, default=None + The number of jobs to run in parallel when training the decision trees + across feature combinations. 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. There is a real speedup only when there are many + feature combinations and/or a large `param_grid` to search — with just + a handful of combinations, thread-dispatch overhead outweighs the gain, + which is why the default stays sequential. + {missing_values} {drop_original} @@ -255,6 +268,7 @@ def __init__( param_grid: Optional[Dict[str, Union[str, int, float, List[int]]]] = None, regression: bool = True, random_state: int = 0, + n_jobs: Optional[int] = None, missing_values: str = "raise", drop_original: bool = False, ) -> None: @@ -283,6 +297,7 @@ def __init__( self.param_grid = param_grid self.regression = regression self.random_state = random_state + self.n_jobs = n_jobs self.missing_values = missing_values self.drop_original = drop_original @@ -337,10 +352,8 @@ def fit(self, X: IntoDataFrame, y: IntoSeries): is_pandas = nwd.is_pandas_dataframe(X) nw_X = nw.from_native(X, eager_only=True) - estimators_ = [] + X_subs = [] for features in input_features: - estimator = self._make_decision_tree(param_grid=param_grid) - # single feature models if isinstance(features, (str, int)): X_sub = nw_X.get_column(features).to_frame().to_native() @@ -349,9 +362,11 @@ def fit(self, X: IntoDataFrame, y: IntoSeries): X_sub = X[features] else: X_sub = nw_X.select(features).to_native() + X_subs.append(X_sub) - estimator.fit(X_sub, y) - estimators_.append(estimator) + estimators_ = Parallel(n_jobs=self.n_jobs, prefer="threads")( + delayed(self._fit_one_tree)(X_sub, y, param_grid) for X_sub in X_subs + ) self.variables_ = variables_ self.input_features_ = input_features @@ -464,6 +479,12 @@ def _make_decision_tree(self, param_grid: Dict): return tree_model + def _fit_one_tree(self, X_sub: IntoDataFrame, y: IntoSeries, param_grid: Dict): + """Instantiate and fit one decision tree on one feature combination.""" + estimator = self._make_decision_tree(param_grid=param_grid) + estimator.fit(X_sub, y) + return estimator + def _create_variable_combinations( self, variables: List, diff --git a/tests/test_creation/test_decision_tree_features.py b/tests/test_creation/test_decision_tree_features.py index 80f4b27a9..b8f1c31cc 100644 --- a/tests/test_creation/test_decision_tree_features.py +++ b/tests/test_creation/test_decision_tree_features.py @@ -10,44 +10,87 @@ from feature_engine.creation import DecisionTreeFeatures from tests.estimator_checks.fit_functionality_checks import check_return_empty - -@pytest.fixture(scope="module") -def df_creation(): - data = { - "Name": [ - "tom", - "nick", - "krish", - "megan", - "peter", - "jordan", - "fred", - "sam", - "alexa", - "brittany", - ], - "Age": [20, 44, 19, 33, 51, 40, 41, 37, 30, 54], - "Height": [164, 150, 178, 158, 188, 190, 168, 174, 176, 171], - "Marks": [1.0, 0.8, 0.6, 0.1, 0.3, 0.4, 0.8, 0.6, 0.5, 0.2], - } - - df = pd.DataFrame(data) - return df - - -@pytest.fixture(scope="module") -def regression_target(): - return pd.Series([4.1, 5.8, 3.9, 6.2, 4.3, 4.5, 7.2, 4.4, 4.1, 6.7]) - - -@pytest.fixture(scope="module") -def classification_target(): - return pd.Series([1, 1, 1, 0, 0, 1, 0, 1, 0, 0]) - - -@pytest.fixture(scope="module") -def multiclass_target(): - return pd.Series([1, 1, 2, 2, 0, 1, 0, 1, 0, 0]) +DATA = { + "Name": [ + "tom", + "nick", + "krish", + "megan", + "peter", + "jordan", + "fred", + "sam", + "alexa", + "brittany", + ], + "Age": [20, 44, 19, 33, 51, 40, 41, 37, 30, 54], + "Height": [164, 150, 178, 158, 188, 190, 168, 174, 176, 171], + "Marks": [1.0, 0.8, 0.6, 0.1, 0.3, 0.4, 0.8, 0.6, 0.5, 0.2], +} +REGRESSION_Y = [4.1, 5.8, 3.9, 6.2, 4.3, 4.5, 7.2, 4.4, 4.1, 6.7] +BINARY_Y = [1, 1, 1, 0, 0, 1, 0, 1, 0, 0] +MULTICLASS_Y = [1, 1, 2, 2, 0, 1, 0, 1, 0, 0] + +COMBOS = [ + "Age", + "Height", + "Marks", + ["Age", "Height"], + ["Age", "Marks"], + ["Height", "Marks"], + ["Age", "Height", "Marks"], +] + + +def _select(X, combo): + cols = combo if isinstance(combo, list) else [combo] + return nw.from_native(X, eager_only=True).select(cols).to_native() + + +def _expected_tree_predictions( + X, + y, + scoring, + random_state, + regression=True, + binary=False, + precision=None, + param_grid=None, +): + # Fits a fresh GridSearchCV per combo on the same backend as X, so this + # works as the reference for both pandas and polars input alike. + if param_grid is None: + param_grid = {"max_depth": [1, 2, 3, 4]} + if regression is True: + est = DecisionTreeRegressor(random_state=random_state) + else: + est = DecisionTreeClassifier(random_state=random_state) + tree = GridSearchCV(est, cv=3, scoring=scoring, param_grid=param_grid) + + expected = {} + for combo in COMBOS: + X_sub = _select(X, combo) + tree.fit(X_sub, y) + if regression is True: + preds = tree.predict(X_sub) + elif binary is True: + preds = tree.predict_proba(X_sub)[:, 1] + else: + preds = tree.predict(X_sub) + if precision is not None: + preds = np.round(preds, precision) + expected[f"tree({combo})"] = list(preds) + return expected + + +def assert_df_equal(X, expected: dict) -> None: + result = nw.from_native(X, eager_only=True).to_dict(as_series=False) + assert list(result.keys()) == list(expected.keys()) + for col, values in expected.items(): + if all(isinstance(v, (int, float, np.integer, np.floating)) for v in values): + assert result[col] == pytest.approx(values, abs=1e-6) + else: + assert result[col] == values @pytest.mark.parametrize("precision", ["string", 0.1, -1, np.nan]) @@ -206,482 +249,198 @@ def test_create_variable_combinations_when_tuple(input_features, expected): assert combos == expected -def test_feature_creation_regression(df_creation, regression_target): - X = df_creation.copy() - y = regression_target.copy() - +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_feature_creation_regression(make_df): + X = make_df(DATA) scoring = "neg_mean_squared_error" rs = 0 tr = DecisionTreeFeatures(scoring=scoring, random_state=rs) - Xt = tr.fit_transform(X, y) - - # get expected - est = DecisionTreeRegressor(random_state=rs) - tree = GridSearchCV( - est, - cv=3, - scoring=scoring, - param_grid={"max_depth": [1, 2, 3, 4]}, - ) - - combos = [ - "Age", - "Height", - "Marks", - ["Age", "Height"], - ["Age", "Marks"], - ["Height", "Marks"], - ["Age", "Height", "Marks"], - ] - var_names = [f"tree({item})" for item in combos] - - X_exp = df_creation.copy() - for i in range(len(combos)): - varn = var_names[i] - combon = combos[i] - if isinstance(combon, str): - tree.fit(X[combon].to_frame(), y) - X_exp[varn] = tree.predict(X[combon].to_frame()) - else: - tree.fit(X[combon], y) - X_exp[varn] = tree.predict(X[combon]) + Xt = tr.fit_transform(X, REGRESSION_Y) - pd.testing.assert_frame_equal(Xt, X_exp) + expected = dict(DATA) + expected.update(_expected_tree_predictions(X, REGRESSION_Y, scoring, rs)) + assert_df_equal(Xt, expected) -def test_feature_creation_regression_and_precision(df_creation, regression_target): - X = df_creation.copy() - y = regression_target.copy() - +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_feature_creation_regression_and_precision(make_df): + X = make_df(DATA) scoring = "neg_mean_squared_error" rs = 0 tr = DecisionTreeFeatures(scoring=scoring, random_state=rs, precision=1) - Xt = tr.fit_transform(X, y) - - # get expected - est = DecisionTreeRegressor(random_state=rs) - tree = GridSearchCV( - est, - cv=3, - scoring=scoring, - param_grid={"max_depth": [1, 2, 3, 4]}, - ) - - combos = [ - "Age", - "Height", - "Marks", - ["Age", "Height"], - ["Age", "Marks"], - ["Height", "Marks"], - ["Age", "Height", "Marks"], - ] - var_names = [f"tree({item})" for item in combos] - - X_exp = df_creation.copy() - for i in range(len(combos)): - varn = var_names[i] - combon = combos[i] - if isinstance(combon, str): - tree.fit(X[combon].to_frame(), y) - preds = tree.predict(X[combon].to_frame()) - X_exp[varn] = np.round(preds, 1) - else: - tree.fit(X[combon], y) - preds = tree.predict(X[combon]) - X_exp[varn] = np.round(preds, 1) - - pd.testing.assert_frame_equal(Xt, X_exp) + Xt = tr.fit_transform(X, REGRESSION_Y) + expected = dict(DATA) + expected.update( + _expected_tree_predictions(X, REGRESSION_Y, scoring, rs, precision=1) + ) + assert_df_equal(Xt, expected) -def test_feature_creation_regression_drop_original(df_creation, regression_target): - X = df_creation.copy() - y = regression_target.copy() +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_feature_creation_regression_drop_original(make_df): + X = make_df(DATA) scoring = "neg_mean_squared_error" rs = 0 tr = DecisionTreeFeatures(scoring=scoring, random_state=rs, drop_original=True) - Xt = tr.fit_transform(X, y) - - # get expected - est = DecisionTreeRegressor(random_state=rs) - tree = GridSearchCV( - est, - cv=3, - scoring=scoring, - param_grid={"max_depth": [1, 2, 3, 4]}, - ) + Xt = tr.fit_transform(X, REGRESSION_Y) - combos = [ - "Age", - "Height", - "Marks", - ["Age", "Height"], - ["Age", "Marks"], - ["Height", "Marks"], - ["Age", "Height", "Marks"], - ] - var_names = [f"tree({item})" for item in combos] - - X_exp = df_creation.copy() - for i in range(len(combos)): - varn = var_names[i] - combon = combos[i] - if isinstance(combon, str): - tree.fit(X[combon].to_frame(), y) - X_exp[varn] = tree.predict(X[combon].to_frame()) - else: - tree.fit(X[combon], y) - X_exp[varn] = tree.predict(X[combon]) - X_exp.drop(["Age", "Height", "Marks"], axis=1, inplace=True) + expected = {"Name": DATA["Name"]} + expected.update(_expected_tree_predictions(X, REGRESSION_Y, scoring, rs)) + assert_df_equal(Xt, expected) - pd.testing.assert_frame_equal(Xt, X_exp) - - -def test_feature_creation_binary_classif(df_creation, classification_target): - X = df_creation.copy() - y = classification_target.copy() +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_feature_creation_binary_classif(make_df): + X = make_df(DATA) scoring = "roc_auc" rs = 0 tr = DecisionTreeFeatures(scoring=scoring, random_state=rs, regression=False) - Xt = tr.fit_transform(X, y) - - # get expected - est = DecisionTreeClassifier(random_state=rs) - tree = GridSearchCV( - est, - cv=3, - scoring=scoring, - param_grid={"max_depth": [1, 2, 3, 4]}, - ) - - combos = [ - "Age", - "Height", - "Marks", - ["Age", "Height"], - ["Age", "Marks"], - ["Height", "Marks"], - ["Age", "Height", "Marks"], - ] - var_names = [f"tree({item})" for item in combos] - - X_exp = df_creation.copy() - for i in range(len(combos)): - varn = var_names[i] - combon = combos[i] - if isinstance(combon, str): - tree.fit(X[combon].to_frame(), y) - preds = tree.predict_proba(X[combon].to_frame()) - X_exp[varn] = preds[:, 1] - else: - tree.fit(X[combon], y) - preds = tree.predict_proba(X[combon]) - X_exp[varn] = preds[:, 1] - - pd.testing.assert_frame_equal(Xt, X_exp) + Xt = tr.fit_transform(X, BINARY_Y) + expected = dict(DATA) + expected.update( + _expected_tree_predictions( + X, BINARY_Y, scoring, rs, regression=False, binary=True + ) + ) + assert_df_equal(Xt, expected) -def test_feature_creation_binary_classif_w_precision( - df_creation, classification_target -): - X = df_creation.copy() - y = classification_target.copy() +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_feature_creation_binary_classif_w_precision(make_df): + X = make_df(DATA) scoring = "roc_auc" rs = 0 tr = DecisionTreeFeatures( scoring=scoring, random_state=rs, regression=False, precision=2 ) - Xt = tr.fit_transform(X, y) - - # get expected - est = DecisionTreeClassifier(random_state=rs) - tree = GridSearchCV( - est, - cv=3, - scoring=scoring, - param_grid={"max_depth": [1, 2, 3, 4]}, - ) - - combos = [ - "Age", - "Height", - "Marks", - ["Age", "Height"], - ["Age", "Marks"], - ["Height", "Marks"], - ["Age", "Height", "Marks"], - ] - var_names = [f"tree({item})" for item in combos] - - X_exp = df_creation.copy() - for i in range(len(combos)): - varn = var_names[i] - combon = combos[i] - if isinstance(combon, str): - tree.fit(X[combon].to_frame(), y) - preds = tree.predict_proba(X[combon].to_frame()) - X_exp[varn] = np.round(preds[:, 1], 2) - else: - tree.fit(X[combon], y) - preds = tree.predict_proba(X[combon]) - X_exp[varn] = np.round(preds[:, 1], 2) - - pd.testing.assert_frame_equal(Xt, X_exp) + Xt = tr.fit_transform(X, BINARY_Y) + expected = dict(DATA) + expected.update( + _expected_tree_predictions( + X, BINARY_Y, scoring, rs, regression=False, binary=True, precision=2 + ) + ) + assert_df_equal(Xt, expected) -def test_feature_creation_binary_multiclass(df_creation, multiclass_target): - X = df_creation.copy() - y = multiclass_target.copy() +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_feature_creation_binary_multiclass(make_df): + X = make_df(DATA) scoring = "roc_auc" rs = 0 tr = DecisionTreeFeatures(scoring=scoring, random_state=rs, regression=False) - Xt = tr.fit_transform(X, y) - - # get expected - est = DecisionTreeClassifier(random_state=rs) - tree = GridSearchCV( - est, - cv=3, - scoring=scoring, - param_grid={"max_depth": [1, 2, 3, 4]}, - ) + Xt = tr.fit_transform(X, MULTICLASS_Y) - combos = [ - "Age", - "Height", - "Marks", - ["Age", "Height"], - ["Age", "Marks"], - ["Height", "Marks"], - ["Age", "Height", "Marks"], - ] - var_names = [f"tree({item})" for item in combos] - - X_exp = df_creation.copy() - for i in range(len(combos)): - varn = var_names[i] - combon = combos[i] - if isinstance(combon, str): - tree.fit(X[combon].to_frame(), y) - preds = tree.predict(X[combon].to_frame()) - X_exp[varn] = preds - else: - tree.fit(X[combon], y) - preds = tree.predict(X[combon]) - X_exp[varn] = preds - - pd.testing.assert_frame_equal(Xt, X_exp) - - -def test_get_feature_names_out(df_creation, regression_target): - X = df_creation.copy() - y = regression_target.copy() - - tr = DecisionTreeFeatures( - variables=["Age", "Marks"], + expected = dict(DATA) + expected.update( + _expected_tree_predictions( + X, MULTICLASS_Y, scoring, rs, regression=False, binary=False + ) ) - - Xt = tr.fit_transform(X, y) - feat_out = Xt.columns.to_list() - assert tr.get_feature_names_out() == feat_out - assert tr.get_feature_names_out(X.columns.to_list()) == feat_out + assert_df_equal(Xt, expected) -def test_get_feature_names_out_from_pipeline(df_creation, regression_target): - X = df_creation.copy() - y = regression_target.copy() +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_get_feature_names_out(make_df): + X = make_df(DATA) + tr = DecisionTreeFeatures(variables=["Age", "Marks"]) + Xt = tr.fit_transform(X, REGRESSION_Y) + feat_out = list(nw.from_native(Xt, eager_only=True).columns) + assert tr.get_feature_names_out() == feat_out + assert tr.get_feature_names_out(list(DATA.keys())) == feat_out - # set up transformer - tr = DecisionTreeFeatures( - variables=["Age", "Marks"], - ) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_get_feature_names_out_from_pipeline(make_df): + X = make_df(DATA) + tr = DecisionTreeFeatures(variables=["Age", "Marks"]) pipe = Pipeline([("transformer", tr)]) - - Xt = pipe.fit_transform(X, y) - feat_out = Xt.columns.to_list() + Xt = pipe.fit_transform(X, REGRESSION_Y) + feat_out = list(nw.from_native(Xt, eager_only=True).columns) assert pipe.get_feature_names_out(input_features=None) == feat_out - assert pipe.get_feature_names_out(input_features=X.columns.to_list()) == feat_out + assert pipe.get_feature_names_out(input_features=list(DATA.keys())) == feat_out +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize("_input_features", ["hola", ["Age", "Marks"]]) -def test_get_feature_names_out_raises_error_when_wrong_param( - _input_features, df_creation, regression_target -): - X = df_creation.copy() - y = regression_target.copy() - - tr = DecisionTreeFeatures( - variables=["Age", "Marks"], - ) - tr.fit(X, y) - +def test_get_feature_names_out_raises_error_when_wrong_param(make_df, _input_features): + X = make_df(DATA) + tr = DecisionTreeFeatures(variables=["Age", "Marks"]) + tr.fit(X, REGRESSION_Y) with pytest.raises(ValueError): tr.get_feature_names_out(input_features=_input_features) -def test_error_when_regression_true_and_target_binary( - df_creation, classification_target -): - X = df_creation.copy() - y = classification_target.copy() +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_error_when_regression_true_and_target_binary(make_df): + X = make_df(DATA) tr = DecisionTreeFeatures(regression=True) msg = ( "Trying to fit a regression to a binary target is not " - + "allowed by this transformer. Check the target values " - + "or set regression to False." + "allowed by this transformer. Check the target values " + "or set regression to False." ) with pytest.raises(ValueError, match=msg): - tr.fit(X, y) + tr.fit(X, BINARY_Y) -def test_user_enter_param_grid(df_creation, classification_target): - X = df_creation.copy() - y = classification_target.copy() +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_user_enter_param_grid(make_df): + X = make_df(DATA) scoring = "roc_auc" rs = 0 grid = {"max_depth": [1, 2, 3, 4]} tr = DecisionTreeFeatures( scoring=scoring, random_state=rs, regression=False, param_grid=grid ) - Xt = tr.fit_transform(X, y) - - # get expected - est = DecisionTreeClassifier(random_state=rs) - tree = GridSearchCV( - est, - cv=3, - scoring=scoring, - param_grid={"max_depth": [1, 2, 3, 4]}, - ) + Xt = tr.fit_transform(X, BINARY_Y) - combos = [ - "Age", - "Height", - "Marks", - ["Age", "Height"], - ["Age", "Marks"], - ["Height", "Marks"], - ["Age", "Height", "Marks"], - ] - var_names = [f"tree({item})" for item in combos] - - X_exp = df_creation.copy() - for i in range(len(combos)): - varn = var_names[i] - combon = combos[i] - if isinstance(combon, str): - tree.fit(X[combon].to_frame(), y) - preds = tree.predict_proba(X[combon].to_frame()) - X_exp[varn] = preds[:, 1] - else: - tree.fit(X[combon], y) - preds = tree.predict_proba(X[combon]) - X_exp[varn] = preds[:, 1] - - pd.testing.assert_frame_equal(Xt, X_exp) + expected = dict(DATA) + expected.update( + _expected_tree_predictions( + X, BINARY_Y, scoring, rs, regression=False, binary=True, param_grid=grid + ) + ) + assert_df_equal(Xt, expected) def test_check_return_empty(): # DecisionTreeFeatures is not part of the check_feature_engine_estimator # pipeline (test_check_estimator_creation.py only feeds MathFeatures, # RelativeFeatures and CyclicalFeatures into it), so return_empty is - # tested directly here instead. + # tested directly here instead. check_return_empty is a shared, + # pandas-only estimator-check helper used across the library. check_return_empty(DecisionTreeFeatures(regression=False)) -# ============================ polars support ============================ - -NUMERIC_DATA = { - "Age": [20, 44, 19, 33, 51, 40, 41, 37, 30, 54], - "Height": [164, 150, 178, 158, 188, 190, 168, 174, 176, 171], -} -REGRESSION_Y = [4.1, 5.8, 3.9, 6.2, 4.3, 4.5, 7.2, 4.4, 4.1, 6.7] -BINARY_Y = [0, 1, 0, 1, 1, 0, 1, 0, 0, 1] -MULTICLASS_Y = [0, 1, 2, 0, 1, 2, 0, 1, 2, 0] - - -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_regression_both_backends_match(make_df): - df = make_df(NUMERIC_DATA) - transformer = DecisionTreeFeatures(features_to_combine=2, random_state=0) - transformer.fit(df, REGRESSION_Y) - Xt = transformer.transform(df) - - result = nw.from_native(Xt, eager_only=True).to_dict(as_series=False) - assert result["tree(['Age', 'Height'])"] == pytest.approx( - [4.1, 6.475, 4.0, 6.475, 4.4, 4.4, 6.475, 4.4, 4.0, 6.475] - ) - - @pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_binary_classification_both_backends_match(make_df): - df = make_df(NUMERIC_DATA) - transformer = DecisionTreeFeatures(regression=False, random_state=0) - transformer.fit(df, BINARY_Y) - Xt = transformer.transform(df) +def test_n_jobs_parallel_matches_sequential(make_df): + # core correctness check for n_jobs: parallelizing tree training across + # feature combinations must produce identical trees, and therefore + # identical predictions, to sequential training (n_jobs=None). + X = make_df(DATA) + tr_seq = DecisionTreeFeatures(n_jobs=None, random_state=0) + tr_seq.fit(X, REGRESSION_Y) + tr_par = DecisionTreeFeatures(n_jobs=2, random_state=0) + tr_par.fit(X, REGRESSION_Y) - pd_df = pd.DataFrame(NUMERIC_DATA) - pd_transformer = DecisionTreeFeatures(regression=False, random_state=0) - pd_transformer.fit(pd_df, BINARY_Y) - expected = pd_transformer.transform(pd_df)["tree(['Age', 'Height'])"].tolist() + Xt_seq = tr_seq.transform(X) + Xt_par = tr_par.transform(X) - result = nw.from_native(Xt, eager_only=True).get_column( - "tree(['Age', 'Height'])" - ) - assert result.to_list() == pytest.approx(expected) - - -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_multiclass_classification_both_backends_match(make_df): - df = make_df(NUMERIC_DATA) - transformer = DecisionTreeFeatures(regression=False, random_state=0) - transformer.fit(df, MULTICLASS_Y) - Xt = transformer.transform(df) - - pd_df = pd.DataFrame(NUMERIC_DATA) - pd_transformer = DecisionTreeFeatures(regression=False, random_state=0) - pd_transformer.fit(pd_df, MULTICLASS_Y) - expected = pd_transformer.transform(pd_df)["tree(['Age', 'Height'])"].tolist() - - result = nw.from_native(Xt, eager_only=True).get_column( - "tree(['Age', 'Height'])" - ) - assert result.to_list() == pytest.approx(expected) - - -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_drop_original_both_backends(make_df): - df = make_df(NUMERIC_DATA) - transformer = DecisionTreeFeatures(features_to_combine=1, drop_original=True) - transformer.fit(df, REGRESSION_Y) - Xt = transformer.transform(df) - assert list(nw.from_native(Xt, eager_only=True).columns) == [ - "tree(Age)", - "tree(Height)", - ] - - -@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) -def test_get_feature_names_out_both_backends(make_df): - df = make_df(NUMERIC_DATA) - transformer = DecisionTreeFeatures(features_to_combine=1) - transformer.fit(df, REGRESSION_Y) - Xt = transformer.transform(df) - feat_out = list(nw.from_native(Xt, eager_only=True).columns) - assert transformer.get_feature_names_out(input_features=None) == feat_out + expected = nw.from_native(Xt_seq, eager_only=True).to_dict(as_series=False) + assert_df_equal(Xt_par, expected) -def test_single_int_named_feature_combo_both_backends(): +def test_single_int_named_feature_combo(): # regression test: a single-variable combo with an integer column name # used to crash (isinstance(features, str) missed the int case), since # X[features] for a bare int returns a 1D Series, not the 2D input # sklearn requires - fixed to check isinstance(features, (str, int)). + # Integer column names are pandas-only - polars requires string columns. df = pd.DataFrame({0: [1.0, 2, 3, 4, 5, 6, 7, 8], 1: [2.0, 3, 4, 5, 6, 7, 8, 9]}) y = [1.0, 2, 3, 4, 5, 6, 7, 8] transformer = DecisionTreeFeatures(features_to_combine=1, random_state=0) From ac184135cac68a04d80a36009b60ac79c8bce3a2 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Mon, 24 Aug 2026 21:06:58 +0200 Subject: [PATCH 6/7] fix: avoid pandas fragmentation warning in DecisionTreeFeatures.transform transform() assigned one new tree-prediction column at a time (X[col_name] = preds), which triggers pandas' "DataFrame is highly fragmented" PerformanceWarning once there are enough feature combinations - confirmed with 10 vars/features_to_combine=3 (175 new columns). .assign(**kwargs) does NOT fix this: it inserts columns one at a time internally too, same warning. The actual fix is building all new columns into one DataFrame and joining once (single insertion). Verified: output is byte-identical to the old behavior (pd.testing.assert_frame_equal on a 3000-row/9-var/129-tree case), drop_original still works, and a new regression test confirms the warning is gone (and fails against the old code, confirming it actually catches the regression). Co-Authored-By: Claude Sonnet 5 --- .../creation/decision_tree_features.py | 7 ++++- .../test_decision_tree_features.py | 26 +++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/feature_engine/creation/decision_tree_features.py b/feature_engine/creation/decision_tree_features.py index 0f9648c1f..1e31be965 100644 --- a/feature_engine/creation/decision_tree_features.py +++ b/feature_engine/creation/decision_tree_features.py @@ -428,6 +428,7 @@ def get_x_sub(features): return nw_X.select(features).to_native() new_series = [] + new_columns = {} # if regression or multiclass, we return the output of predict(); # if binary classification, we return the probability for features, estimator in zip(self.input_features_, self.estimators_): @@ -446,13 +447,17 @@ def get_x_sub(features): preds = estimator.predict(X_sub) if is_pandas is True: - X[col_name] = preds + new_columns[col_name] = preds else: new_series.append( nw.new_series(col_name, preds, backend=nw_X.implementation) ) if is_pandas is True: + # assign() still inserts columns one at a time internally, so it + # doesn't avoid fragmentation with many feature combinations; + # building one DataFrame and joining it does (single insertion). + X = X.join(type(X)(new_columns, index=X.index)) if self.drop_original is True: X = X.drop(columns=self.variables_) else: diff --git a/tests/test_creation/test_decision_tree_features.py b/tests/test_creation/test_decision_tree_features.py index b8f1c31cc..a97f68475 100644 --- a/tests/test_creation/test_decision_tree_features.py +++ b/tests/test_creation/test_decision_tree_features.py @@ -1,3 +1,5 @@ +import warnings + import narwhals as nw import numpy as np import pandas as pd @@ -435,6 +437,30 @@ def test_n_jobs_parallel_matches_sequential(make_df): assert_df_equal(Xt_par, expected) +def test_transform_does_not_fragment_pandas_output(): + # regression test: transform() used to assign one new tree column at a + # time (X[col_name] = preds), which triggers pandas' "DataFrame is + # highly fragmented" PerformanceWarning once there are enough feature + # combinations - fixed by building all new columns in one DataFrame + # and joining once. Needs enough variables to cross pandas' internal + # fragmentation threshold (a handful of combos won't trigger it). + rng = np.random.RandomState(0) + n_vars = 9 + X = pd.DataFrame( + rng.rand(200, n_vars), columns=[f"v{i}" for i in range(n_vars)] + ) + y = rng.rand(200) + + tr = DecisionTreeFeatures( + features_to_combine=3, param_grid={"max_depth": [1, 2]}, random_state=0 + ) + tr.fit(X, y) + + with warnings.catch_warnings(): + warnings.simplefilter("error", pd.errors.PerformanceWarning) + tr.transform(X) + + def test_single_int_named_feature_combo(): # regression test: a single-variable combo with an integer column name # used to crash (isinstance(features, str) missed the int case), since From c41a12692c65d3a8710a7eb1872d331971453f9d Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Mon, 24 Aug 2026 21:27:30 +0200 Subject: [PATCH 7/7] shorten docstring --- feature_engine/creation/decision_tree_features.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/feature_engine/creation/decision_tree_features.py b/feature_engine/creation/decision_tree_features.py index 1e31be965..135a62815 100644 --- a/feature_engine/creation/decision_tree_features.py +++ b/feature_engine/creation/decision_tree_features.py @@ -141,10 +141,7 @@ class DecisionTreeFeatures(TransformerMixin, BaseEstimator, GetFeatureNamesOutMi 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. There is a real speedup only when there are many - feature combinations and/or a large `param_grid` to search — with just - a handful of combinations, thread-dispatch overhead outweighs the gain, - which is why the default stays sequential. + available processors. {missing_values}