基于itertools combinations的配对代码大数据场景性能优化求助

问题背景 #

需从50万+行的DataFrame中生成同一招标项目下的企业ID配对组合，原代码执行效率极低（耗时数小时），需在保证结果一致的前提下提速。

原代码核心低效点：

• 遍历所有tender_product条目（含大量重复值），重复计算同一招标的配对
• 多次执行dataframe[dataframe['tender_product'] == item]筛选，IO开销大
• 后期冗余的query、drop_duplicates操作额外消耗资源

优化方案1：按招标分组批量生成配对 #

利用pandas.groupby直接按招标项目分组，避免重复遍历与筛选，结合itertools.combinations批量生成配对。

import pandas as pd
from itertools import combinations
from itertools import chain

def get_all_pairs_optimized(df):
    # 按招标项目分组，获取每个项目下的企业ID列表
    grouped = df.groupby('tender_product')['ID_name']
    
    # 批量生成所有配对的迭代器
    pairs_iter = chain.from_iterable(
        ((c1, c2, tender) for c1, c2 in combinations(ids, 2))
        for tender, ids in grouped
    )
    
    # 转换为目标DataFrame
    result_df = pd.DataFrame(pairs_iter, columns=['ID_name1', 'ID_name2', 'tender_product'])
    
    # 生成pair与反向pair（修正原代码的笔误）
    result_df['pair'] = result_df['ID_name1'].astype(str) + '_' + result_df['ID_name2'].astype(str)
    result_df['reversed_pair'] = result_df['ID_name2'].astype(str) + '_' + result_df['ID_name1'].astype(str)
    
    return result_df

优化说明 #

• groupby一次性完成分组，避免重复筛选同一招标的企业ID
• chain.from_iterable合并分组后的配对迭代器，减少内存占用
• 移除冗余操作：combinations本身不会生成相同ID的配对，无需query('ID_name1 != ID_name2')；分组生成的配对天然无重复，无需drop_duplicates

优化方案2：用numpy向量化操作进一步提速 #

针对超大规模数据，利用numpy的广播机制生成配对，替代Python层面循环，进一步压缩耗时。

import pandas as pd
import numpy as np

def get_all_pairs_vectorized(df):
    result_list = []
    
    for tender, group in df.groupby('tender_product')['ID_name']:
        ids = group.values
        n = len(ids)
        if n < 2:
            continue
        
        # 生成所有i<j的索引对（numpy底层C实现，速度远快于Python循环）
        idx = np.triu_indices(n, k=1)
        id1 = ids[idx[0]]
        id2 = ids[idx[1]]
        
        # 组装临时DataFrame并加入结果列表
        temp_df = pd.DataFrame({
            'ID_name1': id1,
            'ID_name2': id2,
            'tender_product': tender
        })
        result_list.append(temp_df)
    
    # 合并所有结果
    result_df = pd.concat(result_list, ignore_index=True)
    
    # 生成pair列
    result_df['pair'] = result_df['ID_name1'].astype(str) + '_' + result_df['ID_name2'].astype(str)
    result_df['reversed_pair'] = result_df['ID_name2'].astype(str) + '_' + result_df['ID_name1'].astype(str)
    
    return result_df

优化说明 #

• numpy.triu_indices直接生成索引对，底层为C实现，比itertools.combinations更快
• 批量生成ID对，减少Python循环的开销
• 跳过企业数量<2的招标项目，避免无效运算

优化方案3：内存预处理（可选） #

若DataFrame内存占用过高，可先转换ID列类型，减少内存开销间接提升运算速度：

# 将ID列从int64转为int32（若ID数值范围允许）
df['ID_name'] = df['ID_name'].astype('int32')
# 若ID为字符串，统一为string类型（比object类型更节省内存）
df['ID_name'] = df['ID_name'].astype('string')

结果一致性验证 #

用样例数据测试，优化后代码输出与原代码完全一致：

初始数据 #

输出结果 #

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