Skip to content

Commit a1e176b

Browse files
committed
Add test for chunked dataarrays
1 parent 3148209 commit a1e176b

File tree

1 file changed

+43
-0
lines changed

1 file changed

+43
-0
lines changed

xarray/tests/test_groupby.py

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -245,6 +245,49 @@ def test_da_groupby_empty() -> None:
245245
empty_array.groupby("dim")
246246

247247

248+
def test_dask_da_groupby_quantile() -> None:
249+
# Only works when the grouped reduction can run blockwise
250+
# Scalar quantile
251+
expected = xr.DataArray(
252+
data=[2, 5], coords={"x": [1, 2], "quantile": 0.5}, dims="x"
253+
)
254+
array = xr.DataArray(
255+
data=[1, 2, 3, 4, 5, 6], coords={"x": [1, 1, 1, 2, 2, 2]}, dims="x"
256+
)
257+
with pytest.raises(ValueError):
258+
array.chunk(x=1).groupby("x").quantile(0.5)
259+
260+
# will work blockwise with flox
261+
actual = array.chunk(x=3).groupby("x").quantile(0.5)
262+
assert_identical(expected, actual)
263+
264+
# will work blockwise with flox
265+
actual = array.chunk(x=-1).groupby("x").quantile(0.5)
266+
assert_identical(expected, actual)
267+
268+
269+
def test_dask_da_groupby_median() -> None:
270+
expected = xr.DataArray(data=[2, 5], coords={"x": [1, 2]}, dims="x")
271+
array = xr.DataArray(
272+
data=[1, 2, 3, 4, 5, 6], coords={"x": [1, 1, 1, 2, 2, 2]}, dims="x"
273+
)
274+
with xr.set_options(use_flox=False):
275+
actual = array.chunk(x=1).groupby("x").median()
276+
assert_identical(expected, actual)
277+
278+
with xr.set_options(use_flox=True):
279+
actual = array.chunk(x=1).groupby("x").median()
280+
assert_identical(expected, actual)
281+
282+
# will work blockwise with flox
283+
actual = array.chunk(x=3).groupby("x").median()
284+
assert_identical(expected, actual)
285+
286+
# will work blockwise with flox
287+
actual = array.chunk(x=-1).groupby("x").median()
288+
assert_identical(expected, actual)
289+
290+
248291
def test_da_groupby_quantile() -> None:
249292
array = xr.DataArray(
250293
data=[1, 2, 3, 4, 5, 6], coords={"x": [1, 1, 1, 2, 2, 2]}, dims="x"

0 commit comments

Comments
 (0)