Support tuple unpacking of DataOps

CHANGES.rst

@@ -12,6 +12,10 @@ Ongoing development
 
 New Features
 ------------
+- It is now possible to unpack a :class:`DataOp` that evaluates to several
+  values, for example ``train, test = data_op``. Each target becomes a DataOp
+  that extracts one of the values.
+  :pr:`2243` by :user:`Elias Strauss <e-strauss>`.
 
 Changes
 -------

doc/modules/data_ops/basics/control_flow.rst

@@ -106,25 +106,30 @@ Result:
 Unpacking multiple outputs from deferred functions
 --------------------------------------------------
 
-When a deferred function returns more than one value, you cannot unpack the
-result directly because unpacking iterates over the result. Iteration is not
-supported on DataOps until evaluation.
-
-In general, it is recommended that deferred functions return a single
-value whenever possible. Returning multiple outputs should be avoided unless
-strictly necessary, as it makes downstream usage more complex.
-
-Instead, keep the result as a single DataOp and index into it:
+When a deferred function returns more than one value, we can unpack the result
+as usual:
 
 >>> test = skrub.var("test", [1, 2])
 >>> @skrub.deferred
 ... def process_test_data(test):
-...     left = test[0]
-...     right = test[1]
-...     return left, right
->>> res = test.skb.apply_func(process_test_data)
->>> left = res[0]
->>> right = res[1]
+...     return test[0], test[1]
+>>> left, right = process_test_data(test)
+>>> left
+<GetItem 0>
+Result:
+―――――――
+1
+
+Each target becomes a DataOp that extracts one of the values, once the
+computation runs. To do so, skrub needs to know how many values we are
+unpacking: it finds out by inspecting the bytecode of the assignment. This
+works on CPython, but it is not guaranteed by the language, and it does not
+cover starred targets such as ``first, *rest = ...`` (the number of values is
+not known in that case). When in doubt, we can always keep the result as a
+single DataOp and index into it, which is equivalent and always works:
+
+>>> res = process_test_data(test)
+>>> left, right = res[0], res[1]
 
 :func:`deferred` is useful not only for our own functions, but also when we
 need to call module-level functions from a library. For example, to delay the

skrub/_data_ops/_data_ops.py

@@ -34,6 +34,7 @@
 import operator
 import pathlib
 import re
+import sys
 import textwrap
 import traceback
 import types
@@ -197,6 +198,28 @@ def _format_data_op_creation_stack():
     return traceback.format_list(stack)
 
 
+def _unpack_arity():
+    """Number of targets in the unpacking assignment being executed, if any.
+
+    Read from the caller's ``UNPACK_SEQUENCE`` instruction, or None if not found.
+    """
+    if (getframe := getattr(sys, "_getframe", None)) is None:
+        return None
+    try:
+        # skip the frames of this function and of DataOp.__iter__
+        frame = getframe(2)
+        for instruction in dis.get_instructions(frame.f_code):
+            if instruction.offset == frame.f_lasti:
+                if instruction.opname == "UNPACK_SEQUENCE":
+                    return instruction.arg
+                return None
+    except Exception:
+        # this is best-effort introspection: anything unexpected must fall back
+        # on refusing to iterate, not raise something else.
+        pass
+    return None
+
+
 class DataOpImpl:
     """Base class for all kinds of DataOps (computation graph nodes).
 
@@ -700,6 +723,13 @@ def __bool__(self):
         )
 
     def __iter__(self):
+        # Unpacking (`a, b = data_op`) is supported: we know how many values are
+        # expected, so we can create a node for each of them. Any other kind of
+        # iteration would need the length of the result, which is unknown until
+        # the DataOp is evaluated.
+        if (arity := _unpack_arity()) is not None:
+            values = unpack(self, arity)
+            return iter([values[i] for i in range(arity)])
         raise TypeError(
             "This object is a DataOp that will be evaluated later, "
             "when your learner runs. So it is not possible to eagerly "
@@ -1662,6 +1692,33 @@ def pretty_repr(self):
         return f"[{_get_preview(self.key)!r}]"
 
 
+class AsTuple(DataOpImpl):
+    """Node created by unpacking a DataOp, e.g. ``a, b = data_op``."""
+
+    _fields = ["iterable", "expected_length"]
+
+    def compute(self, e, mode, environment):
+        # converting to a tuple (rather than indexing into the result directly)
+        # allows unpacking any iterable, and makes sure an iterator is consumed
+        # only once even though each target indexes into this node.
+        result = tuple(e.iterable)
+        expected, got = e.expected_length, len(result)
+        if got != expected:
+            problem = "not enough" if got < expected else "too many"
+            raise ValueError(
+                f"{problem} values to unpack (expected {expected}, got {got})"
+            )
+        return result
+
+    def __repr__(self):
+        return f"<{self.__class__.__name__} {short_repr(self.iterable)}>"
+
+
+@checked_data_op_constructor
+def unpack(iterable, expected_length):
+    return DataOp(AsTuple(iterable, expected_length))
+
+
 class Call(DataOpImpl):
     _fields = [
         "func",

skrub/_data_ops/tests/test_data_ops.py

@@ -46,6 +46,33 @@ def test_slice():
     assert c.skb.eval({"a": list(range(10, 20)), "b": 5}) == [10, 11, 12, 13, 14]
 
 
+def test_unpacking():
+    a = skrub.var("a", [10, 20, 30])
+    first, second, third = a
+    assert repr(first).startswith("<GetItem 0>")
+    assert (first.skb.eval(), second.skb.eval(), third.skb.eval()) == (10, 20, 30)
+    assert (first + second).skb.eval({"a": [1, 2, 3]}) == 3
+
+
+def test_unpacking_iterable():
+    # the result does not need to be a sequence, and an iterator must be
+    # consumed only once even though each target indexes into it.
+    @skrub.deferred
+    def f():
+        yield 1
+        yield 2
+        yield 3
+
+    a, b, c = f()
+    assert (a + b + c).skb.eval() == 6
+
+
+def test_unpacking_nested():
+    a = skrub.var("a", [1, [2, 3]])
+    first, (second, third) = a
+    assert (first.skb.eval(), second.skb.eval(), third.skb.eval()) == (1, 2, 3)
+
+
 def test_estimator():
     c = skrub.var("c")
     e = skrub.as_data_op(LogisticRegression(C=c))

skrub/_data_ops/tests/test_errors.py

@@ -2,6 +2,7 @@
 import re
 import sys
 import traceback
+import types
 
 import numpy as np
 import pytest
@@ -11,6 +12,7 @@
 from sklearn.preprocessing import StandardScaler
 
 import skrub
+from skrub._data_ops import _data_ops
 from skrub._utils import PassThrough
 from skrub.conftest import skip_polars_installed_without_pyarrow
 
@@ -28,6 +30,83 @@ def test_for():
             pass
 
 
+def test_star_unpacking():
+    # unlike `a, b = data_op`, iteration with a starred target does not tell us
+    # how many values are expected so it remains unsupported.
+    a = skrub.var("a", [1, 2, 3])
+    with pytest.raises(
+        TypeError, match=".*it is not possible to eagerly iterate over it"
+    ):
+        _first, *_rest = a
+    with pytest.raises(
+        TypeError, match=".*it is not possible to eagerly iterate over it"
+    ):
+        (lambda *args: None)(*a)
+
+
+def test_unpacking_wrong_number_of_targets():
+    a = skrub.var("a", [1, 2, 3])
+    with pytest.raises(
+        RuntimeError,
+        match=r"(?s)Evaluation of 'unpack\(\)' failed"
+        r".*too many values to unpack \(expected 2, got 3\)",
+    ):
+        _first, _second = a
+    with pytest.raises(
+        RuntimeError,
+        match=r"(?s)Evaluation of 'unpack\(\)' failed"
+        r".*not enough values to unpack \(expected 4, got 3\)",
+    ):
+        _first, _second, _third, _fourth = a
+
+
+def _bytecode_inspection_failure(code):
+    raise RuntimeError("cannot inspect bytecode")
+
+
+@pytest.mark.parametrize(
+    "module_name, replacement",
+    [
+        # a Python implementation that does not provide sys._getframe
+        pytest.param("sys", types.SimpleNamespace(), id="no_getframe"),
+        # inspecting the bytecode fails
+        pytest.param(
+            "dis",
+            types.SimpleNamespace(get_instructions=_bytecode_inspection_failure),
+            id="inspection_error",
+        ),
+        # the instruction being executed is not found in the bytecode
+        pytest.param(
+            "dis",
+            types.SimpleNamespace(get_instructions=lambda code: iter(())),
+            id="instruction_not_found",
+        ),
+    ],
+)
+def test_unpacking_without_bytecode_inspection(monkeypatch, module_name, replacement):
+    # `a, b = data_op` relies on finding the UNPACK_SEQUENCE instruction that is
+    # being executed. When that is not possible we fall back on refusing to
+    # iterate, as users can always index into the result instead.
+    monkeypatch.setattr(_data_ops, module_name, replacement)
+    a = skrub.var("a", [1, 2])
+    with pytest.raises(
+        TypeError, match=".*it is not possible to eagerly iterate over it"
+    ):
+        _first, _second = a
+
+
+def test_unpacking_wrong_number_of_targets_at_runtime():
+    # without a value for 'a' the length is only known when the plan runs
+    a = skrub.var("a")
+    first, _second = a
+    if sys.version_info < (3, 11):
+        err_t, err_msg = RuntimeError, "Evaluation of node <AsTuple"
+    else:
+        err_t, err_msg = ValueError, "too many values to unpack"
+    with pytest.raises(err_t, match=err_msg):
+        first.skb.eval({"a": [1, 2, 3]})
+
+
 def test_if():
     a = skrub.var("a", True)
     with pytest.raises(