如何用math.isclose对Pandas DataFrame浮点列排序？解决TypeError

需要按b列（使用math.isclose判断浮点数相等）、再按a列对DataFrame排序。尝试用functools.cmp_to_key转换自定义比较函数后调用df.sort_values，触发错误：

TypeError: object of type 'functools.KeyWrapper' has no len()

已知sort_values的key参数要求接收Series且函数需向量化，不知道如何实现。

原始DataFrame #

尝试的代码 #

import pandas as pd
from functools import cmp_to_key
from math import isclose
import numpy as np

my_list = [
[0, 1.2],
[2, 0.07076863960397785],
[1, 0.07076863960397783],
[4, 0.02],
[3, 0.07076863960397784]
]

df = pd.DataFrame(my_list,columns=['a','b'])

def compare(a,b):
  if isclose(a,b):
    return 0
  elif a-b<0:
    return -1
  else:
    return 1

df.sort_values(by=['b','a'],key= cmp_to_key(compare))

期望排序结果 #

解决方案 #

sort_values的key参数不支持cmp_to_key转换后的两两比较函数——它要求输入一个向量化函数，接收Series并输出同长度的Series（为每个元素生成排序键）。我们可以通过给b列的相似浮点数分配相同分组键的方式实现需求：

实现代码

import pandas as pd
from math import isclose

my_list = [
    [0, 1.2],
    [2, 0.07076863960397785],
    [1, 0.07076863960397783],
    [4, 0.02],
    [3, 0.07076863960397784]
]

df = pd.DataFrame(my_list, columns=['a','b'])

# 为b列生成分组键：isclose判定为相等的元素分配相同键
def generate_b_group(s):
    group_keys = [0]
    for idx in range(1, len(s)):
        if isclose(s.iloc[idx], s.iloc[idx-1]):
            group_keys.append(group_keys[-1])
        else:
            group_keys.append(group_keys[-1] + 1)
    return pd.Series(group_keys, index=s.index)

# 添加临时分组键列，按分组键+ a列排序后删除临时列
df['b_group'] = generate_b_group(df['b'])
sorted_df = df.sort_values(by=['b_group', 'a']).drop('b_group', axis=1)

print(sorted_df)

原理说明

1. 遍历b列，用isclose判断当前元素与前一个元素是否属于同一组，生成连续的分组键
2. 按b_group（保证相似浮点数排在一起）和a列排序，最后移除临时列即可得到期望结果

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