如何在Pandas中按id匹配并依据日期区间关联phase列

需求说明 #

现有两个DataFrame：

• foo1：存储每个id对应不同phase的日期区间（date_start到date_end）
• foo2：存储每个id的具体日期记录

需要为foo2新增phase列，规则为：

• 匹配相同id
• 当date落在对应id的某一phase日期区间内时，填充该phase值
• 若date不在任何对应id的日期区间内，填充NaN

实现方案 #

步骤1：转换日期类型 #

首先将所有日期列从字符串转换为datetime类型，才能进行区间比较：

import pandas as pd

foo1 = pd.DataFrame({'id':[1,1,2,2],
              'phase':['Pre','Post','Pre','Post'],
              'date_start': ['2022-07-24', '2022-12-25', '2022-09-30', '2022-12-25'],
               'date_end': ['2022-07-30', '2023-03-07', '2022-10-05', '2023-03-04']})

foo2 = pd.DataFrame({'id': [1,1,1,1,
                     2,2,2,2],
              'date': ['2022-07-24', '2022-07-25', '2022-12-26', '2023-01-01',
                      '2022-10-04', '2022-11-25', '2022-12-26', '2023-03-01']})

# 转换日期列类型
foo1['date_start'] = pd.to_datetime(foo1['date_start'])
foo1['date_end'] = pd.to_datetime(foo1['date_end'])
foo2['date'] = pd.to_datetime(foo2['date'])

步骤2：逐行匹配phase #

定义一个函数，为foo2的每一行匹配对应的phase，然后通过apply批量处理：

def match_phase(row):
    # 筛选当前id下，日期在区间内的phase记录
    matched = foo1[(foo1['id'] == row['id']) & 
                   (foo1['date_start'] <= row['date']) & 
                   (foo1['date_end'] >= row['date'])]['phase']
    # 有匹配结果则返回第一个，否则返回NaN
    return matched.iloc[0] if not matched.empty else pd.NA

# 为foo2添加phase列
foo2['phase'] = foo2.apply(match_phase, axis=1)

步骤3：查看结果 #

执行后foo2的输出如下：

id       date phase
0   1 2022-07-24   Pre
1   1 2022-07-25   Pre
2   1 2022-12-26  Post
3   1 2023-01-01  Post
4   2 2022-10-04   Pre
5   2 2022-11-25  <NA>
6   2 2022-12-26  Post
7   2 2023-03-01  Post

高效优化方案（适用于大数据量） #

如果数据量较大，逐行apply效率较低，可以通过笛卡尔积合并+布尔索引的方式实现：

# 按id合并两个DataFrame
merged = foo2.merge(foo1, on='id', how='left')
# 筛选日期在区间内的记录
mask = (merged['date'] >= merged['date_start']) & (merged['date'] <= merged['date_end'])
# 将匹配到的phase映射回原foo2，未匹配的填充NaN
foo2['phase'] = merged.loc[mask, 'phase'].reindex(foo2.index)

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