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__(
@@ -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:

feature_engine/discretisation/base_discretiser.py

@@ -1,7 +1,11 @@
 # Authors: Morgan Sell <morganpsell@gmail.com>
 # 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