amend after review

Author: makbigc

doc/source/whatsnew/v0.25.0.rst

@@ -252,7 +252,7 @@ Other API Changes
 - The ``arg`` argument in :meth:`pandas.core.groupby.DataFrameGroupBy.agg` has been renamed to ``func`` (:issue:`26089`)
 - The ``arg`` argument in :meth:`pandas.core.window._Window.aggregate` has been renamed to ``func`` (:issue:`26372`)
 - Most Pandas classes had a ``__bytes__`` method, which was used for getting a python2-style bytestring representation of the object. This method has been removed as a part of dropping Python2 (:issue:`26447`)
-- :meth:`ExtensionArray.argsort` places the nan at the end of the sorted array (:issue:`21801`)
+- :meth:`ExtensionArray.argsort` places the ``nan`` at the end of the sorted array (:issue:`21801`)
 
 .. _whatsnew_0250.deprecations:
 

pandas/core/arrays/base.py

@@ -373,6 +373,8 @@ def _values_for_argsort(self) -> Tuple[np.ndarray, np.ndarray]:
         ndarray
             The mask which indicates the NA values.
 
+            .. versionadded:: 0.25.0
+
         See Also
         --------
         ExtensionArray.argsort
@@ -410,17 +412,15 @@ def argsort(self, ascending=True, kind='quicksort', *args, **kwargs):
         ascending = nv.validate_argsort_with_ascending(ascending, args, kwargs)
         values, mask = self._values_for_argsort()
 
-        def permutation(mask):
-            # Return a permutation which maps the indices of the
-            # subarray without nan to the indices of the original array.
-            permu = np.arange(len(mask))
-            permu = permu[~mask]
-            return permu
-
         if mask.any():
             notmask = ~mask
             notnull = np.argsort(values[notmask], kind=kind, **kwargs)
-            permu = permutation(mask)
+
+            # permu maps the indices of the subarray
+            # without nan to the indices of the original array.
+            permu = np.arange(len(mask))
+            permu = permu[~mask]
+
             notnull = permu[notnull]
             allnan = np.arange(len(self))[mask]
             result = np.append(notnull, allnan)

pandas/core/arrays/integer.py

@@ -523,6 +523,8 @@ def _values_for_argsort(self) -> Tuple[np.ndarray, np.ndarray]:
         ndarray
             The mask which indicates the NA values.
 
+            .. versionadded:: 0.25.0
+
         See Also
         --------
         ExtensionArray.argsort