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

[max-sixty]

Is a single group limited to a single chunk? Assuming so, if we get one giant chuck, could that present any performance problems?

[dcherian]

The other chunks get "auto" reshaped. This is controlled by the chunks kwarg, which only takes "auto" at the moment.

https://docs.dask.org/en/latest/generated/dask.array.shuffle.html

[max-sixty]

I'm sure I'm being slow here — but doesn't .shuffle claim to sort the array?

Or this means "when resampling with a dimension other than the chunking dimension"?

[dcherian]

Yes this is confusing.

Resampling only works with sorted axes, so shuffling is basically identical to rechunking. Maybe I should just delete this line?

[max-sixty]

ATM shuffle_by returns the shuffled Dataset or DataArray. groupby.shuffle() returns a new Groupby object that uses the shuffled data.

I am not very attached to this. FYO Stephan mentioned sortby as a possible alternate entrypoint instead.

Yes, this makes sense.

Is there a case for automatically running the equivalent of .shuffle_by on a chunked array when we run .sortby? And then we wouldn't need .shuffle_by? Or there are enough times when changing the chunk sizes would be confusing that it's not worthwhile?

(For context: I'm keen to keep the interface of our main objects really small & orthogonal. This seems useful enough that it is worth adding a method for, but wanted to see if there's any way to get the benefits without even adding a method...)

[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!