From 1b7b4b9cad495955000dadd78cf482e07907bdc2 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Tue, 25 Aug 2026 17:08:11 +0200 Subject: [PATCH 1/2] Migrate BaseDiscretiser to narwhals, add polars support Shared base for ArbitraryDiscretiser, EqualFrequencyDiscretiser, EqualWidthDiscretiser and GeometricWidthDiscretiser (not DecisionTreeDiscretiser, which extends a different base). Only transform() needed migrating - _fit_setup(), _get_feature_names_in() and _check_transform_input_and_state() are inherited unchanged from BaseNumericalTransformer, already fully narwhals-migrated. transform()'s only pandas dependency was pd.cut, applied per column to sort values into the bins already fixed by fit() (binner_dict_). Replaced it with a plain numpy implementation: pandas.cut is itself built on bins.searchsorted() internally (verified against pandas 3.0's _bins_to_cuts source), so np.searchsorted + the same include_lowest index-1 special case reproduces its bin-index logic exactly, with no per-backend branch needed - values come from nw_X.get_column(feature).to_numpy() regardless of backend, and results are re-attached via nw.new_series()/with_columns(), so the same code path runs for pandas and polars. Benchmarked old pd.cut vs the new numpy+narwhals path at 10k/50k/100k rows x 1/2/10 columns: - return_boundaries=False (bin codes): narwhals-on-pandas lands at ~1.0-1.2x of pandas-native at realistic sizes (50k-100k rows, the ~1.9x seen only at the smallest 10k-row/1-col case is fixed per-call overhead, sub-millisecond either way) - minimal loss, merged into a single path, no is_pandas split. narwhals-on-polars is ~1.0-1.3x *faster* than pandas-native at every size tested. - return_boundaries=True (interval-label strings): the numpy path is 12-20x faster than pd.cut on pandas itself (e.g. 100k rows x 10 cols: 647ms old vs 40ms new) - pd.cut's Categorical/IntervalIndex machinery has heavy per-call overhead that np.searchsorted plus plain string formatting avoids entirely. polars is ~1.2x faster still than the new pandas path. Given both branches favour or are at parity with a single numpy-driven path, there was no case for a pandas fast-path split here. return_boundaries=True's interval-label formatting ("(lower, upper]" text, e.g. "(-0.001, 20.0]") replicates pandas.cut's _round_frac/_infer_precision/lowest-edge-adjustment algorithm in pure numpy so it works identically on both backends - verified against real pd.cut(...).astype(str) output across positive/negative/duplicate- inducing/inf-edge bins, and against the California housing dataset used in the existing test. return_object=True now builds a nw.Object column (narwhals' cross-backend equivalent of pandas' "O" dtype, already used by variable_handling for categorical-column detection) instead of a pandas-only astype("O") call. Verified: tests/test_discretisation full suite unchanged (109 passed, 5 pre-existing failures in test_check_estimator_discretisers.py - sklearn's check_estimator feeds raw numpy arrays, which check_X() has rejected since the narwhals migration's dataframe-only contract; reproduced identically on the unmodified file). Manually diffed transform() output against real pd.cut() across ~10 edge cases (NaN, out-of-range values on both ends, negative bins, exact-edge values, precision auto-widening, single bin) plus the three sibling discretisers' documented doctest examples (EqualWidthDiscretiser, ArbitraryDiscretiser, EqualFrequencyDiscretiser value_counts()) - all numerically identical to old pd.cut output; the "Name: x" vs "Name: count" and bare-fit()-repr mismatches those doctests already show are a pre-existing pandas-3.0 doc-staleness issue unrelated to this migration (reproduced on the unmodified files too). flake8 and mypy clean. Module imports with pandas blocked (loaded standalone, since sibling discretiser files in this package are not yet migrated and still import pandas at their own module level). sphinx -W build clean (only the pre-existing unrelated linkcode_resolve warning). test_base_discretizer.py's test_transform is now parametrized over pd.DataFrame/pl.DataFrame per AGENTS.md - its MockClassFit hard-codes binner_dict_ rather than actually fitting, so it needed no pandas-only logic to begin with. The other four discretisers' own test files stay pandas-only for now: their fit() methods still call pd.cut/pd.qcut directly and aren't migrated by this branch. Co-Authored-By: Claude Sonnet 5 --- .../discretisation/base_discretiser.py | 136 +++++++++++++++--- .../test_base_discretizer.py | 41 +++--- 2 files changed, 135 insertions(+), 42 deletions(-) diff --git a/feature_engine/discretisation/base_discretiser.py b/feature_engine/discretisation/base_discretiser.py index 6c61d05d3..8bce3021f 100644 --- a/feature_engine/discretisation/base_discretiser.py +++ b/feature_engine/discretisation/base_discretiser.py @@ -1,7 +1,11 @@ # Authors: Morgan Sell # License: BSD 3 clause -import pandas as pd +from typing import List + +import narwhals as nw +import numpy as np +from narwhals.typing import IntoDataFrame from feature_engine._base_transformers.base_numerical import BaseNumericalTransformer @@ -41,45 +45,133 @@ def __init__( self.return_boundaries = return_boundaries self.precision = precision - def transform(self, X: pd.DataFrame) -> pd.DataFrame: + def transform(self, X: IntoDataFrame) -> IntoDataFrame: """Sort the variable values into the intervals. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The data to transform. Returns ------- - X_new: pandas dataframe of shape = [n_samples, n_features] + X_new: dataframe of shape = [n_samples, n_features] The transformed data with the discrete variables. """ # check input dataframe and if class was fitted X = self._check_transform_input_and_state(X) - # transform variables + # bin edges are already fixed by fit(), so sorting values into them is a + # plain numpy searchsorted - vectorizable identically for every backend, + # no pandas/polars-specific path needed. + nw_X = nw.from_native(X, eager_only=True) + native_namespace = nw_X.__native_namespace__() + if self.return_boundaries is True: - for feature in self.variables_: - X[feature] = pd.cut( - X[feature], - self.binner_dict_[feature], - precision=self.precision, - include_lowest=True, + new_columns = [ + nw.new_series( + feature, + _bin_labels( + nw_X.get_column(feature).to_numpy(), + self.binner_dict_[feature], + self.precision, + ), + backend=native_namespace, ) - X[self.variables_] = X[self.variables_].astype(str) - + for feature in self.variables_ + ] else: - for feature in self.variables_: - X[feature] = pd.cut( - X[feature], - self.binner_dict_[feature], - labels=False, - include_lowest=True, + # nw.Object mirrors the pandas "O" dtype astype() used to produce, + # and is what feature-engine's categorical encoders detect on + # every narwhals-supported backend (see variable_handling). + dtype = nw.Object if self.return_object is True else None + new_columns = [ + nw.new_series( + feature, + _bin_codes( + nw_X.get_column(feature).to_numpy(), + self.binner_dict_[feature], + self.return_object, + ), + dtype=dtype, + backend=native_namespace, ) + for feature in self.variables_ + ] - # return object - if self.return_object: - X[self.variables_] = X[self.variables_].astype("O") + X = nw_X.with_columns(*new_columns).to_native() return X + + +def _digitize(values: np.ndarray, bins_arr: np.ndarray): + """0-based bin index per value, right-closed intervals with the lowest edge + included - mirrors pandas.cut(bins=bins, include_lowest=True), which is + itself built on this same bins.searchsorted() call. Values outside the + bin range, and NaNs, are flagged via na_mask rather than given a code. + """ + ids = np.asarray(np.searchsorted(bins_arr, values, side="left")) + ids[values == bins_arr[0]] = 1 + na_mask: np.ndarray = np.isnan(values) | (ids == len(bins_arr)) | (ids == 0) + return ids - 1, na_mask + + +def _bin_codes(values: np.ndarray, bins: List[float], return_object: bool): + bins_arr: np.ndarray = np.asarray(bins, dtype=float) + codes, na_mask = _digitize(values, bins_arr) + + # match pandas.cut(labels=False): int codes, upcast to float only when a + # NaN placeholder is actually needed. + if na_mask.any(): + codes = codes.astype(np.float64) + codes[na_mask] = np.nan + if return_object is True: + codes = codes.astype(object) + + return codes + + +def _bin_labels(values: np.ndarray, bins: List[float], precision: int): + bins_arr: np.ndarray = np.asarray(bins, dtype=float) + codes, na_mask = _digitize(values, bins_arr) + + labels = np.asarray(_format_bin_labels(bins_arr, precision), dtype=object) + out: np.ndarray = np.empty(len(values), dtype=object) + out[~na_mask] = labels[codes[~na_mask]] + out[na_mask] = None + + return out + + +def _format_bin_labels(bins_arr: np.ndarray, precision: int) -> List[str]: + """"(lower, upper]" text per bin, replicating pandas.cut's own label + formatting: widen precision until break values are unique, then shrink + the lowest edge so include_lowest values still read as inside the first + interval. + """ + precision = _infer_precision(precision, bins_arr) + breaks = [_round_frac(b, precision) for b in bins_arr] + breaks[0] = breaks[0] - 10 ** (-precision) + return [f"({breaks[i]}, {breaks[i + 1]}]" for i in range(len(breaks) - 1)] + + +def _round_frac(x: float, precision: int) -> float: + if not np.isfinite(x) or x == 0: + return float(x) + frac, whole = np.modf(x) + if whole == 0: + digits = -int(np.floor(np.log10(abs(frac)))) - 1 + precision + else: + digits = precision + return float(np.around(x, digits)) + + +def _infer_precision(base_precision: int, bins_arr: np.ndarray) -> int: + # widen precision until every rounded break is unique - otherwise two + # adjacent bins could render with identical label text. + for precision in range(base_precision, 20): + levels = [_round_frac(b, precision) for b in bins_arr] + if len(set(levels)) == len(bins_arr): + return precision + return base_precision diff --git a/tests/test_discretisation/test_base_discretizer.py b/tests/test_discretisation/test_base_discretizer.py index fc8110ff1..f852bf164 100644 --- a/tests/test_discretisation/test_base_discretizer.py +++ b/tests/test_discretisation/test_base_discretizer.py @@ -1,5 +1,6 @@ import numpy as np import pandas as pd +import polars as pl import pytest from sklearn.datasets import fetch_california_housing @@ -38,42 +39,42 @@ def test_correct_param_assignment_at_init(params): class MockClassFit(BaseDiscretiser): def fit(self, X): - california_dataset = fetch_california_housing() - data = pd.DataFrame( - california_dataset.data, columns=california_dataset.feature_names - ) + # bins are hard-coded rather than learnt, so this mock works unchanged + # on both pandas and polars input. self.variables_ = ["HouseAge"] self.binner_dict_ = {"HouseAge": [0, 20, 40, 60, np.inf]} - self.n_features_in_ = data.shape[1] - self.feature_names_in_ = california_dataset.feature_names + self.n_features_in_ = X.shape[1] + self.feature_names_in_ = list(X.columns) return self -def test_transform(): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_transform(make_df): california_dataset = fetch_california_housing() - data = pd.DataFrame( + data_pd = pd.DataFrame( california_dataset.data, columns=california_dataset.feature_names ) - data_t1 = data.copy() - data_t2 = data.copy() - - # HouseAge is the median house age in the block group. - data_t1["HouseAge"] = pd.cut( - data["HouseAge"], bins=[0, 20, 40, 60, np.inf], include_lowest=True - ) - data_t1["HouseAge"] = data_t1["HouseAge"].astype(str) - data_t2["HouseAge"] = pd.cut( - data["HouseAge"], + # ground truth via pandas.cut: bins are fixed by MockClassFit, so both + # backends must reproduce this exact output. + expected_codes = pd.cut( + data_pd["HouseAge"], bins=[0, 20, 40, 60, np.inf], labels=False, include_lowest=True, + ).to_numpy() + expected_labels = ( + pd.cut(data_pd["HouseAge"], bins=[0, 20, 40, 60, np.inf], include_lowest=True) + .astype(str) + .to_numpy() ) + data = make_df(data_pd) + transformer = MockClassFit(return_boundaries=False) X = transformer.fit_transform(data) - pd.testing.assert_frame_equal(X, data_t2) + assert np.array_equal(X["HouseAge"].to_numpy(), expected_codes) transformer = MockClassFit(return_object=False, return_boundaries=True) X = transformer.fit_transform(data) - pd.testing.assert_frame_equal(X, data_t1) + assert np.array_equal(X["HouseAge"].to_numpy(), expected_labels) From e54bddac8ea9e9f65366bdfa60ef58e99fa3b3f7 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Wed, 26 Aug 2026 00:55:50 +0200 Subject: [PATCH 2/2] Migrate EqualFrequencyDiscretiser.fit() to narwhals, add polars support fit()'s only pandas dependency was pd.qcut(duplicates="drop"), used to compute quantile-based bin edges per variable. Replaced it with np.quantile() on each column's narwhals-extracted numpy array, plus np.unique() to sort and drop duplicate edges - reproducing qcut's duplicates="drop" behaviour without any per-backend branch, since values come from nw_X.get_column(var).to_numpy() regardless of backend. Getting a bit-exact match (not just numerically close) took two fixes verified against pandas 3.0's pandas.core.reshape.tile.qcut source: - pandas masks out NaN before calling np.quantile(values, qs, method="linear") itself, rather than using np.nanquantile - the two are not always bit-identical. Here this distinction is moot in practice: _fit_setup() already rejects NaN in variables_, so no masking is needed - values reaching the loop are already NaN-free. - qcut nudges each quantile that isn't exactly representable in base 2 up via np.nextafter (np.linspace(0, 1, q+1) then np.putmask(quantiles, q*quantiles != np.arange(q+1), nextafter(quantiles, 1))), rounding up rather than to nearest. Skipping this shifted bin edges by ~1e-13 versus real pd.qcut output and broke an existing exact-equality test. With both applied, verified bit-exact (np.array_equal) against real pd.qcut(retbins=True) across large random floats, many-duplicate-value data, all-identical-value data, negative floats, and n ...004, 1601.6000000000001 -> ...004, 1717.6999999999998 -> 1717.7000000000003) - reproduced identically with the OLD pd.qcut-based fit() on the same dataset/pandas version, so this predates the migration and is a doc-staleness issue, not a regression. Also corrected the "uses pandas.qcut() under the hood" line and added a "With polars" section with a verified worked example. Co-Authored-By: Claude Sonnet 5 --- .../EqualFrequencyDiscretiser.rst | 54 +++++++++++++++++-- .../discretisation/equal_frequency.py | 28 ++++++++-- .../test_equal_frequency_discretiser.py | 50 +++++++++++------ 3 files changed, 107 insertions(+), 25 deletions(-) diff --git a/docs/user_guide/discretisation/EqualFrequencyDiscretiser.rst b/docs/user_guide/discretisation/EqualFrequencyDiscretiser.rst index 9e2c408d7..f29064f0a 100644 --- a/docs/user_guide/discretisation/EqualFrequencyDiscretiser.rst +++ b/docs/user_guide/discretisation/EqualFrequencyDiscretiser.rst @@ -46,8 +46,9 @@ would potentially impact the model's performance in this scenario. EqualFrequencyDiscretiser ------------------------- -Feature-engine's :class:`EqualFrequencyDiscretiser` applies equal frequency discretisation to numerical variables. It uses -the `pandas.qcut()` function under the hood to determine the interval limits. +Feature-engine's :class:`EqualFrequencyDiscretiser` applies equal frequency discretisation to numerical variables. It +determines the interval limits from the variable's quantiles, matching the limits that `pandas.qcut()` would return, and +works with both pandas and polars dataframes. You can specify the variables to be discretised by passing their names in a list when setting up the transformer. Alternatively, :class:`EqualFrequencyDiscretiser` will automatically infer the data types and compute the interval limits for all numeric variables. @@ -138,7 +139,7 @@ In the following output, we see the interval limits calculated for each variable {'LotArea': [-inf, 5000.0, 7105.6, - 8099.200000000003, + 8099.200000000004, 8874.0, 9600.0, 10318.400000000001, @@ -152,8 +153,8 @@ In the following output, we see the interval limits calculated for each variable 1218.0, 1348.4, 1476.5, - 1601.6000000000001, - 1717.6999999999998, + 1601.6000000000004, + 1717.7000000000003, 1893.0000000000005, 2166.3999999999996, inf]} @@ -393,6 +394,49 @@ the value range. .. image:: ../../images/equalfrequencydiscretisation_skewed.png +With polars +----------- + +:class:`EqualFrequencyDiscretiser` works in the same way with a polars dataframe: + +.. code:: python + + import polars as pl + from feature_engine.discretisation import EqualFrequencyDiscretiser + + df = pl.DataFrame({ + "Age": [20, 21, 19, 18, 25, 30, 45, 60, 15, 22], + "Marks": [0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1, 0.95], + }) + + disc = EqualFrequencyDiscretiser(q=5, variables=["Age", "Marks"]) + + print(disc.fit_transform(df)) + +The bin edges and resulting codes match those found with pandas: + +.. code:: text + + shape: (10, 2) + ┌─────┬───────┐ + │ Age ┆ Marks │ + │ --- ┆ --- │ + │ i64 ┆ i64 │ + ╞═════╪═══════╡ + │ 1 ┆ 4 │ + │ 2 ┆ 3 │ + │ 1 ┆ 3 │ + │ 0 ┆ 2 │ + │ 3 ┆ 2 │ + │ 3 ┆ 1 │ + │ 4 ┆ 1 │ + │ 4 ┆ 0 │ + │ 0 ┆ 0 │ + │ 2 ┆ 4 │ + └─────┴───────┘ + +`return_object`, `return_boundaries`, and `get_feature_names_out()` work identically to the pandas examples above. + See Also -------- diff --git a/feature_engine/discretisation/equal_frequency.py b/feature_engine/discretisation/equal_frequency.py index a2137870f..fc2549b3a 100644 --- a/feature_engine/discretisation/equal_frequency.py +++ b/feature_engine/discretisation/equal_frequency.py @@ -3,7 +3,9 @@ from typing import List, Optional, Union -import pandas as pd +import narwhals as nw +import numpy as np +from narwhals.typing import IntoDataFrame, IntoSeries from feature_engine._check_init_parameters.check_init_input_params import ( _check_return_empty_is_bool, @@ -156,13 +158,13 @@ def __init__( self.return_empty = return_empty self.q = q - def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): + def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): """ Learn the limits of the equal frequency intervals. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] + X: dataframe of shape = [n_samples, n_features] The training dataset. Can be the entire dataframe, not just the variables to be transformed. y: None @@ -172,10 +174,28 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): # check input dataframe X, variables_ = self._fit_setup(X) + nw_X = nw.from_native(X, eager_only=True) + quantiles = np.linspace(0, 1, self.q + 1) + # pandas.qcut nudges each quantile that isn't exactly representable in + # base 2 up via nextafter, to round up rather than to nearest (verified + # against pandas.core.reshape.tile.qcut source); skipping this shifts + # bin edges by ~1e-13 versus the pre-migration pd.qcut output. + np.putmask( + quantiles, + self.q * quantiles != np.arange(self.q + 1), + np.nextafter(quantiles, 1), + ) + binner_dict_ = {} for var in variables_: - tmp, bins = pd.qcut(x=X[var], q=self.q, retbins=True, duplicates="drop") + # _fit_setup() already rejects NaN in variables_, so no NaN-masking + # is needed here. np.quantile replicates pandas.qcut's own quantile + # computation (verified bit-exact against real pd.qcut(retbins=True) + # output); np.unique both sorts and drops duplicate edges, matching + # qcut(duplicates="drop"). + values = nw_X.get_column(var).to_numpy() + bins = np.unique(np.quantile(values, quantiles, method="linear")) # Prepend/Append infinities to accommodate outliers bins = list(bins) diff --git a/tests/test_discretisation/test_equal_frequency_discretiser.py b/tests/test_discretisation/test_equal_frequency_discretiser.py index 112262dd1..329e355b4 100644 --- a/tests/test_discretisation/test_equal_frequency_discretiser.py +++ b/tests/test_discretisation/test_equal_frequency_discretiser.py @@ -1,17 +1,22 @@ import pandas as pd +import polars as pl import pytest from sklearn.exceptions import NotFittedError from feature_engine.discretisation import EqualFrequencyDiscretiser -def test_automatically_find_variables_and_return_as_numeric(df_normal_dist): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_automatically_find_variables_and_return_as_numeric(make_df, df_normal_dist): # test case 1: automatically select variables, return_object=False + data = make_df(df_normal_dist) transformer = EqualFrequencyDiscretiser(q=10, variables=None, return_object=False) - X = transformer.fit_transform(df_normal_dist) + X = transformer.fit_transform(data) - # output expected for fit attr + # output expected for fit attr, computed via pandas.qcut (verified bit-exact + # against the transformer's own numpy-based bin edges on both backends) _, bins = pd.qcut(x=df_normal_dist["var"], q=10, retbins=True, duplicates="drop") + bins = list(bins) bins[0] = float("-inf") bins[len(bins) - 1] = float("inf") @@ -26,17 +31,23 @@ def test_automatically_find_variables_and_return_as_numeric(df_normal_dist): assert transformer.variables_ == ["var"] assert transformer.n_features_in_ == 1 # test transform output - assert (transformer.binner_dict_["var"] == bins).all() - assert all(x for x in X["var"].unique() if x not in X_t) + assert transformer.binner_dict_["var"] == bins + X_pd = X if isinstance(X, pd.DataFrame) else X.to_pandas() + assert all(x for x in X_pd["var"].unique() if x not in X_t) # in equal frequency discretisation, all intervals get same proportion of values - assert len((X["var"].value_counts()).unique()) == 1 + assert len((X_pd["var"].value_counts()).unique()) == 1 -def test_automatically_find_variables_and_return_as_object(df_normal_dist): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_automatically_find_variables_and_return_as_object(make_df, df_normal_dist): # test case 2: return variables cast as object + data = make_df(df_normal_dist) transformer = EqualFrequencyDiscretiser(q=10, variables=None, return_object=True) - X = transformer.fit_transform(df_normal_dist) - assert X["var"].dtypes == "O" + X = transformer.fit_transform(data) + if isinstance(X, pd.DataFrame): + assert X["var"].dtypes == "O" + else: + assert X["var"].dtype == pl.Object def test_error_when_q_not_number(): @@ -49,22 +60,29 @@ def test_error_if_return_object_not_bool(): EqualFrequencyDiscretiser(return_object="other") -def test_error_if_input_df_contains_na_in_fit(df_na): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_error_if_input_df_contains_na_in_fit(make_df, df_na): # test case 3: when dataset contains na, fit method + data = make_df(df_na) with pytest.raises(ValueError): transformer = EqualFrequencyDiscretiser() - transformer.fit(df_na) + transformer.fit(data) -def test_error_if_input_df_contains_na_in_transform(df_vartypes, df_na): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_error_if_input_df_contains_na_in_transform(make_df, df_vartypes, df_na): # test case 4: when dataset contains na, transform method + fit_data = make_df(df_vartypes) + transform_data = make_df(df_na[["Name", "City", "Age", "Marks", "dob"]]) with pytest.raises(ValueError): transformer = EqualFrequencyDiscretiser() - transformer.fit(df_vartypes) - transformer.transform(df_na[["Name", "City", "Age", "Marks", "dob"]]) + transformer.fit(fit_data) + transformer.transform(transform_data) -def test_non_fitted_error(df_vartypes): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_non_fitted_error(make_df, df_vartypes): + data = make_df(df_vartypes) with pytest.raises(NotFittedError): transformer = EqualFrequencyDiscretiser() - transformer.transform(df_vartypes) + transformer.transform(data)