Add GroupBy.shuffle_to_chunks()

[dcherian]

This adds some new API to shuffle an Xarray object. Shuffling means we sort so that members of a group occur in the same chunk, with the possibility of multiple groups in a single chunk.

gb = ds.groupby(..)
shuffled : Dataset = gb.distributed_shuffle()
shuffled.groupby(...).quantile()

I've also added shuffle_by to DataArray and Dataset. This generalizes sortby, and lets you persist a shuffled Xarray object to disk.

ds.shuffle_by(grouper)

• xref Groupby-map is slow with out of order indices #9220, xref Using the shuffle primitive in Xarray #9546
• Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst
• New functions/methods are listed in api.rst
• wait on dask 2024.08.1
• upstream flox release
• add chunks to signature.
• error if group is chunked

cc @phofl

[dcherian]

We could also just have GroupBy.shuffle() -> Dataset | DataArray.

Then

# Save the groupby kwargs since we need it twice
groupers = {"label": UniqueGrouper()}

# get the shuffled dataset
shuffled: Dataset = ds.groupby(groupers).shuffle()

# Now group the shuffled dataset the same way and 
# map a UDF that expects all members of a group in memory at the same time.
shuffled.groupby(groupers).map(UDF)

Alternative names could be GroupBy.as_sorted() -> Dataset | DataArray, GroupBy.sorted_object() -> Dataset | DataArray.

---

After thinking about it a while, sortby is similar but different w.r.t chunking. It's vectorized indexing with an index you get from calling argsort, and dask/cubed is free to rechunk as they please. For example sortby(dim=UniqueGrouper()) when all(chunksizes["dim"] == 1) is basically free in terms of rechunking, but with a shuffle, you'd get the same labels in the same chunk so that chunking would have to change (assuming duplicated dim labels). We could have sortby a grouper object basically provide np.concatenate(_group_indices) as the indexer.

This is all fairly advanced stuff, so we can start with the smallest possible API surface, which would be GroupBy.shuffle IMO (or any of the other names above)

cc @shoyer

[max-sixty]

Very much agree with all of your message — sortby being a bit different, starting with GroupBy.shuffle

We could also just have GroupBy.shuffle() -> Dataset | DataArray.

(If it were GroupBy.shuffle() -> GroupBy, then I guess .shuffle_by would be equivalent to .groupby(groupers).shuffle().map(lambda x: x)...)

[shoyer]

Would it be too verbose to call this distributed_shuffle? Shuffling just has a very specific meaningful in distributed computing that is quite different from the generic meaning of "shuffle" for rearranging data.

[max-sixty]

Would it be too verbose to call this distributed_shuffle? Shuffling just has a very specific meaningful in distributed computing that is quite different from the generic meaning of "shuffle" for rearranging data.

For the method is on the GroupBy object, is shuffle less ambiguous? FWIW I had correctly presumed it meant its actual meaning given it was associated with group by

Otherwise we could workshop distributed_shuffle to something terser :)

[dcherian]

OK I went with distributed_shuffle returning a shuffled Xarray object. Any further groupby operations will need to repeat the groupby statement:
e.g.

gb = array.groupby_bins("dim_0", bins=bins, **cut_kwargs)

shuffled: DataArray = gb.distributed_shuffle()

gb = shuffle.groupby_bins("dim_0", bins=bins, **cut_kwargs)

I couldn't think of a shorter name, but happy to change it.

[dcherian]

how about shuffle_to_chunks?

[max-sixty]

Looks good! (No strong view on the exact name...)

[dcherian]

OK I'm going with shuffle_to_chunks. Will merge in ~24 hours. It's been around a while now :)

[shoyer]

Sounds good!

…

On Wed, Nov 20, 2024 at 3:58 PM Deepak Cherian ***@***.***> wrote: OK I'm going with shuffle_to_chunks. Will merge in ~12 hours. It's been around a while now :) — Reply to this email directly, view it on GitHub <#9320 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJJFVXAWGTXX62ULVGF6L32BUOZVAVCNFSM6AAAAABMDO3MWGVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDIOBZG44DEMBZGY> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[max-sixty]

Really like the new name!

[dcherian]

Thanks for the input!