Pandas/Pyspark中对比两个DataFrame并高效打标拆分行的方案

针对百万级数据的性能瓶颈，核心思路是利用pandas的矢量化操作、分组累积量计算和近似匹配，彻底避免嵌套循环。以下是分步实现方案：

1. 预处理与排序 #

先统一两个DataFrame的结构，计算组内累积数量（用于区间匹配），同时标记数据来源：

import pandas as pd
import numpy as np

# 对expected_orders预处理
expected_orders = expected_orders.sort_values(['country', 'product', 'date'])
expected_orders['start_cum'] = expected_orders.groupby(['country', 'product'])['qty'].cumsum().shift(fill_value=0)
expected_orders['end_cum'] = expected_orders.groupby(['country', 'product'])['qty'].cumsum()
expected_orders['source'] = 'expected'

# 对actual_orders预处理
actual_orders = actual_orders.sort_values(['country', 'product', 'date'])
actual_orders['start_cum'] = actual_orders.groupby(['country', 'product'])['qty'].cumsum().shift(fill_value=0)
actual_orders['end_cum'] = actual_orders.groupby(['country', 'product'])['qty'].cumsum()
actual_orders['source'] = 'actual'

• start_cum/end_cum：标记当前行的数量在组内的累积区间（比如前3行总数量是100，当前行数量是50，则start_cum=100，end_cum=150）
• 排序确保后续分组匹配的顺序正确

2. 分组区间匹配与打标 #

定义分组处理函数，用merge_asof实现矢量化的区间匹配，替代循环拆分：

def process_group(group):
    # 拆分当前组的expected和actual数据
    exp = group[group['source'] == 'expected']
    act = group[group['source'] == 'actual']
    
    # 提取所有累积量端点，生成匹配用的数量片段
    all_cum_points = np.sort(np.unique(np.concatenate([exp[['start_cum', 'end_cum']].values.flatten(),
                                                      act[['start_cum', 'end_cum']].values.flatten()])))
    segments = pd.DataFrame({
        'start': all_cum_points[:-1],
        'end': all_cum_points[1:]
    })
    segments['qty'] = segments['end'] - segments['start']
    segments['country'] = group['country'].iloc[0]
    segments['product'] = group['product'].iloc[0]
    
    # 匹配片段对应的expected日期
    segments = pd.merge_asof(segments, exp[['start_cum', 'date']],
                             left_on='start', right_on='start_cum', direction='backward')
    segments.rename(columns={'date': 'expected_date'}, inplace=True)
    
    # 匹配片段对应的actual日期
    segments = pd.merge_asof(segments, act[['start_cum', 'date']],
                             left_on='start', right_on='start_cum', direction='backward')
    segments.rename(columns={'date': 'actual_date'}, inplace=True)
    
    # 按规则打标签
    segments['label'] = np.select(
        [
            (segments['expected_date'].isna()) | (segments['actual_date'].isna()),
            segments['expected_date'] == segments['actual_date'],
            segments['expected_date'] < segments['actual_date'],
            segments['expected_date'] > segments['actual_date']
        ],
        [
            '无对应日期',
            '相同日期',
            '早于对应日期',
            '晚于对应日期'
        ]
    )
    
    # 整理输出列
    return segments[['country', 'product', 'expected_date', 'actual_date', 'qty', 'label']]

# 合并数据后分组处理
combined = pd.concat([expected_orders, actual_orders], ignore_index=True)
final_result = combined.groupby(['country', 'product'], group_keys=False).apply(process_group).reset_index(drop=True)

3. 性能优化说明 #

• 矢量化替代循环：用merge_asof（排序后的近似匹配）和np.select（矢量化条件判断）替代嵌套循环，时间复杂度从O(n*m)降至O(n log n + m log m)，百万级数据可在数分钟内完成处理
• 内存控制：如果数据量过大导致内存不足，可以按country分块处理，每次仅加载部分数据到内存
• 避免冗余计算：预处理阶段一次性计算累积量，避免分组内重复计算

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