Add is_supported_aggregation (#455)

Author: dcherian

flox/__init__.py

@@ -3,7 +3,7 @@
 """Top-level module for flox ."""
 
 from . import cache
-from .aggregations import Aggregation, Scan
+from .aggregations import Aggregation, Scan, is_supported_aggregation
 from .core import (
     groupby_reduce,
     groupby_scan,
@@ -36,4 +36,5 @@ def _get_version():
     "set_options",
     "ReindexStrategy",
     "ReindexArrayType",
+    "is_supported_aggregation",
 ]

flox/aggregations.py

@@ -14,7 +14,7 @@
 
 from . import aggregate_flox, aggregate_npg, xrutils
 from . import xrdtypes as dtypes
-from .lib import sparse_array_type
+from .lib import dask_array_type, sparse_array_type
 
 if TYPE_CHECKING:
     FuncTuple = tuple[Callable | str, ...]
@@ -895,3 +895,20 @@ def _initialize_aggregation(
     agg.simple_combine = tuple(simple_combine)
 
     return agg
+
+
+def is_supported_aggregation(array, func: str) -> bool:
+    if isinstance(array, dask_array_type):
+        array = array._meta
+
+    if isinstance(array, sparse_array_type):
+        from flox.core import _is_sparse_supported_reduction
+
+        return _is_sparse_supported_reduction(func)
+
+    module, *_ = type(array).__module__.split(".")
+
+    if module in ["numpy", "cubed"]:
+        return func in AGGREGATIONS
+    else:
+        return False

flox/core.py

@@ -209,7 +209,10 @@ def _postprocess_numbagg(result, *, func, fill_value, size, seen_groups):
     if needs_masking:
         mask = np.isin(groups, seen_groups, assume_unique=True, invert=True)
         if mask.any():
-            result[..., groups[mask]] = fill_value
+            if isinstance(result, sparse_array_type):
+                result.fill_value = fill_value
+            else:
+                result[..., groups[mask]] = fill_value
     return result