Migrate ArbitraryDiscretiser to narwhals, add polars support

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 <https://nbviewer.org/github/feature-engine/feature-engine-examples/blob/main/discretisation/ArbitraryDiscretiser_plus_MeanEncoder.ipynb>`_.
 
+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
 --------------------
 

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__(
@@ -127,7 +150,7 @@ def __init__(
                 f"variable. Got {binning_dict} instead."
             )
 
-        if errors not in ["ignore", "raise"]:
+        if not isinstance(errors, str) or errors not in ["ignore", "raise"]:
             raise ValueError(
                 "errors only takes values 'ignore' and 'raise'. "
                 f"Got {errors} instead."
@@ -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:

tests/test_discretisation/test_arbitrary_discretiser.py

@@ -1,94 +1,118 @@
+import re
+
 import numpy as np
 import pandas as pd
 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
+from tests.backend_helpers import frame_to_dict
+
+BINS = [0, 20, 40, 60, np.inf]
+
+MSG_NA = (
+    "Some of the variables in the dataset contain NaN. Check and "
+    "remove those before using this transformer."
+)
 
 
-def test_arbitrary_discretiser():
-    california_dataset = fetch_california_housing()
-    data = pd.DataFrame(
-        california_dataset.data, columns=california_dataset.feature_names
+# init parameters
+@pytest.mark.parametrize("binning_dict", ["HOLA", 1, False, None, [0, 10, 20]])
+def test_error_if_binning_dict_not_dict_type(binning_dict):
+    msg = (
+        "binning_dict must be a dictionary with the interval limits per "
+        f"variable. Got {binning_dict} instead."
     )
-    user_dict = {"HouseAge": [0, 20, 40, 60, np.inf]}
+    with pytest.raises(ValueError, match=re.escape(msg)):
+        ArbitraryDiscretiser(binning_dict=binning_dict)
 
-    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
+@pytest.mark.parametrize(
+    "errors", ["medialuna", "Ignore", "", 1, None, ["ignore"], ("raise",)]
+)
+def test_error_if_errors_not_permitted_value(errors):
+    msg = f"errors only takes values 'ignore' and 'raise'. Got {errors} instead."
+    with pytest.raises(ValueError, match=re.escape(msg)):
+        ArbitraryDiscretiser(binning_dict={"Age": BINS}, errors=errors)
+
+
+@pytest.mark.parametrize(
+    "binning_dict, return_object, return_boundaries, precision, errors",
+    [
+        ({"HouseAge": BINS}, False, False, 3, "ignore"),
+        ({"HouseAge": BINS, "MedInc": [0, 5, np.inf]}, True, False, 1, "raise"),
+        ({"Age": [0, 10, np.inf]}, False, True, 10, "raise"),
+    ],
+)
+def test_init_param_assignment(
+    binning_dict, return_object, return_boundaries, precision, errors
+):
+    transformer = ArbitraryDiscretiser(
+        binning_dict=binning_dict,
+        return_object=return_object,
+        return_boundaries=return_boundaries,
+        precision=precision,
+        errors=errors,
     )
-    data_t1["HouseAge"] = data_t1["HouseAge"].astype(str)
-    data_t2["HouseAge"] = pd.cut(
-        data["HouseAge"],
-        bins=[0, 20, 40, 60, np.inf],
-        labels=False,
-        include_lowest=True,
+    assert transformer.binning_dict == binning_dict
+    assert transformer.return_object is return_object
+    assert transformer.return_boundaries is return_boundaries
+    assert transformer.precision == precision
+    assert transformer.errors == errors
+
+
+# fit and transform
+def test_arbitrary_discretiser(make_df, data_california):
+    user_dict = {"HouseAge": BINS}
+
+    # ground truth via pandas.cut - bins are user-supplied and fixed, so both
+    # backends must reproduce this exact output.
+    house_age = pd.Series(data_california["HouseAge"])
+    expected_codes = pd.cut(
+        house_age, bins=BINS, labels=False, include_lowest=True
+    ).tolist()
+    expected_labels = (
+        pd.cut(house_age, bins=BINS, include_lowest=True).astype(str).tolist()
     )
 
+    data = make_df(data_california)
+
     transformer = ArbitraryDiscretiser(
         binning_dict=user_dict, return_object=False, return_boundaries=False
     )
     X = transformer.fit_transform(data)
 
-    # init params
-    assert transformer.return_object is False
-    assert transformer.return_boundaries is False
     # fit params
     assert transformer.variables_ == ["HouseAge"]
     assert transformer.binner_dict_ == user_dict
     # transform params
-    pd.testing.assert_frame_equal(X, data_t2)
+    assert isinstance(X, make_df)
+    assert frame_to_dict(X)["HouseAge"] == 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)
+    assert isinstance(X, make_df)
+    assert frame_to_dict(X)["HouseAge"] == 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
+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()
+    transformer = ArbitraryDiscretiser(binning_dict=age_dict)
+    transformer.fit(data)
+    with pytest.raises(ValueError, match=re.escape(MSG_NA)):
+        transformer.transform(data_na)
 
-    # 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"]
-
-    # 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("return_object", [False, True])
+def test_error_when_nan_introduced_during_transform(make_df, return_object):
+    # values outside the bin edges learned in fit become NaN, which warns or raises
+    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,40 +122,17 @@ 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.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):
-        ArbitraryDiscretiser(binning_dict=age_dict, errors="medialuna")
-
-
-@pytest.mark.parametrize("binning_dict", ["HOLA", 1, False])
-def test_error_if_binning_dict_not_dict_type(binning_dict):
-    msg = (
-        "binning_dict must be a dictionary with the interval limits per "
-        f"variable. Got {binning_dict} instead."
+    transformer = ArbitraryDiscretiser(
+        binning_dict=limits_dict, return_object=return_object, errors="raise"
     )
-    with pytest.raises(ValueError) as record:
-        ArbitraryDiscretiser(binning_dict=binning_dict)
-
-    # check that error message matches
-    assert str(record.value) == msg
+    transformer.fit(train)
+    with pytest.raises(ValueError, match=re.escape(msg)):
+        transformer.transform(test)