如何高效实现两个DataFrame基于时间窗口与质量的最近值匹配

原实现性能问题分析 #

• 每次调用匹配函数都对全量df_ID做两次排序，属于完全重复的冗余计算，是性能差的核心原因
• 逐行apply的Python循环效率极低，无法处理大规模数据
• 未处理searchsorted返回索引越界、时间窗口内无匹配的边界情况，容易触发索引报错

向量化优化方案 #

核心思路是把所有可预计算的操作提到循环外，用numpy向量化操作批量计算时间窗口边界，大幅减少冗余计算：

import pandas as pd
import numpy as np

# 预排序df_ID，全局仅执行1次
df_ID_sorted = df_ID.sort_values("time", ignore_index=True)
id_times = df_ID_sorted["time"].values
id_masses = df_ID_sorted["mass"].values
id_ids = df_ID_sorted["ID"].values

# 批量提取df_quant的匹配字段
quant_times = df_quant["time"].values
quant_masses = df_quant["mass"].values

# 向量化计算每个样本对应的时间窗口左右边界索引
left_bounds = np.searchsorted(id_times, quant_times - 9, side="left")
right_bounds = np.searchsorted(id_times, quant_times + 9, side="right")

# 匹配最接近的mass对应的ID
result_ids = []
for i in range(len(df_quant)):
    l, r = left_bounds[i], right_bounds[i]
    # 处理窗口无匹配的情况，可根据需求调整默认值
    if l >= r:
        result_ids.append(None)
        continue
    # 取当前窗口内的mass和ID数组
    window_masses = id_masses[l:r]
    window_ids = id_ids[l:r]
    # 找最小差对应的索引
    min_idx = np.argmin(np.abs(window_masses - quant_masses[i]))
    result_ids.append(window_ids[min_idx])

df_quant["ID"] = result_ids

超大规模数据可选加速方案 #

如果数据量达到百万级以上，可引入numba编译核心循环，速度还能再提升10~100倍：

from numba import njit

# 用numba编译循环为机器码，关闭GIL提升并行效率
@njit
def match_id_numba(left_bounds, right_bounds, id_masses, id_ids, quant_masses):
    n = len(quant_masses)
    res = [None] * n
    for i in range(n):
        l, r = left_bounds[i], right_bounds[i]
        if l >= r:
            continue
        min_diff = 1e18
        best_idx = l
        for j in range(l, r):
            curr_diff = abs(id_masses[j] - quant_masses[i])
            if curr_diff < min_diff:
                min_diff = curr_diff
                best_idx = j
        res[i] = id_ids[best_idx]
    return res

df_quant["ID"] = match_id_numba(left_bounds, right_bounds, id_masses, id_ids, quant_masses)

性能提升效果 #

原有实现时间复杂度为 O(N * M log M)（N为df_quant长度，M为df_ID长度），优化后时间复杂度为 O(M log M + N log M + N*K)（K为平均时间窗口内的样本数），常规场景下性能提升可达几十到上百倍。

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