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 2df003c941e33a9fcdbaa61ab85423da4a34a412 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Wed, 26 Aug 2026 00:55:57 +0200 Subject: [PATCH 2/2] Migrate ArbitraryDiscretiser to narwhals, add polars support fit() needed no changes: it already delegates entirely to the already-migrated FitFromDictMixin._fit_from_dict(). The pandas dependency was in transform()'s post-hoc NaN-introduced check, which used X[...].isnull().sum().sum() / .columns / .any() / .tolist() - pandas-only calls that broke outright on polars input coming back from the now-migrated BaseDiscretiser.transform(). Replaced it with a narwhals-based per-column check, branched on return_boundaries rather than dtype: labels (return_boundaries=True) use None for missing values, which narwhals' is_null() detects correctly on both backends. Codes (return_boundaries=False) are numeric, so a numpy float cast + np.isnan is used instead of is_null()/is_nan() directly. That numeric-cast branch isn't just style - narwhals' is_null() (and polars' own null semantics) do NOT see a boxed np.nan sitting inside a polars Object-dtype column (return_object=True's output dtype): verified with a direct repro, is_null().any() returns False on a polars Object series holding all-NaN values, silently swallowing the warning/error this method exists to raise. is_nan() isn't usable there either - narwhals raises "is_nan only supported for numeric dtype, not Object". The numpy-float-cast approach sidesteps both issues and was confirmed to raise/warn correctly across all pandas/polars x return_object x return_boundaries combinations. Benchmarked old (pandas-only) vs new (narwhals) transform() at 10k/50k/100k rows x 1/2/10 cols on pandas input: return_object=False lands at parity (0.9-1.05x, within noise); return_object=True is 1.15-1.3x slower (e.g. 100k rows x 10 cols: 36.2ms old vs 44.9ms new) since the per-variable numpy float-cast replaces one vectorized pandas isnull().sum().sum() call. This falls within the "minimal loss" band used to decide against a pandas/polars split elsewhere in this migration, so a single narwhals-driven path was kept - no is_pandas branch was added. narwhals-on-polars is faster than narwhals-on-pandas at every size tested, consistent with the base branch's own findings. Verified: tests/test_discretisation full suite (114 passed, same 5 pre-existing check_estimator failures as the unmodified base branch - reproduced there too, predates this change). Rewrote test_arbitrary_discretiser.py per AGENTS.md: one parametrized test per behavior over pd.DataFrame/pl.DataFrame (previously pandas-only), switched pytest.raises()/pytest.warns() to the match= form instead of capturing and asserting on the record. flake8 and mypy clean. Module imports with pandas blocked. sphinx -W build clean (only the pre-existing unrelated linkcode_resolve warning). Verified the existing docstring/rst examples against real output before touching: the "Name: x" vs "Name: count" and bare-fit()-repr doctest mismatches are the same pre-existing pandas-3.0 doc-staleness noted in the base branch commit (reproduced on the unmodified file too) - left alone, out of scope here. Added a "With polars" example to both the class docstring and ArbitraryDiscretiser.rst, output verified against a real run. Co-Authored-By: Claude Sonnet 5 --- .../discretisation/ArbitraryDiscretiser.rst | 36 ++++++ feature_engine/discretisation/arbitrary.py | 63 +++++++-- .../test_arbitrary_discretiser.py | 120 ++++++++---------- 3 files changed, 137 insertions(+), 82 deletions(-) diff --git a/docs/user_guide/discretisation/ArbitraryDiscretiser.rst b/docs/user_guide/discretisation/ArbitraryDiscretiser.rst index b4d81e604..9c42a2cd1 100644 --- a/docs/user_guide/discretisation/ArbitraryDiscretiser.rst +++ b/docs/user_guide/discretisation/ArbitraryDiscretiser.rst @@ -110,6 +110,42 @@ obtain monotonic relationships between the variable and the target, you can do s seamlessly by setting `return_object` to True. You can find an example of discretisation followed by encoding to obtain monotonic releationships `here `_. +With polars +----------- + +:class:`ArbitraryDiscretiser()` also works with polars dataframes. + +.. code:: python + + import polars as pl + import numpy as np + from feature_engine.discretisation import ArbitraryDiscretiser + + X = pl.DataFrame({ + "MedInc": [1.5, 3.0, 5.0, 8.0, 0.8], + }) + + user_dict = {"MedInc": [0, 2, 4, 6, np.inf]} + + transformer = ArbitraryDiscretiser(binning_dict=user_dict, return_boundaries=False) + X_t = transformer.fit_transform(X) + print(X_t) + +.. code:: text + + shape: (5, 1) + ┌────────┐ + │ MedInc │ + │ --- │ + │ i64 │ + ╞════════╡ + │ 0 │ + │ 1 │ + │ 2 │ + │ 3 │ + │ 0 │ + └────────┘ + Additional resources -------------------- diff --git a/feature_engine/discretisation/arbitrary.py b/feature_engine/discretisation/arbitrary.py index 5776cd71d..4a6c75fd8 100644 --- a/feature_engine/discretisation/arbitrary.py +++ b/feature_engine/discretisation/arbitrary.py @@ -4,7 +4,9 @@ import warnings from typing import Dict, List, Optional, Union -import pandas as pd +import narwhals as nw +import numpy as np +from narwhals.typing import IntoDataFrame, IntoSeries from feature_engine._base_transformers.mixins import FitFromDictMixin from feature_engine._docstrings.fit_attributes import ( @@ -110,6 +112,27 @@ class ArbitraryDiscretiser(BaseDiscretiser, FitFromDictMixin): 3 25 1 17 Name: x, dtype: int64 + + With polars: + + >>> import polars as pl + >>> from feature_engine.discretisation import ArbitraryDiscretiser + >>> X = pl.DataFrame({"x": [10, 30, 60, 90]}) + >>> bins = dict(x=[0, 25, 50, 75, 100]) + >>> ad = ArbitraryDiscretiser(binning_dict=bins) + >>> ad.fit(X) + >>> ad.transform(X) + shape: (4, 1) + ┌─────┐ + │ x │ + │ --- │ + │ i64 │ + ╞═════╡ + │ 0 │ + │ 1 │ + │ 2 │ + │ 3 │ + └─────┘ """ def __init__( @@ -138,13 +161,13 @@ def __init__( self.binning_dict = binning_dict self.errors = errors - def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): + def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): """ This transformer does not learn any parameter. 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. @@ -161,30 +184,44 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): return self - 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. """ X = super().transform(X) - # check if NaN values were introduced by the discretisation procedure. - if X[self.variables_].isnull().sum().sum() > 0: - - # obtain the name(s) of the columns with null values - nan_columns = ( - X[self.variables_].columns[X[self.variables_].isnull().any()].tolist() - ) + # check if NaN values were introduced by the discretisation procedure. + nw_X = nw.from_native(X, eager_only=True) + if self.return_boundaries is True: + # missing labels are set to None by _bin_labels() + nan_columns = [ + var + for var in self.variables_ + if nw_X.get_column(var).is_null().any() + ] + else: + # codes are numeric; when return_object=True they're boxed as + # python floats in an Object column, where polars' is_null()/ + # is_nan() can't see NaN - a numpy float cast is reliable on + # both backends. + nan_columns = [ + var + for var in self.variables_ + if np.isnan(nw_X.get_column(var).to_numpy().astype(float)).any() + ] + + if len(nan_columns) > 0: if len(nan_columns) > 1: nan_columns_str = ", ".join(nan_columns) else: diff --git a/tests/test_discretisation/test_arbitrary_discretiser.py b/tests/test_discretisation/test_arbitrary_discretiser.py index f1b2db712..6827c2c48 100644 --- a/tests/test_discretisation/test_arbitrary_discretiser.py +++ b/tests/test_discretisation/test_arbitrary_discretiser.py @@ -1,35 +1,39 @@ +import re + +import narwhals as nw import numpy as np import pandas as pd +import polars as pl import pytest -from numpy.random import default_rng -from scipy.stats import skewnorm from sklearn.datasets import fetch_california_housing from feature_engine.discretisation import ArbitraryDiscretiser -def test_arbitrary_discretiser(): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_arbitrary_discretiser(make_df): california_dataset = fetch_california_housing() - data = pd.DataFrame( + data_pd = pd.DataFrame( california_dataset.data, columns=california_dataset.feature_names ) user_dict = {"HouseAge": [0, 20, 40, 60, np.inf]} - 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 user-supplied and fixed, 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 = ArbitraryDiscretiser( binning_dict=user_dict, return_object=False, return_boundaries=False ) @@ -42,53 +46,38 @@ def test_arbitrary_discretiser(): assert transformer.variables_ == ["HouseAge"] assert transformer.binner_dict_ == user_dict # transform params - pd.testing.assert_frame_equal(X, data_t2) + result_codes = nw.from_native(X, eager_only=True).get_column("HouseAge").to_numpy() + assert np.array_equal(result_codes, expected_codes) transformer = ArbitraryDiscretiser( binning_dict=user_dict, return_object=False, return_boundaries=True ) X = transformer.fit_transform(data) - pd.testing.assert_frame_equal(X, data_t1) + result_labels = nw.from_native(X, eager_only=True).get_column("HouseAge").to_numpy() + assert np.array_equal(result_labels, expected_labels) -def test_error_if_input_df_contains_na_in_transform(df_vartypes, df_na): - # test case 1: when dataset contains na, transform method +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_error_if_input_df_contains_na_in_transform(make_df): + # test case 1: when dataset contains na, transform method raises age_dict = {"Age": [0, 10, 20, 30, np.inf]} + data = make_df({"Age": [20.0, 21.0, 19.0, 18.0]}) + data_na = make_df({"Age": [20.0, 21.0, None, 18.0]}) - with pytest.raises(ValueError): - transformer = ArbitraryDiscretiser(binning_dict=age_dict) - transformer.fit(df_vartypes) - transformer.transform(df_na[["Name", "City", "Age", "Marks", "dob"]]) - - -def test_error_when_nan_introduced_during_transform(): - # test error when NA are introduced during the discretisation. - rng = default_rng() - - # create dataframe with 2 variables, 1 normal and 1 skewed - random = skewnorm.rvs(a=-50, loc=4, size=100) - random = random - min(random) # Shift so the minimum value is equal to zero. - - train = pd.concat( - [ - pd.Series(rng.standard_normal(100)), - pd.Series(random), - ], - axis=1, - ) - - train.columns = ["var_a", "var_b"] + transformer = ArbitraryDiscretiser(binning_dict=age_dict) + transformer.fit(data) + with pytest.raises(ValueError, match="Some of the variables in the dataset"): + transformer.transform(data_na) - # create a dataframe with 2 variables normally distributed - test = pd.concat( - [ - pd.Series(rng.standard_normal(100)), - pd.Series(rng.standard_normal(100)), - ], - axis=1, - ) - test.columns = ["var_a", "var_b"] +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +@pytest.mark.parametrize("return_object", [False, True]) +def test_error_when_nan_introduced_during_transform(make_df, return_object): + # test warning/error when NA are introduced during the discretisation, + # i.e. when a value in the data to transform falls outside the bin edges + # fitted on the training data. + train = make_df({"var_a": [-4.0, -1.0, 1.0, 4.0], "var_b": [1.0, 2.0, 3.0, 4.0]}) + test = make_df({"var_a": [-4.0, -1.0, 1.0, 4.0], "var_b": [10.0, 20.0, 30.0, 40.0]}) msg = ( "During the discretisation, NaN values were introduced " @@ -98,29 +87,25 @@ def test_error_when_nan_introduced_during_transform(): limits_dict = {"var_a": [-5, -2, 0, 2, 5], "var_b": [0, 2, 5]} # check for warning when errors equals 'ignore' - with pytest.warns(UserWarning) as record: - transformer = ArbitraryDiscretiser(binning_dict=limits_dict, errors="ignore") - transformer.fit(train) + transformer = ArbitraryDiscretiser( + binning_dict=limits_dict, return_object=return_object, errors="ignore" + ) + transformer.fit(train) + with pytest.warns(UserWarning, match=re.escape(msg)): transformer.transform(test) - # check that only one warning was returned - assert len(record) == 1 - # check that message matches - assert record[0].message.args[0] == msg - # check for error when errors equals 'raise' - with pytest.raises(ValueError) as record: - transformer = ArbitraryDiscretiser(binning_dict=limits_dict, errors="raise") - transformer.fit(train) + transformer = ArbitraryDiscretiser( + binning_dict=limits_dict, return_object=return_object, errors="raise" + ) + transformer.fit(train) + with pytest.raises(ValueError, match=re.escape(msg)): transformer.transform(test) - # check that error message matches - assert str(record.value) == msg - def test_error_if_not_permitted_value_is_errors(): age_dict = {"Age": [0, 10, 20, 30, np.inf]} - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="errors only takes values"): ArbitraryDiscretiser(binning_dict=age_dict, errors="medialuna") @@ -130,8 +115,5 @@ def test_error_if_binning_dict_not_dict_type(binning_dict): "binning_dict must be a dictionary with the interval limits per " f"variable. Got {binning_dict} instead." ) - with pytest.raises(ValueError) as record: + with pytest.raises(ValueError, match=msg): ArbitraryDiscretiser(binning_dict=binning_dict) - - # check that error message matches - assert str(record.value) == msg