Bump skrub 0.8.0 -> 0.11.0

pyproject.toml

@@ -25,7 +25,7 @@ requires-python = ">=3.11"
 dependencies = [
     "scikit-learn==1.8",
     "pandas==3.0.2",
-    "skrub==0.8.0",
+    "skrub==0.11.0",
     "polars",
     "graphviz",
     "pygments>=2.20.0",

stratum/_api.py

@@ -1,10 +1,11 @@
 import pandas as pd
 from dataclasses import dataclass
 from skrub import DataOp
+from skrub._data_ops._evaluation import find_scoring_node
 
 from stratum._config import FLAGS
 from stratum.optimizer._optimize import SearchConfig, optimize
-from stratum.optimizer.logical._scoring import resolve_scoring
+from stratum.optimizer.logical._scoring import resolve_declared_scoring, resolve_scoring
 from stratum.runtime._scheduler import SequentialScheduler
 from stratum.frontend._skrub_graph import get_data
 from time import perf_counter
@@ -20,8 +21,15 @@ class RunTimings:
 
 #TODO: Rename this file
 def grid_search(dag: DataOp, cv=None, scoring=None, return_predictions=False, env=None):
-    """Perform grid search with cross-validation on a DataOp DAG. ``scoring`` is required."""
-    if scoring is None:
+    """Perform grid search with cross-validation on a DataOp DAG.
+
+    ``scoring`` is required unless the DAG declares its metric with
+    ``.skb.with_scoring()``; an explicit ``scoring`` overrides that, as in skrub.
+    """
+    declared = scoring is None and find_scoring_node(dag) is not None
+    if declared:
+        metric = resolve_declared_scoring(dag)
+    elif scoring is None:
         # A search ranks a set, so every candidate has to be measured the same way.
         # `.skb.make_grid_search()` still honours scoring=None, by handing the call to
         # skrub. See docs/adr/0004-a-search-always-names-its-metric.md.
@@ -39,8 +47,8 @@ def grid_search(dag: DataOp, cv=None, scoring=None, return_predictions=False, en
         env[k] = v
     # The scorer is plan-time state: it decides what the plan's last operator computes,
     # and an unusable `scoring=` fails here rather than after a fold has been fitted.
-    search = SearchConfig(metric=resolve_scoring(scoring),
-                          return_predictions=return_predictions)
+    search = SearchConfig(metric=metric if declared else resolve_scoring(scoring),
+                          return_predictions=return_predictions, declared_kwargs=declared)
     # Resolve variables to constants at compile time, so the scheduler runs
     # without an environment.
     optimize_start = perf_counter()

stratum/adapters/table_vectorizer.py

@@ -26,6 +26,7 @@
 from skrub._clean_categories import CleanCategories
 from skrub._clean_null_strings import CleanNullStrings
 from skrub._drop_uninformative import DropUninformative
+from skrub._duration_to_float import DurationToFloat
 from skrub._join_utils import pick_column_names
 from skrub._select_cols import Drop
 from skrub._table_vectorizer import (
@@ -181,6 +182,7 @@ def _preprocessors(self):
                 drop_if_unique=self.drop_if_unique,
             ),
             ToDatetime(format=self.datetime_format),
+            DurationToFloat(),
             ToFloat(),
             CleanCategories(),
             ToStr(),

stratum/optimizer/_optimize.py

@@ -1,14 +1,15 @@
 from skrub._data_ops._evaluation import _Graph
 from skrub._data_ops import DataOp
+from skrub._data_ops._data_ops import Apply
 from skrub._data_ops._subsampling import SubsamplePreviews
 from collections import deque, defaultdict
 from dataclasses import dataclass
 from typing import Any
 from .logical._op_cse import apply_op_cse
 from .logical._dataframe_ops import extract_dataframe_op, add_splitting_op
 from .logical._numeric_ops import extract_numeric_op
-from .logical._candidate_ops import CollectCandidatesOp, ScoreCandidatesOp
-from .logical._ops import BaseEstimatorOp, ChoiceOp, Op, OperandRef, as_op, check_choices_not_shared
+from .logical._candidate_ops import CollectCandidatesOp, DeclaredScoringOp, ScoreCandidatesOp
+from .logical._ops import BaseEstimatorOp, ChoiceOp, Op, OperandRef, as_op, check_choices_not_shared, remap_operand_refs
 from .logical._split_ops import SplitOutput
 from ._op_utils import clone_sub_dag, find_choice_naive, replace_op_in_outputs, show_graph, topological_iterator, validate_dag
 from ._explain import explain_linear_plan
@@ -63,6 +64,9 @@ class SearchConfig:
 
     metric: Any
     return_predictions: bool = False
+    #: The metric is the one the plan declares with `.skb.with_scoring()`, so the
+    #: scorer also reads the kwargs declared with it, computed on each test fold.
+    declared_kwargs: bool = False
 
 
 class OptConfig():
@@ -311,6 +315,7 @@ def convert_to_ops(dag: DataOp, env: dict = None) -> Op:
     check_choices_not_shared(nodes.values())
     order = topological_traverse(nodes, parents, children)
     root_id = order[-1]
+    feeds_estimator = _apply_fed_nodes(nodes, parents, order)
 
     # id(DataOp) -> Op. Keyed by DataOp identity (not the graph's node keys) so
     # as_op's operand binder can resolve inputs found in the impl fields directly.
@@ -325,14 +330,33 @@ def convert_to_ops(dag: DataOp, env: dict = None) -> Op:
             input_key = children.get(node_key, [])[0]
             ids_to_ops[id(skrub_op)] = ids_to_ops[id(nodes[input_key])]
             continue
-        ids_to_ops[id(skrub_op)] = as_op(skrub_op, ids_to_ops, env)
+        ids_to_ops[id(skrub_op)] = as_op(skrub_op, ids_to_ops, env,
+                                         feeds_estimator=node_key in feeds_estimator)
 
     root = ids_to_ops[id(nodes[root_id])]
     log_time("conversion took", start)
     _debug_show_graph(root, "conversion")
     return root
 
 
+def _apply_fed_nodes(nodes: dict, parents: dict, order: list) -> set:
+    """Keys of the nodes whose value some `.skb.apply()` consumes, directly or not.
+
+    skrub decides at evaluation time whether an Apply is the last estimator: if
+    another Apply is still waiting on its output, it transforms (or predicts) in
+    every response pass instead of answering the scorer. The plan settles the same
+    question once, from the graph: any downstream Apply, through any of its fields,
+    counts. skrub's answer can depend on the order its traversal first reaches a
+    node shared by an Apply and a non-Apply path; this static rule ignores that.
+    """
+    fed = set()
+    for key in reversed(order):
+        if any(isinstance(nodes[p]._skrub_impl, Apply) or p in fed
+               for p in parents.get(key, ())):
+            fed.add(key)
+    return fed
+
+
 def install_candidate_set(root: Op, search: SearchConfig | None) -> Op:
     """Terminate the plan in a candidate-set operator.
 
@@ -350,6 +374,10 @@ def install_candidate_set(root: Op, search: SearchConfig | None) -> Op:
     replacing = isinstance(root, ChoiceOp)
     names = root.make_outcome_names() if replacing else ["default"]
     candidates = list(root.inputs) if replacing else [root]
+    # A `.skb.with_scoring()` tail is metadata for the scorer, not part of what a
+    # candidate produces: the candidate is the value it passes through.
+    declared = [c if isinstance(c, DeclaredScoringOp) else None for c in candidates]
+    candidates = [d.inputs[0] if d is not None else c for c, d in zip(candidates, declared)]
 
     if search is None or search.metric is None:
         node = CollectCandidatesOp(names)
@@ -364,23 +392,59 @@ def install_candidate_set(root: Op, search: SearchConfig | None) -> Op:
         node = ScoreCandidatesOp(names, search.metric, response_mode=mode,
                                  emit_predictions=search.return_predictions)
     node.inputs = list(candidates)
-    for candidate in candidates:
+    for candidate, d in zip(candidates, declared):
         if replacing:
             # The choice is the root, so it is the only consumer being redirected. A
             # rebuild rather than `replace_output` because one operator may fill two
             # candidate slots and therefore be redirected twice.
-            candidate.outputs = [node if out is root else out for out in candidate.outputs]
+            candidate.outputs = [node if out is root or out is d else out
+                                 for out in candidate.outputs]
         else:
+            candidate.outputs = [out for out in candidate.outputs if out is not d]
             candidate.add_output(node)
 
     if isinstance(node, ScoreCandidatesOp):
         _, y_op = fold
         node.y_ref = OperandRef(node.add_input(y_op))
         y_op.add_output(node)
+        if search.declared_kwargs:
+            if any(d is None for d in declared):
+                raise ValueError(
+                    "not every candidate ends in `.skb.with_scoring()`, so the declared"
+                    " scorer cannot score them all. Pass `scoring=` instead.")
+            node.kwargs_refs = [_rebind(d.scorers[0]["kwargs"], d.inputs, node)
+                                for d in declared]
+    for d in declared:
+        if d is not None:
+            # The kwargs the scorer now reads have it as their consumer; the rest of
+            # the declaration (and a search scoring with another metric) is dead.
+            d.detach_and_prune()
     log_time("installing the candidate set took", start)
     return node
 
 
+def _rebind(value, operands: list[Op], consumer: Op):
+    """``value`` with its OperandRefs into ``operands`` re-pointed at inputs of ``consumer``."""
+    mapping = {}
+    for k in _referenced(value):
+        if k not in mapping:
+            mapping[k] = consumer.add_input(operands[k])
+            operands[k].add_output(consumer)
+    return remap_operand_refs(value, mapping)
+
+
+def _referenced(value):
+    """The operand indices ``value`` refers to, through nested tuples, lists and dicts."""
+    if isinstance(value, OperandRef):
+        yield value.k
+    elif isinstance(value, (list, tuple)):
+        for v in value:
+            yield from _referenced(v)
+    elif isinstance(value, dict):
+        for v in value.values():
+            yield from _referenced(v)
+
+
 def _response_mode(metric, names: list[str], candidates: list[Op]) -> str:
     """The response method the scoring pass runs, chosen once here.
 

stratum/optimizer/logical/_aggregation_ops.py

@@ -475,10 +475,7 @@ def _detach_and_rewire(new_op: AggregateOp, df: Op, folder: _Folder,
     replaced. Downstream consumers are rewired by the caller.
     """
     for node in folder.absorbed:
-        for inp in node.inputs:
-            inp.outputs = [o for o in inp.outputs if o is not node]
-        node.inputs = []
-        node.outputs = []
+        node.detach()
     replaced_ids = {id(node) for node in replaced if node is not None}
     for producer in (df, *folder.leaf_ops):
         producer.outputs = [o for o in producer.outputs

stratum/optimizer/logical/_base.py

@@ -366,6 +366,32 @@ def replace_output_of_inputs(self, new_output):
         for in_ in self.inputs:
             in_.replace_output(self, new_output)
 
+    def detach(self) -> None:
+        """Take this node out of the graph: drop it from its producers' outputs and
+        clear its own edges. Its consumers are the caller's to rewire, before or after.
+
+        Producers left without a consumer stay as they are, which is what a rewrite
+        wants while it is about to hand them a new consumer; see ``detach_and_prune``.
+        """
+        for in_ in self.inputs:
+            in_.outputs = [out for out in in_.outputs if out is not self]
+        self.inputs = []
+        self.outputs = []
+
+    def detach_and_prune(self) -> None:
+        """``detach`` this node, then every producer that leaves without a consumer.
+
+        A producer nobody consumes is unreachable from the root, but it is still in
+        its own producers' outputs, and the topological walk rejects that edge. So the
+        dead cone goes with the node. Call it only once every surviving producer has
+        its new consumer: one whose last edge is momentarily gone would be pruned too.
+        """
+        producers = list(dict.fromkeys(self.inputs))
+        self.detach()
+        for in_ in producers:
+            if not in_.outputs:
+                in_.detach_and_prune()
+
     def clone(self):
         if getattr(self.__class__, "fields", None) is None:
             raise NotImplementedError(f"Cloning of {self.__class__.__name__} objects is not implemented yet. Please implement it.")

stratum/optimizer/logical/_candidate_ops.py

@@ -20,7 +20,7 @@
 import polars as pl
 
 from stratum.optimizer.logical import _schema
-from stratum.optimizer.logical._base import OutputType
+from stratum.optimizer.logical._base import OutputType, _resolve_operand
 from stratum.optimizer.logical._ops import FITTING_MODE, Op, OperandRef
 from stratum.optimizer.logical._scoring import Metric, read_response
 
@@ -30,6 +30,27 @@
 _MAX_LABEL = 50
 
 
+class DeclaredScoringOp(Op):
+    """``pred.skb.with_scoring(...)``: ``pred`` unchanged, plus the scorers it declares.
+
+    The scorers keep their skrub shape (``{"scoring", "kwargs", "name"}``), with every
+    DataOp in them bound to an operand of this op, so a ``sample_weight`` computed from
+    the data is part of the plan. Like skrub's ``Scoring`` node it passes ``pred``
+    through; the candidate set strips it and, when the search scores with the declared
+    metric, reads the kwargs operands per fold (see ``install_candidate_set``).
+    """
+
+    logical_family = "DeclaredScoring"
+    fields = ["scorers"]
+
+    def __init__(self, scorers: list, inputs: list | None = None):
+        super().__init__(inputs=list(inputs) if inputs else [])
+        self.scorers = scorers
+
+    def process(self, mode: str, inputs: list):
+        return inputs[0]
+
+
 class CandidateSetOp(Op):
     """Fan-in over every candidate's output, carrying one label per candidate.
 
@@ -100,6 +121,9 @@ def __init__(self, candidate_names: list[str], metric: Metric,
         self.emit_predictions = emit_predictions
         self.output_type = OutputType.FRAME
         self.y_ref: OperandRef | None = None
+        # One entry per candidate: the kwargs its declared scorer reads, with
+        # OperandRefs for the ones the plan computes on the fold; None for none.
+        self.kwargs_refs: list | None = None
 
     def propagate_output_schema(self):
         """The scored table is built column by column in ``process``: one row per
@@ -133,8 +157,10 @@ def process(self, mode: str, inputs: list):
         ids, scores, values = [], [], []
         for i, name in enumerate(self.candidate_names):
             produced = inputs[i]
+            kwargs = _resolve_operand(self.kwargs_refs[i], inputs) if self.kwargs_refs else None
             try:
-                scores.append(self.metric(y_true, read_response(produced, self.response_mode)))
+                scores.append(self.metric(y_true, read_response(produced, self.response_mode),
+                                          **(kwargs or {})))
             except Exception as e:
                 raise RuntimeError(
                     f"[scoring] {self.metric!r} failed on candidate {name!r}"

stratum/optimizer/logical/_dataframe_ops.py

@@ -3,6 +3,7 @@
 from pandas import DataFrame
 from polars import DataFrame as PolarsDataFrame
 from skrub import SelectCols
+from skrub.selectors._base import All
 import operator
 from stratum.optimizer.logical import _schema
 import pandas as pd
@@ -38,7 +39,12 @@
 class ConcatOp(Op):
     """Logical frame concatenation. Pure config -- execution is provided by the
     physical impls in ``physical/_concat_execs.py`` (Pandas/PolarsConcatOp),
-    selected at plan time."""
+    selected at plan time.
+
+    Like skrub's ``.skb.concat`` it also stacks numpy arrays, which the physical
+    impls detect at run time. A graph-fed operand's kind is unknown when planning,
+    so the op is taken to produce a frame unless its first operand is a literal
+    array."""
     logical_family = "Concat"
     fields = ["first", "others", "axis"] # Add more if needed
 
@@ -50,7 +56,8 @@ def __init__(self, first, others: list, axis):
         self.first = first
         self.others = list(others)
         self.axis = axis
-        self.output_type = OutputType.FRAME
+        self.output_type = (OutputType.UNKNOWN if isinstance(first, np.ndarray)
+                            else OutputType.FRAME)
 
     def propagate_output_schema(self):
         """See :func:`_schema.concat_schemas` for the per-axis rules.
@@ -290,8 +297,11 @@ def extract_dataframe_op(op: Op, root: Op, selection_op = True, map_op = True,
             op.output_type = OutputType.FRAME
             # skrub implements `skb.select(cols)` as an Apply of its SelectCols
             # transformer; surface the selector as a dedicated ColumnSelectorOp.
+            # Only while the Apply hands it every column: restricted by `cols` or
+            # `exclude_cols`, skrub passes the other columns through untouched.
             if (isinstance(op, TransformerOp) and isinstance(op.estimator, SelectCols)
-                    and not op.param_refs and not isinstance(op.y, OperandRef)):
+                    and not op.param_refs and not isinstance(op.y, OperandRef)
+                    and isinstance(op.cols, All) and op.exclude_cols is None):
                 new_op = make_column_selector_op(op)
 
         elif isinstance(op, ChoiceOp):

stratum/optimizer/logical/_map_ops.py

@@ -64,10 +64,7 @@ def _detach_absorbed_and_rewire(op: Op, new_op: MapOp, folder: _Folder) -> None:
     Downstream consumers of ``op`` are rewired by the caller.
     """
     for node in folder.absorbed:
-        for inp in node.inputs:
-            inp.outputs = [o for o in inp.outputs if o is not node]
-        node.inputs = []
-        node.outputs = []
+        node.detach()
     for producer in new_op.inputs:
         producer.outputs = [o for o in producer.outputs if o is not op]
         producer.add_output(new_op)