Repository: dask/dask
Pull request: Use blockwise in `single_partition_join`
~Blocked by https://github.com/dask/dask/issues/8338~

- [ ] Closes #8306
- [x] Tests added / passed
- [x] Passes `pre-commit run --all-files`


Problem statement:
Use Blockwise/`map_partitions` in various DataFrame join methods
I noticed that some join methods have things like
```python
        dsk = {
            (name, i): (apply, merge_chunk, [left_key, right_key], kwargs)
            for i, right_key in enumerate(right.__dask_keys__())
        }
```
where we're generating a low-level graph that could just be done with `map_partitions`. Using `map_partitions` in these scenarios would both speed up graph transmission and allow for blockwise fusion across the operations. Refactoring this simple sorts of graphs should be straightforward.

- [x] `single_partition_join`
- [ ] `hash_join`'s `merge_chunk`
- [ ] `stack_partitions` should use `HighLevelGraph.from_collections` instead of merging all of the input graphs

cc @rjzamora @ncclementi @jrbourbeau

Hints from issue discussion:
Thanks for flagging these @gjoseph92. In general, lifting low-level task graphs into high-level graphs is a nice improvement 
I just realized this is slightly more important than I'd originally thought, because now that low-level optimization is turned off for DataFrames, the only way we get task fusion is through HighLevelGraphs. So even simple linear chains won't be fused, exposing us to root task overproduction (https://github.com/dask/distributed/issues/5223).

For example, this means that a `single_partition_join` followed by a `map_partitions` operation may have worse memory performance than doing the join yourself within a `map_partitions`, since lots of extra `single_partition_join` outputs can accumulate in memory.

cc @jrbourbeau @ncclementi 
Ah, thanks for digging deeper here @gjoseph92

Commit message: Add utils for getting HLG layers

Patch to review:
diff --git a/dask/dataframe/multi.py b/dask/dataframe/multi.py
index 3f7f744ec84..ab54ffb91de 100644
--- a/dask/dataframe/multi.py
+++ b/dask/dataframe/multi.py
@@ -61,7 +61,7 @@
 import numpy as np
 import pandas as pd
 from pandas.api.types import is_categorical_dtype, is_dtype_equal
-from tlz import first, merge_sorted, unique
+from tlz import merge_sorted, unique
 
 from ..base import is_dask_collection, tokenize
 from ..highlevelgraph import HighLevelGraph
@@ -412,42 +412,50 @@ def single_partition_join(left, right, **kwargs):
     kwargs["empty_index_dtype"] = meta.index.dtype
     kwargs["categorical_columns"] = meta.select_dtypes(include="category").columns
 
-    name = "merge-" + tokenize(left, right, **kwargs)
-
     if right.npartitions == 1 and kwargs["how"] in allowed_left:
-        right_key = first(right.__dask_keys__())
-        dsk = {
-            (name, i): (apply, merge_chunk, [left_key, right_key], kwargs)
-            for i, left_key in enumerate(left.__dask_keys__())
-        }
+        joined = left.map_partitions(
+            merge_chunk,
+            right,
+            meta=meta,
+            enforce_metadata=False,
+            transform_divisions=False,
+            align_dataframes=False,
+            token="merge",  # NOTE: misleadingly, this is actually the name
+            **kwargs,
+        )
+
         if use_left:
-            divisions = left.divisions
+            joined.divisions = left.divisions
         elif use_right and len(right.divisions) == len(left.divisions):
-            divisions = right.divisions
+            joined.divisions = right.divisions
         else:
-            divisions = [None for _ in left.divisions]
+            joined.divisions = [None for _ in left.divisions]
 
     elif left.npartitions == 1 and kwargs["how"] in allowed_right:
-        left_key = first(left.__dask_keys__())
-        dsk = {
-            (name, i): (apply, merge_chunk, [left_key, right_key], kwargs)
-            for i, right_key in enumerate(right.__dask_keys__())
-        }
+        joined = right.map_partitions(
+            lambda right, left, **kwargs: merge_chunk(left, right, **kwargs),
+            left,
+            meta=meta,
+            enforce_metadata=False,
+            transform_divisions=False,
+            align_dataframes=False,
+            token="merge",  # NOTE: misleadingly, this is actually the name
+            **kwargs,
+        )
 
         if use_right:
-            divisions = right.divisions
+            joined.divisions = right.divisions
         elif use_left and len(left.divisions) == len(right.divisions):
-            divisions = left.divisions
+            joined.divisions = left.divisions
         else:
-            divisions = [None for _ in right.divisions]
+            joined.divisions = [None for _ in right.divisions]
 
     else:
         raise NotImplementedError(
             "single_partition_join has no fallback for invalid calls"
         )
 
-    graph = HighLevelGraph.from_collections(name, dsk, dependencies=[left, right])
-    return new_dd_object(graph, name, meta, divisions)
+    return joined
 
 
 def warn_dtype_mismatch(left, right, left_on, right_on):
diff --git a/dask/dataframe/tests/test_multi.py b/dask/dataframe/tests/test_multi.py
index c6aeb0f934b..fd884868b79 100644
--- a/dask/dataframe/tests/test_multi.py
+++ b/dask/dataframe/tests/test_multi.py
@@ -23,6 +23,7 @@
     has_known_categories,
     make_meta,
 )
+from dask.utils_test import hlg_layer_topological
 
 
 def test_align_partitions():
@@ -1372,6 +1373,7 @@ def test_cheap_inner_merge_with_pandas_object():
     b = pd.DataFrame({"x": [1, 2, 3, 4], "z": list("abda")})
 
     dc = da.merge(b, on="x", how="inner")
+    assert not hlg_layer_topological(dc.dask, -1).is_materialized()
     assert all("shuffle" not in k[0] for k in dc.dask)
 
     list_eq(da.merge(b, on="x", how="inner"), a.merge(b, on="x", how="inner"))
@@ -1387,6 +1389,7 @@ def test_cheap_single_partition_merge():
     bb = dd.from_pandas(b, npartitions=1, sort=False)
 
     cc = aa.merge(bb, on="x", how="inner")
+    assert not hlg_layer_topological(cc.dask, -1).is_materialized()
     assert all("shuffle" not in k[0] for k in cc.dask)
     assert len(cc.dask) == len(aa.dask) * 2 + len(bb.dask)
 
@@ -1403,6 +1406,8 @@ def test_cheap_single_partition_merge_divisions():
     bb = dd.from_pandas(b, npartitions=1, sort=False)
 
     actual = aa.merge(bb, on="x", how="inner")
+    assert not hlg_layer_topological(actual.dask, -1).is_materialized()
+
     assert not actual.known_divisions
     assert_divisions(actual)
 
@@ -1422,11 +1427,13 @@ def test_cheap_single_parition_merge_left_right(how):
     actual = aa.merge(bb, left_index=True, right_on="x", how=how)
     expected = a.merge(b, left_index=True, right_on="x", how=how)
 
+    assert not hlg_layer_topological(actual.dask, -1).is_materialized()
     assert_eq(actual, expected)
 
     actual = aa.merge(bb, left_on="x", right_index=True, how=how)
     expected = a.merge(b, left_on="x", right_index=True, how=how)
 
+    assert not hlg_layer_topological(actual.dask, -1).is_materialized()
     assert_eq(actual, expected)
 
 
@@ -1447,6 +1454,7 @@ def test_cheap_single_partition_merge_on_index():
     # for empty joins.
     expected.index = expected.index.astype("int64")
 
+    assert not hlg_layer_topological(actual.dask, -1).is_materialized()
     assert not actual.known_divisions
     assert_eq(actual, expected)
 
@@ -1454,6 +1462,7 @@ def test_cheap_single_partition_merge_on_index():
     expected = b.merge(a, right_index=True, left_on="x", how="inner")
     expected.index = expected.index.astype("int64")
 
+    assert not hlg_layer_topological(actual.dask, -1).is_materialized()
     assert not actual.known_divisions
     assert_eq(actual, expected)
 
diff --git a/dask/tests/test_utils_test.py b/dask/tests/test_utils_test.py
new file mode 100644
index 00000000000..1738b3adde2
--- /dev/null
+++ b/dask/tests/test_utils_test.py
@@ -0,0 +1,31 @@
+import pytest
+
+from dask import utils_test
+from dask.highlevelgraph import HighLevelGraph
+
+
+def test_hlg_layer():
+    a = {"x": 1}
+    b = {"y": (utils_test.inc, "x")}
+    layers = {"a-layer": a, "bee-layer": b}
+    dependencies = {"a-layer": set(), "bee-layer": {"a-layer"}}
+    hg = HighLevelGraph(layers, dependencies)
+
+    assert utils_test.hlg_layer(hg, "a") is hg.layers["a-layer"]
+    assert utils_test.hlg_layer(hg, "b") is hg.layers["bee-layer"]
+    with pytest.raises(KeyError, match="No layer starts with"):
+        utils_test.hlg_layer(hg, "foo")
+
+
+def test_hlg_layer_topological():
+    a = {"x": 1}
+    b = {"y": (utils_test.inc, "x")}
+    c = {"z": (utils_test.inc, "x")}
+    d = {"r": (sum, ["y", "z"])}
+    layers = {"a": a, "b": b, "c": c, "d": d}
+    dependencies = {"a": set(), "b": {"a"}, "c": {"a"}, "d": {"b", "c"}}
+    hg = HighLevelGraph(layers, dependencies)
+
+    assert utils_test.hlg_layer_topological(hg, -1) is hg.layers["d"]
+    assert utils_test.hlg_layer_topological(hg, 0) is hg.layers["a"]
+    assert utils_test.hlg_layer_topological(hg, 1) in (hg.layers["b"], hg.layers["c"])
diff --git a/dask/utils_test.py b/dask/utils_test.py
index 27949a52275..73724c9012c 100644
--- a/dask/utils_test.py
+++ b/dask/utils_test.py
@@ -1,4 +1,10 @@
+from __future__ import annotations
+
 import importlib
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .highlevelgraph import HighLevelGraph, Layer
 
 
 def inc(x):
@@ -123,3 +129,16 @@ def import_or_none(name):
         return importlib.import_module(name)
     except (ImportError, AttributeError):
         return None
+
+
+def hlg_layer(hlg: HighLevelGraph, prefix: str) -> Layer:
+    "Get the first layer from a HighLevelGraph whose name starts with a prefix"
+    for key, lyr in hlg.layers.items():
+        if key.startswith(prefix):
+            return lyr
+    raise KeyError(f"No layer starts with {prefix!r}: {list(hlg.layers)}")
+
+
+def hlg_layer_topological(hlg: HighLevelGraph, i: int) -> Layer:
+    "Get the layer from a HighLevelGraph at position ``i``, topologically"
+    return hlg.layers[hlg._toposort_layers()[i]]


Code review comments:
diff_hunk: @@ -412,42 +412,50 @@ def single_partition_join(left, right, **kwargs):
     kwargs["empty_index_dtype"] = meta.index.dtype
     kwargs["categorical_columns"] = meta.select_dtypes(include="category").columns
 
-    name = "merge-" + tokenize(left, right, **kwargs)
-
     if right.npartitions == 1 and kwargs["how"] in allowed_left:
-        right_key = first(right.__dask_keys__())
-        dsk = {
-            (name, i): (apply, merge_chunk, [left_key, right_key], kwargs)
-            for i, left_key in enumerate(left.__dask_keys__())
-        }
+        joined = left.map_partitions(
+            merge_chunk,
+            right,
+            meta=meta,
+            enforce_metadata=False,
+            transform_divisions=False,
+            align_dataframes=False,
+            token="merge",  # NOTE: misleadingly, this is actually the name
+            **kwargs,
+        )
+
         if use_left:
-            divisions = left.divisions
+            joined.divisions = left.divisions
         elif use_right and len(right.divisions) == len(left.divisions):
-            divisions = right.divisions
+            joined.divisions = right.divisions
         else:
-            divisions = [None for _ in left.divisions]
+            joined.divisions = [None for _ in left.divisions]
 
     elif left.npartitions == 1 and kwargs["how"] in allowed_right:
-        left_key = first(left.__dask_keys__())
-        dsk = {
-            (name, i): (apply, merge_chunk, [left_key, right_key], kwargs)
-            for i, right_key in enumerate(right.__dask_keys__())
-        }
+        joined = right.map_partitions(
+            lambda right, left, **kwargs: merge_chunk(left, right, **kwargs),
+            left,
path: dask/dataframe/multi.py
text: @user1:
I feel like I've heard that we try to avoid inlining lambda functions in the code base, but I can't remember why. It seems like you could just use `dd.map_partitions` and achieve the same outcome thought right?

```suggestion
        joined = map_partitions(
            merge_chunk,
            left,
            right,
```

@author:
Great point, thanks! Despite reading the top-level function too many times I forgot you could actually just use it 😁