Class DiskShuffle

dask/dataframe/dask_expr/_shuffle.py:499–538

Disk-based shuffle implementation

class DiskShuffle(SimpleShuffle):
    """Disk-based shuffle implementation"""

    @staticmethod
    def _shuffle_group(df, col, _filter, p):
        with ensure_cleanup_on_exception(p):
            g = df.groupby(col)
            d = {i: g.get_group(i) for i in g.groups if i in _filter}
            p.append(d, fsync=True)

    def _layer(self):
        from dask.dataframe.dispatch import partd_encode_dispatch

        column = self.partitioning_index
        df = self.frame

        always_new_token = uuid.uuid1().hex

        p = (f"zpartd-{always_new_token}",)
        encode_cls = partd_encode_dispatch(df._meta)
        dsk1 = {p: (maybe_buffered_partd(encode_cls=encode_cls),)}

        # Partition data on disk
        name = f"shuffle-partition-{always_new_token}"
        dsk2 = {
            (name, i): (self._shuffle_group, key, column, self._partitions, p)
            for i, key in enumerate(df.__dask_keys__())
        }

        # Barrier
        barrier_token = (f"barrier-{always_new_token}",)
        dsk3 = {barrier_token: (barrier, list(dsk2))}

        # Collect groups
        dsk4 = {
            (self._name, j): (collect, p, k, df._meta, barrier_token)
            for j, k in enumerate(self._partitions)
        }

        return toolz.merge(dsk1, dsk2, dsk3, dsk4)


def _shuffle_transfer(