高效实现Pandas中两个DataFrame的匹配与合并需求

针对你遇到的嵌套循环效率低、全量merge内存溢出问题，这里提供两种高效的解决方案，核心思路是避免全量笛卡尔积计算，通过分组缩小处理范围：

方案一：Pandas分组+Numpy广播（适合中等数据量） #

先按chr+family分组，在每个组内完成pos范围匹配和sample过滤，这样既减少了计算量，又控制了单批次内存占用。

步骤与代码： #

1. 生成分组键，将chr和family组合成唯一标识，方便后续分组：

import pandas as pd
import numpy as np

# 为两个DataFrame添加分组键
FR['group_key'] = FR['chr'].astype(str) + '_' + FR['family'].astype(str)
UNPAIRED['group_key'] = UNPAIRED['chr'].astype(str) + '_' + UNPAIRED['family'].astype(str)

2. 定义组内处理函数，用Numpy广播加速范围匹配（比逐行判断快一个数量级）：

def process_single_group(fr_subset, unpaired_subset):
    # 提取当前组的pos和sample数组，用广播生成匹配mask
    fr_pos_array = fr_subset['pos'].values[:, np.newaxis]
    unpaired_pos_array = unpaired_subset['pos'].values
    # 匹配pos在±200范围内的mask
    pos_match = (unpaired_pos_array >= fr_pos_array - 200) & (unpaired_pos_array <= fr_pos_array + 200)
    
    # 匹配sample不同的mask
    fr_sample_array = fr_subset['sample'].values[:, np.newaxis]
    unpaired_sample_array = unpaired_subset['sample'].values
    sample_match = unpaired_sample_array != fr_sample_array
    
    # 合并两个mask，找到所有符合条件的行索引对
    valid_pairs = np.where(pos_match & sample_match)
    fr_rows = fr_subset.iloc[valid_pairs[0]].reset_index(drop=True)
    unpaired_rows = unpaired_subset.iloc[valid_pairs[1]].reset_index(drop=True)
    
    # 合并FR和UNPAIRED的匹配行
    return pd.concat([fr_rows, unpaired_rows], axis=1)

3. 遍历共同分组，收集结果：

# 只处理两个DataFrame都存在的分组，减少无效计算
common_groups = set(FR['group_key']) & set(UNPAIRED['group_key'])

nonref_list = []
for group in common_groups:
    fr_group = FR[FR['group_key'] == group]
    unpaired_group = UNPAIRED[UNPAIRED['group_key'] == group]
    if fr_group.empty or unpaired_group.empty:
        continue
    # 处理当前分组并添加到结果列表
    nonref_list.append(process_single_group(fr_group, unpaired_group))

# 合并所有分组结果得到最终的NONREF_ALL
NONREF_ALL = pd.concat(nonref_list, ignore_index=True)
# 清理临时分组键（如果不需要保留）
NONREF_ALL.drop(columns='group_key', inplace=True)

方案二：Dask并行处理（适合超大数据量） #

如果数据量极大，Pandas单进程内存仍不够，用Dask将数据拆分到多个分区并行处理，避免一次性加载全量数据到内存。

步骤与代码： #

import dask.dataframe as dd

# 将Pandas DataFrame转换为Dask DataFrame，按CPU核心数设置分区数
dd_FR = dd.from_pandas(FR, npartitions=4)
dd_UNPAIRED = dd.from_pandas(UNPAIRED, npartitions=4)

# 添加分组键和来源标记
dd_FR['group_key'] = dd_FR['chr'].astype(str) + '_' + dd_FR['family'].astype(str)
dd_FR['source'] = 'FR'
dd_UNPAIRED['group_key'] = dd_UNPAIRED['chr'].astype(str) + '_' + dd_UNPAIRED['family'].astype(str)
dd_UNPAIRED['source'] = 'UNPAIRED'

# 合并两个Dask DataFrame，按group_key分组处理
combined_dd = dd.concat([dd_FR, dd_UNPAIRED])

# 定义适配Dask的分组处理函数（内部逻辑和Pandas版本一致）
def dask_group_processor(df):
    fr_sub = df[df['source'] == 'FR']
    unpaired_sub = df[df['source'] == 'UNPAIRED']
    if fr_sub.empty or unpaired_sub.empty:
        return pd.DataFrame()
    
    fr_pos = fr_sub['pos'].values[:, np.newaxis]
    unpaired_pos = unpaired_sub['pos'].values
    pos_match = (unpaired_pos >= fr_pos - 200) & (unpaired_pos <= fr_pos + 200)
    
    fr_sample = fr_sub['sample'].values[:, np.newaxis]
    unpaired_sample = unpaired_sub['sample'].values
    sample_match = unpaired_sample != fr_sample
    
    valid_pairs = np.where(pos_match & sample_match)
    fr_rows = fr_sub.iloc[valid_pairs[0]].reset_index(drop=True)
    unpaired_rows = unpaired_sub.iloc[valid_pairs[1]].reset_index(drop=True)
    
    return pd.concat([fr_rows, unpaired_rows], axis=1)

# 执行分组计算并转换为Pandas DataFrame
result_dd = combined_dd.groupby('group_key').apply(
    dask_group_processor,
    meta=NONREF_ALL.dtypes  # 提前指定输出结构，加速计算
).compute()

NONREF_ALL = result_dd.reset_index(drop=True).drop(columns=['group_key', 'source'])

为什么这两个方案更高效？ #

• 避免了全量merge生成笛卡尔积：全量merge会把所有chr+family相同的行两两配对，行数可能暴涨到百万甚至千万级，直接撑爆内存；分组后只在组内处理，单组数据量小，内存可控。
• 用Numpy广播替代逐行循环：Numpy的向量化操作比Python循环快几十倍，大幅降低时间消耗。
• Dask的并行分区处理：把数据拆分成多个小批次，分布式处理，不会一次性占用全部内存。

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