@@ -245,6 +245,49 @@ def test_da_groupby_empty() -> None:
245
245
empty_array .groupby ("dim" )
246
246
247
247
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
+
248
291
def test_da_groupby_quantile () -> None :
249
292
array = xr .DataArray (
250
293
data = [1 , 2 , 3 , 4 , 5 , 6 ], coords = {"x" : [1 , 1 , 1 , 2 , 2 , 2 ]}, dims = "x"
0 commit comments