[Data] Heavy data skew after groupby

[bveeramani]

What happened + What you expected to happen

I called map_groups after a groupby, and observed that most of the tasks finished immediately, but a few tasks continued to run for a very long time. The reason is that there's heavy data skew after the groupby (see repro).

Versions / Dependencies

0702f7c

Reproduction script

import ray

refs = (
    ray.data.read_parquet("data.parquet")
    .sort(["spam", "ham"])

    .get_internal_block_refs()
)
print([len(ray.get(ref)) for ref in refs])
# Blocks either contain ~500 or ~8500 rows
# [620, 8495, 647, 8623, 8570, 314, 8308, 8391, 443, 8579, 395, 8498, 8628, 378, 8522, 8510, 312, 8596, 355, 11]            

data.parquet.zip

Issue Severity

None