API: ExtensionArray.argsort places the missing value at the end (#27137)

Author: makbigc

asv_bench/benchmarks/algorithms.py

@@ -141,4 +141,16 @@ def time_quantile(self, quantile, interpolation, dtype):
         self.idx.quantile(quantile, interpolation=interpolation)
 
 
+class SortIntegerArray:
+    params = [10**3, 10**5]
+
+    def setup(self, N):
+        data = np.arange(N, dtype=float)
+        data[40] = np.nan
+        self.array = pd.array(data, dtype='Int64')
+
+    def time_argsort(self, N):
+        self.array.argsort()
+
+
 from .pandas_vb_common import setup  # noqa: F401 isort:skip

doc/source/whatsnew/v0.25.0.rst

@@ -661,7 +661,7 @@ when both are :class:`Series` (:issue:`23293`).
 
 *Previous behavior*
 
-.. code-block:: python
+.. code-block:: ipython
 
    In [5]: np.power(s1, s2)
    Out[5]:
@@ -684,6 +684,36 @@ applying the ufunc.
 
    np.power(s1, s2.array)
 
+Categorical.argsort now places missing values at the end
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+:meth:`Categorical.argsort` now places missing values at the end of the array, making it
+consistent with NumPy and the rest of pandas (:issue:`21801`).
+
+.. ipython:: python
+
+   cat = pd.Categorical(['b', None, 'a'], categories=['a', 'b'], ordered=True)
+
+*Previous behavior*
+
+.. code-block:: ipython
+
+   In [2]: cat = pd.Categorical(['b', None, 'a'], categories=['a', 'b'], ordered=True)
+
+   In [3]: cat.argsort()
+   Out[3]: array([1, 2, 0])
+
+   In [4]: cat[cat.argsort()]
+   Out[4]:
+   [NaN, a, b]
+   categories (2, object): [a < b]
+
+*New behavior*
+
+.. ipython:: python
+
+   cat.argsort()
+   cat[cat.argsort()]
 
 .. _whatsnew_0250.api_breaking.deps:
 
@@ -767,6 +797,7 @@ Other API changes
 - Removed support of gtk package for clipboards (:issue:`26563`)
 - Using an unsupported version of Beautiful Soup 4 will now raise an ``ImportError`` instead of a ``ValueError`` (:issue:`27063`)
 - :meth:`Series.to_excel` and :meth:`DataFrame.to_excel` will now raise a ``ValueError`` when saving timezone aware data. (:issue:`27008`, :issue:`7056`)
+- :meth:`ExtensionArray.argsort` places NA values at the end of the sorted array. (:issue:`21801`)
 - :meth:`DataFrame.to_hdf` and :meth:`Series.to_hdf` will now raise a ``NotImplementedError`` when saving a :class:`MultiIndex` with extention data types for a ``fixed`` format. (:issue:`7775`)
 - Passing duplicate ``names`` in :meth:`read_csv` will now raise a ``ValueError`` (:issue:`17346`)
 

pandas/core/arrays/base.py

@@ -23,6 +23,7 @@
 
 from pandas._typing import ArrayLike
 from pandas.core import ops
+from pandas.core.sorting import nargsort
 
 _not_implemented_message = "{} does not implement {}."
 
@@ -409,7 +410,8 @@ def argsort(self, ascending=True, kind='quicksort', *args, **kwargs):
         Returns
         -------
         index_array : ndarray
-            Array of indices that sort ``self``.
+            Array of indices that sort ``self``. If NaN values are contained,
+            NaN values are placed at the end.
 
         See Also
         --------
@@ -420,10 +422,9 @@ def argsort(self, ascending=True, kind='quicksort', *args, **kwargs):
         # 1. _values_for_argsort : construct the values passed to np.argsort
         # 2. argsort : total control over sorting.
         ascending = nv.validate_argsort_with_ascending(ascending, args, kwargs)
-        values = self._values_for_argsort()
-        result = np.argsort(values, kind=kind, **kwargs)
-        if not ascending:
-            result = result[::-1]
+
+        result = nargsort(self, kind=kind, ascending=ascending,
+                          na_position='last')
         return result
 
     def fillna(self, value=None, method=None, limit=None):

pandas/core/arrays/categorical.py

@@ -1531,13 +1531,14 @@ def check_for_ordered(self, op):
     def _values_for_argsort(self):
         return self._codes.copy()
 
-    def argsort(self, *args, **kwargs):
-        # TODO(PY2): use correct signature
-        # We have to do *args, **kwargs to avoid a a py2-only signature
-        # issue since np.argsort differs from argsort.
+    def argsort(self, ascending=True, kind='quicksort', *args, **kwargs):
         """
         Return the indices that would sort the Categorical.
 
+        .. versionchanged:: 0.25.0
+
+           Changed to sort missing values at the end.
+
         Parameters
         ----------
         ascending : bool, default True
@@ -1574,9 +1575,14 @@ def argsort(self, *args, **kwargs):
         ...                      ordered=True)
         >>> cat.argsort()
         array([3, 0, 1, 2])
+
+        Missing values are placed at the end
+
+        >>> cat = pd.Categorical([2, None, 1])
+        >>> cat.argsort()
+        array([2, 0, 1])
         """
-        # Keep the implementation here just for the docstring.
-        return super().argsort(*args, **kwargs)
+        return super().argsort(ascending=ascending, kind=kind, *args, **kwargs)
 
     def sort_values(self, inplace=False, ascending=True, na_position='last'):
         """

pandas/core/sorting.py

@@ -240,20 +240,6 @@ def nargsort(items, kind='quicksort', ascending=True, na_position='last'):
 
     items = extract_array(items)
     mask = np.asarray(isna(items))
-    # specially handle Categorical
-    if is_categorical_dtype(items):
-        if na_position not in {'first', 'last'}:
-            raise ValueError('invalid na_position: {!r}'.format(na_position))
-
-        cnt_null = mask.sum()
-        sorted_idx = items.argsort(ascending=ascending, kind=kind)
-        if ascending and na_position == 'last':
-            # NaN is coded as -1 and is listed in front after sorting
-            sorted_idx = np.roll(sorted_idx, -cnt_null)
-        elif not ascending and na_position == 'first':
-            # NaN is coded as -1 and is listed in the end after sorting
-            sorted_idx = np.roll(sorted_idx, cnt_null)
-        return sorted_idx
 
     if is_extension_array_dtype(items):
         items = items._values_for_argsort()

pandas/tests/extension/base/methods.py

@@ -47,6 +47,14 @@ def test_argsort(self, data_for_sorting):
         expected = pd.Series(np.array([2, 0, 1], dtype=np.int64))
         self.assert_series_equal(result, expected)
 
+    def test_argsort_missing_array(self, data_missing_for_sorting):
+        result = data_missing_for_sorting.argsort()
+        expected = np.array([2, 0, 1], dtype=np.dtype("int"))
+        # we don't care whether it's int32 or int64
+        result = result.astype("int64", casting="safe")
+        expected = expected.astype("int64", casting="safe")
+        tm.assert_numpy_array_equal(result, expected)
+
     def test_argsort_missing(self, data_missing_for_sorting):
         result = pd.Series(data_missing_for_sorting).argsort()
         expected = pd.Series(np.array([1, -1, 0], dtype=np.int64))

pandas/tests/frame/test_sorting.py

@@ -539,7 +539,7 @@ def test_sort_index_categorical_index(self):
         assert_frame_equal(result, expected)
 
         result = df.sort_index(ascending=False)
-        expected = df.iloc[[3, 2, 5, 1, 0, 4]]
+        expected = df.iloc[[2, 3, 0, 1, 5, 4]]
         assert_frame_equal(result, expected)
 
     def test_sort_index(self):
@@ -629,7 +629,7 @@ def test_sort_index_na_position_with_categories(self):
 
         reversed_categories = sorted(categories, reverse=True)
         reversed_category_indices = sorted(category_indices, reverse=True)
-        reversed_na_indices = sorted(na_indices, reverse=True)
+        reversed_na_indices = sorted(na_indices)
 
         df = pd.DataFrame({
             column_name: pd.Categorical(['A', np.nan, 'B', np.nan, 'C'],