Skip to content

Commit dc0ef70

Browse files
committed
Use nargsort inside EA.argsort to sort
1 parent 989f912 commit dc0ef70

File tree

5 files changed

+24
-8
lines changed

5 files changed

+24
-8
lines changed

doc/source/whatsnew/v0.25.0.rst

+1
Original file line numberDiff line numberDiff line change
@@ -561,6 +561,7 @@ Other API changes
561561
- The ``.str``-accessor has been disabled for 1-level :class:`MultiIndex`, use :meth:`MultiIndex.to_flat_index` if necessary (:issue:`23679`)
562562
- Removed support of gtk package for clipboards (:issue:`26563`)
563563
- Using an unsupported version of Beautiful Soup 4 will now raise an ``ImportError`` instead of a ``ValueError`` (:issue:`27063`)
564+
- :meth:`ExtensionArray.argsort` places NA values at the end of the sorted array, except :class:`Categorical` (:issue:`21801`)
564565

565566
.. _whatsnew_0250.deprecations:
566567

pandas/core/arrays/base.py

+6-5
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323

2424
from pandas._typing import ArrayLike
2525
from pandas.core import ops
26+
from pandas.core.sorting import nargsort
2627

2728
_not_implemented_message = "{} does not implement {}."
2829

@@ -398,7 +399,8 @@ def argsort(self, ascending=True, kind='quicksort', *args, **kwargs):
398399
Returns
399400
-------
400401
index_array : ndarray
401-
Array of indices that sort ``self``.
402+
Array of indices that sort ``self``. If NaN values are contained,
403+
NaN values are placed at the end.
402404
403405
See Also
404406
--------
@@ -409,10 +411,9 @@ def argsort(self, ascending=True, kind='quicksort', *args, **kwargs):
409411
# 1. _values_for_argsort : construct the values passed to np.argsort
410412
# 2. argsort : total control over sorting.
411413
ascending = nv.validate_argsort_with_ascending(ascending, args, kwargs)
412-
values = self._values_for_argsort()
413-
result = np.argsort(values, kind=kind, **kwargs)
414-
if not ascending:
415-
result = result[::-1]
414+
415+
result = nargsort(self, kind=kind, ascending=ascending,
416+
na_position='last')
416417
return result
417418

418419
def fillna(self, value=None, method=None, limit=None):

pandas/core/arrays/categorical.py

+7-3
Original file line numberDiff line numberDiff line change
@@ -1523,7 +1523,7 @@ def check_for_ordered(self, op):
15231523
def _values_for_argsort(self):
15241524
return self._codes.copy()
15251525

1526-
def argsort(self, *args, **kwargs):
1526+
def argsort(self, ascending=True, kind='quicksort', *args, **kwargs):
15271527
# TODO(PY2): use correct signature
15281528
# We have to do *args, **kwargs to avoid a a py2-only signature
15291529
# issue since np.argsort differs from argsort.
@@ -1567,8 +1567,12 @@ def argsort(self, *args, **kwargs):
15671567
>>> cat.argsort()
15681568
array([3, 0, 1, 2])
15691569
"""
1570-
# Keep the implementation here just for the docstring.
1571-
return super().argsort(*args, **kwargs)
1570+
ascending = nv.validate_argsort_with_ascending(ascending, args, kwargs)
1571+
values = self._values_for_argsort()
1572+
result = np.argsort(values, kind=kind, **kwargs)
1573+
if not ascending:
1574+
result = result[::-1]
1575+
return result
15721576

15731577
def sort_values(self, inplace=False, ascending=True, na_position='last'):
15741578
"""

pandas/tests/extension/base/methods.py

+5
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,11 @@ def test_argsort(self, data_for_sorting):
4747
expected = pd.Series(np.array([2, 0, 1], dtype=np.int64))
4848
self.assert_series_equal(result, expected)
4949

50+
def test_argsort_nan_last(self, data_missing_for_sorting):
51+
# GH 21801
52+
result = data_missing_for_sorting.argsort()
53+
assert result[-1] == 1
54+
5055
def test_argsort_missing(self, data_missing_for_sorting):
5156
result = pd.Series(data_missing_for_sorting).argsort()
5257
expected = pd.Series(np.array([1, -1, 0], dtype=np.int64))

pandas/tests/extension/test_categorical.py

+5
Original file line numberDiff line numberDiff line change
@@ -193,6 +193,11 @@ def test_searchsorted(self, data_for_sorting):
193193
if not data_for_sorting.ordered:
194194
raise pytest.skip(reason="searchsorted requires ordered data.")
195195

196+
def test_argsort_nan_last(self, data_missing_for_sorting):
197+
# GH 21801
198+
# TODO: Categorical.argsort places NA values at the end
199+
pass
200+
196201

197202
class TestCasting(base.BaseCastingTests):
198203
pass

0 commit comments

Comments
 (0)