基于位置区间映射两个DataFrame的性能优化求助

问题背景 #

我有两个大体积DataFrame：

• df1（60万行）：包含染色体、位置及其他附属列

  index  chr     position    other_cols
  1      1          100      ...
  2      1          2100    
  3      1          3300
  4      2          4
  5      2          2200
  6      3          8420
  

• df2（120万行）：包含染色体、区间起始/终止位置

  index          chr       start       stop
      1          1         1           100
      2          1         2000        3000 
      3          1         4000        8000 
      4          2         1           1500
      5          3         20          40000 
  

需要筛选出df1中位置落在df2对应染色体区间内的行，最终输出格式如下：

index  chr     position    matched_start  matched_end  other_cols (from df1)
1      1          100          1           100            ...
2      1          2100         2000        3000  
3      2          4            1           1500
4      3          8420         20          40000 

原代码直接按chr合并后过滤，会产生海量笛卡尔积（比如同染色体的所有df1和df2行强制配对），直接撑爆内存导致笔记本崩溃，分块处理也没效果：

# Merge the DataFrames based on the 'chromosome' column
merged_df = pd.merge(df1, df2, on='chr', how='inner')

# Filter the rows where variant_position lies within the start/end ranges
filtered_df = merged_df[(merged_df['position'] >= merged_df['start']) & (merged_df['position'] <= merged_df['end'])]

高效解决方案 #

方法1：按染色体分组逐处理（降低单次内存负载） #

核心思路：先按染色体分组，每组单独做区间匹配，彻底避免全量笛卡尔积。利用pd.IntervalIndex快速判断位置归属。

import pandas as pd

# 先对df2按染色体分组，每组构建可快速查询的区间索引
df2_interval_map = df2.groupby('chr').apply(
    lambda x: pd.IntervalIndex.from_arrays(x['start'], x['stop'], closed='both')
).to_dict()

# 定义单染色体组的匹配逻辑
def process_chr_group(group):
    chr_num = group['chr'].iloc[0]
    intervals = df2_interval_map.get(chr_num)
    if not intervals:
        return pd.DataFrame()  # 无对应区间的染色体直接返回空
    
    # 找到每个位置对应的区间索引，过滤无匹配的行
    match_indices = group['position'].apply(lambda pos: intervals.get_loc(pos, errors='coerce'))
    valid_rows = group[match_indices.notna()].copy()
    
    if valid_rows.empty:
        return pd.DataFrame()
    
    # 匹配对应的区间起止值
    matched_intervals = intervals[match_indices[valid_rows.index].astype(int)]
    valid_rows['matched_start'] = matched_intervals.left
    valid_rows['matched_end'] = matched_intervals.right
    return valid_rows

# 按染色体分组处理df1，合并结果
final_df = df1.groupby('chr').apply(process_chr_group).reset_index(drop=True)

方法2：用pyjanitor的interval_join（简洁高效） #

pyjanitor库内置了专门的区间合并函数，内部优化了匹配逻辑，无需手动处理分组：

import pandas as pd
import janitor

# 直接执行区间合并，按染色体关联，匹配位置在[start, stop]内的行
final_df = df1.interval_join(
    df2,
    left_on='position',
    right_start='start',
    right_end='stop',
    by='chr',
    how='inner'
).rename(columns={'start': 'matched_start', 'stop': 'matched_end'})

方法3：numpy向量化+分组（极致性能） #

追求最高性能时，用numpy广播减少循环开销，同时控制单组数据量：

import pandas as pd
import numpy as np

final_result = []

# 遍历每个染色体，单独处理对应数据
for chr_num in df1['chr'].unique():
    df1_sub = df1[df1['chr'] == chr_num].reset_index(drop=True)
    df2_sub = df2[df2['chr'] == chr_num].reset_index(drop=True)
    
    if df2_sub.empty:
        continue
    
    # 构建广播矩阵，批量判断位置是否在区间内
    pos_array = df1_sub['position'].values[:, np.newaxis]
    start_array = df2_sub['start'].values[np.newaxis, :]
    stop_array = df2_sub['stop'].values[np.newaxis, :]
    
    # 找到每个位置匹配的第一个区间（多匹配场景可自行调整逻辑）
    match_mask = (pos_array >= start_array) & (pos_array <= stop_array)
    valid_mask = match_mask.any(axis=1)
    match_indices = np.argmax(match_mask, axis=1)[valid_mask]
    
    # 拼接结果
    matched_rows = df1_sub[valid_mask].copy()
    matched_rows['matched_start'] = df2_sub.loc[match_indices, 'start'].values
    matched_rows['matched_end'] = df2_sub.loc[match_indices, 'stop'].values
    final_result.append(matched_rows)

final_df = pd.concat(final_result, ignore_index=True)

核心优化逻辑 #

原方案崩溃的根源是全量笛卡尔积：按chr合并后，同染色体的所有df1、df2行强制配对，临时数据量可能达到千万甚至亿级，直接超出内存上限。

上面的方案都通过分组缩小处理范围，搭配区间索引/向量化操作，既减少了内存占用，又大幅提升了匹配速度。

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