Repository: dask/dask

Issue:
Broadcast join fails when passing a list of columns to merge on
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**Describe the issue**:

The broadcast merge codepath for dataframes throws an error when passing a list of columns in the `on` argument with a `keyError`

**Minimal Complete Verifiable Example**:

```python
import pandas as pd
import numpy as np
from dask import dataframe as dd
from distributed import Client, wait
c = Client()

df1 = pd.DataFrame({"a":np.arange(20),"b":np.arange(20),"c":[1,2]*10})
df2 = df1.copy(deep=True)
df1 = dd.from_pandas(df1,2)
df2 = dd.from_pandas(df2,5)

len(df1.merge(df2,on=["a"], how="inner", shuffle="tasks", broadcast=True))
# KeyError: "('a',)"
# Works with on="a"
```

**Anything else we need to know?**:
Same issue persists when joining on multiple columns

**Environment**:

- Dask version: 2022.12.0
- Python version: 3.9
- Operating System: ubuntu18.04
- Install method (conda, pip, source): pip



Discussion:
I can confirm and can reproduce this. My best guess is that this is a serialization error or we're calling a stringify or smth too eagerly.

basically the list is cast to a tuple and then stringified somewhere, i.e. `["a"]` becomes `"('a',)"`. Passing a tuple directly or the literal works as expected.

If I had to guess, [this](https://github.com/dask/dask/blob/96a72df0aeb4a7d04844e1ebc2fddb30265a0bc7/dask/layers.py#L889-L892) looks suspicious but I don't fully understand what's going on there. Maybe @rjzamora ?
Thanks for raising @ayushdg ! I ran into this bug yesterday, but didn't get a chance to raise and issue and investigate yet. @fjetter is correct that this is likely another HLG-serialization edge case (probably having to do with msgpack not distinguishing lists/tuples). I will try to figure out a fix, but also look forward to something like https://github.com/dask/distributed/pull/6028 avoiding these problems altogether :)

Resolving patch:
diff --git a/dask/layers.py b/dask/layers.py
--- a/dask/layers.py
+++ b/dask/layers.py
@@ -872,6 +872,8 @@ def __init__(
         rhs_npartitions,
         parts_out=None,
         annotations=None,
+        left_on=None,
+        right_on=None,
         **merge_kwargs,
     ):
         super().__init__(annotations=annotations)
@@ -882,14 +884,12 @@ def __init__(
         self.rhs_name = rhs_name
         self.rhs_npartitions = rhs_npartitions
         self.parts_out = parts_out or set(range(self.npartitions))
+        self.left_on = tuple(left_on) if isinstance(left_on, list) else left_on
+        self.right_on = tuple(right_on) if isinstance(right_on, list) else right_on
         self.merge_kwargs = merge_kwargs
         self.how = self.merge_kwargs.get("how")
-        self.left_on = self.merge_kwargs.get("left_on")
-        self.right_on = self.merge_kwargs.get("right_on")
-        if isinstance(self.left_on, list):
-            self.left_on = (list, tuple(self.left_on))
-        if isinstance(self.right_on, list):
-            self.right_on = (list, tuple(self.right_on))
+        self.merge_kwargs["left_on"] = self.left_on
+        self.merge_kwargs["right_on"] = self.right_on
 
     def get_output_keys(self):
         return {(self.name, part) for part in self.parts_out}


Test patch:
diff --git a/dask/tests/test_distributed.py b/dask/tests/test_distributed.py
--- a/dask/tests/test_distributed.py
+++ b/dask/tests/test_distributed.py
@@ -148,6 +148,27 @@ def test_fused_blockwise_dataframe_merge(c, fuse):
     )
 
 
+@pytest.mark.parametrize("on", ["a", ["a"]])
+@pytest.mark.parametrize("broadcast", [True, False])
+def test_dataframe_broadcast_merge(c, on, broadcast):
+    # See: https://github.com/dask/dask/issues/9870
+    pd = pytest.importorskip("pandas")
+    dd = pytest.importorskip("dask.dataframe")
+
+    pdfl = pd.DataFrame({"a": [1, 2] * 2, "b_left": range(4)})
+    pdfr = pd.DataFrame({"a": [2, 1], "b_right": range(2)})
+    dfl = dd.from_pandas(pdfl, npartitions=4)
+    dfr = dd.from_pandas(pdfr, npartitions=2)
+
+    ddfm = dd.merge(dfl, dfr, on=on, broadcast=broadcast, shuffle="tasks")
+    dfm = ddfm.compute()
+    dd.utils.assert_eq(
+        dfm.sort_values("a"),
+        pd.merge(pdfl, pdfr, on=on).sort_values("a"),
+        check_index=False,
+    )
+
+
 @pytest.mark.parametrize(
     "computation",
     [