Repository: dask/dask
Pull request #11907 (2025-04-24 14:05:29)

Map_partitions again accepts delayed objects
With the move to dask-expr we apparently lost the feature that map_partitions accepts delayed objects. Maybe other APIs are also affected but I don' tknow of a good way to test this.

github-actions[bot]: ## Unit Test Results
_See [test report](https://dask.github.io/dask/test_report.html) for an extended history of previous test failures. This is useful for diagnosing flaky tests._

      9 files  ±0        9 suites  ±0   3h 25m 37s ⏱️ + 2m 26s
 18 013 tests +1   16 799 ✅ +1   1 214 💤 ±0  0 ❌ ±0 
161 191 runs  +9  149 086 ✅ +8  12 105 💤 +1  0 ❌ ±0 

Results for commit a5596524. ± Comparison against base commit c68e71d2.

<details>
  <summary>This pull request <b>removes</b> 1 and <b>adds</b> 2 tests. <i>Note that renamed tests count towards both.</i></summary>

```
dask.dataframe.dask_expr.tests.test_distributed ‑ test_future_in_map_partitions
```

```
dask.dataframe.dask_expr.tests.test_distributed ‑ test_future_in_map_partitions[delayed]
dask.dataframe.dask_expr.tests.test_distributed ‑ test_future_in_map_partitions[future]
```
</details>

[test-results]:data:application/gzip;base64,H4sIAGJPCmgC/03MTQ6DIBCG4as0rLtglJlxepnGICSk/jQIq6Z3LxqlLt9nMt9H+TC6VT1ucr+pNYdUY8ixT2GZS0LTtlyoHNN2hk5De/ZzzdZuSCzyx1d4759gqvk+jMV0BRfjEg+Jed6nCUDg6DptRHdUsU5rPO0yvfd12S7TFFIJ1SMKYWMAGL1phFnQdXYYSHsgcsiGrSNR3x95n5eYGAEAAA==


Patch:
diff --git a/dask/_task_spec.py b/dask/_task_spec.py
index bf404f166..5a2aa1888 100644
--- a/dask/_task_spec.py
+++ b/dask/_task_spec.py
@@ -954,15 +954,15 @@ class Dict(NestedContainer, Mapping):
     def __iter__(self):
         yield from self.args[::2]
 
+    def __len__(self):
+        return len(self.args) // 2
+
     def __getitem__(self, key):
         for k, v in batched(self.args, 2, strict=True):
             if k == key:
                 return v
         raise KeyError(key)
 
-    def __len__(self):
-        return len(self.args) // 2
-
     @staticmethod
     def constructor(args):
         return dict(batched(args, 2, strict=True))
diff --git a/dask/dataframe/dask_expr/_collection.py b/dask/dataframe/dask_expr/_collection.py
index 7afcb6898..91fea7840 100644
--- a/dask/dataframe/dask_expr/_collection.py
+++ b/dask/dataframe/dask_expr/_collection.py
@@ -28,6 +28,7 @@ import dask.dataframe.methods as methods
 from dask import compute, get_annotations
 from dask._collections import new_collection
 from dask._expr import OptimizerStage
+from dask._task_spec import Dict, TaskRef
 from dask.array import Array
 from dask.base import DaskMethodsMixin, is_dask_collection, named_schedulers
 from dask.core import flatten
@@ -65,6 +66,7 @@ from dask.dataframe.dask_expr._expr import (
     ToDatetime,
     ToNumeric,
     ToTimedelta,
+    _DelayedExpr,
     no_default,
 )
 from dask.dataframe.dask_expr._merge import JoinRecursive, Merge
@@ -6166,6 +6168,18 @@ def map_partitions(
         # will need to call `Repartition` on operands that are not
         # aligned with `self.expr`.
         raise NotImplementedError()
+    args = [_DelayedExpr(a) if isinstance(a, Delayed) else a for a in args]
+    newkwargs = {}
+    delayed_kwargs = []
+    for k, v in kwargs.items():
+        if isinstance(v, Delayed):
+            dexpr = _DelayedExpr(v)
+            delayed_kwargs.append(dexpr)
+            newkwargs[k] = TaskRef(dexpr.__dask_keys__()[0])
+        else:
+            newkwargs[k] = v
+    del kwargs
+
     new_expr = expr.MapPartitions(
         args[0],
         func,
@@ -6175,9 +6189,11 @@ def map_partitions(
         clear_divisions,
         align_dataframes,
         parent_meta,
-        kwargs.pop("token", None),
-        kwargs,
+        newkwargs.pop("token", None),
+        Dict(newkwargs),
+        len(args) - 1,
         *args[1:],
+        *delayed_kwargs,
     )
     return new_collection(new_expr)
 
diff --git a/dask/dataframe/dask_expr/_expr.py b/dask/dataframe/dask_expr/_expr.py
index 0474830b0..461bc58dc 100644
--- a/dask/dataframe/dask_expr/_expr.py
+++ b/dask/dataframe/dask_expr/_expr.py
@@ -656,12 +656,14 @@ class MapPartitions(Blockwise):
         "parent_meta",
         "token",
         "kwargs",
+        "nargs",
     ]
     _defaults: dict = {
         "kwargs": None,
         "align_dataframes": True,
         "parent_meta": None,
         "token": None,
+        "nargs": 0,
     }
 
     @functools.cached_property
@@ -687,14 +689,20 @@ class MapPartitions(Blockwise):
 
     @functools.cached_property
     def args(self):
-        return [self.frame] + self.operands[len(self._parameters) :]
+        return [self.frame] + self.operands[
+            len(self._parameters) : len(self._parameters) + self.nargs
+        ]
 
     @functools.cached_property
     def _meta(self):
         meta = self.operand("meta")
         return _get_meta_map_partitions(
             self.args,
-            [e for e in self.args if isinstance(e, Expr)],
+            [
+                e
+                for e in self.args
+                if isinstance(e, Expr) and not isinstance(e, _DelayedExpr)
+            ],
             self.func,
             self.kwargs,
             meta,
@@ -725,7 +733,7 @@ class MapPartitions(Blockwise):
 
     def _task(self, name: Key, index: int) -> Task:
         args = [self._blockwise_arg(op, index) for op in self.args]
-        kwargs = (self.kwargs if self.kwargs is not None else {}).copy()
+        kwargs = dict(self.kwargs if self.kwargs is not None else {})
         if self._has_partition_info:
             kwargs["partition_info"] = {
                 "number": index,
@@ -978,6 +986,7 @@ class MapOverlap(MapPartitions):
             None,
             self.token,
             self._kwargs,
+            len(self.args[1:]),
             *self.args[1:],
         )
 
diff --git a/dask/dataframe/dask_expr/_groupby.py b/dask/dataframe/dask_expr/_groupby.py
index 1364dca17..fa543c9e1 100644
--- a/dask/dataframe/dask_expr/_groupby.py
+++ b/dask/dataframe/dask_expr/_groupby.py
@@ -1391,6 +1391,7 @@ class GroupByCumulative(Expr, GroupByBase):
             None,
             None,
             {"chunk": self.chunk, "columns": columns, **dropna, **self.numeric_only},
+            len(self.by),
             *self.by,
         )
         cum_raw = frame
@@ -1420,6 +1421,7 @@ class GroupByCumulative(Expr, GroupByBase):
             None,
             None,
             {"chunk": M.last, "columns": columns, **dropna},
+            len(by),
             *by,
         )
         return GroupByCumulativeFinalizer(
diff --git a/dask/dataframe/dask_expr/_merge_asof.py b/dask/dataframe/dask_expr/_merge_asof.py
index 436eb229c..eb9684cc6 100644
--- a/dask/dataframe/dask_expr/_merge_asof.py
+++ b/dask/dataframe/dask_expr/_merge_asof.py
@@ -132,6 +132,7 @@ class MergeAsof(Merge):
                 None,
                 None,
                 self._kwargs,
+                1,
                 self.right,
             )
 
@@ -154,6 +155,7 @@ class MergeAsof(Merge):
                 None,
                 None,
                 {"left_index": True, "right_index": True},
+                1,
                 right,
             )
 


Test patch:
diff --git a/dask/dataframe/dask_expr/tests/test_distributed.py b/dask/dataframe/dask_expr/tests/test_distributed.py
index a1333c1ea..91128141a 100644
--- a/dask/dataframe/dask_expr/tests/test_distributed.py
+++ b/dask/dataframe/dask_expr/tests/test_distributed.py
@@ -350,8 +350,9 @@ def test_merge_combine_similar_squash_merges(add_repartition):
     )
 
 
+@pytest.mark.parametrize("kind", ["future", "delayed"])
 @gen_cluster(client=True)
-async def test_future_in_map_partitions(c, s, a, b):
+async def test_future_in_map_partitions(c, s, a, b, kind):
     # xgboost uses this pattern
 
     def test_func(n):
@@ -361,11 +362,22 @@ async def test_future_in_map_partitions(c, s, a, b):
 
     df = from_pandas(pd.DataFrame({"a": [1, 2, 3, 4]}), npartitions=2)
 
-    f = c.submit(test_func, 100)
-    q = map_partitions(lambda x, y: y + x.sum(), f, df, meta=df._meta)
+    if kind == "future":
+        f = c.submit(test_func, 100)
+        z = c.submit(test_func, 10)
+    elif kind == "delayed":
+        f = dask.delayed(test_func)(100)
+        z = dask.delayed(test_func)(10)
+    else:
+        raise ValueError("kind must be 'future' or 'delayed'")
+
+    def _foo(x, y, z):
+        return x.sum() + y + z.sum()
+
+    q = map_partitions(_foo, f, df, z=z, meta=df._meta)
     result = c.compute(q)
     result = await result
-    expected = pd.DataFrame({"a": [4951, 4952, 4953, 4954]})
+    expected = pd.DataFrame({"a": [4996, 4997, 4998, 4999]})
     pd.testing.assert_frame_equal(result, expected)
 
 
diff --git a/dask/tests/test_task_spec.py b/dask/tests/test_task_spec.py
index becbefca8..4986c3efb 100644
--- a/dask/tests/test_task_spec.py
+++ b/dask/tests/test_task_spec.py
@@ -4,6 +4,7 @@ import itertools
 import pickle
 import sys
 from collections import namedtuple
+from collections.abc import Mapping
 
 import pytest
 
@@ -1013,6 +1014,21 @@ def test_dict_class():
 
     d = Dict([["columns", ["a", "b"]]])
     assert d() == {"columns": ["a", "b"]}
+    # Can be converted to a dict, e.g. also used as **kwargs
+    assert isinstance(t, Mapping)
+    assert dict(t) == {
+        "k": Task("key-1", func, "a", "b"),
+        "v": Task("key-2", func, "c", "d"),
+    }
+    assert len(t) == len(dict(t)) == 2
+
+    def as_kwargs(**kwargs):
+        return kwargs == {
+            "k": Task("key-1", func, "a", "b"),
+            "v": Task("key-2", func, "c", "d"),
+        }
+
+    assert as_kwargs(**t)
 
 
 def test_block_io_fusion():