窗口函数
窗口函数是具有超能力的表达式。它们允许你在select上下文中对组执行聚合操作。让我们来感受一下这意味着什么。
首先，我们加载一个宝可梦数据集：
import polars as pl
types = (
"Grass Water Fire Normal Ground Electric Psychic Fighting Bug Steel "
"Flying Dragon Dark Ghost Poison Rock Ice Fairy".split()
)
type_enum = pl.Enum(types)
# 然后让我们加载一些关于宝可梦的CSV数据
pokemon = pl.read_csv(
"https://gist.githubusercontent.com/ritchie46/cac6b337ea52281aa23c049250a4ff03/raw/89a957ff3919d90e6ef2d34235e6bf22304f3366/pokemon.csv",
).cast({"Type 1": type_enum, "Type 2": type_enum})
print(pokemon.head())
use polars::prelude::*;
use reqwest::blocking::Client;
let data: Vec<u8> = Client::new()
.get("https://gist.githubusercontent.com/ritchie46/cac6b337ea52281aa23c049250a4ff03/raw/89a957ff3919d90e6ef2d34235e6bf22304f3366/pokemon.csv")
.send()?
.text()?
.bytes()
.collect();
let file = std::io::Cursor::new(data);
let df = CsvReadOptions::default()
.with_has_header(true)
.into_reader_with_file_handle(file)
.finish()?;
println!("{}", df.head(Some(5)));
shape: (5, 13)
┌─────┬───────────────────────┬────────┬────────┬───┬─────────┬───────┬────────────┬───────────┐
│ # ┆ Name ┆ Type 1 ┆ Type 2 ┆ … ┆ Sp. Def ┆ Speed ┆ Generation ┆ Legendary │
│ --- ┆ --- ┆ --- ┆ --- ┆ ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ str ┆ enum ┆ enum ┆ ┆ i64 ┆ i64 ┆ i64 ┆ bool │
╞═════╪═══════════════════════╪════════╪════════╪═══╪═════════╪═══════╪════════════╪═══════════╡
│ 1 ┆ Bulbasaur ┆ Grass ┆ Poison ┆ … ┆ 65 ┆ 45 ┆ 1 ┆ false │
│ 2 ┆ Ivysaur ┆ Grass ┆ Poison ┆ … ┆ 80 ┆ 60 ┆ 1 ┆ false │
│ 3 ┆ Venusaur ┆ Grass ┆ Poison ┆ … ┆ 100 ┆ 80 ┆ 1 ┆ false │
│ 3 ┆ VenusaurMega Venusaur ┆ Grass ┆ Poison ┆ … ┆ 120 ┆ 80 ┆ 1 ┆ false │
│ 4 ┆ Charmander ┆ Fire ┆ null ┆ … ┆ 50 ┆ 65 ┆ 1 ┆ false │
└─────┴───────────────────────┴────────┴────────┴───┴─────────┴───────┴────────────┴───────────┘
每组操作
当我们想要在一个组内执行操作时，窗口函数是理想的选择。例如，假设我们想要根据“速度”列对我们的宝可梦进行排名。然而，我们不是想要一个全局排名，而是想要在每个由“Type 1”列定义的组内对速度进行排名。我们编写表达式以根据“速度”列对数据进行排名，然后我们添加函数over来指定这应该在“Type 1”列的唯一值上进行：
result = pokemon.select(
pl.col("Name", "Type 1"),
pl.col("Speed").rank("dense", descending=True).over("Type 1").alias("Speed rank"),
)
print(result)
let result = df
.clone()
.lazy()
.select([
col("Name"),
col("Type 1"),
col("Speed")
.rank(
RankOptions {
method: RankMethod::Dense,
descending: true,
},
None,
)
.over(["Type 1"])
.alias("Speed rank"),
])
.collect()?;
println!("{}", result);
shape: (163, 3)
┌───────────────────────┬─────────┬────────────┐
│ Name ┆ Type 1 ┆ Speed rank │
│ --- ┆ --- ┆ --- │
│ str ┆ enum ┆ u32 │
╞═══════════════════════╪═════════╪════════════╡
│ Bulbasaur ┆ Grass ┆ 6 │
│ Ivysaur ┆ Grass ┆ 3 │
│ Venusaur ┆ Grass ┆ 1 │
│ VenusaurMega Venusaur ┆ Grass ┆ 1 │
│ Charmander ┆ Fire ┆ 7 │
│ … ┆ … ┆ … │
│ Moltres ┆ Fire ┆ 5 │
│ Dratini ┆ Dragon ┆ 3 │
│ Dragonair ┆ Dragon ┆ 2 │
│ Dragonite ┆ Dragon ┆ 1 │
│ Mewtwo ┆ Psychic ┆ 2 │
└───────────────────────┴─────────┴────────────┘
为了帮助可视化此操作，您可以想象Polars选择了“Type 1”列中具有相同值的数据子集，然后仅对这些值计算排名表达式。然后，该特定组的结果被投影回原始行，Polars对所有现有组执行此操作。下图突出显示了“Type 1”等于“Grass”的宝可梦的排名计算。
请注意，宝可梦“Golbat”的“速度”值为90，这比宝可梦“Venusaur”的值80要大，然而后者排名第一，因为“Golbat”和“Venusar”在“Type 1”列中的值不同。
函数 over 接受任意数量的表达式来指定执行计算的分组。我们可以重复上述排名，但针对列“Type 1”和“Type 2”的组合进行更细粒度的排名：
shape: (163, 4)
┌───────────────────────┬─────────┬────────┬────────────┐
│ Name ┆ Type 1 ┆ Type 2 ┆ Speed rank │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ enum ┆ enum ┆ u32 │
╞═══════════════════════╪═════════╪════════╪════════════╡
│ Bulbasaur ┆ Grass ┆ Poison ┆ 6 │
│ Ivysaur ┆ Grass ┆ Poison ┆ 3 │
│ Venusaur ┆ Grass ┆ Poison ┆ 1 │
│ VenusaurMega Venusaur ┆ Grass ┆ Poison ┆ 1 │
│ Charmander ┆ Fire ┆ null ┆ 7 │
│ … ┆ … ┆ … ┆ … │
│ Moltres ┆ Fire ┆ Flying ┆ 2 │
│ Dratini ┆ Dragon ┆ null ┆ 2 │
│ Dragonair ┆ Dragon ┆ null ┆ 1 │
│ Dragonite ┆ Dragon ┆ Flying ┆ 1 │
│ Mewtwo ┆ Psychic ┆ null ┆ 2 │
└───────────────────────┴─────────┴────────┴────────────┘
通常，使用函数over得到的结果也可以通过聚合后调用函数explode来实现，尽管行的顺序会有所不同：
shape: (163, 3)
┌────────────┬──────────┬────────────┐
│ Name ┆ Type 1 ┆ Speed rank │
│ --- ┆ --- ┆ --- │
│ str ┆ enum ┆ u32 │
╞════════════╪══════════╪════════════╡
│ Mankey ┆ Fighting ┆ 4 │
│ Primeape ┆ Fighting ┆ 1 │
│ Machop ┆ Fighting ┆ 7 │
│ Machoke ┆ Fighting ┆ 6 │
│ Machamp ┆ Fighting ┆ 5 │
│ … ┆ … ┆ … │
│ Weepinbell ┆ Grass ┆ 4 │
│ Victreebel ┆ Grass ┆ 2 │
│ Exeggcute ┆ Grass ┆ 7 │
│ Exeggutor ┆ Grass ┆ 4 │
│ Tangela ┆ Grass ┆ 3 │
└────────────┴──────────┴────────────┘
这表明，通常group_by和over会产生不同形状的结果：
group_by通常生成一个结果数据框，其行数与用于聚合的组数相同；并且over通常生成一个与原始数据框行数相同的数据框。
函数 over 并不总是产生与原始数据框行数相同的结果，这是我们接下来要探讨的内容。
将结果映射到数据框行
函数 over 接受一个参数 mapping_strategy，该参数决定了表达式在组上的结果如何映射回数据框的行。
group_to_rows
默认行为是 "group_to_rows"：表达式在组上的结果应与组的长度相同，并且结果会映射回该组的行。
如果行的顺序不重要，选项"explode"性能更高。Polars不是将结果值映射到原始行，而是创建一个新的数据框，其中来自同一组的值彼此相邻。为了帮助理解这一区别，请考虑以下数据框：
shape: (6, 3)
┌─────────┬─────────┬──────┐
│ athlete ┆ country ┆ rank │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ i64 │
╞═════════╪═════════╪══════╡
│ A ┆ PT ┆ 6 │
│ B ┆ NL ┆ 1 │
│ C ┆ NL ┆ 5 │
│ D ┆ PT ┆ 4 │
│ E ┆ PT ┆ 2 │
│ F ┆ NL ┆ 3 │
└─────────┴─────────┴──────┘
我们可以按运动员在自己国家的排名进行排序。如果我们这样做，荷兰运动员在第二、第三和第六行，他们将保持在那里。将改变的是运动员姓名的顺序，从“B”、“C”和“F”变为“B”、“F”和“C”：
shape: (6, 3)
┌─────────┬──────┬─────────┐
│ athlete ┆ rank ┆ country │
│ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ str │
╞═════════╪══════╪═════════╡
│ E ┆ 2 ┆ PT │
│ B ┆ 1 ┆ NL │
│ F ┆ 3 ┆ NL │
│ D ┆ 4 ┆ PT │
│ A ┆ 6 ┆ PT │
│ C ┆ 5 ┆ NL │
└─────────┴──────┴─────────┘
下图表示此转换：
explode
如果我们将参数mapping_strategy设置为"explode"，那么来自同一国家的运动员将被分组在一起，但行的最终顺序——关于国家的——将不会相同，如图所示：
因为Polars不需要跟踪每个组的行位置，使用"explode"通常比"group_to_rows"更快。然而，使用"explode"也需要更多的注意，因为它意味着我们需要重新排序我们希望保留的其他列。产生此结果的代码如下
shape: (6, 3)
┌─────────┬─────────┬──────┐
│ athlete ┆ country ┆ rank │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ i64 │
╞═════════╪═════════╪══════╡
│ E ┆ PT ┆ 2 │
│ D ┆ PT ┆ 4 │
│ A ┆ PT ┆ 6 │
│ B ┆ NL ┆ 1 │
│ F ┆ NL ┆ 3 │
│ C ┆ NL ┆ 5 │
└─────────┴─────────┴──────┘
join
参数 mapping_strategy 的另一个可能值是 "join"，它将结果值聚合到一个列表中，并在同一组的所有行上重复该列表：
shape: (6, 3)
┌─────────┬─────────┬───────────┐
│ athlete ┆ country ┆ rank │
│ --- ┆ --- ┆ --- │
│ str ┆ str ┆ list[i64] │
╞═════════╪═════════╪═══════════╡
│ A ┆ PT ┆ [2, 4, 6] │
│ B ┆ NL ┆ [1, 3, 5] │
│ C ┆ NL ┆ [1, 3, 5] │
│ D ┆ PT ┆ [2, 4, 6] │
│ E ┆ PT ┆ [2, 4, 6] │
│ F ┆ NL ┆ [1, 3, 5] │
└─────────┴─────────┴───────────┘
窗口聚合表达式
如果应用于组值的表达式产生标量值，则该标量将在组的行中广播：
result = pokemon.select(
pl.col("Name", "Type 1", "Speed"),
pl.col("Speed").mean().over(pl.col("Type 1")).alias("Mean speed in group"),
)
print(result)
let result = df
.clone()
.lazy()
.select([
col("Name"),
col("Type 1"),
col("Speed"),
col("Speed")
.mean()
.over(["Type 1"])
.alias("Mean speed in group"),
])
.collect()?;
println!("{}", result);
shape: (163, 4)
┌───────────────────────┬─────────┬───────┬─────────────────────┐
│ Name ┆ Type 1 ┆ Speed ┆ Mean speed in group │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ enum ┆ i64 ┆ f64 │
╞═══════════════════════╪═════════╪═══════╪═════════════════════╡
│ Bulbasaur ┆ Grass ┆ 45 ┆ 54.230769 │
│ Ivysaur ┆ Grass ┆ 60 ┆ 54.230769 │
│ Venusaur ┆ Grass ┆ 80 ┆ 54.230769 │
│ VenusaurMega Venusaur ┆ Grass ┆ 80 ┆ 54.230769 │
│ Charmander ┆ Fire ┆ 65 ┆ 86.285714 │
│ … ┆ … ┆ … ┆ … │
│ Moltres ┆ Fire ┆ 90 ┆ 86.285714 │
│ Dratini ┆ Dragon ┆ 50 ┆ 66.666667 │
│ Dragonair ┆ Dragon ┆ 70 ┆ 66.666667 │
│ Dragonite ┆ Dragon ┆ 80 ┆ 66.666667 │
│ Mewtwo ┆ Psychic ┆ 130 ┆ 99.25 │
└───────────────────────┴─────────┴───────┴─────────────────────┘
更多示例
更多练习，以下是一些供我们计算的窗口函数：
- 按类型对所有宝可梦进行排序；
- 选择每种类型的第一个
3宝可梦作为"Type 1"； - 按速度降序排列某一类型中的宝可梦，并选择前
3个作为"fastest/group"； - 按攻击力降序排列同一类型的宝可梦，并选择前
3个作为"最强/组"；以及 - 按名称对类型中的宝可梦进行排序，并选择前
3个作为"sorted_by_alphabet"。
result = pokemon.sort("Type 1").select(
pl.col("Type 1").head(3).over("Type 1", mapping_strategy="explode"),
pl.col("Name")
.sort_by(pl.col("Speed"), descending=True)
.head(3)
.over("Type 1", mapping_strategy="explode")
.alias("fastest/group"),
pl.col("Name")
.sort_by(pl.col("Attack"), descending=True)
.head(3)
.over("Type 1", mapping_strategy="explode")
.alias("strongest/group"),
pl.col("Name")
.sort()
.head(3)
.over("Type 1", mapping_strategy="explode")
.alias("sorted_by_alphabet"),
)
print(result)
let result = df
.clone()
.lazy()
.select([
col("Type 1")
.head(Some(3))
.over_with_options(["Type 1"], None, WindowMapping::Explode)
.flatten(),
col("Name")
.sort_by(
["Speed"],
SortMultipleOptions::default().with_order_descending(true),
)
.head(Some(3))
.over_with_options(["Type 1"], None, WindowMapping::Explode)
.flatten()
.alias("fastest/group"),
col("Name")
.sort_by(
["Attack"],
SortMultipleOptions::default().with_order_descending(true),
)
.head(Some(3))
.over_with_options(["Type 1"], None, WindowMapping::Explode)
.flatten()
.alias("strongest/group"),
col("Name")
.sort(Default::default())
.head(Some(3))
.over_with_options(["Type 1"], None, WindowMapping::Explode)
.flatten()
.alias("sorted_by_alphabet"),
])
.collect()?;
println!("{:?}", result);
shape: (43, 4)
┌────────┬───────────────────────┬───────────────────────┬─────────────────────────┐
│ Type 1 ┆ fastest/group ┆ strongest/group ┆ sorted_by_alphabet │
│ --- ┆ --- ┆ --- ┆ --- │
│ enum ┆ str ┆ str ┆ str │
╞════════╪═══════════════════════╪═══════════════════════╪═════════════════════════╡
│ Grass ┆ Venusaur ┆ Victreebel ┆ Bellsprout │
│ Grass ┆ VenusaurMega Venusaur ┆ VenusaurMega Venusaur ┆ Bulbasaur │
│ Grass ┆ Victreebel ┆ Exeggutor ┆ Exeggcute │
│ Water ┆ Starmie ┆ GyaradosMega Gyarados ┆ Blastoise │
│ Water ┆ Tentacruel ┆ Kingler ┆ BlastoiseMega Blastoise │
│ … ┆ … ┆ … ┆ … │
│ Rock ┆ Kabutops ┆ Kabutops ┆ Geodude │
│ Ice ┆ Jynx ┆ Articuno ┆ Articuno │
│ Ice ┆ Articuno ┆ Jynx ┆ Jynx │
│ Fairy ┆ Clefable ┆ Clefable ┆ Clefable │
│ Fairy ┆ Clefairy ┆ Clefairy ┆ Clefairy │
└────────┴───────────────────────┴───────────────────────┴─────────────────────────┘