如何在Polars中按Symbol分组实现基于下一个可用日期的asof join

需求：将events表中的每个Earnings_Date，匹配到fr表中同Symbol下相同或之后的第一个可用Date，并关联对应的Event数据。

问题分析 #

原代码未按Symbol分组处理日期匹配，多Symbol场景下会出现匹配错误；且join_asof默认匹配的是<=目标日期的最近行，不符合"相同或之后首个日期"的需求。

解决方案代码 #

1. 构造多Symbol测试数据

import polars as pl

# fr表：包含多个Symbol的日期序列
fr = pl.DataFrame(
    {
        'Symbol': ['A']*5 + ['B']*4,
        'Date': ['2010-08-29', '2010-09-01', '2010-09-05', '2010-11-30', '2010-12-02'] +
                ['2010-07-01', '2010-08-01', '2010-09-01', '2010-10-01'],
    }
).with_columns(pl.col('Date').str.to_date('%Y-%m-%d')).sort(['Symbol', 'Date'])

# events表：包含多个Symbol的事件日期与事件值
events = pl.DataFrame(
    {
         'Symbol': ['A']*3 + ['B']*2,
         'Earnings_Date': ['2010-06-01', '2010-09-01', '2010-12-01'] +
                          ['2010-07-15', '2010-09-05'],
         'Event': [1, 4, 7, 2, 5],
     }
).with_columns(pl.col('Earnings_Date').str.to_date('%Y-%m-%d')).sort(['Symbol', 'Earnings_Date'])

2. 按Symbol分组匹配日期并关联

# 提取每个Symbol对应的fr日期列表
fr_symbol_dates = fr.group_by('Symbol').agg(pl.col('Date').alias('fr_dates'))

# 为每个事件日期找到对应Symbol下首个符合条件的fr日期
events_matched = events.join(fr_symbol_dates, on='Symbol').with_columns(
    # 使用search_sorted的left模式，定位第一个>=Earnings_Date的日期索引
    pl.col('fr_dates').list.search_sorted(pl.col('Earnings_Date'), side='left').alias('match_index')
).with_columns(
    # 根据索引提取匹配日期，索引等于列表长度则表示无符合条件的日期（设为None）
    pl.when(pl.col('match_index') < pl.col('fr_dates').list.len())
      .then(pl.col('fr_dates').list.get(pl.col('match_index')))
      .alias('matched_fr_date')
).drop(['fr_dates', 'match_index'])

# 将匹配结果关联回fr表，得到最终数据
final_result = fr.join(events_matched, 
                       left_on=['Symbol', 'Date'], 
                       right_on=['Symbol', 'matched_fr_date'], 
                       how='left')

print(final_result)

结果说明 #

• 对于events中Symbol=A、Earnings_Date=2010-12-01，成功匹配到fr中2010-12-02；
• 对于Symbol=B、Earnings_Date=2010-09-05，匹配到fr中2010-10-01；
• 无符合条件日期的事件会显示None，保留fr表的所有行。

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