Polars如何原生实现列表列中指定字符串的先后顺序判断

你不需要借助map_elements跳出原生执行链路，Polars 原生表达式完全支持这类索引比较需求，且性能远高于Python侧循环实现。

基于已聚合列表列的实现 #

直接在你现有的df_groups结果上做过滤即可，核心是通过list.eval在列表内部计算两个目标字符串首次出现的索引，再比较大小：

pre = 'A'
succ = 'B'

df_groups_filtered = df_groups.filter(
    # 先过滤掉两个目标字符串不全的分组
    pl.col('str_list').list.contains(pre)
    & pl.col('str_list').list.contains(succ)
    # 原生比较两个字符串首次出现的索引先后
    & pl.col('str_list').list.eval(
        pl.arg_where(pl.element() == pre).first()
        < pl.arg_where(pl.element() == succ).first()
    ).first()
)

运行结果和你预期完全一致：

shape: (1, 2)
┌───────┬─────────────────┐
│ group ┆ str_list        │
│ ---   ┆ ---             │
│ i64   ┆ list[str]       │
╞═══════╪═════════════════╡
│ 1     ┆ ["A", "B", "B"] │
└───────┴─────────────────┘

说明：list.eval会直接在Polars Rust引擎内部执行列表元素遍历和索引计算，没有Python侧函数调用开销，大数据量下性能比map_elements实现高10~100倍。

更高性能的实现（跳过中间列表构造） #

如果你不需要保留聚合后的str_list列做其他计算，可以直接在分组聚合阶段完成顺序判断，省去构造字符串列表的内存和计算开销：

pre = 'A'
succ = 'B'

df_result = df.group_by('group').agg(
    pl.col('str').alias('str_list'),
    # 布尔列的arg_max会直接返回第一个True的位置，即元素首次出现索引
    (pl.col('str').eq(pre).arg_max() < pl.col('str').eq(succ).arg_max())
    .alias('pre_before_succ')
).filter(
    pl.col('pre_before_succ')
    & pl.col('str_list').list.contains(pre)
    & pl.col('str_list').list.contains(succ)
)

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