pd.testing.assert_frame_equal对比DataFrame触发AssertionError，双方dtype显示一致却不匹配

这个问题我之前也碰到过，明明两个DataFrame的dtype显示完全一样，结果断言就是失败，确实挺让人困惑的😅

先帮你梳理下问题场景：你用pyarrow作为存储后端生成了带特定dtype的DataFrame，然后转成pyarrow Table再转回pandas，虽然看起来dtype都是string[pyarrow]，但pd.testing.assert_frame_equal却抛出了断言错误，提示dtype不同。

问题根源 #

其实问题出在pandas的断言默认会严格检查dtype对象的「同一性」（也就是用Python的is运算符比较），而不仅仅是它们的字符串表示或者逻辑等价性。

你自定义的types_mapper函数每次都会新建一个pd.StringDtype("pyarrow")实例，而原DataFrame里的dtype是之前convert_dtypes(dtype_backend="pyarrow")生成的另一个实例——虽然它们的显示结果一模一样，但在Python里是两个完全独立的对象，所以断言就判定它们「不同」了。

解决方案 #

给你几个可行的解决办法，按推荐程度排序：

1. 直接用dtype_backend替代自定义types_mapper（最推荐）

其实pa.Table.to_pandas()方法已经原生支持dtype_backend参数了，完全不需要自己写类型映射函数。用这个参数指定pyarrow后端，就能自动生成和原DataFrame一致的dtype实例：

import pandas as pd
import pyarrow as pa
import numpy as np
from datetime import datetime, timedelta

df = pd.DataFrame({
   "product_id": pd.Series(
       ["PROD_" + str(np.random.randint(1000, 9999)) for _ in range(100)], 
       dtype=pd.StringDtype(storage="pyarrow")
   ),
   "transaction_timestamp": pd.date_range(
       start=datetime.now() - timedelta(days=30), 
       periods=100, 
       freq='1H'
   ),
   "sales_amount": pd.Series(
       np.round(np.random.normal(500, 150, 100), 2),
       dtype=pd.Float64Dtype()
   ),
   "customer_segment": pd.Series(
       np.random.choice(['Premium', 'Standard', 'Basic'], 100),
       dtype=pd.StringDtype(storage="pyarrow")
   ),
   "is_repeat_customer": pd.Series(
       np.random.choice([True, False], 100, p=[0.3, 0.7])
   )
})

df = df.convert_dtypes(dtype_backend="pyarrow")
# 这里直接用dtype_backend参数
df_pa = pa.Table.from_pandas(df).to_pandas(dtype_backend="pyarrow")
pd.testing.assert_frame_equal(df, df_pa)  # 现在可以正常通过断言了

2. 修改types_mapper，复用原DataFrame的dtype实例

如果你确实需要自定义类型映射逻辑，可以提前保存原DataFrame的dtype，在映射函数里直接返回原实例，而不是新建：

# 先保存原df的所有dtype
original_dtypes = df.dtypes.to_dict()

def types_mapper(pa_type):
    if pa_type == pa.string():
        # 返回原df中对应类型的dtype实例，而不是新建
        return original_dtypes["product_id"]
    # 其他类型可以同理处理，比如Float64等
    elif pa_type == pa.float64():
        return original_dtypes["sales_amount"]
    # 剩下的类型用默认处理
    else:
        return None

df_pa = pa.Table.from_pandas(df).to_pandas(types_mapper=types_mapper)
pd.testing.assert_frame_equal(df, df_pa)

3. 放宽断言的dtype检查（不推荐）

如果只是想临时跳过dtype的同一性检查，可以在断言时设置check_dtype=False，但这个方法可能会掩盖真正的dtype不匹配问题，所以只建议在确定不需要严格检查dtype的场景下使用：

pd.testing.assert_frame_equal(df, df_pa, check_dtype=False)

备注：内容来源于stack exchange，提问作者itstoocold