优化百万级DataFrame基于最近日期匹配填充Visit_occurrence_id的方案

你的核心问题在于原代码用了循环遍历每个person+逐行日期计算，这种方法在70万行数据下时间复杂度极高（近似O(N*M)）。我们可以用Pandas的merge_asof函数来实现向量化匹配，时间复杂度降到O(N log N + M log M)，效率会提升几个数量级。

关键步骤说明 #

• 统一日期类型：先把visit_start_date和measurement_date从object转为datetime，这是高效日期操作的基础。
• 排序数据：merge_asof要求左右表都按匹配键（person_id）+ 日期列排序，这样才能快速定位到最近的匹配项。
• 用merge_asof关联：按person_id分组，匹配measurement_date对应的最近visit_start_date，直接关联对应的visit_occurrence_id。

完整优化代码 #

import pandas as pd
import numpy as np

# 构造原始数据（和你的代码一致）
measurement_data = {'measurement_date':['2017-09-04', '2018-04-24', '2018-05-22', '2019-02-02', '2019-01-28', '2019-05-07', '2018-12-11','2017-04-28'], 
                    'person_id':[1, 2, 2, 1, 3, 1, 3, 3],
                    'visit_occurrence_id':[np.nan,np.nan,np.nan,np.nan,np.nan,np.nan,np.nan,np.nan]}
visit_data = {'visit_occurrence_id':[1, 2, 3, 4, 5], 
              'visit_start_date':['2016-06-01', '2019-05-01', '2016-01-22', '2017-02-14', '2018-05-11'], 
              'person_id':[1, 2, 1, 2, 3]}

measurement = pd.DataFrame(measurement_data)
visit = pd.DataFrame(visit_data)

# ---------------------- 优化核心部分 ----------------------
# 1. 转换日期列为datetime类型
measurement['measurement_date'] = pd.to_datetime(measurement['measurement_date'])
visit['visit_start_date'] = pd.to_datetime(visit['visit_start_date'])

# 2. 对两个表按person_id和日期排序
visit_sorted = visit.sort_values(by=['person_id', 'visit_start_date'])
measurement_sorted = measurement.sort_values(by=['person_id', 'measurement_date'])

# 3. 使用merge_asof匹配最近的visit
# direction='nearest'表示找最接近的日期（无论前后）
# 若只想要measurement_date之前的最近记录，改成direction='backward'
result = pd.merge_asof(
    measurement_sorted,
    visit_sorted[['person_id', 'visit_start_date', 'visit_occurrence_id']],
    left_on='measurement_date',
    right_on='visit_start_date',
    by='person_id',
    direction='nearest'
)

# 4. 恢复原表的列顺序和索引（可选，按需调整）
result = result[measurement.columns].sort_index()
print(result)

为什么这个方法更快？ #

• 原代码的循环+apply是逐行处理，每个measurement行都要遍历当前person的所有visit日期，70万行数据的计算量会呈指数级增长。
• merge_asof是向量化操作，先通过排序把数据整理成有序结构，再用类似二分查找的逻辑快速匹配最近项，整个过程是Pandas内部优化的C级操作，速度比Python循环快几十到上百倍。

可选配置说明 #

• 如果只允许匹配measurement_date之前的最近visit，把direction='nearest'改为direction='backward'。
• 若某个person没有任何visit记录，visit_occurrence_id会保留NaN，完全符合需求。
• 最后一步的sort_index()可以帮你恢复原measurement表的索引顺序，不需要的话可以直接去掉。

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