如何在Polars中无需中间分组DataFrame标记组内最大值行？

需求说明 #

现有包含group和value列的Polars DataFrame，需要新增一列group_max，用于标记每组中value为最大值的记录，要求无需创建单独的分组DataFrame再与原表连接。

你目前已实现的方式需要生成中间分组DataFrame，代码如下：

grouped = df.groupby('group').agg(pl.col('value').max()).with_columns(pl.lit(1).alias('value'))
df = df.join(grouped, on = ['group', 'value'], how='outer')

更简洁的实现方式 #

Polars支持窗口函数（over子句），可以直接在原DataFrame上完成计算，完全不需要中间表和连接操作：

import polars as pl

# 示例DataFrame
df = pl.DataFrame({
    'group': ['A', 'A', 'B', 'B', 'B'],
    'value': [10, 20, 5, 15, 15]
})

# 新增group_max列，标记每组最大值记录（转成1/0格式）
df = df.with_columns(
    pl.when(pl.col('value') == pl.col('value').max().over('group'))
      .then(1)
      .otherwise(0)
      .alias('group_max')
)

# 若不需要数字标记，也可直接保留布尔值
# df = df.with_columns(
#     (pl.col('value') == pl.col('value').max().over('group')).alias('group_max')
# )

运行后结果示例：

shape: (5, 3)
┌───────┬───────┬───────────┐
│ group ┆ value ┆ group_max │
│ ---   ┆ ---   ┆ ---       │
│ str   ┆ i64   ┆ i32       │
╞═══════╪═══════╪═══════════╡
│ A     ┆ 10    ┆ 0         │
│ A     ┆ 20    ┆ 1         │
│ B     ┆ 5     ┆ 0         │
│ B     ┆ 15    ┆ 1         │
│ B     ┆ 15    ┆ 1         │
└───────┴───────┴───────────┘

原理说明 #

pl.col('value').max().over('group')会对每个group分组计算value的最大值，并将该最大值广播到对应分组的每一行。之后通过比较当前行的value是否等于分组最大值，即可标记出每组的最大值记录。这种方式完全在原DataFrame的计算流程中完成，避免了额外的分组和连接操作，代码更简洁高效。

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