@@ -284,16 +284,13 @@ def test_da_groupby_empty() -> None:
284
284
285
285
@requires_dask
286
286
def test_dask_da_groupby_quantile () -> None :
287
- # Only works when the grouped reduction can run blockwise
288
287
# Scalar quantile
289
288
expected = xr .DataArray (
290
289
data = [2 , 5 ], coords = {"x" : [1 , 2 ], "quantile" : 0.5 }, dims = "x"
291
290
)
292
291
array = xr .DataArray (
293
292
data = [1 , 2 , 3 , 4 , 5 , 6 ], coords = {"x" : [1 , 1 , 1 , 2 , 2 , 2 ]}, dims = "x"
294
293
)
295
- with pytest .raises (ValueError ):
296
- array .chunk (x = 1 ).groupby ("x" ).quantile (0.5 )
297
294
298
295
# will work blockwise with flox
299
296
actual = array .chunk (x = 3 ).groupby ("x" ).quantile (0.5 )
@@ -327,7 +324,8 @@ def test_dask_da_groupby_median() -> None:
327
324
assert_identical (expected , actual )
328
325
329
326
330
- def test_da_groupby_quantile () -> None :
327
+ @pytest .mark .parametrize ("use_flox" , [pytest .param (True , marks = requires_flox ), False ])
328
+ def test_da_groupby_quantile (use_flox : bool ) -> None :
331
329
array = xr .DataArray (
332
330
data = [1 , 2 , 3 , 4 , 5 , 6 ], coords = {"x" : [1 , 1 , 1 , 2 , 2 , 2 ]}, dims = "x"
333
331
)
@@ -336,16 +334,19 @@ def test_da_groupby_quantile() -> None:
336
334
expected = xr .DataArray (
337
335
data = [2 , 5 ], coords = {"x" : [1 , 2 ], "quantile" : 0.5 }, dims = "x"
338
336
)
339
- actual = array .groupby ("x" ).quantile (0.5 )
340
- assert_identical (expected , actual )
337
+
338
+ with xr .set_options (use_flox = use_flox ):
339
+ actual = array .groupby ("x" ).quantile (0.5 )
340
+ assert_identical (expected , actual )
341
341
342
342
# Vector quantile
343
343
expected = xr .DataArray (
344
344
data = [[1 , 3 ], [4 , 6 ]],
345
345
coords = {"x" : [1 , 2 ], "quantile" : [0 , 1 ]},
346
346
dims = ("x" , "quantile" ),
347
347
)
348
- actual = array .groupby ("x" ).quantile ([0 , 1 ])
348
+ with xr .set_options (use_flox = use_flox ):
349
+ actual = array .groupby ("x" ).quantile ([0 , 1 ])
349
350
assert_identical (expected , actual )
350
351
351
352
array = xr .DataArray (
@@ -356,7 +357,8 @@ def test_da_groupby_quantile() -> None:
356
357
e = [np .nan , 5 ] if skipna is False else [2.5 , 5 ]
357
358
358
359
expected = xr .DataArray (data = e , coords = {"x" : [1 , 2 ], "quantile" : 0.5 }, dims = "x" )
359
- actual = array .groupby ("x" ).quantile (0.5 , skipna = skipna )
360
+ with xr .set_options (use_flox = use_flox ):
361
+ actual = array .groupby ("x" ).quantile (0.5 , skipna = skipna )
360
362
assert_identical (expected , actual )
361
363
362
364
# Multiple dimensions
0 commit comments