表达式展开:看似打破规则,其实遵守规则

我们曾在一则社交媒体帖子中提出一个问题:.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 模块的功能,以及通过程序生成表达式。

原文:Breaking the rules with expression expansion;作者 Rodrigo Girão Serrão;2024年12月5日。来源 Polars。依转载授权汉化,保留原文代码与示例输出;输出为原文所载结果,未在本环境执行。原文未标明具体 Polars 版本。

© 版权声明
THE END
喜欢就支持一下吧
点赞0 分享
评论 抢沙发

请登录后发表评论

    暂无评论内容