Support nanargmin, nanargmax

Author: dcherian

flox/aggregations.py

@@ -421,7 +421,7 @@ def _pick_second(*x):
     chunk=("nanmax", "nanargmax"),  # order is important
     combine=("max", "argmax"),
     reduction_type="argreduce",
-    fill_value=(dtypes.NINF, -1),
+    fill_value=(dtypes.NINF, 0),
     final_fill_value=-1,
     finalize=_pick_second,
     dtypes=(None, np.intp),
@@ -434,7 +434,7 @@ def _pick_second(*x):
     chunk=("nanmin", "nanargmin"),  # order is important
     combine=("min", "argmin"),
     reduction_type="argreduce",
-    fill_value=(dtypes.INF, -1),
+    fill_value=(dtypes.INF, 0),
     final_fill_value=-1,
     finalize=_pick_second,
     dtypes=(None, np.intp),

tests/conftest.py

@@ -1,7 +1,7 @@
 import pytest
 
 
-@pytest.fixture(scope="module", params=["flox"])
+@pytest.fixture(scope="module", params=["flox", "numpy", "numba"])
 def engine(request):
     if request.param == "numba":
         try:

tests/test_core.py

@@ -55,7 +55,7 @@ def dask_array_ones(*args):
     "nansum",
     "argmax",
     "nanfirst",
-    pytest.param("nanargmax", marks=(pytest.mark.skip,)),
+    "nanargmax",
     "prod",
     "nanprod",
     "mean",
@@ -233,8 +233,9 @@ def test_groupby_reduce_all(nby, size, chunks, func, add_nan_by, engine):
                 # computing silences a bunch of dask warnings
                 array_ = array.compute() if chunks is not None else array
                 if "arg" in func and add_nan_by:
-                    array_[..., nanmask] = np.nan
-                    expected = getattr(np, "nan" + func)(array_, axis=-1, **kwargs)
+                    func_ = f"nan{func}" if "nan" not in func else func
+                    array[..., nanmask] = np.nan
+                    expected = getattr(np, func_)(array, axis=-1, **kwargs)
                 # elif func in ["first", "last"]:
                 #    expected = getattr(xrutils, f"nan{func}")(array_[..., ~nanmask], axis=-1, **kwargs)
                 elif func in ["nanfirst", "nanlast"]: