Repository: dask/distributed

Issue:
Culling in HashJoinP2PLayer can blow up memory
<!-- 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**:

``_cull_dependencies`` in ``HashJoinP2PLayer`` creates a defaultdict that is huge as the number of partitions grows. we have 10476 parts out and 10476 partitions, which will create relatively big sets.  The memory footprint of deps is over 10GB

**Minimal Complete Verifiable Example**:

```python
def read_data(filename):

    path = "s3://coiled-runtime-ci/tpch-scale-1000/" + filename + "/"
    return dd.read_parquet(path, engine="pyarrow")


if __name__ == "__main__":
    cluster = LocalCluster()
    client = cluster.get_client()

    lineitem_ds = read_data("lineitem")
    orders_ds = read_data("orders")
    cutomer_ds = read_data("customer")

    jn1 = cutomer_ds.merge(orders_ds, left_on="c_custkey", right_on="o_custkey")
    jn2 = jn1.merge(lineitem_ds, left_on="o_orderkey", right_on="l_orderkey")
    jn2["revenue"] = jn2.l_extendedprice * (1 - jn2.l_discount)
    total = jn2.groupby(["l_orderkey", "o_orderdate", "o_shippriority"])[
        "revenue"
    ].sum()
    total.reset_index().head(10)[
        ["l_orderkey", "revenue", "o_orderdate", "o_shippriority"]
    ]
```

sorry for the non-ideal reproducer

**Anything else we need to know?**:

Can we reduce the memory footprint here somehow?

cc @fjetter 

**Environment**:

- Dask version: main
- Python version: 3.11
- Operating System: mac
- Install method (conda, pip, source):



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
@@ -1,7 +1,6 @@
 # mypy: ignore-errors
 from __future__ import annotations
 
-from collections import defaultdict
 from collections.abc import Iterable, Sequence
 from typing import TYPE_CHECKING, Any
 
@@ -243,15 +242,14 @@ def _cull_dependencies(
         all input partitions. This method does not require graph
         materialization.
         """
-        deps = defaultdict(set)
+        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)}
+        # Protect against mutations later on with frozenset
+        keys = frozenset(keys)
         for part in parts_out:
-            deps[(self.name, part)] |= {
-                (self.name_input_left, i) for i in range(self.npartitions)
-            }
-            deps[(self.name, part)] |= {
-                (self.name_input_right, i) for i in range(self.npartitions)
-            }
+            deps[(self.name, part)] = keys
         return deps
 
     def _keys_to_parts(self, keys: Iterable[str]) -> set[str]:
diff --git a/distributed/shuffle/_shuffle.py b/distributed/shuffle/_shuffle.py
--- a/distributed/shuffle/_shuffle.py
+++ b/distributed/shuffle/_shuffle.py
@@ -227,8 +227,11 @@ def cull(
         parameter.
         """
         parts_out = self._keys_to_parts(keys)
-        input_parts = {(self.name_input, i) for i in range(self.npartitions_input)}
-        culled_deps = {(self.name, part): input_parts.copy() for part in parts_out}
+        # Protect against mutations later on with frozenset
+        input_parts = frozenset(
+            {(self.name_input, i) for i in range(self.npartitions_input)}
+        )
+        culled_deps = {(self.name, part): input_parts for part in parts_out}
 
         if parts_out != set(self.parts_out):
             culled_layer = self._cull(parts_out)


Test patch: