如何为Pandas MultiIndex多级别应用自定义键函数排序？

需求说明 #

现有带三级MultiIndex的DataFrame，需按以下规则排序：

1. 按number级别升序排列
2. 按size级别自定义顺序（small→medium→tall）排列
3. 按letter级别升序排列

同时解决调用sort_index(level="size", key=custom_key)时出现的KeyError: 'Level size not found'问题。

问题原因与解决 #

1. KeyError问题解析 #

出现KeyError是因为pandas 1.5.0版本之前，sort_index的key参数仅支持通过位置索引指定级别，不支持使用级别名称。解决方式二选一：

• 升级pandas至1.5.0及以上版本，即可直接使用级别名称绑定key函数
• 旧版中改用级别位置（size是第2级，索引为1）指定

2. 多级别排序+自定义规则实现 #

针对多级别排序需求，我们可以通过sort_index的key参数传入字典/列表，为指定级别绑定自定义排序键，其余级别使用默认排序逻辑。

完整代码（pandas 1.5.0+）

import pandas as pd
import numpy as np

# 构建原始DataFrame
my_index = pd.MultiIndex.from_product(
    [(3,1,2), ("small", "tall", "medium"), ("B", "A", "C")],
    names=["number", "size", "letter"]
)
df_0 = pd.DataFrame(np.random.rand(27, 2), columns=["x", "y"], index=my_index)

# 定义size的自定义排序映射与键函数
size_order = {"small": 0, "medium": 1, "tall": 2}
custom_key = np.vectorize(lambda x: size_order[x])

# 执行多级别排序
df_sorted = df_0.sort_index(
    level=["number", "size", "letter"],
    key={"size": custom_key},  # 仅为size级别应用自定义键
    inplace=False
)

print(df_sorted)

旧版pandas兼容写法

# 旧版需用位置索引指定级别，key参数传入对应位置的键函数（None表示默认排序）
df_sorted = df_0.sort_index(
    level=[0, 1, 2],  # 对应number、size、letter的位置
    key=[None, custom_key, None],
    inplace=False
)

输出结果 #

执行代码后，DataFrame将按照number(1→2→3)→size(small→medium→tall)→letter(A→B→C)的顺序排列，与预期输出一致。

内容的提问来源于stack exchange，提问作者TheGreenHeptagon