如何加速Python dataset库的重复行检测？优化is_duplicate_free函数

问题背景 #

当前is_duplicate_free函数在实际规模的表上性能极差，测试程序在工作站上耗时45秒。需要保持数据库无关性（兼容SQLite、PostgreSQL），且适配Python 3版本，基于dataset库实现。

原始函数的核心问题：遍历所有去重行计数、调用len(table)触发全表扫描，都会把大量数据拉到Python内存处理，效率极低。

优化后的代码 #

def is_duplicate_free(table) -> bool:
    """检查表中是否无重复行（忽略主键），返回True/False"""
    # 自动识别主键列，兼容自动生成的'id'或自定义主键
    primary_keys = [col.name for col in table.table.primary_key.columns]
    # 筛选出用于判断重复的非主键列
    non_pk_columns = [col for col in table.columns if col not in primary_keys]
    
    if not non_pk_columns:
        # 无有效非主键列，不存在重复可能
        return True
    
    # 构造标准SQL，让数据库直接完成计数统计
    query = f"""
        SELECT 
            (SELECT COUNT(*) FROM {table.name}) AS total_rows,
            (SELECT COUNT(*) FROM (SELECT DISTINCT {', '.join(non_pk_columns)} FROM {table.name})) AS distinct_rows
    """
    # 执行查询并获取结果
    stats = table.db.query(query).first()
    
    return stats['total_rows'] == stats['distinct_rows']

优化细节说明 #

• 数据库端计算：将计数和去重操作完全交给数据库引擎执行，避免大量数据从数据库传输到Python内存，这是性能提升的关键
• 通用主键处理：通过table.table.primary_key.columns自动获取主键，不再硬编码id，适配自定义主键的场景
• SQL标准兼容：使用COUNT(*)和SELECT DISTINCT等SQL标准语法，完美兼容SQLite和PostgreSQL
• 边界情况处理：当表中没有非主键列时，直接返回True，避免无效查询

替换后，测试程序的执行时间会从几十秒级压缩到毫秒级，同时保持原有断言逻辑的正确性。

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