Python对比遍历两个有序生成器时如何判断哪个生成器已耗尽

原代码存在两个核心问题：

1. 分支判断逻辑笔误：第二个elif重复写了A.on[i] < B.on[i]，永远不会触发推进B生成器的逻辑，正确判断条件应为A.on[i] > B.on[i]
2. 未做异常来源区分：所有next()调用放在同一个try块中，触发StopIteration时无法判断是哪个生成器耗尽，也没有实现差异元素收集、B耗尽后剩余A元素的追加逻辑

核心实现思路 #

• 用next(迭代器, 哨兵值)替代无默认参数的next()调用，从根源避免StopIteration异常，通过判断返回值是否为哨兵即可明确知道哪个生成器已耗尽。哨兵用自定义的object()实例实现，不会和SQL返回的真实行数据冲突。
• 采用双指针遍历逻辑：因为两个生成器的数据均按对比键有序，当A的当前键小于B的当前键时，说明该条A数据在B中不存在，直接收入差集；当A的当前键大于B的当前键时，说明该条B数据在A中不存在，直接推进B的指针即可；键相等时说明两边都存在，同时推进两个指针。
• 当检测到B生成器耗尽时，因为数据有序，A中剩余的所有数据都不可能在B中存在，直接把当前持有的A数据、以及A迭代器后续所有数据全部追加到差集即可，终止遍历。
• 差集数据先存在普通列表中，最后一次性转为DataFrame，避免逐行拼接DataFrame带来的性能损耗，适配海量数据场景。

完整实现代码 #

import pandas as pd

def get_a_only_diff(tableAEnum, tableBEnum, compare_key_idx=0):
    """
    对比两个有序生成器的差集，返回仅在tableAEnum中存在的元素组成的DataFrame
    :param tableAEnum: 按对比键升序排列的A表数据生成器
    :param tableBEnum: 按对比键升序排列的B表数据生成器
    :param compare_key_idx: 用于对比的on字段索引，默认取第0位
    """
    # 自定义哨兵值，标记迭代器耗尽，不会和真实行数据冲突
    END = object()
    diff_rows = []

    # 初始化两个指针的当前值，不抛异常
    current_a = next(tableAEnum, END)
    current_b = next(tableBEnum, END)

    while current_a is not END:
        # B已经耗尽，剩余所有A的元素都属于差集
        if current_b is END:
            diff_rows.append(current_a)
            # 批量追加A剩余的所有元素
            diff_rows.extend(tableAEnum)
            break
        
        a_key = current_a.on[compare_key_idx]
        b_key = current_b.on[compare_key_idx]

        if a_key < b_key:
            # A当前元素B中不存在，收入差集，推进A指针
            diff_rows.append(current_a)
            current_a = next(tableAEnum, END)
        elif a_key > b_key:
            # B当前元素A中不存在，只推进B指针
            current_b = next(tableBEnum, END)
        else:
            # 两边元素匹配，同时推进两个指针
            current_a = next(tableAEnum, END)
            current_b = next(tableBEnum, END)
    
    # 一次性转成DataFrame返回
    return pd.DataFrame(diff_rows)

注意事项 #

该实现全程只在内存中保留当前遍历的两个行对象、以及最终的差集结果，不需要全量加载两个SQL查询的所有数据，内存占用稳定，完全适配海量有序数据的差集计算场景。如果对比的是联合键，只需要把a_key/b_key的取值逻辑改成对应元组即可，整体逻辑不需要调整。

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