allow passing kwargs in .skb.apply() (skrub-data#1642)

Author: jeromedockes

CHANGES.rst

@@ -11,6 +11,9 @@ Ongoing Development
 
 New features
 ------------
+- :meth:`DataOp.skb.apply` now allows passing extra named arguments to the
+  estimator's methods through the parameters ``fit_kwargs``, ``predict_kwargs``
+  etc. :pr:`1642` by :user:`Jérôme Dockès <jeromedockes>`.
 - TableReport now displays the mean statistic for boolean columns.
   :pr:`1647` by :user:`Abdelhakim Benechehab <abenechehab>`.
 - :meth:`DataOp.skb.iter_cv_splits` iterates over the training and testing

skrub/_apply_to_cols.py

@@ -73,9 +73,12 @@ class SingleColumnTransformer(BaseEstimator):
 
     __single_column_transformer__ = True
 
-    def fit(self, column, y=None):
+    def fit(self, column, y=None, **kwargs):
         """Fit the transformer.
 
+        This default implementation simply calls ``fit_transform()`` and
+        returns ``self``.
+
         Subclasses should implement ``fit_transform`` and ``transform``.
 
         Parameters
@@ -87,12 +90,15 @@ def fit(self, column, y=None):
         y : column or dataframe
             Prediction targets.
 
+        **kwargs
+            Extra named arguments are passed to ``self.fit_transform()``.
+
         Returns
         -------
         self
             The fitted transformer.
         """
-        self.fit_transform(column, y=y)
+        self.fit_transform(column, y=y, **kwargs)
         return self
 
     def _check_single_column(self, column, function_name):
@@ -132,35 +138,35 @@ def _wrap_add_check_single_column(f):
     if f.__name__ == "fit":
 
         @functools.wraps(f)
-        def fit(self, X, y=None):
+        def fit(self, X, y=None, **kwargs):
             self._check_single_column(X, f.__name__)
-            return f(self, X, y=y)
+            return f(self, X, y=y, **kwargs)
 
         return fit
     elif f.__name__ == "partial_fit":
 
         @functools.wraps(f)
-        def partial_fit(self, X, y=None):
+        def partial_fit(self, X, y=None, **kwargs):
             self._check_single_column(X, f.__name__)
-            return f(self, X, y=y)
+            return f(self, X, y=y, **kwargs)
 
         return partial_fit
 
     elif f.__name__ == "fit_transform":
 
         @functools.wraps(f)
-        def fit_transform(self, X, y=None):
+        def fit_transform(self, X, y=None, **kwargs):
             self._check_single_column(X, f.__name__)
-            return f(self, X, y=y)
+            return f(self, X, y=y, **kwargs)
 
         return fit_transform
     else:
         assert f.__name__ == "transform", f.__name__
 
         @functools.wraps(f)
-        def transform(self, X):
+        def transform(self, X, **kwargs):
             self._check_single_column(X, f.__name__)
-            return f(self, X)
+            return f(self, X, **kwargs)
 
         return transform
 
@@ -438,7 +444,7 @@ def __init__(
         self.rename_columns = rename_columns
         self.n_jobs = n_jobs
 
-    def fit(self, X, y=None):
+    def fit(self, X, y=None, **kwargs):
         """Fit the transformer on each column independently.
 
         Parameters
@@ -449,15 +455,19 @@ def fit(self, X, y=None):
         y : Pandas or Polars Series or DataFrame, default=None
             The target data.
 
+        **kwargs
+            Extra named arguments are passed to the ``fit_transform()`` method of
+            the individual column transformers (the clones of ``self.transformer``).
+
         Returns
         -------
         ApplyToCols
             The transformer itself.
         """
-        self.fit_transform(X, y)
+        self.fit_transform(X, y, **kwargs)
         return self
 
-    def fit_transform(self, X, y=None):
+    def fit_transform(self, X, y=None, **kwargs):
         """Fit the transformer on each column independently and transform X.
 
         Parameters
@@ -468,6 +478,10 @@ def fit_transform(self, X, y=None):
         y : Pandas or Polars Series or DataFrame, default=None
             The target data.
 
+        **kwargs
+            Extra named arguments are passed to the ``fit_transform()`` method of
+            the individual column transformers (the clones of ``self.transformer``).
+
         Returns
         -------
         result : Pandas or Polars DataFrame
@@ -485,19 +499,25 @@ def fit_transform(self, X, y=None):
                 self._columns,
                 self.transformer,
                 self.allow_reject,
+                kwargs,
             )
             for col_name in all_columns
         )
         return self._process_fit_transform_results(results, X)
 
-    def transform(self, X):
+    def transform(self, X, **kwargs):
         """Transform a dataframe.
 
         Parameters
         ----------
         X : Pandas or Polars DataFrame
             The column to transform.
 
+        **kwargs
+            Extra named arguments are passed to the ``transform()`` method of
+            the fitted individual column transformers (the values of
+            ``self.transformers_``, which are clones of ``self.transformer``).
+
         Returns
         -------
         result : Pandas or Polars DataFrame
@@ -507,10 +527,7 @@ def transform(self, X):
         parallel = Parallel(n_jobs=self.n_jobs)
         func = delayed(_transform_column)
         outputs = parallel(
-            func(
-                sbd.col(X, col_name),
-                self.transformers_.get(col_name),
-            )
+            func(sbd.col(X, col_name), self.transformers_.get(col_name), kwargs)
             for col_name in sbd.column_names(X)
         )
         transformed_columns = []
@@ -586,7 +603,9 @@ def _prepare_transformer_input(transformer, column):
     return sbd.make_dataframe_like(column, [column])
 
 
-def _fit_transform_column(column, y, columns_to_handle, transformer, allow_reject):
+def _fit_transform_column(
+    column, y, columns_to_handle, transformer, allow_reject, kwargs
+):
     col_name = sbd.name(column)
     if col_name not in columns_to_handle:
         return col_name, [column], None
@@ -595,7 +614,7 @@ def _fit_transform_column(column, y, columns_to_handle, transformer, allow_rejec
     transformer_input = _prepare_transformer_input(transformer, column)
     allowed = (RejectColumn,) if allow_reject else ()
     try:
-        output = transformer.fit_transform(transformer_input, y=y)
+        output = transformer.fit_transform(transformer_input, y=y, **kwargs)
     except allowed:
         return col_name, [column], None
     except Exception as e:
@@ -608,12 +627,12 @@ def _fit_transform_column(column, y, columns_to_handle, transformer, allow_rejec
     return col_name, output_cols, transformer
 
 
-def _transform_column(column, transformer):
+def _transform_column(column, transformer, kwargs):
     if transformer is None:
         return [column]
     transformer_input = _prepare_transformer_input(transformer, column)
     try:
-        output = transformer.transform(transformer_input)
+        output = transformer.transform(transformer_input, **kwargs)
     except Exception as e:
         raise ValueError(
             f"Transformer {transformer.__class__.__name__}.transform "

skrub/_apply_to_frame.py

@@ -137,7 +137,7 @@ def __init__(
         self.keep_original = keep_original
         self.rename_columns = rename_columns
 
-    def fit(self, X, y=None):
+    def fit(self, X, y=None, **kwargs):
         """Fit the transformer on all columns jointly.
 
         Parameters
@@ -148,15 +148,19 @@ def fit(self, X, y=None):
         y : Pandas or Polars Series or DataFrame, default=None
             The target data.
 
+        **kwargs
+            Extra named arguments are passed to the ``fit_transform()`` method
+            of ``self.transformer``.
+
         Returns
         -------
         ApplyToFrame
             The transformer itself.
         """
-        self.fit_transform(X, y)
+        self.fit_transform(X, y, **kwargs)
         return self
 
-    def fit_transform(self, X, y=None):
+    def fit_transform(self, X, y=None, **kwargs):
         """Fit the transformer on all columns jointly and transform X.
 
         Parameters
@@ -167,6 +171,10 @@ def fit_transform(self, X, y=None):
         y : Pandas or Polars Series or DataFrame, default=None
             The target data.
 
+        **kwargs
+            Extra named arguments are passed to the ``fit_transform()`` method
+            of ``self.transformer``.
+
         Returns
         -------
         result : Pandas or Polars DataFrame
@@ -183,7 +191,7 @@ def fit_transform(self, X, y=None):
         if self._columns:
             self.transformer_ = clone(self.transformer)
             _utils.set_output(self.transformer_, X)
-            transformed = self.transformer_.fit_transform(to_transform, y)
+            transformed = self.transformer_.fit_transform(to_transform, y, **kwargs)
             transformed = _utils.check_output(
                 self.transformer_, to_transform, transformed, allow_column_list=False
             )
@@ -212,14 +220,18 @@ def fit_transform(self, X, y=None):
         result = sbd.copy_index(X, result)
         return result
 
-    def transform(self, X):
+    def transform(self, X, **kwargs):
         """Transform a dataframe.
 
         Parameters
         ----------
         X : Pandas or Polars DataFrame
             The column to transform.
 
+        **kwargs
+            Extra named arguments are passed to the ``transform()`` method
+            of ``self.transformer_``.
+
         Returns
         -------
         result : Pandas or Polars DataFrame
@@ -236,7 +248,7 @@ def transform(self, X):
             passthrough = selectors.select(X, selectors.inv(self._columns))
         if not self._columns:
             return passthrough
-        transformed = self.transformer_.transform(to_transform)
+        transformed = self.transformer_.transform(to_transform, **kwargs)
         # we do not call `_utils.check_output` here, assuming that if the output
         # had a correct type (e.g. polars dataframe) in `fit_transform` it will
         # have the same (correct) type in `transform`.

skrub/_data_ops/_data_ops.py

@@ -1253,7 +1253,16 @@ def check_subsampled_X_y_shape(X_op, y_op, X_value, y_value, mode, environment,
 class Apply(DataOpImpl):
     """.skb.apply() nodes."""
 
-    _fields = ["X", "estimator", "y", "cols", "how", "allow_reject", "unsupervised"]
+    _fields = [
+        "X",
+        "estimator",
+        "y",
+        "cols",
+        "how",
+        "allow_reject",
+        "unsupervised",
+        "kwargs",
+    ]
 
     # We define `eval()` rather than `compute` because some children may not
     # need to be evaluated depending on the mode. For example in "predict" mode
@@ -1299,8 +1308,10 @@ def eval(self, *, mode, environment):
             # `.skb.preview()` or `.skb.eval()`). We replace `.transform()`
             # with `.predict()`
             if method_name == "fit_transform":
-                self.estimator_.fit(X, y)
-            pred = self.estimator_.predict(X)
+                fit_kwargs = yield from self._eval_kwargs("fit")
+                self.estimator_.fit(X, y, **fit_kwargs)
+            predict_kwargs = yield from self._eval_kwargs("predict")
+            pred = self.estimator_.predict(X, **predict_kwargs)
             # In `(fit_)transform` mode only, format the predictions as a
             # dataframe or column if y was one during `fit()`
             return self._format_predictions(X, pred)
@@ -1316,7 +1327,8 @@ def eval(self, *, mode, environment):
             y_arg = (y,)
         else:
             y_arg = ()
-        return getattr(self.estimator_, method_name)(X, *y_arg)
+        method_kwargs = yield from self._eval_kwargs(method_name)
+        return getattr(self.estimator_, method_name)(X, *y_arg, **method_kwargs)
 
     def _store_y_format(self, y):
         if sbd.is_dataframe(y):
@@ -1355,6 +1367,26 @@ def _format_predictions(self, X, pred):
             return sbd.copy_index(X, pred)
         return pred
 
+    def _eval_kwargs(self, method_name):
+        """
+        Evaluate the kwargs we need to pass to the given method.
+
+        The values in ``self.kwargs`` can be (or contain) DataOps or choices.
+        This looks up the kwargs for ``method_name``, yields it for evaluation,
+        and checks that the result is actually a dictionary before returning it.
+        """
+        kwargs = yield self.kwargs.get(method_name, None)
+        if kwargs is None:
+            # We check if kwargs is None _after_ evaluation
+            kwargs = {}
+        if not isinstance(kwargs, dict):
+            raise TypeError(
+                f"The `{method_name}_kwargs` passed to `.skb.apply()` should be a dict"
+                " of named arguments. Got an object of type"
+                f" {type(kwargs).__name__!r} instead: {kwargs!r}"
+            )
+        return kwargs
+
     def supported_modes(self):
         """
         Used by SkrubLearner and param search to decide if they have the