Polars 将于下一个可用日期加入

Question

我有一个框架（事件），我想将其加入到另一个框架（fr）中，加入日期和符号。不一定有任何日期重叠。事件中的日期仅与 fr 中相同或较晚日期的第一次出现匹配，因此，如果事件日期是 2010-08-29，它将在同一日期加入，或者如果不存在，则加入下一个可用日期（2010 年） -09-01).

我尝试使用 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",
        strategy="backward",
    )
)

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()

构造。

def null_if_duplicate(col: str, by: str) -> pl.Expr:
    return pl.when(
        (pl.col(col) != pl.col(col).shift().over(by)).fill_null(True)
    ).then(
        pl.col(col)
    )

(
    fr
    .join_asof(
        events,
        left_on="Date",
        right_on="Earnings_Date",
        by="Symbol",
        strategy="backward",
    )
    .with_columns(
        null_if_duplicate("Earnings_Date", by="Symbol"),
        null_if_duplicate("Event", 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 ┆ null          ┆ null  │
│ 3     ┆ A      ┆ 2010-11-30 ┆ null          ┆ null  │
│ 4     ┆ A      ┆ 2010-12-02 ┆ 2010-12-01    ┆ 7     │
└───────┴────────┴────────────┴───────────────┴───────┘