Repository: dask/dask

Issue:
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 


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 

Resolving patch:
diff --git a/dask/dataframe/multi.py b/dask/dataframe/multi.py
--- 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,14 +412,7 @@ 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__())
-        }
         if use_left:
             divisions = left.divisions
         elif use_right and len(right.divisions) == len(left.divisions):
@@ -428,26 +421,29 @@ def single_partition_join(left, right, **kwargs):
             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__())
-        }
-
         if use_right:
             divisions = right.divisions
         elif use_left and len(left.divisions) == len(right.divisions):
             divisions = left.divisions
         else:
             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)
+    joined = map_partitions(
+        merge_chunk,
+        left,
+        right,
+        meta=meta,
+        enforce_metadata=False,
+        transform_divisions=False,
+        align_dataframes=False,
+        **kwargs,
+    )
+    joined.divisions = divisions
+    return joined
 
 
 def warn_dtype_mismatch(left, right, left_on, right_on):


Test patch:
diff --git a/dask/dataframe/tests/test_multi.py b/dask/dataframe/tests/test_multi.py
--- 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,12 +1373,14 @@ 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"))
 
 
-def test_cheap_single_partition_merge():
+@pytest.mark.parametrize("flip", [False, True])
+def test_cheap_single_partition_merge(flip):
     a = pd.DataFrame(
         {"x": [1, 2, 3, 4, 5, 6], "y": list("abdabd")}, index=[10, 20, 30, 40, 50, 60]
     )
@@ -1386,11 +1389,15 @@ def test_cheap_single_partition_merge():
     b = pd.DataFrame({"x": [1, 2, 3, 4], "z": list("abda")})
     bb = dd.from_pandas(b, npartitions=1, sort=False)
 
-    cc = aa.merge(bb, on="x", how="inner")
+    pd_inputs = (b, a) if flip else (a, b)
+    inputs = (bb, aa) if flip else (aa, bb)
+
+    cc = dd.merge(*inputs, 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)
 
-    list_eq(aa.merge(bb, on="x", how="inner"), a.merge(b, on="x", how="inner"))
+    list_eq(cc, pd.merge(*pd_inputs, on="x", how="inner"))
 
 
 def test_cheap_single_partition_merge_divisions():
@@ -1403,6 +1410,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)
 
@@ -1412,21 +1421,27 @@ def test_cheap_single_partition_merge_divisions():
 
 
 @pytest.mark.parametrize("how", ["left", "right"])
-def test_cheap_single_parition_merge_left_right(how):
+@pytest.mark.parametrize("flip", [False, True])
+def test_cheap_single_parition_merge_left_right(how, flip):
     a = pd.DataFrame({"x": range(8), "z": list("ababbdda")}, index=range(8))
     aa = dd.from_pandas(a, npartitions=1)
 
     b = pd.DataFrame({"x": [1, 2, 3, 4], "z": list("abda")}, index=range(4))
     bb = dd.from_pandas(b, npartitions=1)
 
-    actual = aa.merge(bb, left_index=True, right_on="x", how=how)
-    expected = a.merge(b, left_index=True, right_on="x", how=how)
+    pd_inputs = (b, a) if flip else (a, b)
+    inputs = (bb, aa) if flip else (aa, bb)
+
+    actual = dd.merge(*inputs, left_index=True, right_on="x", how=how)
+    expected = pd.merge(*pd_inputs, 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)
+    actual = dd.merge(*inputs, left_on="x", right_index=True, how=how)
+    expected = pd.merge(*pd_inputs, left_on="x", right_index=True, how=how)
 
+    assert not hlg_layer_topological(actual.dask, -1).is_materialized()
     assert_eq(actual, expected)
 
 
@@ -1447,6 +1462,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 +1470,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
--- /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
--- 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]]