dask support

Author: dcherian

flox/aggregations.py

@@ -586,8 +586,8 @@ def topk_new_dims_func(k) -> tuple[Dim]:
 topk = Aggregation(
     name="topk",
     fill_value=dtypes.NINF,
-    chunk=None,
-    combine=None,
+    chunk="topk",
+    combine=xrutils.topk,
     final_dtype=None,
     new_dims_func=topk_new_dims_func,
 )
@@ -890,10 +890,7 @@ def _initialize_aggregation(
     simple_combine: list[Callable | None] = []
     for combine in agg.combine:
         if isinstance(combine, str):
-            if combine in ["nanfirst", "nanlast"]:
-                simple_combine.append(getattr(xrutils, combine))
-            else:
-                simple_combine.append(getattr(np, combine))
+            simple_combine.append(getattr(np, combine))
         else:
             simple_combine.append(combine)
 

flox/xrutils.py

@@ -378,3 +378,21 @@ def nanlast(values, axis, keepdims=False):
         return np.expand_dims(result, axis=axis)
     else:
         return result
+
+
+def topk(a, k, axis, keepdims):
+    """Chunk and combine function of topk
+
+    Extract the k largest elements from a on the given axis.
+    If k is negative, extract the -k smallest elements instead.
+    Note that, unlike in the parent function, the returned elements
+    are not sorted internally.
+    """
+    assert keepdims is True
+    axis = axis[0]
+    if abs(k) >= a.shape[axis]:
+        return a
+
+    a = np.partition(a, -k, axis=axis)
+    k_slice = slice(-k, None) if k > 0 else slice(-k)
+    return a[tuple(k_slice if i == axis else slice(None) for i in range(a.ndim))]