使用pandas Series sort_index的key参数时排序结果不符的问题

问题场景 #

尝试用sort_index(key=lambda idx: foo(idx))将pandas Series的第一个索引元素移到末尾，代码如下：

自定义排序函数 #

def foo(idx):
    print("pre",idx)
    if idx.name == "pca_n":
        ret = pd.Index(list(idx[1:]) + list(idx[:1]),name=idx.name)
    else:
        ret = idx.copy()
    print("post",ret)
    return ret

调用代码 #

print("index before sort",byHyp.index)
byHyp = byHyp.sort_index(key = lambda x: foo(x))
print("index after sort",byHyp.index)

运行输出 #

index before sort Int64Index([-1, 2, 5, 10, 20], dtype='int64', name='pca_n')
pre Int64Index([-1, 2, 5, 10, 20], dtype='int64', name='pca_n')
post Int64Index([2, 5, 10, 20, -1], dtype='int64', name='pca_n')
index after sort Int64Index([20, -1, 2, 5, 10], dtype='int64', name='pca_n')

可见foo返回了预期的索引顺序，但最终Series的索引排序结果不符。

核心误解 #

你完全误解了sort_index的key参数用法：key函数不是直接返回最终的索引顺序，而是返回一组与原索引元素一一对应的排序键，pandas会根据这些键的大小对原索引元素进行排序。

为什么会得到这个结果 #

你的foo函数返回的[2,5,10,20,-1]是一组排序键，每个键对应原索引的一个元素：

• 原索引-1 → 键2
• 原索引2 → 键5
• 原索引5 → 键10
• 原索引10 → 键20
• 原索引20 → 键-1

pandas会按照键从小到大排序原索引元素：

1. 键-1对应原索引20 → 排第一
2. 键2对应原索引-1 → 排第二
3. 键5对应原索引2 → 排第三
4. 键10对应原索引5 → 排第四
5. 键20对应原索引10 → 排第五

这就是最终得到[20,-1,2,5,10]的原因。

正确实现方法 #

方法1：直接构造新索引（推荐） #

不需要用sort_index，直接用reindex构造目标顺序：

# 跳过第一个元素，再将第一个元素追加到末尾
new_index = list(byHyp.index[1:]) + list(byHyp.index[:1])
byHyp = byHyp.reindex(new_index)

方法2：正确使用sort_index的key参数 #

给每个索引元素分配排序键，让第一个元素的键最大，其他元素按原顺序递增：

byHyp = byHyp.sort_index(
    key=lambda idx: [len(idx) if i == 0 else i for i in range(len(idx))]
)

这里第一个元素的键是5（索引长度），其余元素的键为1,2,3,4，排序后第一个元素会被排到末尾，得到预期的[2,5,10,20,-1]。

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