From f9051f727841a6ef39d3c3b15404c67a70a0c11b Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Fri, 18 Sep 2026 14:35:05 +0200 Subject: [PATCH 1/2] Support integer column names in DecisionTreeDiscretiser and EqualWidthDiscretiser Co-Authored-By: Claude Opus 5 --- feature_engine/discretisation/decision_tree.py | 14 +++++--------- feature_engine/discretisation/equal_width.py | 3 ++- .../test_decision_tree_discretiser.py | 15 +++++++++++++++ .../test_equal_width_discretiser.py | 14 ++++++++++++++ 4 files changed, 36 insertions(+), 10 deletions(-) diff --git a/feature_engine/discretisation/decision_tree.py b/feature_engine/discretisation/decision_tree.py index d7a3a45e8..a4f09c890 100644 --- a/feature_engine/discretisation/decision_tree.py +++ b/feature_engine/discretisation/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 @@ -374,14 +373,11 @@ def transform(self, X: IntoDataFrame) -> IntoDataFrame: values = nw_X.get_column(feature).to_numpy() new_columns[feature] = self._bin_index(values, thresholds) - if nwd.is_pandas_dataframe(X) is True: - X = X.assign(**new_columns) - else: - new_series = [ - nw.new_series(name, values, backend=nw_X.implementation) - for name, values in new_columns.items() - ] - X = nw_X.with_columns(*new_series).to_native() + new_series = [ + nw.new_series(name, values, backend=nw_X.implementation) + for name, values in new_columns.items() + ] + X = nw_X.with_columns(*new_series).to_native() return X diff --git a/feature_engine/discretisation/equal_width.py b/feature_engine/discretisation/equal_width.py index eaf09ddb3..8930fd337 100644 --- a/feature_engine/discretisation/equal_width.py +++ b/feature_engine/discretisation/equal_width.py @@ -189,7 +189,8 @@ def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): if len(variables_) > 0: # one narwhals call for every variable at once, instead of a # get_column() round-trip per variable. - arr = nw.from_native(X, eager_only=True).select(variables_).to_numpy() + nw_X = nw.from_native(X, eager_only=True) + arr = nw_X.select(nw.col(variables_)).to_numpy() mins = arr.min(axis=0) maxs = arr.max(axis=0) for var, mn, mx in zip(variables_, mins, maxs): diff --git a/tests/test_discretisation/test_decision_tree_discretiser.py b/tests/test_discretisation/test_decision_tree_discretiser.py index 04b4bb892..c380384ca 100644 --- a/tests/test_discretisation/test_decision_tree_discretiser.py +++ b/tests/test_discretisation/test_decision_tree_discretiser.py @@ -1,6 +1,7 @@ import re import numpy as np +import pandas as pd import pytest from sklearn.exceptions import NotFittedError @@ -370,6 +371,20 @@ def test_regression_rounds_predictions(make_df, data_normal_dist, params): assert sorted(set(frame_to_dict(Xt)["var"])) == sorted(params[1]) +@pytest.mark.parametrize("bin_output", ["prediction", "bin_number", "boundaries"]) +def test_integer_column_names(bin_output): + # integer column names are pandas-only + X = pd.DataFrame({0: DATA_TWO_VARS["var_A"], 1: DATA_TWO_VARS["var_B"]}) + y = pd.Series(TARGET_TWO_VARS) + params = dict(bin_output=bin_output, precision=3, regression=False, random_state=0) + + Xt = DecisionTreeDiscretiser(**params).fit_transform(X, y) + expected = DecisionTreeDiscretiser(**params).fit_transform(X.rename(columns=str), y) + + assert list(Xt.columns) == [0, 1] + assert Xt.to_numpy().tolist() == expected.to_numpy().tolist() + + def test_non_fitted_error(make_df, data_normal_dist): transformer = DecisionTreeDiscretiser() msg = ( diff --git a/tests/test_discretisation/test_equal_width_discretiser.py b/tests/test_discretisation/test_equal_width_discretiser.py index ebfc3bb3d..7a85f1b94 100644 --- a/tests/test_discretisation/test_equal_width_discretiser.py +++ b/tests/test_discretisation/test_equal_width_discretiser.py @@ -115,6 +115,20 @@ def test_error_if_input_df_contains_na_in_transform(make_df, data_vartypes, data transformer.transform(transform_data) +def test_integer_column_names(data_normal_dist): + # integer column names are pandas-only + values = data_normal_dist["var"] + X = pd.DataFrame({0: values, 1: [2 * v for v in values]}) + + transformer = EqualWidthDiscretiser(bins=10) + Xt = transformer.fit_transform(X) + expected = EqualWidthDiscretiser(bins=10).fit_transform(X.rename(columns=str)) + + assert list(transformer.binner_dict_) == [0, 1] + assert list(Xt.columns) == [0, 1] + assert Xt.to_numpy().tolist() == expected.to_numpy().tolist() + + def test_non_fitted_error(make_df, data_vartypes): transformer = EqualWidthDiscretiser() msg = ( From 202fe06fd459082bbc323a2d057eeb742fe49dc0 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Fri, 18 Sep 2026 14:56:38 +0200 Subject: [PATCH 2/2] Keep the pandas fast path in EqualWidthDiscretiser.fit Co-Authored-By: Claude Opus 5 --- feature_engine/discretisation/equal_width.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/feature_engine/discretisation/equal_width.py b/feature_engine/discretisation/equal_width.py index 8930fd337..d8344091e 100644 --- a/feature_engine/discretisation/equal_width.py +++ b/feature_engine/discretisation/equal_width.py @@ -4,6 +4,7 @@ from typing import List, Optional, Union import narwhals as nw +import narwhals.dependencies as nwd import numpy as np from narwhals.typing import IntoDataFrame, IntoSeries @@ -187,10 +188,12 @@ def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): binner_dict_ = {} if len(variables_) > 0: - # one narwhals call for every variable at once, instead of a - # get_column() round-trip per variable. - nw_X = nw.from_native(X, eager_only=True) - arr = nw_X.select(nw.col(variables_)).to_numpy() + # one call for all variables; pandas is faster than narwhals here + if nwd.is_pandas_dataframe(X) is True: + arr = X[variables_].to_numpy() + else: + nw_X = nw.from_native(X, eager_only=True) + arr = nw_X.select(nw.col(variables_)).to_numpy() mins = arr.min(axis=0) maxs = arr.max(axis=0) for var, mn, mx in zip(variables_, mins, maxs):