argmin/argmax implementation based on _values_for_argsort

pandas/core/arrays/base.py

@@ -28,7 +28,7 @@
 from pandas.core import ops
 from pandas.core.algorithms import _factorize_array, unique
 from pandas.core.missing import backfill_1d, pad_1d
-from pandas.core.sorting import nargsort
+from pandas.core.sorting import nargminmax, nargsort
 
 _extension_array_shared_docs: Dict[str, str] = dict()
 
@@ -520,10 +520,7 @@ def argmin(self):
         --------
         ExtensionArray.argmax
         """
-        if len(self) == 0:
-            raise ValueError("attempt to get argmin of an empty sequence")
-
-        return self.argsort()[0]
+        return nargminmax(self, "argmin")
 
     def argmax(self):
         """
@@ -538,12 +535,7 @@ def argmax(self):
         --------
         ExtensionArray.argmin
         """
-
-        if len(self) == 0:
-            raise ValueError("attempt to get argmax of an empty sequence")
-
-        no_nan = self.isna().sum()
-        return self.argsort()[-1 - no_nan]
+        return nargminmax(self, "argmax")
 
     def fillna(self, value=None, method=None, limit=None):
         """

pandas/core/sorting.py

@@ -319,6 +319,33 @@ def nargsort(
     return indexer
 
 
+def nargminmax(values, method: str):
+    """
+    Implementation of np.argmin/argmax but for ExtensionArray and which
+    handles missing values.
+
+    Parameters
+    ----------
+    values : ExtensionArray
+    method : {"argmax", "argmin"}
+
+    Returns
+    -------
+    int
+    """
+    assert method in {"argmax", "argmin"}
+    func = np.argmax if method == "argmax" else np.argmin
+
+    mask = np.asarray(isna(values))
+    values = values._values_for_argsort()
+
+    idx = np.arange(len(values))
+    non_nans = values[~mask]
+    non_nan_idx = idx[~mask]
+
+    return non_nan_idx[func(non_nans)]
+
+
 def ensure_key_mapped_multiindex(index, key: Callable, level=None):
     """
     Returns a new MultiIndex in which key has been applied

pandas/tests/extension/base/methods.py

@@ -75,26 +75,41 @@ def test_argsort_missing(self, data_missing_for_sorting):
         expected = pd.Series(np.array([1, -1, 0], dtype=np.int64))
         self.assert_series_equal(result, expected)
 
-    def test_argmax(self, data_missing_for_sorting):
+    def test_argmin_argmax(self, data_for_sorting, data_missing_for_sorting, na_value):
         # GH 24382
-        result = data_missing_for_sorting.argmax()
-        expected = 0
-        assert result == expected
 
-    def test_argmin(self, data_missing_for_sorting):
-        # GH 24382
-        result = data_missing_for_sorting.argmin()
-        expected = 2
-        assert result == expected
+        # data_for_sorting -> [B, C, A] with A < B < C
+        assert data_for_sorting.argmax() == 1
+        assert data_for_sorting.argmin() == 2
+
+        # with repeated values -> first occurence
+        data = data_for_sorting.take([2, 0, 0, 1, 1, 2])
+        assert data.argmax() == 3
+        assert data.argmin() == 0
+
+        # with missing values
+        # data_missing_for_sorting -> [B, NA, A] with A < B and NA missing.
+        assert data_missing_for_sorting.argmax() == 0
+        assert data_missing_for_sorting.argmin() == 2
 
     @pytest.mark.parametrize(
         "method", ["argmax", "argmin"],
     )
-    def test_extremize_empty_array(self, method, data_missing_for_sorting):
+    def test_argmin_argmax_empty_array(self, method, data):
         # GH 24382
         err_msg = "attempt to get"
         with pytest.raises(ValueError, match=err_msg):
-            getattr(data_missing_for_sorting[:0], method)()
+            getattr(data[:0], method)()
+
+    @pytest.mark.parametrize(
+        "method", ["argmax", "argmin"],
+    )
+    def test_argmin_argmax_all_na(self, method, data, na_value):
+        # all missing with skipna=True is the same as emtpy
+        err_msg = "attempt to get"
+        data_na = type(data)._from_sequence([na_value, na_value], dtype=data.dtype)
+        with pytest.raises(ValueError, match=err_msg):
+            getattr(data_na, method)()
 
     @pytest.mark.parametrize(
         "na_position, expected",

pandas/tests/extension/test_boolean.py

@@ -235,6 +235,23 @@ def test_searchsorted(self, data_for_sorting, as_series):
     def test_value_counts(self, all_data, dropna):
         return super().test_value_counts(all_data, dropna)
 
+    def test_argmin_argmax(self, data_for_sorting, data_missing_for_sorting):
+        # override because there are only 2 unique values
+
+        # data_for_sorting -> [B, C, A] with A < B < C -> here True, True, False
+        assert data_for_sorting.argmax() == 0
+        assert data_for_sorting.argmin() == 2
+
+        # with repeated values -> first occurence
+        data = data_for_sorting.take([2, 0, 0, 1, 1, 2])
+        assert data.argmax() == 1
+        assert data.argmin() == 0
+
+        # with missing values
+        # data_missing_for_sorting -> [B, NA, A] with A < B and NA missing.
+        assert data_missing_for_sorting.argmax() == 0
+        assert data_missing_for_sorting.argmin() == 2
+
 
 class TestCasting(base.BaseCastingTests):
     pass

pandas/tests/extension/test_sparse.py

@@ -316,6 +316,19 @@ def test_shift_0_periods(self, data):
         data._sparse_values[0] = data._sparse_values[1]
         assert result._sparse_values[0] != result._sparse_values[1]
 
+    @pytest.mark.parametrize(
+        "method", ["argmax", "argmin"],
+    )
+    def test_argmin_argmax_all_na(self, method, data, na_value):
+        # overriding because Sparse[int64, 0] cannot handle na_value
+        if data.dtype.fill_value == 0:
+            pytest.skip("missing values not supported with Sparse[int64, 0]")
+
+        err_msg = "attempt to get"
+        data_na = type(data)._from_sequence([na_value, na_value], dtype=data.dtype)
+        with pytest.raises(ValueError, match=err_msg):
+            getattr(data_na, method)()
+
 
 class TestCasting(BaseSparseTests, base.BaseCastingTests):
     def test_astype_object_series(self, all_data):