Polars分组后筛选seq_grp匹配项及非匹配项的最高分

需求背景 #

已有初始Polars DataFrame，按seq列分组得到grouped_df，已完成需求1，现需实现需求2：

• 需求1：筛选出seq_grp与match列表元素匹配且对应score为该组最高分的记录
• 需求2：对存在match元素不等于seq_grp的分组（如baz），获取非seq_grp匹配项的最高分、对应match值及该元素在match列表中的索引

初始数据与分组代码 #

初始DataFrame代码 #

import polars as pl

df = pl.from_repr("""
┌─────┬─────────┬───────┬───────┐
│ seq ┆ seq_grp ┆ match ┆ score │
│ --- ┆ ---     ┆ ---   ┆ ---   │
│ str ┆ str     ┆ str   ┆ i64   │
╞═════╪═════════╪═══════╪═══════╡
│ foo ┆ aa      ┆ aa    ┆ 10    │
│ bar ┆ bb      ┆ cc    ┆ 8     │
│ bar ┆ bb      ┆ bb    ┆ 20    │
│ duk ┆ dd      ┆ dd    ┆ 8     │
│ duk ┆ dd      ┆ dd    ┆ 7     │
│ baz ┆ cc      ┆ ff    ┆ 5     │
│ baz ┆ cc      ┆ cc    ┆ 6     │
│ baz ┆ cc      ┆ cc    ┆ 4     │
│ zed ┆ zz      ┆ yy    ┆ 6     │
└─────┴─────────┴───────┴───────┘
""")

分组代码及结果 #

分组代码：

grouped_df = (
    df.group_by("seq", maintain_order=True)
      .all()
      .with_columns(pl.col("seq_grp").list.get(0))
)

分组后结果：

┌─────┬─────────┬────────────────────┬───────────┐
│ seq ┆ seq_grp ┆ match              ┆ score     │
│ --- ┆ ---     ┆ ---                ┆ ---       │
│ str ┆ str     ┆ list[str]          ┆ list[i64] │
╞═════╪═════════╪════════════════════╪═══════════╡
│ foo ┆ aa      ┆ ["aa"]             ┆ [10]      │
│ bar ┆ bb      ┆ ["cc", "bb"]       ┆ [8, 20]   │
│ duk ┆ dd      ┆ ["dd", "dd"]       ┆ [8, 7]    │
│ baz ┆ cc      ┆ ["ff", "cc", "cc"] ┆ [5, 6, 4] │
│ zed ┆ zz      ┆ ["yy"]             ┆ [6]       │
└─────┴─────────┴────────────────────┴───────────┘

需求2实现方案 #

实现代码 #

result = grouped_df.with_columns(
    # 生成过滤后的非seq_grp匹配项：包含score、match值、列表索引
    pl.struct(["match", "score"])
      .list.eval(
          pl.when(pl.element().match != pl.col("seq_grp"))
            .then(pl.struct(
                score=pl.element().score,
                match_val=pl.element().match,
                index=pl.int_range(0, pl.len())
            ))
            .drop_nulls()
      )
      .alias("non_match_items"),
    # 按score降序排序后取第一个，即最高分的非匹配项
    pl.col("non_match_items")
      .list.sort_by("score", descending=True)
      .list.get(0)
      .alias("top_non_match"),
).with_columns(
    # 拆分结构体为单独列
    pl.col("top_non_match").struct.field("score").alias("top_non_match_score"),
    pl.col("top_non_match").struct.field("match_val").alias("top_non_match_match"),
    pl.col("top_non_match").struct.field("index").alias("top_non_match_index"),
).drop(["non_match_items", "top_non_match"])

# 过滤出存在非匹配项的分组
result = result.filter(pl.col("top_non_match_score").is_not_null())

执行结果 #

┌─────┬─────────┬────────────────────┬───────────┬─────────────────────┬──────────────────────┬──────────────────────┐
│ seq ┆ seq_grp ┆ match              ┆ score     ┆ top_non_match_score ┆ top_non_match_match  ┆ top_non_match_index  │
│ --- ┆ ---     ┆ ---                ┆ ---       ┆ ---                 ┆ ---                  ┆ ---                  │
│ str ┆ str     ┆ list[str]          ┆ list[i64] ┆ i64                 ┆ str                  ┆ i64                  │
╞═════╪═════════╪════════════════════╪═══════════╪═════════════════════╪══════════════════════╪══════════════════════╡
│ bar ┆ bb      ┆ ["cc", "bb"]       ┆ [8, 20]   ┆ 8                   ┆ cc                   ┆ 0                    │
│ baz ┆ cc      ┆ ["ff", "cc", "cc"] ┆ [5, 6, 4] ┆ 5                   ┆ ff                   ┆ 0                    │
│ zed ┆ zz      ┆ ["yy"]             ┆ [6]       ┆ 6                   ┆ yy                   ┆ 0                    │
└─────┴─────────┴────────────────────┴───────────┴─────────────────────┴──────────────────────┴──────────────────────┘

代码说明 #

1. 利用list.eval结合when条件，筛选出match不等于seq_grp的条目，同时通过int_range生成列表索引，保留所需字段
2. 对过滤后的列表按score降序排序，取第一个元素即为该组非匹配项的最高分条目
3. 将结构体字段拆分为独立列，最后过滤掉无匹配项的分组（如foo、duk）

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