我们曾在一则社交媒体帖子中提出一个问题:.struct.unnest 会不会打破 Polars 最基本的原则之一——一个表达式必须恰好输出一列?本文会详细回答这个问题,先透露结论:最后一切都说得通。
简短回答
.struct.unnest 使用了表达式展开。它会为结构体中的每个字段展开出一个表达式,因此看上去像是一个表达式输出了多列。
问题在哪里?
Polars 查询遵守一条简单规则:查询中的一个表达式输出一列,而且只输出一列。
import polars as pl
df = pl.DataFrame(
{
"name": ["Anne", "Abe", "Anastasia", "Anton"],
"age": [16, 23, 62, 8],
}
)
print(
df.select(
pl.col("name").str.len_chars().alias("name_length"), # 1
(pl.col("age") > 18).alias("is_adult"), # 2
(2024 - pl.col("age")).alias("birthyear"), # 3
(pl.col("name").str.len_chars() * pl.col("age")).alias("hun?"), # 4
)
)
shape: (4, 4)
# 1 # 2 # 3 # 4
┌─────────────┬──────────┬───────────┬──────┐
│ name_length ┆ is_adult ┆ birthyear ┆ hun? │
│ --- ┆ --- ┆ --- ┆ --- │
│ u32 ┆ bool ┆ i64 ┆ i64 │
╞═════════════╪══════════╪═══════════╪══════╡
│ 4 ┆ false ┆ 2008 ┆ 64 │
│ 3 ┆ true ┆ 2001 ┆ 69 │
│ 9 ┆ true ┆ 1962 ┆ 558 │
│ 5 ┆ false ┆ 2016 ┆ 40 │
└─────────────┴──────────┴───────────┴──────┘
“一个表达式产生一列”是一项设计决定,让复杂查询的结果 DataFrame 模式更容易推理。然而,使用 .struct.unnest 时,一个表达式似乎产生了多列:
df = pl.DataFrame(
{
"structs": [
{"a": 1, "b": 2, "c": 3},
{"a": 4, "b": 5, "c": 6},
{"a": 7, "b": 8, "c": 9},
]
}
)
print(
df.select(
pl.col("structs").struct.unnest(), # 1
)
)
shape: (3, 3)
# 1 # 2 # 3
┌─────┬─────┬─────┐
│ a ┆ b ┆ c │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ 1 ┆ 2 ┆ 3 │
│ 4 ┆ 5 ┆ 6 │
│ 7 ┆ 8 ┆ 9 │
└─────┴─────┴─────┘
这个例子似乎与规则相悖,但实际上并没有。理解原因,需要先了解表达式展开。
表达式展开
表达式展开是 Polars 的一项功能,可以用简短而强大的表达式避免结构重复。当需要对多列执行相同计算时,可以只写一个表达式,让 Polars 把它展开成多个并行表达式。
下面两个表达式的结构完全相同:
df = pl.DataFrame(
{
"first_name": ["Anne", "Abe", "Anastasia", "Anton"],
"last_name": ["Holmes", "Watson", "Kent", "Wayne"],
}
)
df.select(
pl.col("first_name").str.len_chars(), # 1.
pl.col("last_name").str.len_chars(), # 2.
)
通过函数 pl.col 的表达式展开,就能避免这种重复:
df.select(
pl.col("first_name", "last_name").str.len_chars(),
)
在这个例子中,我们提前知道 Polars 会创建多少个并行表达式,因为 pl.col 明确列出了所有需要计算字符串长度的列。
表达式展开也可以让并行表达式的数量取决于应用它的 DataFrame 模式。例如:
import polars.selectors as cs
plus_one = cs.integer() + 1
plus_one 使用列选择器引用所有整数列。不过,表达式只是计算的表示方式;它需要一个 Polars 上下文才能判断具体有哪些整数列。在下面的上下文中,它没有引用任何列:
df = pl.DataFrame(
{
"first_name": ["Anne", "Abe", "Anastasia", "Anton"],
"last_name": ["Holmes", "Watson", "Kent", "Wayne"],
}
)
print(df.select(plus_one))
shape: (0, 0)
┌┐
╞╡
└┘
回到本文最初的 DataFrame,plus_one 则展开为一列:
df = pl.DataFrame(
{
"name": ["Anne", "Abe", "Anastasia", "Anton"],
"age": [16, 23, 62, 8],
}
)
print(df.select(plus_one))
shape: (4, 1)
┌─────┐
│ age │
│ --- │
│ i64 │
╞═════╡
│ 17 │
│ 24 │
│ 63 │
│ 9 │
└─────┘
再换一个 DataFrame,可以看到 plus_one 展开为五个并行表达式:
df = pl.DataFrame({"one": 1, "two": 2, "three": 3, "four": 4, "five": 5})
print(df.select(plus_one))
shape: (1, 5)
┌─────┬─────┬───────┬──────┬──────┐
│ one ┆ two ┆ three ┆ four ┆ five │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═══════╪══════╪══════╡
│ 2 ┆ 3 ┆ 4 ┆ 5 ┆ 6 │
└─────┴─────┴───────┴──────┴──────┘
结构体的表达式展开
我们已经看到了 pl.col 函数与 selectors 模块如何展开表达式。这只是它们功能的表面,接下来把注意力转向结构体。
结构体数据类型可以想象成类似 Python 字典,由任意数量的命名字段组成。下面,value_counts 表达式生成一个含两个字段的结构体:
df = pl.DataFrame(
{
"ballot_id": [6145, 176723, 345623, 77234, 75246],
"best_movie": ["Cars", "Toy Story", "Toy Story", "Cars", "Toy Story"],
}
)
votes = df.select(pl.col("best_movie").value_counts())
print(votes)
shape: (2, 1)
┌─────────────────┐
│ best_movie │
│ --- │
│ struct[2] │
╞═════════════════╡
│ {"Cars",2} │
│ {"Toy Story",3} │
└─────────────────┘
value_counts 返回结构体,是因为票数需要和电影名配对;它本身是单个表达式,所以必须只输出一列。
可以用 .struct.field 访问结构体字段,该表达式接受要提取的字段名。票数所在字段名为 "count",因此可以这样提取:
print(votes.select(pl.col("best_movie").struct.field("count")))
shape: (2, 1)
┌───────┐
│ count │
│ --- │
│ u32 │
╞═══════╡
│ 2 │
│ 3 │
└───────┘
不过,与 pl.col 一样,.struct.field 也接受任意数量的待提取字段名:
print(votes.select(pl.col("best_movie").struct.field("best_movie", "count")))
shape: (2, 2)
┌────────────┬───────┐
│ best_movie ┆ count │
│ --- ┆ --- │
│ str ┆ u32 │
╞════════════╪═══════╡
│ Cars ┆ 2 │
│ Toy Story ┆ 3 │
└────────────┴───────┘
pl.col 和 .struct.field 还支持通过正则表达式模式进行展开。它们通常支持正则表达式语法,也支持特殊通配符参数 "*",匹配全部名称:
df = pl.DataFrame(
{
"ballot_id": [6145, 176723, 345623, 77234, 75246],
"best_movie": ["Cars", "Toy Story", "Toy Story", "Cars", "Toy Story"],
}
)
print(df.select(pl.col("*").name.to_uppercase()))
shape: (5, 2)
┌───────────┬────────────┐
│ BALLOT_ID ┆ BEST_MOVIE │
│ --- ┆ --- │
│ i64 ┆ str │
╞═══════════╪════════════╡
│ 6145 ┆ Cars │
│ 176723 ┆ Toy Story │
│ 345623 ┆ Toy Story │
│ 77234 ┆ Cars │
│ 75246 ┆ Toy Story │
└───────────┴────────────┘
print(votes.select(pl.col("best_movie").struct.field("*")))
shape: (2, 2)
┌────────────┬───────┐
│ best_movie ┆ count │
│ --- ┆ --- │
│ str ┆ u32 │
╞════════════╪═══════╡
│ Cars ┆ 2 │
│ Toy Story ┆ 3 │
└────────────┴───────┘
最后一块拼图是 Polars 为 pl.col("*") 和 .struct.field("*") 提供了别名:前者推荐写成更易读的 pl.all(),后者的别名则是 .struct.unnest()。因此,.struct.unnest 看似以一个表达式产生多列,其实是通过通配符匹配展开了表达式。
当它应用于结构体列时,Polars 会根据字段数量展开出同样数量的并行表达式,分别将每个字段提取为一列。
进一步阅读
表达式展开还有许多本文未介绍的内容,欢迎阅读用户指南的表达式展开章节:其中包括 pl.col 的更多展开方式、从模式中排除列、在展开过程中重命名列、selectors 模块的功能,以及通过程序生成表达式。











暂无评论内容