Move _postprocess_numbagg to lib.py

Move the numbagg postprocessing function to lib.py for better
organization.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

Author: dcherian

flox/core.py

@@ -48,7 +48,7 @@
     quantile_new_dims_func,
 )
 from .cache import memoize
-from .lib import ArrayLayer, dask_array_type, sparse_array_type
+from .lib import ArrayLayer, _postprocess_numbagg, dask_array_type, sparse_array_type
 from .options import OPTIONS
 from .xrutils import (
     _contains_cftime_datetimes,
@@ -196,27 +196,6 @@ class FactorizeKwargs(TypedDict, total=False):
     sort: bool
 
 
-def _postprocess_numbagg(result, *, func, fill_value, size, seen_groups):
-    """Account for numbagg not providing a fill_value kwarg."""
-    from .aggregate_numbagg import DEFAULT_FILL_VALUE
-
-    if not isinstance(func, str) or func not in DEFAULT_FILL_VALUE:
-        return result
-    # The condition needs to be
-    # len(found_groups) < size; if so we mask with fill_value (?)
-    default_fv = DEFAULT_FILL_VALUE[func]
-    needs_masking = fill_value is not None and not np.array_equal(fill_value, default_fv, equal_nan=True)
-    groups = np.arange(size)
-    if needs_masking:
-        mask = np.isin(groups, seen_groups, assume_unique=True, invert=True)
-        if mask.any():
-            if isinstance(result, sparse_array_type):
-                result.fill_value = fill_value
-            else:
-                result[..., groups[mask]] = fill_value
-    return result
-
-
 def identity(x: T) -> T:
     return x
 

flox/lib.py

@@ -1,5 +1,7 @@
 from dataclasses import dataclass
 
+import numpy as np
+
 from .types import DaskArray, Graph
 
 try:
@@ -29,3 +31,25 @@ def to_array(self, dep: DaskArray) -> DaskArray:
 
         graph = HighLevelGraph.from_collections(self.name, self.layer, dependencies=[dep])
         return Array(graph, self.name, self.chunks, meta=dep._meta)
+
+
+def _postprocess_numbagg(result, *, func, fill_value, size, seen_groups):
+    """Account for numbagg not providing a fill_value kwarg."""
+
+    from .aggregate_numbagg import DEFAULT_FILL_VALUE
+
+    if not isinstance(func, str) or func not in DEFAULT_FILL_VALUE:
+        return result
+    # The condition needs to be
+    # len(found_groups) < size; if so we mask with fill_value (?)
+    default_fv = DEFAULT_FILL_VALUE[func]
+    needs_masking = fill_value is not None and not np.array_equal(fill_value, default_fv, equal_nan=True)
+    groups = np.arange(size)
+    if needs_masking:
+        mask = np.isin(groups, seen_groups, assume_unique=True, invert=True)
+        if mask.any():
+            if isinstance(result, sparse_array_type):
+                result.fill_value = fill_value
+            else:
+                result[..., groups[mask]] = fill_value
+    return result