Repository: dask/dask

Issue:
`partition_info` in `map_partitions` materializes the graph unnecessarily
Using a `partition_info` kwarg in your `map_partitions` function forces the entire blockwise graph to be materialized. This can be slow, lead to slower scheduling, and eliminates optimization opportunities.

The way the injection of `partition_info` is [currently implemented](https://github.com/dask/dask/blob/3b61950cbdf033d53ae98937d19a9ff0b2a38f40/dask/dataframe/core.py#L5919-L5938) is a bit convoluted, involving manipulation of individual tasks within the materialized graph.

Since we now have the [`BlockwiseDep`](https://github.com/dask/dask/blob/fd1b02b61e546692ddf80350e30d0c789f40f2d9/dask/blockwise.py#L35-L55) and [`BlockwiseDepDict`](https://github.com/dask/dask/blob/fd1b02b61e546692ddf80350e30d0c789f40f2d9/dask/blockwise.py#L101-L155) interface for just this sort of purpose (including auxiliary information into a blockwise graph; concept described in https://github.com/dask/dask/issues/7513), this could be done more simply. The main challenge is just that it needs to be given as a kwarg, and we don't currently have the infrastructure in `map_partitions` to pass blockwise-y things as kwargs (xref https://github.com/dask/dask/issues/8308), so it may require an odd dance with a wrapper function.

cc @rjzamora @ian-r-rose 


Discussion:

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -1,4 +1,3 @@
-import copy
 import operator
 import warnings
 from collections.abc import Iterator, Sequence
@@ -22,11 +21,10 @@
 from .. import core, threaded
 from ..array.core import Array, normalize_arg
 from ..base import DaskMethodsMixin, dont_optimize, is_dask_collection, tokenize
-from ..blockwise import Blockwise, blockwise, subs
+from ..blockwise import Blockwise, BlockwiseDep, BlockwiseDepDict, blockwise
 from ..context import globalmethod
 from ..delayed import Delayed, delayed, unpack_collections
 from ..highlevelgraph import HighLevelGraph
-from ..optimization import SubgraphCallable
 from ..utils import (
     IndexCallable,
     M,
@@ -5901,9 +5899,6 @@ def map_partitions(
     name = kwargs.pop("token", None)
     parent_meta = kwargs.pop("parent_meta", None)
 
-    if has_keyword(func, "partition_info"):
-        kwargs["partition_info"] = {"number": -1, "divisions": None}
-
     assert callable(func)
     if name is not None:
         token = tokenize(meta, *args, **kwargs)
@@ -5930,9 +5925,6 @@ def map_partitions(
     else:
         meta = make_meta(meta, index=meta_index, parent_meta=parent_meta)
 
-    if has_keyword(func, "partition_info"):
-        kwargs["partition_info"] = "__dummy__"
-
     if all(isinstance(arg, Scalar) for arg in args):
         layer = {
             (name, 0): (apply, func, (tuple, [(arg._name, 0) for arg in args]), kwargs)
@@ -5972,22 +5964,6 @@ def map_partitions(
         if collections:
             simple = False
 
-    if enforce_metadata:
-        dsk = partitionwise_graph(
-            apply_and_enforce,
-            name,
-            *args2,
-            dependencies=dependencies,
-            _func=func,
-            _meta=meta,
-            **kwargs3,
-        )
-    else:
-        kwargs4 = kwargs if simple else kwargs3
-        dsk = partitionwise_graph(
-            func, name, *args2, **kwargs4, dependencies=dependencies
-        )
-
     divisions = dfs[0].divisions
     if transform_divisions and isinstance(dfs[0], Index) and len(dfs) == 1:
         try:
@@ -6003,25 +5979,32 @@ def map_partitions(
                 divisions = [None] * (dfs[0].npartitions + 1)
 
     if has_keyword(func, "partition_info"):
-        dsk = dict(dsk)
-
-        for k, v in dsk.items():
-            subgraph = v[0]
-            number = k[-1]
-            assert isinstance(number, int)
-            info = {"number": number, "division": divisions[number]}
-            # Replace the __dummy__ keyword argument with `info`
-            subgraph_dsk = copy.copy(subgraph.dsk)
-            [(key, task)] = subgraph_dsk.items()
-            subgraph_dsk[key] = subs(task, {"__dummy__": info})
-            dsk[k] = (
-                SubgraphCallable(
-                    dsk=subgraph_dsk,
-                    outkey=subgraph.outkey,
-                    inkeys=subgraph.inkeys,
-                    name=f"{subgraph.name}-info-{tokenize(info)}",
-                ),
-            ) + v[1:]
+        partition_info = {
+            (i,): {"number": i, "division": division}
+            for i, division in enumerate(divisions[:-1])
+        }
+
+        args2.insert(0, BlockwiseDepDict(partition_info))
+        orig_func = func
+        func = lambda partition_info, *args, **kwargs: orig_func(
+            *args, **kwargs, partition_info=partition_info
+        )
+
+    if enforce_metadata:
+        dsk = partitionwise_graph(
+            apply_and_enforce,
+            name,
+            *args2,
+            dependencies=dependencies,
+            _func=func,
+            _meta=meta,
+            **kwargs3,
+        )
+    else:
+        kwargs4 = kwargs if simple else kwargs3
+        dsk = partitionwise_graph(
+            func, name, *args2, **kwargs4, dependencies=dependencies
+        )
 
     graph = HighLevelGraph.from_collections(name, dsk, dependencies=dependencies)
     return new_dd_object(graph, name, meta, divisions)
@@ -7146,6 +7129,15 @@ def partitionwise_graph(func, name, *args, **kwargs):
             else:
                 raise ValueError("Can't add multi-dimensional array to dataframes")
             numblocks[arg._name] = arg.numblocks
+        elif isinstance(arg, BlockwiseDep):
+            if len(arg.numblocks) == 1:
+                pairs.extend([arg, "i"])
+            elif len(arg.numblocks) == 2:
+                pairs.extend([arg, "ij"])
+            else:
+                raise ValueError(
+                    f"BlockwiseDep arg {arg!r} has {len(arg.numblocks)} dimensions; only 1 or 2 are supported."
+                )
         else:
             pairs.extend([arg, None])
     return blockwise(


Test patch:
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -30,6 +30,7 @@
 from dask.dataframe.utils import assert_eq, assert_max_deps, make_meta
 from dask.datasets import timeseries
 from dask.utils import M, is_dataframe_like, is_series_like, put_lines
+from dask.utils_test import hlg_layer
 
 dsk = {
     ("x", 0): pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}, index=[0, 1, 3]),
@@ -907,6 +908,8 @@ def test_map_partitions():
     assert_eq(d.map_partitions(lambda df: df, meta=d), full)
     assert_eq(d.map_partitions(lambda df: df), full)
     result = d.map_partitions(lambda df: df.sum(axis=1))
+    layer = hlg_layer(result.dask, "lambda-")
+    assert not layer.is_materialized(), layer
     assert_eq(result, full.sum(axis=1))
 
     assert_eq(
@@ -936,7 +939,11 @@ def f(df, partition_info=None):
         assert dsk[("x", d.divisions.index(partition_info["division"]))].equals(df)
         return df
 
-    result = d.map_partitions(f, meta=d).compute(scheduler="single-threaded")
+    df = d.map_partitions(f, meta=d)
+    layer = hlg_layer(df.dask, "f-")
+    assert not layer.is_materialized()
+    df.dask.validate()
+    result = df.compute(scheduler="single-threaded")
     assert type(result) == pd.DataFrame