Polars按ID和Action分组取最大值的实现问题及优化

问题场景 #

现有包含ID、Action、Where、Value字段的Polars DataFrame，示例数据如下：

data = {
    "ID" : [1, 1, 2,2,3,3],
    "Action" : ["A", "A", "B", "B", "A", "A"],
    "Where" : ["Office", "Home", "Home", "Office", "Home", "Home"],
    "Value" : [1, 2, 3, 4, 5, 6]
}

df = pl.DataFrame(data)

需求：按ID和Action分组，筛选每组中Value最大的记录，以此获取用户偏好的操作地点。

现有实现的问题 #

当前实现通过窗口函数添加TOP列后，按所有字段降序排序再去重，但输出不符合预期：ID=1的Where字段错误显示为Office（预期应为Home）。

现有代码：

(
    df
    .select(
        pl.col("ID"),
        pl.col("Action"),
        pl.col("Where"),
        TOP = pl.col("Value").max().over(["ID", "Action"]))
    .sort(
        pl.col("*"), descending =True
    )
    .unique(
        subset = ["ID", "Action"],
        maintain_order = True,
        keep = "first"
    )
)

问题根源：排序时使用pl.col("*")会包含Where字段，字符串"Office"字典序大于"Home"，导致ID=1的组中，Value更小的Office记录排在了前面，去重时被错误保留。

正确且高效的实现方法 #

方法1：窗口函数直接过滤（推荐） #

直接筛选每组中Value等于组内最大值的行，逻辑清晰且性能最优：

result = (
    df
    .filter(pl.col("Value") == pl.col("Value").max().over(["ID", "Action"]))
    .with_columns(TOP=pl.col("Value"))
)
print(result)

输出符合预期：

shape: (3, 4)
┌─────┬────────┬────────┬─────┐
│ ID  ┆ Action ┆ Where  ┆ TOP │
│ --- ┆ ---    ┆ ---    ┆ --- │
│ i64 ┆ str    ┆ str    ┆ i64 │
╞═════╪════════╪════════╪═════╡
│ 1   ┆ A      ┆ Home   ┆ 2   │
│ 2   ┆ B      ┆ Office ┆ 4   │
│ 3   ┆ A      ┆ Home   ┆ 6   │
└─────┴────────┴────────┴─────┘

方法2：分组聚合 #

如果组内存在多个相同最大值的记录，可通过分组聚合指定保留第一条匹配的Where：

result = (
    df
    .group_by(["ID", "Action"], maintain_order=True)
    .agg(
        # 筛选组内Value最大的Where，取第一条
        Where=pl.col("Where").filter(pl.col("Value") == pl.col("Value").max()).first(),
        TOP=pl.col("Value").max()
    )
)
print(result)

方法3：修正排序逻辑的去重法 #

若坚持使用排序去重，需调整排序优先级，确保Value降序优先级高于Where：

result = (
    df
    .with_columns(TOP=pl.col("Value").max().over(["ID", "Action"]))
    # 先按ID、Action分组，再按Value降序排序
    .sort(["ID", "Action", "Value"], descending=[False, False, True])
    .unique(subset=["ID", "Action"], maintain_order=True, keep="first")
)
print(result)

性能对比 #

• 方法1（窗口过滤）：无需全量排序，仅做组内比较，性能最优，适合大数据集。
• 方法2（分组聚合）：逻辑清晰，适合需要对聚合结果做更多自定义处理的场景。
• 方法3（排序去重）：性能略逊于前两者，仅在特定场景下适用。

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