如何在Polars DataFrame中按row_nr组获取count最大值记录？

首先还原示例DataFrame：

import polars as pl

df = pl.DataFrame(
[{'row_nr': 0, 'middle_name_or_initial': 'R.', 'count': 8},
 {'row_nr': 0, 'middle_name_or_initial': 'B.', 'count': 1},
 {'row_nr': 1, 'middle_name_or_initial': 'D.', 'count': 1},
 {'row_nr': 1, 'middle_name_or_initial': 'J.', 'count': 11},
 {'row_nr': 2, 'middle_name_or_initial': 'Micha', 'count': 1},
 {'row_nr': 2, 'middle_name_or_initial': 'M.', 'count': 1},
 {'row_nr': 2, 'middle_name_or_initial': 'J.', 'count': 1},
 {'row_nr': 3, 'middle_name_or_initial': 'E.', 'count': 1},
 {'row_nr': 3, 'middle_name_or_initial': 'S.', 'count': 1},
 {'row_nr': 4, 'middle_name_or_initial': 'M.', 'count': 1},
 {'row_nr': 4, 'middle_name_or_initial': 'P.', 'count': 1},
 {'row_nr': 5, 'middle_name_or_initial': 'Christopher', 'count': 15},
 {'row_nr': 5, 'middle_name_or_initial': 'Robert', 'count': 1},
 {'row_nr': 6, 'middle_name_or_initial': 'Lusi', 'count': 1},
 {'row_nr': 6, 'middle_name_or_initial': 'Luis', 'count': 1}])

接下来实现需求逻辑：

# 按row_nr分组，计算每组的最大count，以及是否所有count都是1
group_stats = df.group_by("row_nr").agg(
    pl.col("count").max().alias("max_count"),
    (pl.col("count").max() == 1).alias("all_ones")
)

# 连接原表与分组统计结果，筛选符合条件的记录
result = df.join(group_stats, on="row_nr").filter(
    (pl.col("count") == pl.col("max_count")) | pl.col("all_ones")
).drop("max_count", "all_ones")

print(result)

逻辑说明 #

1. 分组统计：通过group_by("row_nr")分组后，计算每组的最大count值；同时用max(count) == 1判断该组所有count是否均为1（最大值为1则所有值必然是1）。
2. 连接筛选：将统计结果与原表按row_nr关联，然后过滤两种情况：
   • 当前记录的count等于组内最大count（保留每组的最高count记录）
   • 组内所有count都是1（直接保留该组全部记录）
3. 清理列：移除用于判断的临时列max_count和all_ones，得到最终结果。

输出结果 #

shape: (12, 3)
┌────────┬──────────────────────┬───────┐
│ row_nr ┆ middle_name_or_initial ┆ count │
│ ---    ┆ ---                  ┆ ---   │
│ i64    ┆ str                  ┆ i64   │
╞════════╪══════════════════════╪═══════╡
│ 0      ┆ R.                   ┆ 8     │
│ 1      ┆ J.                   ┆ 11    │
│ 2      ┆ Micha                ┆ 1     │
│ 2      ┆ M.                   ┆ 1     │
│ 2      ┆ J.                   ┆ 1     │
│ 3      ┆ E.                   ┆ 1     │
│ 3      ┆ S.                   ┆ 1     │
│ 4      ┆ M.                   ┆ 1     │
│ 4      ┆ P.                   ┆ 1     │
│ 5      ┆ Christopher          ┆ 15    │
│ 6      ┆ Lusi                 ┆ 1     │
│ 6      ┆ Luis                 ┆ 1     │
└────────┴──────────────────────┴───────┘

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