Polars：根据另一DataFrame查找首个前置匹配的索引/日期

主DataFrame定义 #

import polars as pl

test = pl.DataFrame(
    {
        "date": ["2020-01-01", "2020-01-02", "2020-01-03", "2020-01-04", "2020-01-05"],
        "value": [1, 2, 3, 4, 5],
    }
).with_row_index().with_columns(pl.col("date").cast(pl.Date))
print(test)

输出：

shape: (5, 3)
┌───────┬────────────┬───────┐
│ index ┆ date       ┆ value │
│ ---   ┆ ---        ┆ ---   │
│ u32   ┆ date       ┆ i64   │
╞═══════╪════════════╪═══════╡
│ 0     ┆ 2020-01-01 ┆ 1     │
│ 1     ┆ 2020-01-02 ┆ 2     │
│ 2     ┆ 2020-01-03 ┆ 3     │
│ 3     ┆ 2020-01-04 ┆ 4     │
│ 4     ┆ 2020-01-05 ┆ 5     │
└───────┴────────────┴───────┘

目标日期DataFrame #

shape: (1, 1)
┌────────────┐
│ date       │
│ ---        │
│ date       │
╞════════════╡
│ 2020-01-04 │
└────────────┘

需求 #

从主DataFrametest中，找到早于2020-01-04的最后一条记录，获取其索引（预期输出：2）或对应日期（预期输出："2020-01-03"）。

解决方案 #

1. 提取目标日期 #

先从目标DataFrame中取出日期值：

target_date = pl.DataFrame({"date": ["2020-01-04"]}).cast(pl.Date).item()

2. 获取目标索引 #

筛选出所有早于目标日期的记录，取最后一条的index：

result_index = test.filter(pl.col("date") < target_date).select("index").last().item()
print(result_index)  # 输出：2

3. 获取目标日期 #

同样筛选后取最后一条的date：

result_date = test.filter(pl.col("date") < target_date).select("date").last().item()
print(result_date)  # 输出：2020-01-03

高效优化写法（适合大数据量） #

如果数据量较大，可通过arg_max直接定位符合条件的最大日期对应记录，避免全表筛选后遍历：

# 获取索引
result_index = test.filter(pl.col("date") < target_date).select(pl.col("index").arg_max()).item()

# 获取日期
result_date = test.filter(pl.col("date") < target_date).select(pl.col("date").max()).item()

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