Fix grouping by dask array (#322)

Author: dcherian

flox/core.py

@@ -1378,9 +1378,7 @@ def _extract_unknown_groups(reduced, dtype) -> tuple[DaskArray]:
 
     groups_token = f"group-{reduced.name}"
     first_block = reduced.ndim * (0,)
-    layer: Graph = {
-        (groups_token, *first_block): (operator.getitem, (reduced.name, *first_block), "groups")
-    }
+    layer: Graph = {(groups_token, 0): (operator.getitem, (reduced.name, *first_block), "groups")}
     groups: tuple[DaskArray] = (
         dask.array.Array(
             HighLevelGraph.from_collections(groups_token, layer, dependencies=[reduced]),

tests/test_core.py

@@ -169,8 +169,6 @@ def test_groupby_reduce(
 ) -> None:
     array = array.astype(dtype)
     if chunk:
-        if expected_groups is None:
-            pytest.skip()
         array = da.from_array(array, chunks=(3,) if array.ndim == 1 else (1, 3))
         by = da.from_array(by, chunks=(3,) if by.ndim == 1 else (1, 3))