Refactor out _prepare_for_flox

Author: dcherian

flox/aggregate_flox.py

@@ -4,6 +4,20 @@
 
 from .xrutils import isnull
 
+def _prepare_for_flox(group_idx, array):
+    """
+    Sort the input array once to save time.
+    """
+    assert array.shape[-1] == group_idx.shape[0]
+    issorted = (group_idx[:-1] <= group_idx[1:]).all()
+    if issorted:
+        ordered_array = array
+    else:
+        perm = group_idx.argsort(kind="stable")
+        group_idx = group_idx[..., perm]
+        ordered_array = array[..., perm]
+    return group_idx, ordered_array
+
 
 def _np_grouped_op(group_idx, array, op, axis=-1, size=None, fill_value=None, dtype=None, out=None):
     """

flox/core.py

@@ -19,6 +19,7 @@
     _initialize_aggregation,
     generic_aggregate,
 )
+from .aggregate_flox import _prepare_for_flox
 from .cache import memoize
 from .xrutils import is_duck_array, is_duck_dask_array, isnull
 
@@ -44,21 +45,6 @@ def _is_arg_reduction(func: str | Aggregation) -> bool:
     return False
 
 
-def _prepare_for_flox(group_idx, array):
-    """
-    Sort the input array once to save time.
-    """
-    assert array.shape[-1] == group_idx.shape[0]
-    issorted = (group_idx[:-1] <= group_idx[1:]).all()
-    if issorted:
-        ordered_array = array
-    else:
-        perm = group_idx.argsort(kind="stable")
-        group_idx = group_idx[..., perm]
-        ordered_array = array[..., perm]
-    return group_idx, ordered_array
-
-
 def _get_expected_groups(by, sort, *, raise_if_dask=True) -> pd.Index | None:
     if is_duck_dask_array(by):
         if raise_if_dask: