根据Python Polars中的不同列在列表列中查找索引

Question

我的最小示例是一个具有 2 列的 Dataframe - 第一列是日期，第二列是日期列表，其中每行中的日期都是唯一的。第二列中的日期列表已排序（按升序）。

目标是找到第 2 列的最小索引 (i)，其中第 1 列中的日期小于每行中的第 i 个列表元素。

import polars as pl
from datetime import date

pl.Config(fmt_table_cell_list_len=10, fmt_str_lengths=80) # repr

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) ]
    ]
})
df

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

在此示例中，结果 Dataframe 应该是：

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

表达式），类似于为此 SO 问题提出的解决方案。

另一种可行的方法使用

explode

，但很高兴看到基于

list.eval

的解决方案。

Answer 1

这应该可以。

(df
 .with_row_index('j')
 .with_columns(i=pl.int_ranges(pl.col('date_list').list.len()))
 .explode('date_list','i')
 .group_by('j',maintain_order=True)
 .agg(pl.col('ref_date').first(),
      pl.col('date_list'),
      pl.col('i').filter(pl.col('date_list')>pl.col('ref_date')).min())
 .drop('j')
)

思考过程始于这样的想法：你想要爆炸而不是试图处理

list.eval

。其次，您需要 df 开始时的索引以及前两个操作

date_list

列的索引。第三个当然是爆炸本身。接下来是按原始行索引进行分组并获取

i

的最小值，其中

date_list

>

ref_date

.