结构体

数据类型Struct是一种复合类型, 可以在单个列中存储多个字段

这一章节中我们会解释为什么要有Struct, 以及如何处理Struct类型的值

我们有一个DataFrame, 显示了美国一些州的几部电影的平均分

1import polars as pl
2
3ratings = pl.DataFrame(
4    {
5        "Movie": ["Cars", "IT", "ET", "Cars", "Up", "IT", "Cars", "ET", "Up", "Cars"],
6        "Theatre": ["NE", "ME", "IL", "ND", "NE", "SD", "NE", "IL", "IL", "NE"],
7        "Avg_Rating": [4.5, 4.4, 4.6, 4.3, 4.8, 4.7, 4.5, 4.9, 4.7, 4.6],
8        "Count": [30, 27, 26, 29, 31, 28, 28, 26, 33, 28],
9    }
10)
11print(ratings)
1shape: (10, 4)
2┌───────┬─────────┬────────────┬───────┐
3│ Movie ┆ Theatre ┆ Avg_Rating ┆ Count │
4│ ---   ┆ ---     ┆ ---        ┆ ---   │
5│ str   ┆ str     ┆ f64        ┆ i64   │
6╞═══════╪═════════╪════════════╪═══════╡
7│ Cars  ┆ NE      ┆ 4.5        ┆ 30    │
8│ IT    ┆ ME      ┆ 4.4        ┆ 27    │
9│ ET    ┆ IL      ┆ 4.6        ┆ 26    │
10│ Cars  ┆ ND      ┆ 4.3        ┆ 29    │
11│ Up    ┆ NE      ┆ 4.8        ┆ 31    │
12│ IT    ┆ SD      ┆ 4.7        ┆ 28    │
13│ Cars  ┆ NE      ┆ 4.5        ┆ 28    │
14│ ET    ┆ IL      ┆ 4.9        ┆ 26    │
15│ Up    ┆ IL      ┆ 4.7        ┆ 33    │
16│ Cars  ┆ NE      ┆ 4.6        ┆ 28    │
17└───────┴─────────┴────────────┴───────┘

何时遇到Struct

在下面的代码中, 我们统计了每个 Theatre 的电影数量, 观察下返回结果, 此时只输出了一列, 类型为struct[2]

1result = ratings.select(pl.col("Theatre").value_counts(sort=True))
2print(result)
1shape: (5, 1)
2┌───────────┐
3│ Theatre   │
4│ ---       │
5│ struct[2] │
6╞═══════════╡
7│ {"NE",4}  │
8│ {"IL",3}  │
9│ {"ME",1}  │
10│ {"ND",1}  │
11│ {"SD",1}  │
12└───────────┘

如果我们想看到之前更熟悉的那种输出, 可以使用unnest, unnest将结构列分解为每个字段的单独列, 会插入一个新的列, 可以看到多了一列"count"

1result = ratings.select(pl.col("Theatre").value_counts(sort=True)).unnest("Theatre")
2print(result)
1shape: (5, 2)
2┌─────────┬───────┐
3│ Theatre ┆ count │
4│ ---     ┆ ---   │
5│ str     ┆ u32   │
6╞═════════╪═══════╡
7│ NE      ┆ 4     │
8│ IL      ┆ 3     │
9│ ME      ┆ 1     │
10│ ND      ┆ 1     │
11│ SD      ┆ 1     │
12└─────────┴───────┘
为什么value_counts返回Struct

Polars表达式使用作用于单个Series, 并返回另一个Series. Struct是一种数据类型, 允许我们将多列作为表达式输入, 或从表达式中输出多列.

从字典推断数据类型

当构建SeriesDataFrame时, Polars会将字典转换为数据类型Struct

1rating_series = pl.Series(
2    "ratings",
3    [
4        {"Movie": "Cars", "Theatre": "NE", "Avg_Rating": 4.5},
5        {"Movie": "Toy Story", "Theatre": "ME", "Avg_Rating": 4.9},
6    ],
7)
8print(rating_series)
1shape: (2,)
2Series: 'ratings' [struct[3]]
3[
4    {"Cars","NE",4.5}
5    {"Toy Story","ME",4.9}
6]

字段数量, 名称和类型是根据第一个看到的字典推断出来的. 后续的不一致可能会导致null值或错误

1null_rating_series = pl.Series(
2    "ratings",
3    [
4        {"Movie": "Cars", "Theatre": "NE", "Avg_Rating": 4.5},
5        {"Mov": "Toy Story", "Theatre": "ME", "Avg_Rating": 4.9},
6        {"Movie": "Snow White", "Theatre": "IL", "Avg_Rating": "4.7"},
7    ],
8    strict=False,  # To show the final structs with `null` values.
9)
10print(null_rating_series)
1shape: (3,)
2Series: 'ratings' [struct[4]]
3[
4    {"Cars","NE","4.5",null}
5    {null,"ME","4.9","Toy Story"}
6    {"Snow White","IL","4.7",null}
7]

提取结构体的三个值

我们可以使用函数field来从中获取字段

1result = rating_series.struct.field("Movie")
2print(result)
1shape: (2,)
2Series: 'Movie' [str]
3[
4    "Cars"
5    "Toy Story"
6]

重命名Struct中的字段

我们可以使用rename_fields方法重命名字段

1import polars as pl
2
3s = pl.Series([{"a": 1, "b": 2}, {"a": 3, "b": 4}])
4print(s.struct.fields)
5s = s.struct.rename_fields(["c", "d"])
6print(s.struct.fields)

可以看到初始字段是ab, 重命名之后是cd

1['a', 'b']
2['c', 'd']

Struct的实际用例

识别重复行

我们可以识别"电影"和"剧院"存在重复的情况

1import polars as pl
2
3ratings = pl.DataFrame(
4    {
5        "Movie": ["Cars", "IT", "ET", "Cars", "Up", "IT", "Cars", "ET", "Up", "Cars"],
6        "Theatre": ["NE", "ME", "IL", "ND", "NE", "SD", "NE", "IL", "IL", "NE"],
7        "Avg_Rating": [4.5, 4.4, 4.6, 4.3, 4.8, 4.7, 4.5, 4.9, 4.7, 4.6],
8        "Count": [30, 27, 26, 29, 31, 28, 28, 26, 33, 28],
9    }
10)
11result = ratings.filter(pl.struct("Movie", "Theatre").is_duplicated())
12print(result)

仔细观察结果, 上述数据中, 存在多个Cars+NE, ET+IL, is_duplicated返回的是组合成的Struct重复的行

1shape: (5, 4)
2┌───────┬─────────┬────────────┬───────┐
3│ Movie ┆ Theatre ┆ Avg_Rating ┆ Count │
4│ ---   ┆ ---     ┆ ---        ┆ ---   │
5│ str   ┆ str     ┆ f64        ┆ i64   │
6╞═══════╪═════════╪════════════╪═══════╡
7│ Cars  ┆ NE      ┆ 4.5        ┆ 30    │
8│ ET    ┆ IL      ┆ 4.6        ┆ 26    │
9│ Cars  ┆ NE      ┆ 4.5        ┆ 28    │
10│ ET    ┆ IL      ┆ 4.9        ┆ 26    │
11│ Cars  ┆ NE      ┆ 4.6        ┆ 28    │
12└───────┴─────────┴────────────┴───────┘

当然也可以来识别独特的案例, 可以使用is_unique

多列排名

假设已知存在重复项, 我们想选择哪个评分具有更高的优先级. 我们假设Count这一列更重要, 如果Count列相同, 那么再考虑Avg_Rating

1# 先增加一列"Rank", 然后再过滤
2result = ratings.with_columns(
3    pl.struct("Count", "Avg_Rating")
4        .rank("dense", descending=True)
5        .over("Movie", "Theatre")
6        .alias("Rank")
7    ).filter(pl.struct("Movie", "Theatre").is_duplicated())
8
9print(result)

有关rank的详细信息, 可以查看后面的窗口函数部分

1shape: (5, 5)
2┌───────┬─────────┬────────────┬───────┬──────┐
3│ Movie ┆ Theatre ┆ Avg_Rating ┆ Count ┆ Rank │
4│ ---   ┆ ---     ┆ ---        ┆ ---   ┆ ---  │
5│ str   ┆ str     ┆ f64        ┆ i64   ┆ u32  │
6╞═══════╪═════════╪════════════╪═══════╪══════╡
7│ Cars  ┆ NE      ┆ 4.5        ┆ 30    ┆ 1    │
8│ ET    ┆ IL      ┆ 4.6        ┆ 26    ┆ 2    │
9│ Cars  ┆ NE      ┆ 4.5        ┆ 28    ┆ 3    │
10│ ET    ┆ IL      ┆ 4.9        ┆ 26    ┆ 1    │
11│ Cars  ┆ NE      ┆ 4.6        ┆ 28    ┆ 2    │
12└───────┴─────────┴────────────┴───────┴──────┘

在单个表达式中使用多个列

当我们想将多列作为表达式的输入的时候, Struct也很有用. 我们可以使用map_ elements

INFO

下面的add只是一个非常简单的示例, 实际上这个场景是一些稍微复杂的计算, 无法直接使用表达式来计算的, 比如无法通过pl.col("m")+pl.col("n")就能实现的

1def add(a: int, b: int)-> int:
2    return a+b
3
4values = pl.DataFrame(
5    {
6        "m": [0, 0, 0, 1, 1, 1, 2],
7        "n": [2, 3, 4, 1, 2, 3, 1],
8    }
9)
10result = values.with_columns(
11    pl.struct(["m", "n"])
12        .map_elements(function=lambda s: add(s["m"], s["n"]), return_dtype=pl.Int64)
13        .alias("add")
14)
15
16print(result)