列表和数组

Polars对两种同类容器数据类型提供一流的支持:ListArray

Polars 支持对这两种数据类型进行许多操作, 并且它们的 API 有所重叠, 因此本用户指南的这一部分旨在阐明何时应选择一种数据类型而不是另一种

列表与数组

数据类型List

List: 容纳值为不同长度的但是数据类型一致的一维容器

我们看下面代码, 一共三列, 每一列的元素都是一个列表, 虽然每一列中每个元素的长度可能不同, 但是数据类型是一样的

1from datetime import datetime
2import polars as pl
3
4df = pl.DataFrame(
5    {
6        "names": [
7            ["Anne", "Averill", "Adams"],
8            ["Brandon", "Brooke", "Borden", "Branson"],
9            ["Camila", "Campbell"],
10            ["Dennis", "Doyle"],
11        ],
12        "children_ages": [
13            [5, 7],
14            [],
15            [],
16            [8, 11, 18],
17        ],
18        "medical_appointments": [
19            [],
20            [],
21            [],
22            [datetime(2022, 5, 22, 16, 30)],
23        ],
24    }
25)
26
27print(df)
28print(df.dtypes)
1shape: (4, 3)
2┌─────────────────────────────────┬───────────────┬───────────────────────┐
3│ names                           ┆ children_ages ┆ medical_appointments  │
4│ ---                             ┆ ---           ┆ ---                   │
5│ list[str]                       ┆ list[i64]     ┆ list[datetime[μs]]    │
6╞═════════════════════════════════╪═══════════════╪═══════════════════════╡
7│ ["Anne", "Averill", "Adams"]    ┆ [5, 7]        ┆ []                    │
8│ ["Brandon", "Brooke", … "Brans… ┆ []            ┆ []                    │
9│ ["Camila", "Campbell"]          ┆ []            ┆ []                    │
10│ ["Dennis", "Doyle"]             ┆ [8, 11, 18]   ┆ [2022-05-22 16:30:00] │
11└─────────────────────────────────┴───────────────┴───────────────────────┘
12[List(String), List(Int64), List(Datetime(time_unit='us', time_zone=None))]
NOTE

数据类型List与Python的list不同, Python的list可以是任意类型.

数据类型Array

Array: 适用于值为已知和固定形状的任意维度的同质容器

下面代码中, 一共两列, 对于每一列来说, 要求每行的元素长度相同, 数据元素类型相同

1df = pl.DataFrame(
2    {
3        "bit_flags": [
4            [True, True, True, True, False],
5            [False, True, True, True, True],
6        ],
7        "tic_tac_toe": [
8            [
9                [" ", "x", "o"],
10                [" ", "x", " "],
11                ["o", "x", " "],
12            ],
13            [
14                ["o", "x", "x"],
15                [" ", "o", "x"],
16                [" ", " ", "o"],
17            ],
18        ],
19    },
20    schema={
21        "bit_flags": pl.Array(pl.Boolean, 5),
22        "tic_tac_toe": pl.Array(pl.String, (3, 3)),
23    },
24)
25
26print(df)
1shape: (2, 2)
2┌───────────────────────┬─────────────────────────────────┐
3│ bit_flags             ┆ tic_tac_toe                     │
4│ ---                   ┆ ---                             │
5│ array[bool, 5]        ┆ array[str, (3, 3)]              │
6╞═══════════════════════╪═════════════════════════════════╡
7│ [true, true, … false] ┆ [[" ", "x", "o"], [" ", "x", "… │
8│ [false, true, … true] ┆ [["o", "x", "x"], [" ", "o", "… │
9└───────────────────────┴─────────────────────────────────┘
10[Array(Boolean, shape=(5,)), Array(String, shape=(3, 3))]

出于性能原因, Polars不会推断Array, 而是默认使用List, 在Python中, 此规则的例外是使用Numpy数组来构建列时. 此时Polars会从Numpy获得所有子数组具有相同形状的保证, 因此n+1维度将生成一列n维数组.

1import numpy as np
2
3array = np.arange(0, 120).reshape((5, 2, 3, 4))  # 4D array
4
5print(pl.Series(array).dtype)  # Column with the 3D subarrays

Array(Int64, shape=(2, 3, 4))

如何选择

优先选择Array, 因为ArrayList内存效率更高, 性能也更好. 如果不能使用Array, 就使用List(列中的值没有固定形状, 或者需要使用只有List才有的功能)

使用List

命名空间list

Polars提供了很多函数去处理List数据类型, 这些函数在命名空间list

下面看官方文档给的示例, 有一列表示气象站的名称, 另外一列表示该站点的观测数据

1weather = pl.DataFrame(
2    {
3        "station": [f"Station {idx}" for idx in range(1, 6)],
4        "temperatures": [
5            "20 5 5 E1 7 13 19 9 6 20",
6            "18 8 16 11 23 E2 8 E2 E2 E2 90 70 40",
7            "19 24 E9 16 6 12 10 22",
8            "E2 E0 15 7 8 10 E1 24 17 13 6",
9            "14 8 E0 16 22 24 E1",
10        ],
11    }
12)
13
14print(weather)
1shape: (5, 2)
2┌───────────┬─────────────────────────────────┐
3│ station   ┆ temperatures                    │
4│ ---       ┆ ---                             │
5│ str       ┆ str                             │
6╞═══════════╪═════════════════════════════════╡
7│ Station 1 ┆ 20 5 5 E1 7 13 19 9 6 20        │
8│ Station 2 ┆ 18 8 16 11 23 E2 8 E2 E2 E2 90… │
9│ Station 3 ┆ 19 24 E9 16 6 12 10 22          │
10│ Station 4 ┆ E2 E0 15 7 8 10 E1 24 17 13 6   │
11│ Station 5 ┆ 14 8 E0 16 22 24 E1             │
12└───────────┴─────────────────────────────────┘

以编程方式创建列表

当我们想对每个站点捕获的温度进行分析时, 可以先把temperatures这一列中的每一行按照空格进行分割组织成列表

1weather = weather.with_columns(
2    pl.col("temperatures").str.split(" "),
3)
4print(weather)
1shape: (5, 2)
2┌───────────┬──────────────────────┐
3│ station   ┆ temperatures         │
4│ ---       ┆ ---                  │
5│ str       ┆ list[str]            │
6╞═══════════╪══════════════════════╡
7│ Station 1 ┆ ["20", "5", … "20"]  │
8│ Station 2 ┆ ["18", "8", … "40"]  │
9│ Station 3 ┆ ["19", "24", … "22"] │
10│ Station 4 ┆ ["E2", "E0", … "6"]  │
11│ Station 5 ┆ ["14", "8", … "E1"]  │
12└───────────┴──────────────────────┘

如果我们想展开温度列表, explode方法可以很方便的做到这一点

1result = weather.explode("temperatures")
2print(result)
1shape: (49, 2)
2┌───────────┬──────────────┐
3│ station   ┆ temperatures │
4│ ---       ┆ ---          │
5│ str       ┆ str          │
6╞═══════════╪══════════════╡
7│ Station 1 ┆ 20           │
8│ Station 1 ┆ 5            │
9│ Station 1 ┆ 5            │
10│ Station 1 ┆ E1           │
11│ Station 1 ┆ 7            │
12│ …         ┆ …            │
13│ Station 5 ┆ E0           │
14│ Station 5 ┆ 16           │
15│ Station 5 ┆ 22           │
16│ Station 5 ┆ 24           │
17│ Station 5 ┆ E1           │
18└───────────┴──────────────┘

对列表进行操作

与处理字符串相似, Polars提供了很多针对列表的函数, 在命名空间list中, 我们直接看代码

1result = weather.with_columns(
2    pl.col("temperatures").list.head(3).alias("head"),  # 获取前3个元素
3    pl.col("temperatures").list.tail(3).alias("tail"),  # 获取最后3个元素
4    pl.col("temperatures").list.slice(-3, 2).alias("two_next_to_last"),
5)
6print(result)
1shape: (5, 5)
2┌───────────┬──────────────────────┬────────────────────┬────────────────────┬──────────────────┐
3│ station   ┆ temperatures         ┆ head               ┆ tail               ┆ two_next_to_last │
4│ ---       ┆ ---                  ┆ ---                ┆ ---                ┆ ---              │
5│ str       ┆ list[str]            ┆ list[str]          ┆ list[str]          ┆ list[str]        │
6╞═══════════╪══════════════════════╪════════════════════╪════════════════════╪══════════════════╡
7│ Station 1 ┆ ["20", "5", … "20"]  ┆ ["20", "5", "5"]   ┆ ["9", "6", "20"]   ┆ ["9", "6"]       │
8│ Station 2 ┆ ["18", "8", … "40"]  ┆ ["18", "8", "16"]  ┆ ["90", "70", "40"] ┆ ["90", "70"]     │
9│ Station 3 ┆ ["19", "24", … "22"] ┆ ["19", "24", "E9"] ┆ ["12", "10", "22"] ┆ ["12", "10"]     │
10│ Station 4 ┆ ["E2", "E0", … "6"]  ┆ ["E2", "E0", "15"] ┆ ["17", "13", "6"]  ┆ ["17", "13"]     │
11│ Station 5 ┆ ["14", "8", … "E1"]  ┆ ["14", "8", "E0"]  ┆ ["22", "24", "E1"] ┆ ["22", "24"]     │
12└───────────┴──────────────────────┴────────────────────┴────────────────────┴──────────────────┘

对列表中的元素逐一计算

如果我们想找到错误最多的站点(以E开头的元素), 我们需要这样做:

  1. 尝试将测量结果转化为数字, 对于失败的会变为null(使用cast, 指定strict=False)
  2. 按行计算列表中非数字值的数量
  3. 将此列输出重命名为"error"

eval函数用于对列表元素执行操作, 我们可以使用上下文element分别引用列表中的每个元素, 然后就可以对元素使用任何Polars表达式了 直接看代码注释

1result = weather.with_columns(
2    pl.col("temperatures")  # 选中temperatures这一列
3        .list               # 获取命名空间 list
4        .eval(              # 使用eval方法, 针对列表中的每个元素进行操作
5            pl.element()    # 获取列表中的元素
6                .cast(pl.Int64, strict=False)   # 转换为数字, 失败的会变为null
7                .is_null()                      # 获取是否为null
8        )
9        .list.sum()                             # 对列表中的元素求和, True为1, False为0
10        .alias("errors"),
11)
12print(result)
1shape: (5, 3)
2┌───────────┬──────────────────────┬────────┐
3│ station   ┆ temperatures         ┆ errors │
4│ ---       ┆ ---                  ┆ ---    │
5│ str       ┆ list[str]            ┆ u32    │
6╞═══════════╪══════════════════════╪════════╡
7│ Station 1 ┆ ["20", "5", … "20"]  ┆ 1      │
8│ Station 2 ┆ ["18", "8", … "40"]  ┆ 4      │
9│ Station 3 ┆ ["19", "24", … "22"] ┆ 1      │
10│ Station 4 ┆ ["E2", "E0", … "6"]  ┆ 3      │
11│ Station 5 ┆ ["14", "8", … "E1"]  ┆ 2      │
12└───────────┴──────────────────────┴────────┘

另一种方法是使用正则表达式来检查测量值是否以字母开头

1result2 = weather.with_columns(
2    pl.col("temperatures")
3        .list.eval(pl.element().str.contains("(?i)[a-z]"))
4        .list.sum()
5        .alias("errors"),
6)
7print(result.equals(result2))

True

逐行计算

我们可以使用pl.all()来引用列表中的所有元素

我们先创建一个包含更多天气数据的DataFrame

1weather_by_day = pl.DataFrame(
2    {
3        "station": [f"Station {idx}" for idx in range(1, 11)],
4        "day_1": [17, 11, 8, 22, 9, 21, 20, 8, 8, 17],
5        "day_2": [15, 11, 10, 8, 7, 14, 18, 21, 15, 13],
6        "day_3": [16, 15, 24, 24, 8, 23, 19, 23, 16, 10],
7    }
8)
9print(weather_by_day)
1shape: (10, 4)
2┌────────────┬───────┬───────┬───────┐
3│ station    ┆ day_1 ┆ day_2 ┆ day_3 │
4│ ---        ┆ ---   ┆ ---   ┆ ---   │
5│ str        ┆ i64   ┆ i64   ┆ i64   │
6╞════════════╪═══════╪═══════╪═══════╡
7│ Station 1  ┆ 17    ┆ 15    ┆ 16    │
8│ Station 2  ┆ 11    ┆ 11    ┆ 15    │
9│ Station 3  ┆ 8     ┆ 10    ┆ 24    │
10│ Station 4  ┆ 22    ┆ 8     ┆ 24    │
11│ Station 5  ┆ 9     ┆ 7     ┆ 8     │
12│ Station 6  ┆ 21    ┆ 14    ┆ 23    │
13│ Station 7  ┆ 20    ┆ 18    ┆ 19    │
14│ Station 8  ┆ 8     ┆ 21    ┆ 23    │
15│ Station 9  ┆ 8     ┆ 15    ┆ 16    │
16│ Station 10 ┆ 17    ┆ 13    ┆ 10    │
17└────────────┴───────┴───────┴───────┘

现在我们计算各个站点每日温度的百分比排名

1rank_pct = (pl.element().rank(descending=True) / pl.element().count()).round(2)
2
3result = weather_by_day.with_columns(
4    # concat_list将每一行中的所有元素合并成一个列表, 用于生成一个临时的中间列 "all_temps"
5    pl.concat_list(pl.all().exclude("station")).alias("all_temps")
6).select(
7    # 选择除了 all_temps 这一列之外的所有列
8    pl.all().exclude("all_temps"),
9    # compute the rank by calling `list.eval`
10    pl.col("all_temps").list.eval(rank_pct, parallel=True).alias("temps_rank"),
11)
12
13print(result)
1shape: (10, 5)
2┌────────────┬───────┬───────┬───────┬────────────────────┐
3│ station    ┆ day_1 ┆ day_2 ┆ day_3 ┆ temps_rank         │
4│ ---        ┆ ---   ┆ ---   ┆ ---   ┆ ---                │
5│ str        ┆ i64   ┆ i64   ┆ i64   ┆ list[f64]          │
6╞════════════╪═══════╪═══════╪═══════╪════════════════════╡
7│ Station 1  ┆ 17    ┆ 15    ┆ 16    ┆ [0.33, 1.0, 0.67]  │
8│ Station 2  ┆ 11    ┆ 11    ┆ 15    ┆ [0.83, 0.83, 0.33] │
9│ Station 3  ┆ 8     ┆ 10    ┆ 24    ┆ [1.0, 0.67, 0.33]  │
10│ Station 4  ┆ 22    ┆ 8     ┆ 24    ┆ [0.67, 1.0, 0.33]  │
11│ Station 5  ┆ 9     ┆ 7     ┆ 8     ┆ [0.33, 1.0, 0.67]  │
12│ Station 6  ┆ 21    ┆ 14    ┆ 23    ┆ [0.67, 1.0, 0.33]  │
13│ Station 7  ┆ 20    ┆ 18    ┆ 19    ┆ [0.33, 1.0, 0.67]  │
14│ Station 8  ┆ 8     ┆ 21    ┆ 23    ┆ [1.0, 0.67, 0.33]  │
15│ Station 9  ┆ 8     ┆ 15    ┆ 16    ┆ [1.0, 0.67, 0.33]  │
16│ Station 10 ┆ 17    ┆ 13    ┆ 10    ┆ [0.33, 0.67, 1.0]  │
17└────────────┴───────┴───────┴───────┴────────────────────┘

使用 Array

Polars通常不会对Array做类型推导, 必须在创建Series/DataFrame是显式指定数据类型, 或者显式地转换列

命名空间arr

下面是一些基本的示例, 展示了几个常见的函数用法

1import polars as pl
2
3df = pl.DataFrame(
4    {
5        "first_last": [
6            ["Anne", "Adams"],
7            ["Brandon", "Branson"],
8            ["Camila", "Campbell"],
9            ["Dennis", "Doyle"],
10        ],
11        "fav_numbers": [
12            [42, 0, 1],
13            [2, 3, 5],
14            [13, 21, 34],
15            [73, 3, 7],
16        ],
17    },
18    schema={
19        "first_last": pl.Array(pl.String, 2),
20        "fav_numbers": pl.Array(pl.Int32, 3),
21    },
22)
23
24result = df.select(
25    pl.col("first_last").arr.join(" ").alias("name"),   # 对name列按照空格连接
26    pl.col("fav_numbers").arr.sort(),                   # 对fav_numbers列进行排序
27    pl.col("fav_numbers").arr.max().alias("largest_fav"),   # 获取fav_numbers列中每行元素的最大值
28    pl.col("fav_numbers").arr.sum().alias("summed"),        # 对fav_numbers列的每行元素分别进行求和
29    pl.col("fav_numbers").arr.contains(3).alias("likes_3"), # 对fav_numbers列的每行元素进行判断, 是否包含3
30)
31print(result)
1shape: (4, 5)
2┌─────────────────┬───────────────┬─────────────┬────────┬─────────┐
3│ name            ┆ fav_numbers   ┆ largest_fav ┆ summed ┆ likes_3 │
4│ ---             ┆ ---           ┆ ---         ┆ ---    ┆ ---     │
5│ str             ┆ array[i32, 3] ┆ i32         ┆ i32    ┆ bool    │
6╞═════════════════╪═══════════════╪═════════════╪════════╪═════════╡
7│ Anne Adams      ┆ [0, 1, 42]    ┆ 42          ┆ 43     ┆ false   │
8│ Brandon Branson ┆ [2, 3, 5]     ┆ 5           ┆ 10     ┆ true    │
9│ Camila Campbell ┆ [13, 21, 34]  ┆ 34          ┆ 68     ┆ false   │
10│ Dennis Doyle    ┆ [3, 7, 73]    ┆ 73          ┆ 83     ┆ true    │
11└─────────────────┴───────────────┴─────────────┴────────┴─────────┘