diff --git a/pandas/core/sorting.py b/pandas/core/sorting.py index 930704e6f62f4..3349a13d7f0bc 100644 --- a/pandas/core/sorting.py +++ b/pandas/core/sorting.py @@ -57,7 +57,7 @@ def get_indexer_indexer( target: Index, - level: Level | list[Level] | None, + level: Level | list[Level] | None, # can level actually be a list here? ascending: list[bool] | bool, kind: SortKind, na_position: NaPosition, @@ -531,7 +531,6 @@ def _ensure_key_mapped_multiindex( level_iter = [level] else: level_iter = level - sort_levels: range | set = {index._get_level_number(lev) for lev in level_iter} else: sort_levels = range(index.nlevels) @@ -581,7 +580,10 @@ def ensure_key_mapped( if isinstance( values, Index ): # convert to a new Index subclass, not necessarily the same - result = Index(result, tupleize_cols=False) + # preserve the original name when creating the new Index + result = Index( + result, tupleize_cols=False, name=getattr(values, "name", None) + ) else: # try to revert to original type otherwise type_of_values = type(values) diff --git a/pandas/tests/frame/methods/test_sort_index.py b/pandas/tests/frame/methods/test_sort_index.py index 1a631e760208a..1ee7d3a91f41b 100644 --- a/pandas/tests/frame/methods/test_sort_index.py +++ b/pandas/tests/frame/methods/test_sort_index.py @@ -1,3 +1,4 @@ +from natsort import index_natsorted # should we import this? or change test? import numpy as np import pytest @@ -943,6 +944,95 @@ def test_sort_index_multiindex_sort_remaining(self, ascending): tm.assert_frame_equal(result, expected) + def test_sort_index_multiindex_by_level_name(self): + # GH#62361 + + df = DataFrame( + [[1, 2], [3, 4]], + columns=MultiIndex.from_product( + [["a"], ["top10", "top2"]], names=("A", "B") + ), + ) + + expected = DataFrame( + [[2, 1], [4, 3]], + columns=MultiIndex.from_product( + [["a"], ["top2", "top10"]], names=("A", "B") + ), + ) + + sorted_df = df.sort_index( + axis=1, level="B", key=lambda x: np.argsort(index_natsorted(x)) + ) + tm.assert_frame_equal(sorted_df, expected) + + def test_sort_index_multiindex_by_level_name_2(self): + # GH#62361 + + df = DataFrame( + [[1, 2, 3], [4, 5, 6]], + columns=MultiIndex.from_tuples( + [("a10", "b12"), ("a2", "b17"), ("a2", "b4")], names=("A", "B") + ), + ) + + expected_A = DataFrame( + [[2, 3, 1], [5, 6, 4]], + columns=MultiIndex.from_tuples( + [("a2", "b17"), ("a2", "b4"), ("a10", "b12")], names=("A", "B") + ), + ) + expected_B = DataFrame( + [[3, 1, 2], [6, 4, 5]], + columns=MultiIndex.from_tuples( + [("a2", "b4"), ("a10", "b12"), ("a2", "b17")], names=("A", "B") + ), + ) + + sorted_df = df.sort_index( + axis=1, level=0, key=lambda x: np.argsort(index_natsorted(x)) + ) + tm.assert_frame_equal(sorted_df, expected_A) + + sorted_df = df.sort_index( + axis=1, level="A", key=lambda x: np.argsort(index_natsorted(x)) + ) + tm.assert_frame_equal(sorted_df, expected_A) + + sorted_df = df.sort_index( + axis=1, level=1, key=lambda x: np.argsort(index_natsorted(x)) + ) + tm.assert_frame_equal(sorted_df, expected_B) + + sorted_df = df.sort_index( + axis=1, level="B", key=lambda x: np.argsort(index_natsorted(x)) + ) + tm.assert_frame_equal(sorted_df, expected_B) + + # actually, only 1 element of list matters for sorting (2nd is ignored) + sorted_df = df.sort_index( + axis=1, level=[0, 1], key=lambda x: np.argsort(index_natsorted(x)) + ) + tm.assert_frame_equal(sorted_df, expected_A) + + # sorted_df = df.sort_index( + # axis=1, level=[1, 0], key=lambda x: np.argsort(index_natsorted(x)) + # ) + # tm.assert_frame_equal(sorted_df, expected_B) + + # sorted_df = df.sort_index( + # axis=1, level=[1, "A"], key=lambda x: np.argsort(index_natsorted(x)) + # ) + # tm.assert_frame_equal(sorted_df, expected_B) + + # # repetition does not matter + # sorted_df = df.sort_index( + # axis=1, + # level=["A", "B", 0, 1, "B"], + # key=lambda x: np.argsort(index_natsorted(x)), + # ) + # tm.assert_frame_equal(sorted_df, expected_A) + def test_sort_index_with_sliced_multiindex(): # GH 55379