在Polars中基于另一列查找列表列的最小索引

我们有一个包含两列的Polars DataFrame：第一列ref_date为日期类型，第二列date_list为已按升序排序的唯一日期列表。需求是为每一行找到date_list中首个大于对应ref_date的元素的最小索引(i)。

示例代码 #

import polars as pl
from datetime import date

pl.Config(fmt_table_cell_list_len=10, fmt_str_lengths=80) # 调整输出格式

df = pl.DataFrame({
    'ref_date': [ date(2023, 9, 10), date(2023, 5, 15) ],
    'date_list': [ 
        [ date(2023, 1, 1), date(2023, 5, 1), date(2023, 12, 31) ],
        [ date(2023, 4, 15), date(2023, 6, 18), date(2023, 12, 1) ]
    ]
})

预期输出 #

shape(2, 3)
┌─────────────┬────────────────────────────────────────┬──────┐
│ ref_date    ┆ date_list                              ┆ i    │
│ ---         ┆ ---                                    ┆ ---  │
│ date        ┆ list[date]                             ┆ u32  │
╞═════════════╪════════════════════════════════════════╪══════╡
│ 2023-09-10  ┆ [2023-01-01, 2023-05-01, 2023-12-31]   ┆ 2    │   
│ 2023-05-15  ┆ [2023-04-15, 2023-06-18, 2023-12-01]   ┆ 1    │
└─────────────┴────────────────────────────────────────┴──────┘

已尝试方案 #

• 使用list.eval()时无法直接传入pl.col("ref_date")（仅能通过pl.lit固定单个值）
• 通过explode展开列表后可以实现需求，但希望找到基于list.eval()的解决方案

基于list.eval()的解决方案 #

利用list.eval()的行级上下文特性，结合pl.col("ref_date").first()获取当前行的参考日期，再通过arg_min定位首个符合条件的索引：

df = df.with_columns(
    i=pl.col("date_list").list.eval(
        pl.arg_min(pl.element() > pl.col("ref_date").first())
    ).list.first()
)

代码说明 #

1. pl.col("ref_date").first()：在list.eval的行级上下文里，获取当前行的ref_date值
2. pl.element() > pl.col("ref_date").first()：生成布尔列表，标记date_list中每个元素是否大于参考日期
3. pl.arg_min()：返回布尔列表中第一个True值的索引（因date_list已升序排序，该索引就是首个大于参考日期元素的位置）
4. list.first()：将list.eval返回的单元素列表展开为单个索引值

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