如何基于另一pl.DataFrame的日期列筛选或关联pl.DataFrame

给定两个Polars DataFrame：df1 存储各标的的日期价格数据，df2 存储每个标的需要筛选的多个日期（日期列数量不固定）。需要从df1中筛选出每个标的在df2中所有指定日期的价格数据。

原始数据定义 #

from datetime import date
import polars as pl

df1 = pl.DataFrame(
    {
        "symbol": [
            "sec1", "sec1", "sec1", "sec1", "sec1", "sec1",
            "sec2", "sec2", "sec2", "sec2", "sec2",
        ],
        "date": [
            date(2021, 9, 14),
            date(2021, 9, 15),
            date(2021, 9, 16),
            date(2021, 9, 17),
            date(2021, 8, 31),
            date(2020, 12, 31),
            date(2021, 9, 14),
            date(2021, 9, 15),
            date(2021, 8, 31),
            date(2021, 12, 30),
            date(2020, 12, 31),
        ],
        "price": range(11),
    }
)

df2 = pl.DataFrame(
    {
        "symbol": ["sec1", "sec2"],
        "current_date": [date(2021, 9, 17), date(2021, 9, 15)],
        "mtd": [date(2021, 8, 31), date(2021, 8, 31)],
        "ytd": [date(2020, 12, 31), date(2020, 12, 30)],
    }
)

需求目标 #

从df1中筛选出每个symbol在df2所有日期列对应的价格数据，最终结果如下：

shape: (6, 3)
┌────────┬────────────┬───────┐
│ symbol ┆ date       ┆ price │
│ ---    ┆ ---        ┆ ---   │
│ str    ┆ date       ┆ i64   │
╞════════╪════════════╪═══════╡
│ sec1   ┆ 2021-09-17 ┆ 3     │
│ sec1   ┆ 2021-08-31 ┆ 4     │
│ sec1   ┆ 2020-12-31 ┆ 5     │
│ sec2   ┆ 2021-09-15 ┆ 7     │
│ sec2   ┆ 2021-08-31 ┆ 8     │
│ sec2   ┆ 2020-12-30 ┆ 9     │
└────────┴────────────┴───────┘

优雅解决方案 #

无需循环拼接，通过宽表转长表+内连接即可实现，适配任意数量的日期列：

# 提取df2中所有日期列（排除symbol列）
date_cols = [col for col in df2.columns if col != "symbol"]

# 将df2转为长表结构，保留symbol和对应的目标日期
df2_long = df2.melt(
    id_vars="symbol",
    value_vars=date_cols,
    value_name="date"
).drop("variable")  # 不需要原日期列的名称，删除该列

# 内连接筛选符合条件的行
result = df1.join(df2_long, on=["symbol", "date"], how="inner")

# 输出结果
with pl.Config(tbl_rows=-1):
    print(result)

方案优势 #

• 无需手动处理每个日期列，自动适配df2中日期列的增减
• 利用Polars的向量化操作，性能高效
• 逻辑简洁，代码可维护性强

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