Small cleanup

Author: dcherian

flox/core.py

@@ -820,6 +820,18 @@ def _expand_dims(results: IntermediateDict) -> IntermediateDict:
     return results
 
 
+def _find_unique_groups(x_chunk):
+    from dask.base import flatten
+    from dask.utils import deepmap
+
+    unique_groups = _unique(tuple(flatten(deepmap(listify_groups, x_chunk))))
+    unique_groups = unique_groups[~isnull(unique_groups)]
+
+    if len(unique_groups) == 0:
+        unique_groups = [np.nan]
+    return unique_groups
+
+
 def _simple_combine(
     x_chunk,
     agg: Aggregation,
@@ -839,18 +851,12 @@ def _simple_combine(
     4. At the final agggregate step, we squeeze out DUMMY_AXIS
     """
     from dask.array.core import deepfirst
-    from dask.base import flatten
     from dask.utils import deepmap
 
     if not reindex:
         # We didn't reindex at the blockwise step
         # So now reindex before combining by reducing along DUMMY_AXIS
-        unique_groups = _unique(tuple(flatten(deepmap(listify_groups, x_chunk))))
-        # print(f"unique_groups: {unique_groups}")
-        unique_groups = unique_groups[~isnull(unique_groups)]
-
-        if len(unique_groups) == 0:
-            unique_groups = [np.nan]
+        unique_groups = _find_unique_groups(x_chunk)
         x_chunk = deepmap(
             partial(reindex_intermediates, agg=agg, unique_groups=unique_groups), x_chunk
         )
@@ -912,7 +918,6 @@ def _grouped_combine(
     sort: bool = True,
 ) -> IntermediateDict:
     """Combine intermediates step of tree reduction."""
-    from dask.base import flatten
     from dask.utils import deepmap
 
     if isinstance(x_chunk, dict):
@@ -923,11 +928,7 @@ def _grouped_combine(
         # when there's only a single axis of reduction, we can just concatenate later,
         # reindexing is unnecessary
         # I bet we can minimize the amount of reindexing for mD reductions too, but it's complicated
-        unique_groups = _unique(tuple(flatten(deepmap(listify_groups, x_chunk))))
-        unique_groups = unique_groups[~isnull(unique_groups)]
-        if len(unique_groups) == 0:
-            unique_groups = [np.nan]
-
+        unique_groups = _find_unique_groups(x_chunk)
         x_chunk = deepmap(
             partial(reindex_intermediates, agg=agg, unique_groups=unique_groups), x_chunk
         )