Repository: dask/distributed

Issue:
hash join merge with P2P fails if npartitions is given
<!-- 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**:

Specifying npartitions breaks  the merge

**Minimal Complete Verifiable Example**:

```python
    from distributed import LocalCluster
    cluster = LocalCluster()
    client = cluster.get_client()
    
    df = pd.DataFrame({"a": [1, 2, 3] * 100, "b": 1})

    ddf = dd.from_pandas(df, npartitions=10)
    ddf2 = dd.from_pandas(df, npartitions=5)

    ddf.merge(ddf2, npartitions=4).compute()
```

**Anything else we need to know?**:

**Environment**:

- Dask version: 2023.09.3
- Python version: 3.11
- Operating System: Mac
- Install method (conda, pip, source): condo



Discussion:

Resolving patch:
diff --git a/distributed/shuffle/_merge.py b/distributed/shuffle/_merge.py
--- a/distributed/shuffle/_merge.py
+++ b/distributed/shuffle/_merge.py
@@ -271,8 +271,8 @@ def _cull_dependencies(
         """
         deps = {}
         parts_out = parts_out or self._keys_to_parts(keys)
-        keys = {(self.name_input_left, i) for i in range(self.npartitions)}
-        keys |= {(self.name_input_right, i) for i in range(self.npartitions)}
+        keys = {(self.name_input_left, i) for i in range(self.n_partitions_left)}
+        keys |= {(self.name_input_right, i) for i in range(self.n_partitions_right)}
         # Protect against mutations later on with frozenset
         keys = frozenset(keys)
         for part in parts_out:
@@ -352,6 +352,7 @@ def cull(self, keys: Iterable[str], all_keys: Any) -> tuple[HashJoinP2PLayer, di
         parameter.
         """
         parts_out = self._keys_to_parts(keys)
+
         culled_deps = self._cull_dependencies(keys, parts_out=parts_out)
         if parts_out != set(self.parts_out):
             culled_layer = self._cull(parts_out)


Test patch:
diff --git a/distributed/shuffle/tests/test_merge.py b/distributed/shuffle/tests/test_merge.py
--- a/distributed/shuffle/tests/test_merge.py
+++ b/distributed/shuffle/tests/test_merge.py
@@ -478,6 +478,20 @@ async def test_index_merge_p2p(c, s, a, b, how):
         )
 
 
+@pytest.mark.parametrize("npartitions", [4, 5, 10, 20])
+@gen_cluster(client=True)
+async def test_merge_with_npartitions(c, s, a, b, npartitions):
+    pdf = pd.DataFrame({"a": [1, 2, 3, 4] * 10, "b": 1})
+
+    left = dd.from_pandas(pdf, npartitions=10)
+    right = dd.from_pandas(pdf, npartitions=5)
+
+    expected = pdf.merge(pdf)
+    with dask.config.set({"dataframe.shuffle.method": "p2p"}):
+        result = await c.compute(left.merge(right, npartitions=npartitions))
+    assert_eq(result, expected, check_index=False)
+
+
 class LimitedGetOrCreateShuffleRunManager(_ShuffleRunManager):
     seen: set[ShuffleId]
     block_get_or_create: asyncio.Event