Repository: koaning/scikit-lego
Pull request #694 (2024-08-07 19:25:01)

feat: bump narwhals and adapt to support pyarrow
# Description

Long story short: I knew that some functionalities would have not worked for pyarrow even if using narwhals and we had to adjust accordingly.

For example we cannot create a series using `native_namespace.Series([...])`, because pyarrow doesn't have a `Series` object, therefore the a workaround was needed:

```diff
- series = native_namespace.Series([...])
+ series = nw.from_dict({"_tmp": [...]})["_tmp"]
```
Namely, create a narhwals dataframe for the given namespace and then select the unique column.

A bump of other changes were needed in the tests to assess with pyarrow tables.

## Type of change

- [x] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)


## Checklist:

- [x] My code follows the style guidelines (ruff)
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation (also to the readme.md)
- [ ] I have added tests that prove my fix is effective or that my feature works
- [ ] I have added tests to check whether the new feature adheres to the sklearn convention
- [ ] New and existing unit tests pass locally with my changes

## Other comments

- The commits are because I branched out from another PR 🙈
- Most likely there are other parts of the codebase in which a change may be required to make sure that pyarrow runs (to be continued)
- Adding @MarcoGorelli to the mix, to make sure there is no better way of going about this

koaning: Oh, and one more thing, must we bump narwhals? I usually prefer not to force the user to use the latest and greatest dependency, but it seems like we need it to support more dataframe types? 
FBruzzesi: > Oh, and one more thing, must we bump narwhals? I usually prefer not to force the user to use the latest and greatest dependency, but it seems like we need it to support more dataframe types?

Sadly yes, if we want to have pyarrow support for grouped meta and shift, these functionalities made it only in the latest release (yet narwhals keeps being dependency-free)


FBruzzesi: > LGTM! Might be good to prep another release?

Sounds good! We went a bit silent, but now have some breaking changes from #693 

Patch:
diff --git a/pyproject.toml b/pyproject.toml
index a859406..dbc98f5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -20,7 +20,7 @@ maintainers = [
 ]
 
 dependencies = [
-    "narwhals>=1.0.0",
+    "narwhals>=1.2.0",
     "pandas>=1.1.5",
     "scikit-learn>=1.0",
     "importlib-metadata >= 1.0; python_version < '3.8'",
@@ -62,8 +62,7 @@ docs = [
 ]
 
 test = [
-    "narwhals[polars]",
-    "pyarrow",
+    "narwhals[polars,pyarrow]",
     "pytest>=6.2.5",
     "pytest-xdist>=1.34.0",
     "pytest-cov>=2.6.1",
diff --git a/setup.py b/setup.py
index 2f3e89f..a0e014a 100644
--- a/setup.py
+++ b/setup.py
@@ -2,6 +2,7 @@
 """
 setup.py is used to distribute scikit-lego as a package using setuptools and twine
 """
+
 from setuptools import setup
 
 if __name__ == "__main__":
diff --git a/sklego/meta/_grouped_utils.py b/sklego/meta/_grouped_utils.py
index 32bec84..c873e0c 100644
--- a/sklego/meta/_grouped_utils.py
+++ b/sklego/meta/_grouped_utils.py
@@ -33,10 +33,11 @@ def parse_X_y(X, y, groups, check_X=True, **kwargs) -> nw.DataFrame:
 
     # Convert y and assign it to the frame
     n_samples = X.shape[0]
-    native_space = nw.get_native_namespace(X)
+    y_series = nw.from_dict(
+        data={"tmp": [None] * n_samples if y is None else y}, native_namespace=nw.get_native_namespace(X)
+    )["tmp"]
 
-    y_native = native_space.Series([None] * n_samples) if y is None else native_space.Series(y)
-    return X.with_columns(__sklego_target__=nw.from_native(y_native, allow_series=True))
+    return X.with_columns(__sklego_target__=y_series)
 
 
 def _validate_groups_values(X: nw.DataFrame, groups: List[int] | List[str]) -> None:
diff --git a/sklego/meta/grouped_predictor.py b/sklego/meta/grouped_predictor.py
index a61f331..4087820 100644
--- a/sklego/meta/grouped_predictor.py
+++ b/sklego/meta/grouped_predictor.py
@@ -123,12 +123,13 @@ class GroupedPredictor(ShrinkageMixin, MetaEstimatorMixin, BaseEstimator):
         if columns is None:
             columns = self._groups
 
+        to_drop = list(set(["__sklego_target__", *columns, *as_list(self.groups)]))
         grouped_estimators = {
             # Fit a clone of the estimators to each group
             (group_name[0] if len(group_name) == 1 else group_name): self.__fit_single_group(
                 group=(group_name[0] if len(group_name) == 1 else group_name),
-                X=nw.to_native(X_grp.drop(["__sklego_target__", *columns, *as_list(self.groups)])),
-                y=(nw.to_native(X_grp.select("__sklego_target__")).to_numpy().reshape(-1) if y is not None else None),
+                X=nw.to_native(X_grp.drop(to_drop)),
+                y=(X_grp.select("__sklego_target__").to_numpy().reshape(-1) if y is not None else None),
             )
             for group_name, X_grp in frame.group_by(columns)
         }
@@ -149,12 +150,12 @@ class GroupedPredictor(ShrinkageMixin, MetaEstimatorMixin, BaseEstimator):
 
         if self.shrinkage is not None and self.use_global_model:
             n_samples = frame.shape[0]
-            native_space = nw.get_native_namespace(frame)
 
             frame = frame.select(
-                nw.from_native(native_space.Series([self._global_col_value] * n_samples), allow_series=True).alias(
-                    self._global_col_name
-                ),
+                nw.from_dict(
+                    data={self._global_col_name: np.full(shape=n_samples, fill_value=self._global_col_value)},
+                    native_namespace=nw.get_native_namespace(frame),
+                )[self._global_col_name],
                 nw.all(),
             )
             groups = [self._global_col_name] if groups is None else [self._global_col_name, *groups]
@@ -285,7 +286,7 @@ class GroupedPredictor(ShrinkageMixin, MetaEstimatorMixin, BaseEstimator):
                         (group_value[0] if len(group_value) == 1 else group_value),
                         nw.to_native(X_grp.drop(["__sklego_index__", *groups, *as_list(self.groups)])),
                         method=method,
-                    ).set_index(nw.to_native(X_grp["__sklego_index__"]).to_numpy().reshape(-1).astype(int))
+                    ).set_index(X_grp["__sklego_index__"].to_numpy().reshape(-1).astype(int))
                     for group_value, X_grp in frame.group_by(groups)
                 ],
                 axis=0,
diff --git a/sklego/meta/grouped_transformer.py b/sklego/meta/grouped_transformer.py
index b603bd1..2dfe18e 100644
--- a/sklego/meta/grouped_transformer.py
+++ b/sklego/meta/grouped_transformer.py
@@ -164,7 +164,7 @@ class GroupedTransformer(TransformerMixin, MetaEstimatorMixin, BaseEstimator):
 
         results = [
             (
-                nw.to_native(X_grp.select("__sklego_index__")).to_numpy().squeeze().astype(int),
+                X_grp.select("__sklego_index__").to_numpy().squeeze().astype(int),
                 self.__transform_single_group(
                     group_name, nw.to_native(X_grp.drop(["__sklego_index__", *self.groups_]))
                 ),
diff --git a/sklego/meta/hierarchical_predictor.py b/sklego/meta/hierarchical_predictor.py
index f7a5327..5d71cc5 100644
--- a/sklego/meta/hierarchical_predictor.py
+++ b/sklego/meta/hierarchical_predictor.py
@@ -281,12 +281,15 @@ class HierarchicalPredictor(ShrinkageMixin, MetaEstimatorMixin, BaseEstimator):
             msg = "Found 0 features, while a minimum of 1 if required."
             raise ValueError(msg)
 
-        native_space = nw.get_native_namespace(X)
-
+        native_namespace = nw.get_native_namespace(X)
+        target_series = nw.from_dict({self._TARGET_NAME: y}, native_namespace=native_namespace)[self._TARGET_NAME]
+        global_series = nw.from_dict({self._GLOBAL_NAME: np.ones(n_samples)}, native_namespace=native_namespace)[
+            self._GLOBAL_NAME
+        ]
         frame = X.with_columns(
             **{
-                self._TARGET_NAME: nw.from_native(native_space.Series(y), allow_series=True),
-                self._GLOBAL_NAME: nw.from_native(native_space.Series([1] * n_samples), allow_series=True),
+                self._TARGET_NAME: target_series,
+                self._GLOBAL_NAME: global_series,
             }
         ).pipe(self.__validate_frame)
 
@@ -318,11 +321,14 @@ class HierarchicalPredictor(ShrinkageMixin, MetaEstimatorMixin, BaseEstimator):
             X = nw.from_native(pd.DataFrame(X))
 
         n_samples = X.shape[0]
-        native_space = nw.get_native_namespace(X)
+        native_namespace = nw.get_native_namespace(X)
+        global_series = nw.from_dict({self._GLOBAL_NAME: np.ones(n_samples)}, native_namespace=native_namespace)[
+            self._GLOBAL_NAME
+        ]
 
         frame = X.with_columns(
             **{
-                self._GLOBAL_NAME: nw.from_native(native_space.Series([1] * n_samples), allow_series=True),
+                self._GLOBAL_NAME: global_series,
                 self._INDEX_NAME: np.arange(n_samples),
             }
         ).pipe(self.__validate_frame)
@@ -340,7 +346,7 @@ class HierarchicalPredictor(ShrinkageMixin, MetaEstimatorMixin, BaseEstimator):
 
         for level_idx, grp_names in enumerate(self.fitted_levels_):
             for grp_values, grp_frame in frame.group_by(grp_names):
-                grp_idx = nw.to_native(grp_frame.select(self._INDEX_NAME)).to_numpy().reshape(-1)
+                grp_idx = grp_frame.select(self._INDEX_NAME).to_numpy().reshape(-1)
 
                 _estimator, _level = _get_estimator(
                     estimators=self.estimators_,
diff --git a/sklego/model_selection.py b/sklego/model_selection.py
index fbc4b4a..9fa2613 100644
--- a/sklego/model_selection.py
+++ b/sklego/model_selection.py
@@ -263,11 +263,13 @@ class TimeGapSplit:
 
         j = 0
         for i in self.split(nw.to_native(X)):
-            update_split_info(native_namespace.Series(i[0]), j, "train", summary)
-            update_split_info(native_namespace.Series(i[1]), j, "valid", summary)
+            train_info = nw.to_native(nw.from_dict({"tmp": i[0]}, native_namespace=native_namespace)["tmp"])
+            valid_info = nw.to_native(nw.from_dict({"tmp": i[1]}, native_namespace=native_namespace)["tmp"])
+            update_split_info(train_info, j, "train", summary)
+            update_split_info(valid_info, j, "valid", summary)
             j = j + 1
 
-        result = nw.from_native(native_namespace.DataFrame(summary))
+        result = nw.from_dict(summary, native_namespace=native_namespace)
         result = nw.maybe_set_index(result, ["fold", "part"])
         return nw.to_native(result)
 


Test patch:
diff --git a/tests/test_meta/test_grouped_predictor.py b/tests/test_meta/test_grouped_predictor.py
index 21606a4..93a181f 100644
--- a/tests/test_meta/test_grouped_predictor.py
+++ b/tests/test_meta/test_grouped_predictor.py
@@ -1,8 +1,10 @@
 from contextlib import nullcontext as does_not_raise
 
+import narwhals.stable.v1 as nw
 import numpy as np
 import pandas as pd
 import polars as pl
+import pyarrow as pa
 import pytest
 from sklearn.dummy import DummyRegressor
 from sklearn.impute import SimpleImputer
@@ -60,32 +62,32 @@ def random_xy_grouped_clf_different_classes(request):
 
 
 @pytest.mark.parametrize("groups, expected", [("diet", {1, 2, 3, 4}), ("chick", set(range(1, 50 + 1)))])
-@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame])
+@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame, pa.table])
 def test_chickweight_keys(groups, expected, frame_func):
-    df = frame_func(load_chicken(as_frame=True))
+    df = nw.from_native(frame_func(load_chicken(as_frame=True).to_dict(orient="list")))
     mod = GroupedPredictor(estimator=LinearRegression(), groups=groups)
-    mod.fit(df[["time", groups]], df["weight"])
+    mod.fit(nw.to_native(df.select("time", groups)), nw.to_native(df["weight"]))
     assert set(mod.estimators_.keys()) == expected
 
 
-@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame])
+@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame, pa.table])
 def test_chickweight_can_do_fallback(frame_func):
-    df = frame_func(load_chicken(as_frame=True))
+    df = nw.from_native(frame_func(load_chicken(as_frame=True).to_dict(orient="list")))
     mod = GroupedPredictor(estimator=LinearRegression(), groups="diet")
-    mod.fit(df[["time", "diet"]], df["weight"])
+    mod.fit(nw.to_native(df.select("time", "diet")), nw.to_native(df["weight"]))
     assert set(mod.estimators_.keys()) == {1, 2, 3, 4}
     to_predict = pd.DataFrame({"time": [21, 21], "diet": [5, 6]})
     assert mod.predict(to_predict).shape == (2,)
     assert mod.predict(to_predict)[0] == mod.predict(to_predict)[1]
 
 
-@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame])
+@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame, pa.table])
 def test_chickweight_can_do_fallback_proba(frame_func):
-    df = frame_func(load_chicken(as_frame=True))
+    df = nw.from_native(frame_func(load_chicken(as_frame=True).to_dict(orient="list")))
 
-    y = np.where(df["weight"] > df["weight"].mean(), 1, 0)
+    y = nw.to_native((df["weight"] > df["weight"].mean()).cast(nw.Int32))
     mod = GroupedPredictor(estimator=LogisticRegression(), groups="diet")
-    mod.fit(df[["time", "diet"]], y)
+    mod.fit(nw.to_native(df.select("time", "diet")), y)
     assert set(mod.estimators_.keys()) == {1, 2, 3, 4}
 
     to_predict = pd.DataFrame({"time": [21, 21], "diet": [5, 6]})
diff --git a/tests/test_meta/test_grouped_transformer.py b/tests/test_meta/test_grouped_transformer.py
index 875bf60..b70f539 100644
--- a/tests/test_meta/test_grouped_transformer.py
+++ b/tests/test_meta/test_grouped_transformer.py
@@ -3,6 +3,7 @@ import itertools as it
 import numpy as np
 import pandas as pd
 import polars as pl
+import pyarrow as pa
 import pytest
 import sklearn
 from sklearn import clone
@@ -264,9 +265,9 @@ def test_array_with_strings():
     transformer.fit_transform(X)
 
 
-@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame])
+@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame, pa.table])
 def test_df(penguins_df, frame_func):
-    penguins_df = frame_func(penguins_df)
+    penguins_df = frame_func(penguins_df.to_dict(orient="list"))
     meta = GroupedTransformer(StandardScaler(), groups=["island", "sex"])
 
     transformed = meta.fit_transform(penguins_df)
@@ -275,14 +276,14 @@ def test_df(penguins_df, frame_func):
     assert transformed.shape == (penguins_df.shape[0], penguins_df.shape[1] - 2)
 
 
-@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame])
+@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame, pa.table])
 def test_df_missing_group(penguins_df, frame_func):
     meta = GroupedTransformer(StandardScaler(), groups=["island", "sex"])
 
     # Otherwise the fixture is changed
     X = penguins_df.copy()
     X.loc[0, "island"] = None
-    X = frame_func(X)
+    X = frame_func(X.to_dict(orient="list"))
     with pytest.raises(ValueError):
         meta.fit_transform(X)
 
@@ -309,18 +310,20 @@ def test_grouping_column_not_in_array(penguins):
         meta.fit_transform(X[:, :3])
 
 
-@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame])
+@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame, pa.table])
 def test_grouping_column_not_in_df(penguins_df, frame_func):
     meta = GroupedTransformer(StandardScaler(), groups=["island", "unexisting_column"])
 
     # This should raise ValueError
     with pytest.raises(ValueError):
-        meta.fit_transform(frame_func(penguins_df))
+        meta.fit_transform(frame_func(penguins_df.to_dict(orient="list")))
 
 
-@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame])
+@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame, pa.table])
 def test_no_grouping(penguins_df, frame_func):
-    penguins_numeric = frame_func(penguins_df[["bill_length_mm", "bill_depth_mm", "flipper_length_mm", "body_mass_g"]])
+    penguins_numeric = frame_func(
+        penguins_df[["bill_length_mm", "bill_depth_mm", "flipper_length_mm", "body_mass_g"]].to_dict(orient="list")
+    )
 
     meta = GroupedTransformer(StandardScaler(), groups=None)
     nonmeta = StandardScaler()
@@ -328,10 +331,10 @@ def test_no_grouping(penguins_df, frame_func):
     assert (meta.fit_transform(penguins_numeric) == nonmeta.fit_transform(penguins_numeric)).all()
 
 
-@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame])
+@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame, pa.table])
 def test_with_y(penguins_df, frame_func):
-    X = frame_func(penguins_df.drop(columns=["sex"]))
-    y = penguins_df["sex"]
+    X = frame_func(penguins_df.drop(columns=["sex"]).to_dict(orient="list"))
+    y = penguins_df["sex"].to_numpy()
 
     meta = GroupedTransformer(StandardScaler(), groups="island")
 
diff --git a/tests/test_meta/test_hierarchical_predictor.py b/tests/test_meta/test_hierarchical_predictor.py
index 94bd016..a6bbc46 100644
--- a/tests/test_meta/test_hierarchical_predictor.py
+++ b/tests/test_meta/test_hierarchical_predictor.py
@@ -5,6 +5,7 @@ import narwhals.stable.v1 as nw
 import numpy as np
 import pandas as pd
 import polars as pl
+import pyarrow as pa
 import pytest
 from sklearn import clone
 from sklearn.datasets import make_classification, make_regression
@@ -17,7 +18,7 @@ from sklearn.utils.estimator_checks import parametrize_with_checks
 
 from sklego.meta import HierarchicalClassifier, HierarchicalRegressor
 
-frame_funcs = [pd.DataFrame, pl.DataFrame]
+frame_funcs = [pd.DataFrame, pl.DataFrame, pa.table]
 
 
 @parametrize_with_checks([HierarchicalRegressor(estimator=LinearRegression(), groups=0)])
@@ -60,14 +61,14 @@ def make_hierarchical_dataset(task, frame_func=pd.DataFrame):
     else:
         raise ValueError("Invalid task")
 
-    X_ = (
+    X_ = frame_func(
         pd.DataFrame(X, columns=[f"x_{i}" for i in range(X.shape[1])])
         .assign(
             g_0=1,
             g_1=["A"] * (n_samples // 2) + ["B"] * (n_samples // 2),
             g_2=["X"] * (n_samples // 4) + ["Y"] * (n_samples // 2) + ["Z"] * (n_samples // 4),
         )
-        .pipe(frame_func)
+        .to_dict(orient="list")
     )
     groups = ["g_0", "g_1", "g_2"]
 
@@ -129,7 +130,7 @@ def test_fit_predict(meta_cls, base_estimator, task, fallback_method, shrinkage,
     """Tests that the model can be fit and predict with different configurations of fallback and shrinkage methods if
     X to predict contains same groups as X used to fit.
     """
-    X, y, groups = make_hierarchical_dataset(task, frame_func=frame_funcs[randint(0, 1)])
+    X, y, groups = make_hierarchical_dataset(task, frame_func=frame_funcs[randint(0, 2)])
 
     meta_model = meta_cls(
         estimator=base_estimator,
@@ -159,10 +160,10 @@ def test_fallback(meta_cls, base_estimator, task, fallback_method, context):
     """Tests that the model fails or not when predicting with different fallback methods if X to predict contains
     unseen group values.
     """
-    X, y, groups = make_hierarchical_dataset(task, frame_func=frame_funcs[randint(0, 1)])
+    X, y, groups = make_hierarchical_dataset(task, frame_func=frame_funcs[randint(0, 2)])
 
     meta_model = meta_cls(estimator=base_estimator, groups=groups, fallback_method=fallback_method).fit(X, y)
-    X[groups] = np.ones((X.shape[0], len(groups))) * -1  # Shortcut assignment that works both in pandas and polars
+    X = nw.to_native(nw.from_native(X).with_columns(**{g: nw.lit(-1) for g in groups}))
 
     with context:
         meta_model.predict(X)
@@ -190,7 +191,7 @@ def test_shrinkage(meta_cls, base_estimator, task, metric, shrinkage, kwargs):
     """Tests that the model performance is better than the base estimator when predicting with different shrinkage
     methods.
     """
-    X, y, groups = make_hierarchical_dataset(task, frame_func=frame_funcs[randint(0, 1)])
+    X, y, groups = make_hierarchical_dataset(task, frame_func=frame_funcs[randint(0, 2)])
 
     X_ = nw.from_native(X).drop(groups).pipe(nw.to_native)
     meta_model = meta_cls(
diff --git a/tests/test_meta/test_regression_outlier.py b/tests/test_meta/test_regression_outlier.py
index 2cd1fa8..068a93c 100644
--- a/tests/test_meta/test_regression_outlier.py
+++ b/tests/test_meta/test_regression_outlier.py
@@ -1,6 +1,7 @@
 import numpy as np
 import pandas as pd
 import polars as pl
+import pyarrow as pa
 import pytest
 from sklearn.linear_model import LinearRegression, LogisticRegression
 from sklearn.utils.estimator_checks import parametrize_with_checks
@@ -37,7 +38,7 @@ def test_obvious_example():
         assert preds[i] == -1
 
 
-@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame])
+@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame, pa.table])
 def test_obvious_example_dataframe(frame_func):
     # generate random data for illustrative example
     np.random.seed(42)
@@ -54,7 +55,7 @@ def test_obvious_example_dataframe(frame_func):
         assert preds[i] == -1
 
 
-@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame])
+@pytest.mark.parametrize("frame_func", [pd.DataFrame, pl.DataFrame, pa.table])
 def test_raises_error(frame_func):
     # generate random data for illustrative example
     np.random.seed(42)