@@ -2848,20 +2848,27 @@ def test_multiple_groupers(use_flox) -> None:
2848
2848
xy = UniqueGrouper (labels = ["a" , "b" , "c" ]),
2849
2849
eagerly_compute_group = eagerly_compute_group ,
2850
2850
)
2851
- with raise_if_dask_computes (max_computes = 1 ):
2852
- if eagerly_compute_group :
2853
- with pytest .warns (DeprecationWarning ):
2854
- gb = b .groupby (** kwargs ) # type: ignore[arg-type]
2855
- else :
2856
- gb = b .groupby (** kwargs ) # type: ignore[arg-type]
2857
- assert is_chunked_array (gb .encoded .codes .data )
2858
- assert not gb .encoded .group_indices
2859
2851
expected = xr .DataArray (
2860
2852
[[[1 , 1 , 1 ], [np .nan , 1 , 2 ]]] * 4 ,
2861
2853
dims = ("z" , "x" , "xy" ),
2862
2854
coords = {"xy" : ("xy" , ["a" , "b" , "c" ], {"foo" : "bar" })},
2863
2855
)
2864
- assert_identical (gb .count (), expected )
2856
+ if eagerly_compute_group :
2857
+ with raise_if_dask_computes (max_computes = 1 ):
2858
+ with pytest .warns (DeprecationWarning ):
2859
+ gb = b .groupby (** kwargs ) # type: ignore[arg-type]
2860
+ assert_identical (gb .count (), expected )
2861
+ else :
2862
+ with raise_if_dask_computes (max_computes = 0 ):
2863
+ gb = b .groupby (** kwargs ) # type: ignore[arg-type]
2864
+ assert is_chunked_array (gb .encoded .codes .data )
2865
+ assert not gb .encoded .group_indices
2866
+ if has_flox :
2867
+ with raise_if_dask_computes (max_computes = 1 ):
2868
+ assert_identical (gb .count (), expected )
2869
+ else :
2870
+ with pytest .raises (ValueError , match = "when lazily grouping" ):
2871
+ gb .count ()
2865
2872
2866
2873
2867
2874
@pytest .mark .parametrize ("use_flox" , [True , False ])
@@ -3003,11 +3010,6 @@ def test_lazy_grouping(grouper, expect_index):
3003
3010
pd .testing .assert_index_equal (encoded .full_index , expect_index )
3004
3011
np .testing .assert_array_equal (encoded .unique_coord .values , np .array (expect_index ))
3005
3012
3006
- lazy = (
3007
- xr .Dataset ({"foo" : data }, coords = {"zoo" : data })
3008
- .groupby (zoo = grouper , eagerly_compute_group = False )
3009
- .count ()
3010
- )
3011
3013
eager = (
3012
3014
xr .Dataset ({"foo" : data }, coords = {"zoo" : data .compute ()})
3013
3015
.groupby (zoo = grouper )
@@ -3017,9 +3019,16 @@ def test_lazy_grouping(grouper, expect_index):
3017
3019
{"foo" : (encoded .codes .name , np .ones (encoded .full_index .size ))},
3018
3020
coords = {encoded .codes .name : expect_index },
3019
3021
)
3020
- assert_identical (eager , lazy )
3021
3022
assert_identical (eager , expected )
3022
3023
3024
+ if has_flox :
3025
+ lazy = (
3026
+ xr .Dataset ({"foo" : data }, coords = {"zoo" : data })
3027
+ .groupby (zoo = grouper , eagerly_compute_group = False )
3028
+ .count ()
3029
+ )
3030
+ assert_identical (eager , lazy )
3031
+
3023
3032
3024
3033
@requires_dask
3025
3034
def test_lazy_grouping_errors ():
0 commit comments