sanitize behavior when estimator passed to apply() is a dataop (skrub-data#1671)

Co-authored-by: Riccardo Cappuzzo <7548232+rcap107@users.noreply.github.com>

Author: jeromedockes

CHANGES.rst

@@ -36,11 +36,15 @@ Changes
 
 Bugfixes
 --------
+- Issues occurring when :meth:`DataOp.skb.apply` was passed a DataOp as the
+  estimator have been fixed in :pr:`1671` by :user:`Jérôme Dockès
+  <jeromedockes>`.
 - :class:`TableReport` could raise an error while trying to check if Polars
   columns with some dtypes (lists, structs) are sorted. It would not indicate
   Polars columns sorted in descending order. Fixed in :pr:`1673` by
   :user:`Jérôme Dockès <jeromedockes>`.
 
+
 Release 0.6.2
 =============
 

skrub/_data_ops/_data_ops.py

@@ -1269,15 +1269,21 @@ class Apply(DataOpImpl):
     # we do not evaluate `y`.
 
     def eval(self, *, mode, environment):
-        if mode not in self.supported_modes():
+        # 1. Find method to call and collect the necessary arguments: X and
+        # possibly y and the estimator (depending on the mode).
+
+        if "fit" in mode or mode == "preview":
+            # we need to fit, evaluate the estimator
+            estimator = yield self.estimator
+        else:
+            # for other modes self.estimator_ will be used
+            estimator = None
+        if mode not in self.supported_modes(estimator):
             # We are not the final estimator, e.g. mode is 'predict' and we are
             # a transformer that comes before the predictor.
             mode = "fit_transform" if "fit" in mode else "transform"
         method_name = "fit_transform" if mode == "preview" else mode
 
-        # 1. Collect the necessary arguments: X and possibly y and the estimator
-        # (depending on the mode).
-
         X = yield self.X
         if ("fit" in method_name and not self.unsupervised) or method_name == "score":
             y = yield self.y
@@ -1288,7 +1294,6 @@ def eval(self, *, mode, environment):
         X = _check_column_names(X)
 
         if "fit" in method_name:
-            estimator = yield self.estimator
             cols = yield self.cols
             how = yield self.how
             allow_reject = yield self.allow_reject
@@ -1387,18 +1392,23 @@ def _eval_kwargs(self, method_name):
             )
         return kwargs
 
-    def supported_modes(self):
+    def supported_modes(self, estimator=None):
         """
         Used by SkrubLearner and param search to decide if they have the
         methods `predict`, `predict_proba` etc.
         """
         modes = ["preview", "fit", "fit_transform", "transform"]
-        try:
-            estimator = self.estimator_
-        except AttributeError:
-            estimator = get_chosen_or_default(self.estimator)
+        if estimator is None:
+            try:
+                estimator = self.estimator_
+            except AttributeError:
+                estimator = get_chosen_or_default(self.estimator)
         for name in FITTED_PREDICTOR_METHODS:
-            if hasattr(estimator, name):
+            if isinstance(estimator, DataOp) or hasattr(estimator, name):
+                # if estimator is a DataOp we cannot know yet if it has the
+                # attribute, in this case we assume it does (and risk a
+                # slightly worse error message when a SkrubLearner tries to use
+                # it if it does not).
                 modes.append(name)
         return modes
 

skrub/_data_ops/_estimator.py

@@ -10,7 +10,7 @@
 from .. import _dataframe as sbd
 from .. import _join_utils
 from ._choosing import BaseNumericChoice, get_default
-from ._data_ops import Apply, check_subsampled_X_y_shape
+from ._data_ops import Apply, DataOp, check_subsampled_X_y_shape
 from ._evaluation import (
     choice_graph,
     chosen_or_default_outcomes,
@@ -409,8 +409,6 @@ def describe_params(self):
         )
 
 
-# Xy_pipeline because it is an actual scikit-learn pippeline rather than
-# a skrub learner
 def _to_Xy_pipeline(learner, environment):
     return learner.__skrub_to_Xy_pipeline__(environment)
 
@@ -442,8 +440,11 @@ def _estimator_type(self):
         first = find_first_apply(self.data_op)
         if first is None:
             return "transformer"
+        estimator = get_default(first._skrub_impl.estimator)
+        if isinstance(estimator, DataOp):
+            return "transformer"
         try:
-            return get_default(first._skrub_impl.estimator)._estimator_type
+            return estimator._estimator_type
         except AttributeError:
             return "transformer"
 
@@ -454,7 +455,13 @@ def __sklearn_tags__(self):
             first = find_first_apply(self.data_op)
             if first is None:
                 return _default_sklearn_tags()
-            return get_default(first._skrub_impl.estimator).__sklearn_tags__()
+            estimator = get_default(first._skrub_impl.estimator)
+            if isinstance(estimator, DataOp):
+                return _default_sklearn_tags()
+            try:
+                return estimator.__sklearn_tags__()
+            except AttributeError:
+                return _default_sklearn_tags()
 
     @property
     def classes_(self):

skrub/_data_ops/tests/test_data_ops.py

@@ -8,7 +8,7 @@
 from sklearn.base import BaseEstimator
 from sklearn.datasets import make_classification, make_regression
 from sklearn.dummy import DummyRegressor
-from sklearn.linear_model import LogisticRegression
+from sklearn.linear_model import LogisticRegression, Ridge
 from sklearn.model_selection import train_test_split
 
 import skrub
@@ -625,12 +625,6 @@ def test_concat_non_str_colname():
         )
 
 
-def test_class_skb():
-    from skrub._data_ops._skrub_namespace import SkrubNamespace
-
-    assert skrub.DataOp.skb is SkrubNamespace
-
-
 def test_get_vars():
     a = skrub.var("a")
     b = skrub.var("b")
@@ -639,3 +633,98 @@ def test_get_vars():
     assert list(d.skb.get_vars().keys()) == ["a", "b"]
     assert d.skb.get_vars()["a"] is a
     assert list(d.skb.get_vars(all_named_ops=True).keys()) == ["a", "b", "c"]
+
+
+@pytest.mark.parametrize("needs_data", [False, True])
+@pytest.mark.parametrize("has_preview", [False, True])
+@pytest.mark.parametrize("regression", [False, True])
+@pytest.mark.parametrize("with_scoring", [False, True])
+def test_estimator_is_a_data_op(needs_data, has_preview, regression, with_scoring):
+    # Check that the data_op, learner and search estimators behave well when
+    # the estimator passed to apply is a data op
+    if regression:
+        X_a, y_a = make_regression(random_state=0)
+    else:
+        X_a, y_a = make_classification(random_state=0)
+    X_df = pd.DataFrame(X_a).rename(columns=str)
+    if has_preview:
+        X, y = skrub.X(X_df), skrub.y(y_a)
+    else:
+        X, y = skrub.X(), skrub.y()
+    if needs_data:
+        # In this case the estimator's automated preview cannot be computed
+        # because it needs a value from the environment, the value is not known
+        # until we fit the learner.
+
+        def get_vectorizer(X):
+            return skrub.TableVectorizer()
+
+        vectorizer = X.skb.apply_func(get_vectorizer)
+
+        def get_predictor(X):
+            return Ridge() if regression else LogisticRegression()
+
+        predictor = X.skb.apply_func(get_predictor)
+    else:
+        # In this case the estimator can be evaluated in the automated preview
+        # when the data op is created.
+        vectorizer = skrub.as_data_op(skrub.TableVectorizer())
+        predictor = skrub.as_data_op(Ridge() if regression else LogisticRegression())
+    pred = X.skb.apply(vectorizer).skb.apply(predictor, y=y)
+    # no information about the estimator: we expose all methods and default to
+    # 'transformer' estimator type.
+    learner = pred.skb.make_learner()
+    assert learner.__skrub_to_Xy_pipeline__({})._estimator_type == "transformer"
+    assert hasattr(learner, "predict")
+    assert hasattr(pred.skb.make_randomized_search(), "predict")
+    if has_preview:
+        assert pred.skb.preview().shape == y_a.shape
+    env = {"X": X_df, "y": y_a}
+    assert pred.skb.eval(env).shape == y_a.shape
+    search = pred.skb.make_grid_search(cv=2).fit(env)
+    min_score = 0.3 if regression else 0.7
+    assert search.best_score_ > min_score
+    if with_scoring:
+        scoring = "r2" if regression else "accuracy"
+    else:
+        scoring = None
+    res = skrub.cross_validate(
+        pred.skb.make_grid_search(cv=2, scoring=scoring),
+        environment=env,
+        cv=2,
+        scoring=scoring,
+    )
+    assert res["test_score"].mean() > min_score
+
+
+def test_apply_no_sklearn_tags():
+    # applying an estimator that does not define __sklearn_tags__
+    class Twice:
+        def fit(self, X, y=None):
+            return self
+
+        def fit_transform(self, X, y=None):
+            return X * 2
+
+        def transform(self, X):
+            return X * 2
+
+        def get_params(self, deep=True):
+            return {}
+
+        def set_params(self):
+            return self
+
+    learner = skrub.var("a").skb.apply(Twice()).skb.make_learner()
+    assert learner.fit_transform({"a": 1}) == 2
+    xy_learner = learner.__skrub_to_Xy_pipeline__({})
+    assert xy_learner._estimator_type == "transformer"
+    if hasattr(xy_learner, "__sklearn_tags__"):
+        # Old scikit-learn versiond don't have __sklearn_tags__
+        assert xy_learner.__sklearn_tags__().estimator_type is None
+
+
+def test_class_skb():
+    from skrub._data_ops._skrub_namespace import SkrubNamespace
+
+    assert skrub.DataOp.skb is SkrubNamespace