Fix arg reductions and filling

I changed -1 to 0 for intermediate dask reductions; This *should* have
no effect but allows an unravelling if needed

Author: dcherian

dask_groupby/aggregations.py

@@ -288,7 +288,7 @@ def _zip_index(array_, idx_):
     chunk=("max", "argmax"),  # order is important
     combine=("max", "argmax"),
     reduction_type="argreduce",
-    fill_value=(dtypes.NINF, -1),
+    fill_value=(dtypes.NINF, 0),
     final_fill_value=-1,
     finalize=lambda *x: x[1],
     dtype=np.intp,
@@ -300,7 +300,7 @@ def _zip_index(array_, idx_):
     chunk=("min", "argmin"),  # order is important
     combine=("min", "argmin"),
     reduction_type="argreduce",
-    fill_value=(dtypes.INF, -1),
+    fill_value=(dtypes.INF, 0),
     final_fill_value=-1,
     finalize=lambda *x: x[1],
     dtype=np.intp,

dask_groupby/core.py

@@ -1186,12 +1186,18 @@ def groupby_reduce(
         )  # type: ignore
 
         if reduction.name in ["argmin", "argmax", "nanargmax", "nanargmin"]:
-            if array.ndim > 1 and by.ndim == 1:
+            if array.ndim > 1:
+                # default fill_value is -1; we can't unravel that;
+                # so replace -1 with 0; unravel; then replace 0 with -1
+                # UGH!
+                idx = results["intermediates"][0]
+                mask = idx == -1
+                idx[mask] = 0
                 # Fix npg bug where argmax with nD array, 1D group_idx, axis=-1
                 # will return wrong indices
-                results["intermediates"][0] = np.unravel_index(
-                    results["intermediates"][0], array.shape
-                )[-1]
+                idx = np.unravel_index(idx, array.shape)[-1]
+                idx[mask] = -1
+                results["intermediates"][0] = idx
         elif reduction.name in ["nanvar", "nanstd"]:
             # Fix npg bug where all-NaN rows are 0 instead of NaN
             value, counts = results["intermediates"]

tests/test_core.py

@@ -367,9 +367,8 @@ def test_dask_reduce_axis_subset():
     "axis", [None, (0, 1, 2), (0, 1), (0, 2), (1, 2), 0, 1, 2, (0,), (1,), (2,)]
 )
 def test_groupby_reduce_axis_subset_against_numpy(func, axis, backend):
-    if not isinstance(axis, int):
-        if "arg" in func and (axis is None or len(axis) > 1):
-            pytest.skip()
+    if not isinstance(axis, int) and "arg" in func and (axis is None or len(axis) > 1):
+        pytest.skip()
     if func in ["all", "any"]:
         fill_value = False
     else: