Polars按最接近日期连接

Question

我有一个数据框（叫做events），我想把它和另一个数据框（叫做fr）合并，合并的依据是日期和符号。这里的日期不一定会有重叠。events中的日期只会和fr中同一天或更晚的日期的第一次出现匹配，所以如果事件的日期是2010年12月1日，它会在同一天合并，如果那天没有，就会合并到下一个可用的日期（比如2010年12月2日）。

我试过用search_sorted和join_asof来实现这个，但我还想按符号这一列进行分组，而且这也不是一个标准的合并方式。这种方法对于单个符号来说有点效果。

fr = pl.DataFrame(
    {
        'Symbol': ['A']*5,
        'Date': ['2010-08-29', '2010-09-01', '2010-09-05',
                 '2010-11-30', '2010-12-02'],
    }
).with_columns(pl.col('Date').str.strptime(pl.Date, '%Y-%m-%d')).with_row_index().set_sorted("Date")

events = pl.DataFrame(
    {
         'Symbol': ['A']*3,
         'Earnings_Date': ['2010-06-01', '2010-09-01', '2010-12-01'],
         'Event': [1, 4, 7],
     }
).with_columns(pl.col('Earnings_Date').str.strptime(pl.Date, '%Y-%m-%d')).set_sorted("Earnings_Date")

idx = fr["Date"].search_sorted(events["Earnings_Date"], "left")

fr = fr.with_columns(
    pl.when(
        pl.col("index").is_in(idx)
    )
    .then(True)
    .otherwise(False)
    .alias("Earnings")
)

fr = fr.join_asof(events, by="Symbol", left_on="Date", right_on="Earnings_Date")
fr = fr.with_columns(
    pl.when(
       pl.col("Earnings") == True
    )
    .then(pl.col("Event"))
    .otherwise(False)
    .alias("Event")
)

Answer 1

听起来你在使用 pl.DataFrame.join_asof 的方向上是对的。要按符号分组，可以使用 by 这个参数。

(
    fr
    .join_asof(
        events,
        left_on="Date", right_on="Earnings_Date",
        by="Symbol",
    )
)

shape: (5, 5)
┌───────┬────────┬────────────┬───────────────┬───────┐
│ index ┆ Symbol ┆ Date       ┆ Earnings_Date ┆ Event │
│ ---   ┆ ---    ┆ ---        ┆ ---           ┆ ---   │
│ u32   ┆ str    ┆ date       ┆ date          ┆ i64   │
╞═══════╪════════╪════════════╪═══════════════╪═══════╡
│ 0     ┆ A      ┆ 2010-08-29 ┆ 2010-06-01    ┆ 1     │
│ 1     ┆ A      ┆ 2010-09-01 ┆ 2010-09-01    ┆ 4     │
│ 2     ┆ A      ┆ 2010-09-05 ┆ 2010-09-01    ┆ 4     │
│ 3     ┆ A      ┆ 2010-11-30 ┆ 2010-09-01    ┆ 4     │
│ 4     ┆ A      ┆ 2010-12-02 ┆ 2010-12-01    ┆ 7     │
└───────┴────────┴────────────┴───────────────┴───────┘

现在，我明白你希望每个事件最多只匹配一次。我觉得仅靠 join_asof 是做不到的。不过，我们可以把所有与前一行相等的事件行设置为 Null。为此，可以使用 pl.when().then() 这个结构。

(
    fr
    .join_asof(
        events,
        left_on="Date", right_on="Earnings_Date",
        by="Symbol",
    )
    .with_columns(
        pl.when(
            pl.col("Earnings_Date", "Event").is_first_distinct()
        ).then(
            pl.col("Earnings_Date", "Event")
        ).over("Symbol")
    )
)

shape: (5, 5)
┌───────┬────────┬────────────┬───────────────┬───────┐
│ index ┆ Symbol ┆ Date       ┆ Earnings_Date ┆ Event │
│ ---   ┆ ---    ┆ ---        ┆ ---           ┆ ---   │
│ u32   ┆ str    ┆ date       ┆ date          ┆ i64   │
╞═══════╪════════╪════════════╪═══════════════╪═══════╡
│ 0     ┆ A      ┆ 2010-08-29 ┆ 2010-06-01    ┆ 1     │
│ 1     ┆ A      ┆ 2010-09-01 ┆ 2010-09-01    ┆ 4     │
│ 2     ┆ A      ┆ 2010-09-05 ┆ null          ┆ null  │
│ 3     ┆ A      ┆ 2010-11-30 ┆ null          ┆ null  │
│ 4     ┆ A      ┆ 2010-12-02 ┆ 2010-12-01    ┆ 7     │
└───────┴────────┴────────────┴───────────────┴───────┘