From 59d39d1d40465027381c0e2d81cb6b5b2635557f Mon Sep 17 00:00:00 2001 From: Patrick Hoefler Date: Sun, 26 Feb 2023 19:06:28 +0100 Subject: [PATCH] ENH: Add na_position support to midx.sort_values --- doc/source/whatsnew/v2.1.0.rst | 2 +- pandas/core/indexes/base.py | 2 +- pandas/core/indexes/multi.py | 8 +++++-- pandas/tests/indexes/multi/test_sorting.py | 25 ++++++++++++++++++++++ 4 files changed, 33 insertions(+), 4 deletions(-) diff --git a/doc/source/whatsnew/v2.1.0.rst b/doc/source/whatsnew/v2.1.0.rst index 45b5c16415f9d..42220e9613580 100644 --- a/doc/source/whatsnew/v2.1.0.rst +++ b/doc/source/whatsnew/v2.1.0.rst @@ -28,7 +28,7 @@ enhancement2 Other enhancements ^^^^^^^^^^^^^^^^^^ -- +- :meth:`MultiIndex.sort_values` now supports ``na_position`` (:issue:`51612`) - .. --------------------------------------------------------------------------- diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py index acebe8a498f03..c66e92c28d13a 100644 --- a/pandas/core/indexes/base.py +++ b/pandas/core/indexes/base.py @@ -5586,7 +5586,7 @@ def sort_values( items=idx, ascending=ascending, na_position=na_position, key=key ) else: - _as = idx.argsort() + _as = idx.argsort(na_position=na_position) if not ascending: _as = _as[::-1] diff --git a/pandas/core/indexes/multi.py b/pandas/core/indexes/multi.py index b381752818ba0..19fedd4a59774 100644 --- a/pandas/core/indexes/multi.py +++ b/pandas/core/indexes/multi.py @@ -2133,11 +2133,15 @@ def append(self, other): except (TypeError, IndexError): return Index(new_tuples) - def argsort(self, *args, **kwargs) -> npt.NDArray[np.intp]: + def argsort( + self, *args, na_position: str = "last", **kwargs + ) -> npt.NDArray[np.intp]: if len(args) == 0 and len(kwargs) == 0: # lexsort is significantly faster than self._values.argsort() target = self._sort_levels_monotonic(raise_if_incomparable=True) - return lexsort_indexer(target._get_codes_for_sorting()) + return lexsort_indexer( + target._get_codes_for_sorting(), na_position=na_position + ) return self._values.argsort(*args, **kwargs) @Appender(_index_shared_docs["repeat"] % _index_doc_kwargs) diff --git a/pandas/tests/indexes/multi/test_sorting.py b/pandas/tests/indexes/multi/test_sorting.py index 2746c12b120ef..5ba028025eeb6 100644 --- a/pandas/tests/indexes/multi/test_sorting.py +++ b/pandas/tests/indexes/multi/test_sorting.py @@ -14,6 +14,7 @@ Index, MultiIndex, RangeIndex, + Series, Timestamp, ) import pandas._testing as tm @@ -303,3 +304,27 @@ def test_sort_values_incomparable(): match = "'<' not supported between instances of 'Timestamp' and 'int'" with pytest.raises(TypeError, match=match): mi.sort_values() + + +@pytest.mark.parametrize("na_position", ["first", "last"]) +@pytest.mark.parametrize("dtype", ["float64", "Int64", "Float64"]) +def test_sort_values_with_na_na_position(dtype, na_position): + # 51612 + arrays = [ + Series([1, 1, 2], dtype=dtype), + Series([1, None, 3], dtype=dtype), + ] + index = MultiIndex.from_arrays(arrays) + result = index.sort_values(na_position=na_position) + if na_position == "first": + arrays = [ + Series([1, 1, 2], dtype=dtype), + Series([None, 1, 3], dtype=dtype), + ] + else: + arrays = [ + Series([1, 1, 2], dtype=dtype), + Series([1, None, 3], dtype=dtype), + ] + expected = MultiIndex.from_arrays(arrays) + tm.assert_index_equal(result, expected)