聚合#
Polars的group_by上下文允许将表达式应用于列的子集, 这些子集由数据分组列的唯一值定义. 本章节讨论下如何使用聚合操作
数据准备#
下面的数据由豆包生成
1import polars as pl
2
3df: pl.DataFrame = pl.DataFrame(
4 {
5 'name': ['Alice', 'Bob', 'Charlie', 'David', 'Eve'],
6 'age': [25, 30, 35, 25, 30],
7 'score': [85, 90, 75, 85, 95],
8 'group': ['A', 'B', 'A', 'B', 'A']
9 }
10)
11print(df)1shape: (5, 4)
2┌─────────┬─────┬───────┬───────┐
3│ name ┆ age ┆ score ┆ group │
4│ --- ┆ --- ┆ --- ┆ --- │
5│ str ┆ i64 ┆ i64 ┆ str │
6╞═════════╪═════╪═══════╪═══════╡
7│ Alice ┆ 25 ┆ 85 ┆ A │
8│ Bob ┆ 30 ┆ 90 ┆ B │
9│ Charlie ┆ 35 ┆ 75 ┆ A │
10│ David ┆ 25 ┆ 85 ┆ B │
11│ Eve ┆ 30 ┆ 95 ┆ A │
12└─────────┴─────┴───────┴───────┘基本聚合#
我们可以轻松地将多个表达式应用于聚合, 只需要在函数agg()中列出所需的所有表达式即可.聚合操作的数据量没有上限, 可以任意组合
我们按照group这一列进行分组, 可以看到分为了下面两组
1shape: (3, 4)
2┌─────────┬─────┬───────┬───────┐
3│ name ┆ age ┆ score ┆ group │
4│ --- ┆ --- ┆ --- ┆ --- │
5│ str ┆ i64 ┆ i64 ┆ str │
6╞═════════╪═════╪═══════╪═══════╡
7│ Alice ┆ 25 ┆ 85 ┆ A │
8│ Charlie ┆ 35 ┆ 75 ┆ A │
9│ Eve ┆ 30 ┆ 95 ┆ A │
10└─────────┴─────┴───────┴───────┘1shape: (2, 4)
2┌───────┬─────┬───────┬───────┐
3│ name ┆ age ┆ score ┆ group │
4│ --- ┆ --- ┆ --- ┆ --- │
5│ str ┆ i64 ┆ i64 ┆ str │
6╞═══════╪═════╪═══════╪═══════╡
7│ Bob ┆ 30 ┆ 90 ┆ B │
8│ David ┆ 25 ┆ 85 ┆ B │
9└───────┴─────┴───────┴───────┘我们来对分组做以下操作
- 获取每组的行数
- 获取年龄和分数这一列的最大值
- 将
name这一列组合成一个列表
1res = (df.group_by("group")
2 .agg(
3 pl.len(), # 计算每个分组的行数
4 pl.col("name"), # 将每个分组的name列组合成一个列表
5 pl.max("age").alias("max_age") # 获取每个分组的年龄最大值
6 )
7)
8print(res)1shape: (2, 4)
2┌───────┬─────┬─────────────────────────────┬─────────┐
3│ group ┆ len ┆ name ┆ max_age │
4│ --- ┆ --- ┆ --- ┆ --- │
5│ str ┆ u32 ┆ list[str] ┆ i64 │
6╞═══════╪═════╪═════════════════════════════╪═════════╡
7│ A ┆ 3 ┆ ["Alice", "Charlie", "Eve"] ┆ 35 │
8│ B ┆ 2 ┆ ["Bob", "David"] ┆ 30 │
9└───────┴─────┴─────────────────────────────┴─────────┘条件语句#
我们可以在聚合中使用条件语句直接查询, 比如我们想知道每个组有多少人分数>80和分数>90的
1res = (
2 df.group_by("group").agg(
3 (pl.col("score") > 80).sum().alias("sum_lt_80"),
4 (pl.col("score") > 90).sum().alias("sum_lt_90")
5 )
6)
7print(res)1shape: (2, 3)
2┌───────┬───────────┬───────────┐
3│ group ┆ sum_lt_80 ┆ sum_lt_90 │
4│ --- ┆ --- ┆ --- │
5│ str ┆ u32 ┆ u32 │
6╞═══════╪═══════════╪═══════════╡
7│ A ┆ 2 ┆ 1 │
8│ B ┆ 2 ┆ 0 │
9└───────┴───────────┴───────────┘过滤#
我们可以对组进行过滤, 比如想计算每个组的平均值, 但不想包含该组的所有值, 又不想真正过滤DataFrame中的行, 因为我们需要这些行进行其他聚合
agg()中的每个表达式过滤时是互不影响的, 下面代码计算每个组的总分数和满足age>=30的分数
1# 构建查询
2res = (
3 df
4 .group_by("group") # 按group列分组
5 .agg(
6 (pl.col("score").filter(pl.col("age")>=30)).sum().alias("sum(age>=30)"),
7 pl.col("score").sum().alias("sum()")
8 )
9)
10print(result)1shape: (2, 3)
2┌───────┬──────────────┬───────┐
3│ group ┆ sum(age>=30) ┆ sum() │
4│ --- ┆ --- ┆ --- │
5│ str ┆ i64 ┆ i64 │
6╞═══════╪══════════════╪═══════╡
7│ B ┆ 90 ┆ 170 │
8│ A ┆ 170 ┆ 255 │
9└───────┴──────────────┴───────┘嵌套分组#
我们可以将group和age作为组合键进行分组, 统计每组的平均分数和人数, 意思是group和age相同的行被归为一组
1import polars as pl
2# 这里一共6行数据, group为A的4行中, 有2行的age都是25, 便于观察
3df: pl.DataFrame = pl.DataFrame(
4 {
5 'name': ['Alice', 'Bob', 'Charlie', 'David', 'Eve','John'],
6 'age': [25, 30, 35, 25, 30, 25],
7 'score': [85, 90, 75, 85, 95, 100],
8 'group': ['A', 'B', 'A', 'B', 'A','A']
9 }
10)
11result = (df
12 .group_by(["group", "age"])
13 .agg([
14 pl.col("score").mean().alias("avg_score"),
15 pl.len().alias("count")
16 ]).sort(["group", "age"])
17)
18
19print(result)注意到有两个('A',25),分数分别是85和100, 所以平均值是92.5
1shape: (5, 4)
2┌───────┬─────┬───────────┬───────┐
3│ group ┆ age ┆ avg_score ┆ count │
4│ --- ┆ --- ┆ --- ┆ --- │
5│ str ┆ i64 ┆ f64 ┆ u32 │
6╞═══════╪═════╪═══════════╪═══════╡
7│ A ┆ 25 ┆ 92.5 ┆ 2 │
8│ A ┆ 30 ┆ 95.0 ┆ 1 │
9│ A ┆ 35 ┆ 75.0 ┆ 1 │
10│ B ┆ 25 ┆ 85.0 ┆ 1 │
11│ B ┆ 30 ┆ 90.0 ┆ 1 │
12└───────┴─────┴───────────┴───────┘排序#
使用排序非常简单, 我们可以先排序然后分组, 最后聚合得到我们想要的结果
下面代码获取每组分数最高的人以及分数, 在代码中pl.col("name").first()是根据排序后的结果来计算的, 也就是说排序的结果会影响后续聚合的行为
1import polars as pl
2
3df: pl.DataFrame = pl.DataFrame(
4 {
5 'name': ['Alice', 'Bob', 'Charlie', 'David', 'Eve','John'],
6 'age': [25, 30, 35, 25, 30, 25],
7 'score': [85, 90, 75, 85, 95, 100],
8 'group': ['A', 'B', 'A', 'B', 'A','A']
9 }
10)
11res = (df
12 .sort("score",descending=True)
13 .group_by("group").agg(
14 pl.col("name").first().alias("top_name"),
15 pl.col("score").first().alias("top_score")
16 )
17)
18print(res)1shape: (2, 3)
2┌───────┬──────────┬───────────┐
3│ group ┆ top_name ┆ top_score │
4│ --- ┆ --- ┆ --- │
5│ str ┆ str ┆ i64 │
6╞═══════╪══════════╪═══════════╡
7│ A ┆ John ┆ 100 │
8│ B ┆ Bob ┆ 90 │
9└───────┴──────────┴───────────┘sort会对整体进行排序, 并影响分组后的聚合行为, 然后我们还可以在分组内重新排序, 以实现更灵活的需求
1df: pl.DataFrame = pl.DataFrame(
2 {
3 'name': ['Alice', 'Bob', 'Charlie', 'David', 'Eve','John'],
4 'age': [25, 40, 35, 25, 30, 25],
5 'score': [85, 90, 75, 85, 95, 100],
6 'group': ['A', 'B', 'A', 'B', 'A','A']
7 }
8)
9q = (
10 df.lazy()
11 .sort("score", descending=True) # 按成绩降序排，最高分靠前
12 .group_by("group")
13 .agg([
14 pl.col("name").first().alias("top_scorer"), # 最高分的人名
15 pl.col("score").first().alias("top_score"), # 最高分
16 pl.col("name").sort().first().alias("alphabetical_first"), # 字典序最小的名字
17 pl.col("age").sort_by(pl.col("name")).first().alias("age_of_alphabetical_first"), # 按名字排序后，第一个人的年龄
18 ])
19 .sort("group")
20)
21
22df_result = q.collect()
23print(df_result)1shape: (2, 5)
2┌───────┬────────────┬───────────┬────────────────────┬───────────────────────────┐
3│ group ┆ top_scorer ┆ top_score ┆ alphabetical_first ┆ age_of_alphabetical_first │
4│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
5│ str ┆ str ┆ i64 ┆ str ┆ i64 │
6╞═══════╪════════════╪═══════════╪════════════════════╪═══════════════════════════╡
7│ A ┆ John ┆ 100 ┆ Alice ┆ 25 │
8│ B ┆ Bob ┆ 90 ┆ Bob ┆ 40 │
9└───────┴────────────┴───────────┴────────────────────┴───────────────────────────┘不要扼杀并行#
在Polars中, 请尽量避免使用lambda函数以及自定义Python函数, 相反多使用Polars提供的API
目录