如何用Polars按两列分组提取多极值行（最大/最小值）

数据示例 #

import polars as pl

df = pl.from_repr("""
┌──────┬──────┬──────┬──────┬──────┬──────┐
│ cid1 ┆ cid2 ┆ cid3 ┆ cid4 ┆ cid5 ┆ cid6 │
│ ---  ┆ ---  ┆ ---  ┆ ---  ┆ ---  ┆ ---  │
│ i64  ┆ i64  ┆ f64  ┆ f64  ┆ f64  ┆ f64  │
╞══════╪══════╪══════╪══════╪══════╪══════╡
│ 1    ┆ 5    ┆ 1.0  ┆ 4.0  ┆ 4.0  ┆ 1.0  │
│ 1    ┆ 5    ┆ 2.0  ┆ 5.0  ┆ 5.0  ┆ 9.0  │
│ 1    ┆ 5    ┆ 9.0  ┆ 6.0  ┆ 4.0  ┆ 9.0  │
│ 3    ┆ 7    ┆ 1.0  ┆ 7.0  ┆ 9.0  ┆ 1.0  │
│ 3    ┆ 7    ┆ 3.0  ┆ 7.0  ┆ 9.0  ┆ 1.0  │
│ 3    ┆ 7    ┆ 8.0  ┆ 8.0  ┆ 3.0  ┆ 1.0  │
└──────┴──────┴──────┴──────┴──────┴──────┘
""")

筛选每组前N个最大值 #

使用rank()窗口函数按cid3降序排名，筛选排名≤N的行即可。示例中取前2个最大值，实际需求替换为10即可：

# 筛选每组前2个cid3最大值的行
top_n_max = df.filter(
    pl.col("cid3").rank(descending=True).over("cid1", "cid2") <= 2
)
print(top_n_max)

输出结果：

┌──────┬──────┬──────┬──────┬──────┬──────┐
│ cid1 ┆ cid2 ┆ cid3 ┆ cid4 ┆ cid5 ┆ cid6 │
│ ---  ┆ ---  ┆ ---  ┆ ---  ┆ ---  ┆ ---  │
│ i64  ┆ i64  ┆ f64  ┆ f64  ┆ f64  ┆ f64  │
╞══════╪══════╪══════╪══════╪══════╪══════╡
│ 1    ┆ 5    ┆ 9.0  ┆ 6.0  ┆ 4.0  ┆ 9.0  │
│ 1    ┆ 5    ┆ 2.0  ┆ 5.0  ┆ 5.0  ┆ 9.0  │
│ 3    ┆ 7    ┆ 8.0  ┆ 8.0  ┆ 3.0  ┆ 1.0  │
│ 3    ┆ 7    ┆ 3.0  ┆ 7.0  ┆ 9.0  ┆ 1.0  │
└──────┴──────┴──────┴──────┴──────┴──────┘

筛选每组前M个最小值 #

同样用rank()窗口函数，按cid3升序排名，筛选排名≤M的行。示例中取前1个最小值，实际需求替换为5即可：

# 筛选每组前1个cid3最小值的行
top_m_min = df.filter(
    pl.col("cid3").rank(descending=False).over("cid1", "cid2") <= 1
)
print(top_m_min)

输出结果：

┌──────┬──────┬──────┬──────┬──────┬──────┐
│ cid1 ┆ cid2 ┆ cid3 ┆ cid4 ┆ cid5 ┆ cid6 │
│ ---  ┆ ---  ┆ ---  ┆ ---  ┆ ---  ┆ ---  │
│ i64  ┆ i64  ┆ f64  ┆ f64  ┆ f64  ┆ f64  │
╞══════╪══════╪══════╪══════╪══════╪══════╡
│ 1    ┆ 5    ┆ 1.0  ┆ 4.0  ┆ 4.0  ┆ 1.0  │
│ 3    ┆ 7    ┆ 1.0  ┆ 7.0  ┆ 9.0  ┆ 1.0  │
└──────┴──────┴──────┴──────┴──────┴──────┘

合并两组结果 #

如果需要将最大值和最小值的结果合并为一个DataFrame，使用pl.concat()：

combined_df = pl.concat([top_n_max, top_m_min])
print(combined_df)

关于排名规则的补充 #

rank()默认使用method="average"，相同cid3值会得到相同排名。如果需要严格按行顺序分配名次（即使值相同也不并列），可以指定method="ordinal"：

pl.col("cid3").rank(descending=True, method="ordinal").over("cid1", "cid2")

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