如何将pycountry.db.Country对象用作pd.DataFrame索引并实现.loc查询？

问题背景 #

在创建国家数据集时，使用pycountry.db.Country对象表示国家以避免歧义，但将其设为pd.DataFrame索引后，用.loc[]传入Country对象查询时会抛出KeyError，即使索引中确实存在该对象（已通过==和is验证实例一致性）。

可复现代码：

import pandas as pd
import pycountry

aruba: pycountry.db.Country = pycountry.countries.get(alpha_3="ABW")
belgium: pycountry.db.Country = pycountry.countries.get(alpha_3="BEL")
canada: pycountry.db.Country = pycountry.countries.get(alpha_3="CAN")

data: list[dict] = [
    {"country": aruba, "population": 106_203},
    {"country": belgium, "population": 11_429_336},
    {"country": canada, "population": 37_058_856},
]

df: pd.DataFrame = pd.DataFrame(data)
df.set_index("country", inplace=True)

# 验证索引对象与目标对象为同一实例
assert df.index[1] == belgium
assert df.index[1] is belgium

# 此处抛出KeyError
belgium_data = df.loc[belgium]

问题原因 #

Pandas在处理.loc[]查询时，对于非原生数据类型（如自定义对象）会尝试将其序列化为内部可识别的结构（此处Country对象被展开为键值对元组列表），导致无法匹配索引中存储的原始对象实例，最终触发KeyError。

解决方案 #

方法1：通过index.get_loc()获取位置后查询 #

利用索引的get_loc()方法直接定位目标对象的位置，再用.iloc[]访问对应行：

# 获取目标对象在索引中的位置
pos = df.index.get_loc(belgium)
# 通过位置查询数据
belgium_data = df.iloc[pos]

方法2：改用Country对象的唯一标识作为索引 #

保留Country对象在数据列中，改用其唯一属性（如alpha_3代码）作为索引，既避免歧义，又简化查询逻辑：

# 构建数据时新增唯一标识列
data = [
    {"country": aruba, "country_code": aruba.alpha_3, "population": 106_203},
    {"country": belgium, "country_code": belgium.alpha_3, "population": 11_429_336},
    {"country": canada, "country_code": canada.alpha_3, "population": 37_058_856},
]

df = pd.DataFrame(data).set_index("country_code")

# 通过唯一标识查询
belgium_data = df.loc[belgium.alpha_3]
# 仍可获取对应的Country对象
belgium_obj = belgium_data["country"]

方法3：自定义索引匹配逻辑（进阶） #

若必须保留Country对象作为索引，可继承pd.Index自定义索引类，重写匹配逻辑：

class CountryIndex(pd.Index):
    def get_loc(self, key):
        if isinstance(key, pycountry.db.Country):
            for idx, item in enumerate(self):
                if item is key or item.alpha_3 == key.alpha_3:
                    return idx
        return super().get_loc(key)

# 替换原索引
df.index = CountryIndex(df.index)

# 此时可直接用.loc查询
belgium_data = df.loc[belgium]

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