Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add GroupBy.shuffle_to_chunks() #9320

Merged
merged 61 commits into from
Nov 21, 2024
Merged
Changes from 1 commit
Commits
Show all changes
61 commits
Select commit Hold shift + click to select a range
3bc51bd
Add GroupBy.shuffle()
dcherian Aug 7, 2024
60d7619
Cleanup
dcherian Aug 7, 2024
d1429cd
Cleanup
dcherian Aug 7, 2024
31fc00e
fix
dcherian Aug 7, 2024
4583853
return groupby instance from shuffle
dcherian Aug 13, 2024
abd9dd2
Fix nD by
dcherian Aug 13, 2024
6b820aa
Merge branch 'main' into groupby-shuffle
dcherian Aug 14, 2024
0d70656
Skip if no dask
dcherian Aug 14, 2024
fafb937
fix tests
dcherian Aug 14, 2024
939db9a
Merge branch 'main' into groupby-shuffle
dcherian Aug 14, 2024
a08450e
Add `chunks` to signature
dcherian Aug 14, 2024
d0cd218
FIx self
dcherian Aug 14, 2024
4edc976
Another Self fix
dcherian Aug 14, 2024
0b42be4
Forward chunks too
dcherian Aug 14, 2024
c52734d
[revert]
dcherian Aug 14, 2024
8180625
undo flox limit
dcherian Aug 14, 2024
7897c91
[revert]
dcherian Aug 14, 2024
7773548
fix types
dcherian Aug 14, 2024
51a7723
Add DataArray.shuffle_by, Dataset.shuffle_by
dcherian Aug 15, 2024
cc95513
Add doctest
dcherian Aug 15, 2024
18f4a40
Refactor
dcherian Aug 15, 2024
f489bcf
tweak docstrings
dcherian Aug 15, 2024
ead1bb4
fix typing
dcherian Aug 15, 2024
75115d0
Fix
dcherian Aug 15, 2024
390863a
fix docstring
dcherian Aug 15, 2024
a408cb0
bump min version to dask>=2024.08.1
dcherian Aug 17, 2024
7038f37
Merge branch 'main' into groupby-shuffle
dcherian Aug 17, 2024
05a0fb4
Fix typing
dcherian Aug 17, 2024
b8e7f62
Fix types
dcherian Aug 17, 2024
6d9ed1c
Merge branch 'main' into groupby-shuffle
dcherian Aug 22, 2024
20a8cd9
Merge branch 'main' into groupby-shuffle
dcherian Aug 30, 2024
7a99c8f
remove shuffle_by for now.
dcherian Aug 30, 2024
5e2fdfb
Add tests
dcherian Aug 30, 2024
a22c7ed
Support shuffling with multiple groupers
dcherian Aug 30, 2024
2d48690
Revert "remove shuffle_by for now."
dcherian Sep 11, 2024
0679d2b
Merge branch 'main' into groupby-shuffle
dcherian Sep 12, 2024
63b3e77
Merge branch 'main' into groupby-shuffle
dcherian Sep 17, 2024
7dc5dd1
bad merge
dcherian Sep 17, 2024
bad0744
Merge branch 'main' into groupby-shuffle
dcherian Sep 18, 2024
91e4bd8
Add a test
dcherian Sep 18, 2024
0542944
Merge branch 'main' into groupby-shuffle
dcherian Nov 1, 2024
1e4f805
Add docs
dcherian Nov 1, 2024
ad502aa
bugfix
dcherian Nov 1, 2024
4b0c143
Refactor out Dataset._shuffle
dcherian Nov 2, 2024
2b2c4ab
Merge branch 'main' into groupby-shuffle
dcherian Nov 3, 2024
f624c8f
fix types
dcherian Nov 3, 2024
888e780
fix tests
dcherian Nov 3, 2024
47e5c17
Merge branch 'main' into groupby-shuffle
dcherian Nov 4, 2024
b100fb1
Handle by is chunked
dcherian Nov 4, 2024
978fad9
Merge branch 'main' into groupby-shuffle
dcherian Nov 7, 2024
d1a3fc1
Some refactoring
dcherian Nov 7, 2024
23b0cac
Merge branch 'main' into groupby-shuffle
dcherian Nov 12, 2024
d533638
Merge branch 'main' into groupby-shuffle
dcherian Nov 19, 2024
d467bc6
Remove shuffle_by
dcherian Nov 19, 2024
231533c
shuffle -> distributed_shuffle
dcherian Nov 19, 2024
c77d7c5
return xarray object from distributed_shuffle
dcherian Nov 19, 2024
bccacfe
fix
dcherian Nov 19, 2024
2d4392a
fix doctest
dcherian Nov 19, 2024
003e9f2
fix api
dcherian Nov 19, 2024
0f80c81
Rename to `shuffle_to_chunks`
dcherian Nov 20, 2024
88bef5d
update docs
dcherian Nov 21, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Remove shuffle_by
dcherian committed Nov 19, 2024
commit d467bc60c84dd3976fc6e2a94dd41515ad3289a1
64 changes: 1 addition & 63 deletions xarray/core/common.py
Original file line number Diff line number Diff line change
@@ -52,7 +52,7 @@
T_Variable,
)
from xarray.core.variable import Variable
from xarray.groupers import Grouper, Resampler
from xarray.groupers import Resampler

DTypeMaybeMapping = Union[DTypeLikeSave, Mapping[Any, DTypeLikeSave]]

@@ -891,68 +891,6 @@ def rolling_exp(

return rolling_exp.RollingExp(self, window, window_type)

def shuffle_by(
self,
group: Hashable | DataArray | Mapping[Any, Grouper] | None = None,
chunks: T_Chunks = None,
**groupers: Grouper,
) -> Self:
"""
Sort or "shuffle" this object by a Grouper.
"Shuffle" means the object is sorted so that all group members occur sequentially,
in the same chunk. Multiple groups may occur in the same chunk.
This method is particularly useful for chunked arrays (e.g. dask, cubed).
For chunked array types, the order of appearance is not guaranteed, but will depend on
the input chunking.
Parameters
----------
group : Hashable or DataArray or IndexVariable or mapping of Hashable to Grouper
Array whose unique values should be used to group this array. If a
Hashable, must be the name of a coordinate contained in this dataarray. If a dictionary,
must map an existing variable name to a :py:class:`Grouper` instance.
chunks : int, tuple of int, "auto" or mapping of hashable to int or tuple of int, optional
How to adjust chunks along dimensions not present in the array being grouped by.
**groupers : Grouper
Grouper objects using which to shuffle the data.
Examples
--------
>>> import dask
>>> from xarray.groupers import UniqueGrouper
>>> da = xr.DataArray(
... dims="x",
... data=dask.array.arange(10, chunks=1),
... coords={"x": [1, 2, 3, 1, 2, 3, 1, 2, 3, 0]},
... name="a",
... )
>>> da
<xarray.DataArray 'a' (x: 10)> Size: 80B
dask.array<arange, shape=(10,), dtype=int64, chunksize=(1,), chunktype=numpy.ndarray>
Coordinates:
* x (x) int64 80B 1 2 3 1 2 3 1 2 3 0
>>> da.shuffle_by(x=UniqueGrouper())
<xarray.DataArray 'a' (x: 10)> Size: 80B
dask.array<shuffle, shape=(10,), dtype=int64, chunksize=(3,), chunktype=numpy.ndarray>
Coordinates:
* x (x) int64 80B 0 1 1 1 2 2 2 3 3 3
Returns
-------
DataArray or Dataset
The same type as this object
See Also
--------
DataArrayGroupBy.shuffle
DatasetGroupBy.shuffle
dask.dataframe.DataFrame.shuffle
dask.array.shuffle
"""
return self.groupby(group=group, **groupers)._shuffle_obj(chunks)

def _resample(
self,
resample_cls: type[T_Resample],
26 changes: 0 additions & 26 deletions xarray/tests/test_dask.py
Original file line number Diff line number Diff line change
@@ -25,7 +25,6 @@
assert_identical,
mock,
raise_if_dask_computes,
requires_dask_ge_2024_08_1,
requires_pint,
requires_scipy_or_netCDF4,
)
@@ -1815,28 +1814,3 @@ def test_minimize_graph_size():
# all the other dimensions.
# e.g. previously for 'x', actual == numchunks['y'] * numchunks['z']
assert actual == numchunks[var], (actual, numchunks[var])


@requires_dask_ge_2024_08_1
@pytest.mark.parametrize(
"chunks, expected_chunks",
[
((1,), (1, 3, 3, 3)),
((10,), (10,)),
],
)
def test_shuffle_by(chunks, expected_chunks):
from xarray.groupers import UniqueGrouper

da = xr.DataArray(
dims="x",
data=dask.array.arange(10, chunks=chunks),
coords={"x": [1, 2, 3, 1, 2, 3, 1, 2, 3, 0]},
name="a",
)
ds = da.to_dataset()

for obj in [ds, da]:
actual = obj.shuffle_by(x=UniqueGrouper())
assert_identical(actual, obj.sortby("x"))
assert actual.chunksizes["x"] == expected_chunks