Support integer column names in DecisionTreeDiscretiser and EqualWidthDiscretiser

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
 

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,9 +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.
-            arr = nw.from_native(X, eager_only=True).select(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):

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 = (

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 = (