import dask.dataframe as dd
from dask.distributed import Client, LocalCluster, performance_report

cluster = LocalCluster(local_directory="/tmp")
client = Client(cluster)

df = dd.demo.make_timeseries(
    start="2000-01-01",
    end="2000-01-04",
    dtypes={"x": float, "y": float, "id": int},
    freq="10ms",
    partition_freq="24h",
)

# key of ["x", "y"] gives ~ 25,920,000 groups

# Doesn't finish
with performance_report(filename="ddf-too-many-groups.html"):
    df.groupby(["x", "y"]).sum().compute()

# Finishes
with performance_report(filename="ddf-split-out.html"):
    df.groupby(["x", "y"]).sum(split_out=4).compute()

KilledWorker                              Traceback (most recent call last)
/var/folders/rf/26llfhwd68x7cftb1z3h000w0000gp/T/ipykernel_39651/3854007570.py in <module>
      1 with performance_report(filename=""):
----> 2     df.groupby(["x", "y"]).sum().compute()

~/miniconda3/envs/main/lib/python3.9/site-packages/dask/base.py in compute(self, **kwargs)
    284         dask.base.compute
    285         """
--> 286         (result,) = compute(self, traverse=False, **kwargs)
    287         return result
    288 

~/miniconda3/envs/main/lib/python3.9/site-packages/dask/base.py in compute(*args, **kwargs)
    566         postcomputes.append(x.__dask_postcompute__())
    567 
--> 568     results = schedule(dsk, keys, **kwargs)
    569     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    570 

~/miniconda3/envs/main/lib/python3.9/site-packages/distributed/client.py in get(self, dsk, keys, workers, allow_other_workers, resources, sync, asynchronous, direct, retries, priority, fifo_timeout, actors, **kwargs)
   2746                     should_rejoin = False
   2747             try:
-> 2748                 results = self.gather(packed, asynchronous=asynchronous, direct=direct)
   2749             finally:
   2750                 for f in futures.values():

~/miniconda3/envs/main/lib/python3.9/site-packages/distributed/client.py in gather(self, futures, errors, direct, asynchronous)
   2023             else:
   2024                 local_worker = None
-> 2025             return self.sync(
   2026                 self._gather,
   2027                 futures,

~/miniconda3/envs/main/lib/python3.9/site-packages/distributed/client.py in sync(self, func, asynchronous, callback_timeout, *args, **kwargs)
    864             return future
    865         else:
--> 866             return sync(
    867                 self.loop, func, *args, callback_timeout=callback_timeout, **kwargs
    868             )

~/miniconda3/envs/main/lib/python3.9/site-packages/distributed/utils.py in sync(loop, func, callback_timeout, *args, **kwargs)
    324     if error[0]:
    325         typ, exc, tb = error[0]
--> 326         raise exc.with_traceback(tb)
    327     else:
    328         return result[0]

~/miniconda3/envs/main/lib/python3.9/site-packages/distributed/utils.py in f()
    307             if callback_timeout is not None:
    308                 future = asyncio.wait_for(future, callback_timeout)
--> 309             result[0] = yield future
    310         except Exception:
    311             error[0] = sys.exc_info()

~/miniconda3/envs/main/lib/python3.9/site-packages/tornado/gen.py in run(self)
    760 
    761                     try:
--> 762                         value = future.result()
    763                     except Exception:
    764                         exc_info = sys.exc_info()

~/miniconda3/envs/main/lib/python3.9/site-packages/distributed/client.py in _gather(self, futures, errors, direct, local_worker)
   1888                             exc = CancelledError(key)
   1889                         else:
-> 1890                             raise exception.with_traceback(traceback)
   1891                         raise exc
   1892                     if errors == "skip":

KilledWorker: ("('dataframe-groupby-sum-agg-f6215c9bc267c471d79fe7d859ac9f04', 0)", <WorkerState 'tcp://127.0.0.1:60293', name: 0, memory: 0, processing: 1>)