如何高效优化Pandas DataFrame中产品跨仓库两两对比代码？

问题背景 #

现有50万条记录的数据集，包含约10000个产品，需要按仓库对比任意两个不同产品的pct_total指标（取两者中的较小值）。原实现采用双重for循环，逻辑正确但运行时长超过1天，需优化性能。

原代码的性能瓶颈 #

原代码核心问题：

• 双重循环遍历所有产品对，每次循环内重复执行concat、merge等高开销操作
• 循环中逐步拼接结果DataFrame，频繁触发内存重分配，性能急剧下降
• 每次循环都重复处理仓库集合，存在大量冗余计算

优化方案 #

以下是基于Pandas矢量化操作和预计算的高效实现，核心思路是先构建完整的仓库-产品数据矩阵，再批量处理产品对：

import pandas as pd
from itertools import permutations

# 示例数据
data = {
    'productid': ['hyundai', 'hyundai', 'hyundai', 'kia','kia', 'kia'],
    'warehouse': ['New Jersey', 'New York', 'California', 'New Jersey', 'New York', 'California'],
    'pct_total': [35, 45, 20,65,55,80]
}
df = pd.DataFrame(data)

# 步骤1：构建全量仓库-产品笛卡尔积，确保每个仓库每个产品都有记录
all_warehouses = df['warehouse'].unique()
all_products = df['productid'].unique()
warehouse_product = pd.MultiIndex.from_product(
    [all_warehouses, all_products], 
    names=['warehouse', 'productid']
).to_frame(index=False)

# 左连接原数据，得到完整的仓库-产品-pct_total矩阵（缺失值为NaN）
full_df = warehouse_product.merge(df, on=['warehouse', 'productid'], how='left')

# 步骤2：转换为宽表，仓库为行，产品为列，值为pct_total
wide_df = full_df.pivot(index='warehouse', columns='productid', values='pct_total')

# 步骤3：生成所有不重复的产品对（排除自身对比）
product_pairs = list(permutations(all_products, 2))

# 步骤4：批量处理每个产品对，计算结果
result_list = []
for prod1, prod2 in product_pairs:
    # 提取当前产品对的pct_total数据
    pair_data = wide_df[[prod1, prod2]].reset_index()
    pair_data.columns = ['warehouse', 'pct_total_x', 'pct_total_y']
    # 计算两者中的较小值
    pair_data['lowervalue'] = pair_data[['pct_total_x', 'pct_total_y']].min(axis=1)
    # 添加产品标识列
    pair_data['productid_x'] = prod1
    pair_data['productid_y'] = prod2
    # 加入结果列表
    result_list.append(pair_data)

# 合并所有结果
final_output = pd.concat(result_list, ignore_index=True)
print(final_output)

优化效果说明 #

1. 预计算减少重复操作：提前构建完整的仓库-产品矩阵，避免循环中重复处理仓库集合和执行merge操作
2. 矢量化操作替代循环内的小数据处理：利用Pandas的pivot和列提取操作，将循环内的低效merge转为内存内的快速数据访问
3. 批量合并结果：先将所有子结果存入列表，最后一次性合并，避免循环中逐步拼接DataFrame的内存开销

对于10000个产品的场景，该方案循环次数仍为10000*9999，但每次循环仅为列提取和矢量化计算，性能远优于原循环内的多次IO型操作，运行时长可从天级压缩到小时级甚至更短。

内容的提问来源于stack exchange，提问作者Stumbling Through Data Science