Polars中assert_frame_equal能否在首个数据不匹配时立即终止？

Polars官方的assert_frame_equal方法在处理百万级数据集时，会完成全量行对比才抛出错误，导致耗时过长。要实现首个不匹配项出现即终止并报错，可以自行实现轻量化的检查逻辑，利用Polars的矢量化操作快速定位首个差异：

自定义检查函数 #

import polars as pl

def assert_frame_first_mismatch(src_df: pl.DataFrame, tgt_df: pl.DataFrame, check_dtype: bool = False):
    # 校验列名一致性
    if src_df.columns != tgt_df.columns:
        raise AssertionError(f"列名不匹配：源表列{src_df.columns}，目标表列{tgt_df.columns}")
    
    # 按需校验数据类型
    if check_dtype:
        for col in src_df.columns:
            if src_df[col].dtype != tgt_df[col].dtype:
                raise AssertionError(f"列{col}数据类型不匹配：源表{src_df[col].dtype}，目标表{tgt_df[col].dtype}")
    
    # 生成每行匹配状态的掩码（True表示该行所有列都匹配）
    match_mask = (src_df == tgt_df).all(axis=1)
    
    # 获取首个不匹配行的索引
    first_mismatch_idx = match_mask.arg_false()
    
    if first_mismatch_idx != -1:
        # 提取并展示首个差异行信息
        src_row = src_df[first_mismatch_idx].to_dict()
        tgt_row = tgt_df[first_mismatch_idx].to_dict()
        raise AssertionError(
            f"首个不匹配行位于索引{first_mismatch_idx}\n"
            f"源表行：{src_row}\n"
            f"目标表行：{tgt_row}"
        )

使用方式 #

1. 如果不需要忽略行顺序，直接传入两个DataFrame即可：

assert_frame_first_mismatch(src_df, tgt_df, check_dtype=False)

2. 如果需要忽略行顺序（对应原代码的check_row_order=False），先对两个DataFrame按相同规则排序，再执行检查：

# 按所有列排序（也可指定特定列作为排序键）
sorted_src = src_df.sort(src_df.columns)
sorted_tgt = tgt_df.sort(src_df.columns)
assert_frame_first_mismatch(sorted_src, sorted_tgt, check_dtype=False)

优势说明 #

• 利用Polars矢量化操作实现高效对比，比Python循环逐行检查快得多
• 找到首个不匹配项后立即终止并抛出错误，无需遍历全量数据
• 直接返回首个差异行的具体信息，便于快速定位问题

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