Merge pull request #349 from lvgig/feature/narwhalify_median_imputer

Feature/narwhalify median imputer

Author: davidhopkinson26

Diff for: CHANGELOG.rst

@@ -77,6 +77,7 @@ Changed
 - Refactored OrdinalEncoderTransformer tests in new format `#330 <https://github.com/lvgig/tubular/issues/330>`_
 - Narwhal-ified NullIndicator `#319 <https://github.com/lvgig/tubular/issues/319>_`
 - Narwhal-ified NearestMeanResponseImputer `#320 <https://github.com/lvgig/tubular/issues/320>_`
+- Narwhal-ified MedianImputer `#317 <https://github.com/lvgig/tubular/issues/317>_`
 
 
 1.3.1 (2024-07-18)

Diff for: tests/imputers/test_MedianImputer.py

@@ -1,8 +1,9 @@
+import narwhals as nw
 import numpy as np
-import pandas as pd
-import test_aide as ta
+import pytest
 
 import tests.test_data as d
+from tests import utils as u
 from tests.base_tests import (
     ColumnStrListInitTests,
     GenericFitTests,
@@ -33,68 +34,77 @@ class TestFit(WeightColumnFitMixinTests, GenericFitTests):
     def setup_class(cls):
         cls.transformer_name = "MedianImputer"
 
-    def test_learnt_values(self):
+    @pytest.mark.parametrize("library", ["pandas", "polars"])
+    def test_learnt_values(self, library):
         """Test that the impute values learnt during fit are expected."""
-        df = d.create_df_3()
-        df["d"] = np.nan
-
-        x = MedianImputer(columns=["a", "b", "c", "d"])
-
-        x.fit(df)
-
-        ta.classes.test_object_attributes(
-            obj=x,
-            expected_attributes={
-                "impute_values_": {
-                    "a": df["a"].median(),
-                    "b": df["b"].median(),
-                    "c": df["c"].median(),
-                    "d": np.float64(np.nan),
-                },
-            },
-            msg="impute_values_ attribute",
+        df = d.create_df_3(library=library)
+
+        df = nw.from_native(df)
+        native_namespace = nw.get_native_namespace(df)
+
+        # replace 'a' with all null values to trigger warning
+        df = df.with_columns(
+            nw.new_series(
+                name="d",
+                values=[None] * len(df),
+                native_namespace=native_namespace,
+            ),
         )
 
-    def test_learnt_values_weighted(self):
+        df = df.to_native()
+
+        transformer = MedianImputer(columns=["a", "b", "c", "d"])
+
+        transformer.fit(df)
+
+        assert transformer.impute_values_ == {
+            "a": df["a"].median(),
+            "b": df["b"].median(),
+            "c": df["c"].median(),
+            "d": None,
+        }, "impute_values_ attribute"
+
+    @pytest.mark.parametrize("library", ["pandas", "polars"])
+    def test_learnt_values_weighted(self, library):
         """Test that the impute values learnt during fit are expected - when using weights."""
-        df = d.create_df_9()
-        df["d"] = np.nan
-
-        df = pd.DataFrame(
-            {
-                "a": [1, 2, 4, 6],
-                "c": [3, 2, 4, 6],
-                "d": np.nan,
-            },
+        df = d.create_df_9(library=library)
+
+        df = nw.from_native(df)
+        native_namespace = nw.get_native_namespace(df)
+
+        # replace 'a' with all null values to trigger warning
+        df = df.with_columns(
+            nw.new_series(
+                name="d",
+                values=[None] * len(df),
+                native_namespace=native_namespace,
+            ),
         )
 
-        x = MedianImputer(columns=["a", "d"], weights_column="c")
+        df = df.to_native()
 
-        x.fit(df)
+        transformer = MedianImputer(columns=["a", "d"], weights_column="c")
 
-        ta.classes.test_object_attributes(
-            obj=x,
-            expected_attributes={
-                "impute_values_": {
-                    "a": np.int64(4),
-                    "d": np.nan,
-                },
-            },
-            msg="impute_values_ attribute",
-        )
+        transformer.fit(df)
+
+        assert transformer.impute_values_ == {
+            "a": np.int64(4),
+            "d": None,
+        }, "impute_values_ attribute"
 
-    def test_fit_not_changing_data(self):
+    @pytest.mark.parametrize("library", ["pandas", "polars"])
+    def test_fit_not_changing_data(self, library):
         """Test fit does not change X."""
-        df = d.create_df_1()
+        df = d.create_df_1(library=library)
 
-        x = MedianImputer(columns="a")
+        transformer = MedianImputer(columns="a")
 
-        x.fit(df)
+        transformer.fit(df)
 
-        ta.equality.assert_equal_dispatch(
-            expected=d.create_df_1(),
-            actual=df,
-            msg="Check X not changing during fit",
+        # Check whole dataframes
+        u.assert_frame_equal_dispatch(
+            d.create_df_1(library=library),
+            df,
         )
 
 

Diff for: tubular/imputers.py

@@ -6,7 +6,6 @@
 from typing import TYPE_CHECKING
 
 import narwhals as nw
-import numpy as np
 import pandas as pd
 
 from tubular.base import BaseTransformer
@@ -181,7 +180,7 @@ class attribute, indicates whether transformer has been converted to polars/pand
 
     """
 
-    polars_compatible = False
+    polars_compatible = True
 
     FITS = True
 
@@ -195,51 +194,52 @@ def __init__(
 
         WeightColumnMixin.check_and_set_weight(self, weights_column)
 
-    def fit(self, X: pd.DataFrame, y: pd.Series | None = None) -> pd.DataFrame:
+    @nw.narwhalify
+    def fit(self, X: FrameT, y: nw.Series | None = None) -> FrameT:
         """Calculate median values to impute with from X.
 
         Parameters
         ----------
-        X : pd.DataFrame
+        X : pd/pl.DataFrame
             Data to "learn" the median values from.
 
-        y : None or pd.DataFrame or pd.Series, default = None
+        y : None or pd/pl.Series, default = None
             Not required.
 
         """
         super().fit(X, y)
 
         self.impute_values_ = {}
 
-        if self.weights_column is not None:
-            WeightColumnMixin.check_weights_column(self, X, self.weights_column)
+        for c in self.columns:
+            # filter out null rows so their weight doesn't influence calc
+            filtered = X.filter(~nw.col(c).is_null())
 
-            for c in self.columns:
-                # filter out null rows so their weight doesn't influence calc
-                filtered = X[X[c].notna()]
+            # if column is only nulls, then median is None
+            if len(filtered) <= 0:
+                self.impute_values_[c] = None
 
-                # below algorithm only works for >1 non null values
-                if len(filtered) <= 0:
-                    median = np.nan
+            elif self.weights_column is not None:
+                WeightColumnMixin.check_weights_column(self, X, self.weights_column)
 
-                else:
-                    # first sort df by column to be imputed (order of weight column shouldn't matter for median)
-                    filtered = filtered.sort_values(c)
+                # first sort df by column to be imputed (order of weight column shouldn't matter for median)
+                filtered = filtered.sort(c)
 
-                    # next calculate cumulative weight sums
-                    cumsum = filtered[self.weights_column].cumsum()
+                # next calculate cumulative weight sums
+                cumsum = filtered[self.weights_column].cum_sum()
 
-                    # find midpoint
-                    cutoff = filtered[self.weights_column].sum() / 2.0
+                # find midpoint
+                cutoff = filtered[self.weights_column].sum() / 2.0
 
-                    # find first value >= this point
-                    median = filtered[c][cumsum >= cutoff].iloc[0]
+                # find first value >= this point
+                median = filtered.filter(cumsum >= cutoff).select(c)[0].item()
 
+                # impute value is weighted median
                 self.impute_values_[c] = median
 
-        else:
-            for c in self.columns:
-                self.impute_values_[c] = X[c].median()
+            else:
+                # impute value is median without considering weight
+                self.impute_values_[c] = X.select(nw.col(c).median()).item()
 
         return self