Pandas中基于多条件和时间窗口关联数据的高效方法

import pandas as pd
import janitor # 确保已安装：pip install pyjanitor

# 交易数据DataFrame
trade = pd.DataFrame({
    'date': ['2019-08-31', '2019-09-01', '2019-09-04'],
    'person': [1, 1, 2],
    'code': [123, 123, 456],
    'value1': [1, 2, 3]
})

# 浏览历史DataFrame
view = pd.DataFrame({
    'date': ['2019-08-29', '2019-08-29', '2019-08-30', '2019-08-31', '2019-09-01',
             '2019-09-01', '2019-09-01', '2019-09-02', '2019-09-03'],
    'person': [1, 1, 1, 2, 1, 2, 2, 1, 2],
    'code': [123, 456, 123, 456, 123, 123, 456, 123, 456],
    'value': [1, 2, 3, 4, 5, 6, 7, 8, 9]
})

# 将日期列转换为datetime类型，以便进行时间计算和比较
trade['date'] = pd.to_datetime(trade['date'])
view['date'] = pd.to_datetime(view['date'])

# 使用janitor.conditional_join进行数据关联
out_janitor = (
    trade
    .assign(start_date=lambda d: d['date'].sub(pd.DateOffset(days=7))) # 步骤1: 计算时间窗口起始点
    .conditional_join(
        view.rename(columns={'date': 'view_dates', 'value': 'view_values'}), # 步骤2: 重命名view列
        ('start_date', 'view_dates', '<='), # 条件1: 浏览日期 >= 交易日期 - 7天
        ('date', 'view_dates', '>='),       # 条件2: 浏览日期 <= 交易日期
        ('person', 'person', '=='),         # 条件3: person相等
        ('code', 'code', '=='),             # 条件4: code相等
        right_columns=['view_dates', 'view_values'] # 保留右侧DataFrame的哪些列
    )
    .drop(columns='start_date') # 步骤3: 删除辅助列
    .assign(view_dates=lambda d: d['view_dates'].dt.strftime('%Y-%m-%d')) # 步骤3: 格式化日期
    .groupby(list(trade.columns), as_index=False).agg(list) # 步骤4: 按原始trade列分组并聚合
)

print("使用janitor.conditional_join的结果:")
print(out_janitor)

# 纯Pandas解决方案进行数据关联
out_pandas = (
    trade
    .merge(
        view.rename(columns={'date': 'view_dates', 'value': 'view_values'}), # 步骤1: 重命名view列
        on=['person', 'code'], # 步骤1: 基于person和code进行等值合并
        how='left' # 使用左连接以保留所有交易，即使没有匹配的浏览记录
    )
    .loc[lambda d: (d['date'].gt(d['view_dates'])) & # 步骤2: 浏览日期在交易日期之前
                   (d['date'].sub(pd.DateOffset(days=7)).le(d['view_dates'])) # 步骤2: 浏览日期在窗口起始点之后
         ]
    .assign(view_dates=lambda d: d['view_dates'].dt.strftime('%Y-%m-%d')) # 步骤3: 格式化日期
    .groupby(list(trade.columns), as_index=False).agg(list) # 步骤4: 按原始trade列分组并聚合
)

print("\n纯Pandas解决方案的结果:")
print(out_pandas)