表达式扩展

表达式扩展功能, 能够编写一个可以扩展为多个不同表达式的表达式,具体怎么扩展取决于使用表达式的上下文的模式。

遐想

(x+y+z) * m = (x * m) + (y * m) + (z * m)

数据准备

Python
1df = pl.DataFrame(
2    {  # As of 14th October 2024, ~3pm UTC
3        "ticker": ["AAPL", "NVDA", "MSFT", "GOOG", "AMZN"],
4        "company_name": ["Apple", "NVIDIA", "Microsoft", "Alphabet (Google)", "Amazon"],
5        "price": [229.9, 138.93, 420.56, 166.41, 188.4],
6        "day_high": [231.31, 139.6, 424.04, 167.62, 189.83],
7        "day_low": [228.6, 136.3, 417.52, 164.78, 188.44],
8        "year_high": [237.23, 140.76, 468.35, 193.31, 201.2],
9        "year_low": [164.08, 39.23, 324.39, 121.46, 118.35],
10    }
11)
12
13print(df)
1shape: (5, 7)
2┌────────┬───────────────────┬────────┬──────────┬─────────┬───────────┬──────────┐
3│ ticker ┆ company_name      ┆ price  ┆ day_high ┆ day_low ┆ year_high ┆ year_low │
4│ ---    ┆ ---               ┆ ---    ┆ ---      ┆ ---     ┆ ---       ┆ ---      │
5│ str    ┆ str               ┆ f64    ┆ f64      ┆ f64     ┆ f64       ┆ f64      │
6╞════════╪═══════════════════╪════════╪══════════╪═════════╪═══════════╪══════════╡
7│ AAPL   ┆ Apple             ┆ 229.9  ┆ 231.31   ┆ 228.6   ┆ 237.23    ┆ 164.08   │
8│ NVDA   ┆ NVIDIA            ┆ 138.93 ┆ 139.6    ┆ 136.3   ┆ 140.76    ┆ 39.23    │
9│ MSFT   ┆ Microsoft         ┆ 420.56 ┆ 424.04   ┆ 417.52  ┆ 468.35    ┆ 324.39   │
10│ GOOG   ┆ Alphabet (Google) ┆ 166.41 ┆ 167.62   ┆ 164.78  ┆ 193.31    ┆ 121.46   │
11│ AMZN   ┆ Amazon            ┆ 188.4  ┆ 189.83   ┆ 188.44  ┆ 201.2     ┆ 118.35   │
12└────────┴───────────────────┴────────┴──────────┴─────────┴───────────┴──────────┘

选择列

  • col函数用来引用 DataFrame 中的一列。

常规取列用法

对数值类型的列都进行除法操作:

Python
1eur_usd_rate = 1.09
2
3newdf = df.select(
4    pl.col("price") / eur_usd_rate,
5    pl.col("day_high") / eur_usd_rate,
6    pl.col("day_low") / eur_usd_rate,
7    pl.col("year_high") / eur_usd_rate,
8    pl.col("year_low") / eur_usd_rate
9)
10
11print(newdf)
1shape: (5, 5)
2┌────────────┬────────────┬────────────┬────────────┬────────────┐
3│ price      ┆ day_high   ┆ day_low    ┆ year_high  ┆ year_low   │
4│ ---        ┆ ---        ┆ ---        ┆ ---        ┆ ---        │
5│ f64        ┆ f64        ┆ f64        ┆ f64        ┆ f64        │
6╞════════════╪════════════╪════════════╪════════════╪════════════╡
7│ 210.917431 ┆ 212.211009 ┆ 209.724771 ┆ 217.642202 ┆ 150.53211  │
8│ 127.458716 ┆ 128.073394 ┆ 125.045872 ┆ 129.137615 ┆ 35.990826  │
9│ 385.834862 ┆ 389.027523 ┆ 383.045872 ┆ 429.678899 ┆ 297.605505 │
10│ 152.669725 ┆ 153.779817 ┆ 151.174312 ┆ 177.348624 ┆ 111.431193 │
11│ 172.844037 ┆ 174.155963 ┆ 172.880734 ┆ 184.587156 ┆ 108.577982 │
12└────────────┴────────────┴────────────┴────────────┴────────────┘

上面常规的用法中,可以很明显的看出对于同一个处理逻辑,需要写很多次col函数,显得很冗余,所以可以使用 col函数提供的扩展功能。

按列名选择

col 函数提供多个列名时, 会发生最简单的表达式扩展形式。

Python
1eur_usd_rate = 1.09
2
3result = df.with_columns(
4    (
5            pl.col(
6                "price",
7                "day_high",
8                "day_low",
9                "year_high",
10                "year_low",
11            )
12            / eur_usd_rate
13    ).round(2)
14)
15print(result)
1shape: (5, 7)
2┌────────┬───────────────────┬────────┬──────────┬─────────┬───────────┬──────────┐
3│ ticker ┆ company_name      ┆ price  ┆ day_high ┆ day_low ┆ year_high ┆ year_low │
4│ ---    ┆ ---               ┆ ---    ┆ ---      ┆ ---     ┆ ---       ┆ ---      │
5│ str    ┆ str               ┆ f64    ┆ f64      ┆ f64     ┆ f64       ┆ f64      │
6╞════════╪═══════════════════╪════════╪══════════╪═════════╪═══════════╪══════════╡
7│ AAPL   ┆ Apple             ┆ 210.92 ┆ 212.21   ┆ 209.72  ┆ 217.64    ┆ 150.53   │
8│ NVDA   ┆ NVIDIA            ┆ 127.46 ┆ 128.07   ┆ 125.05  ┆ 129.14    ┆ 35.99    │
9│ MSFT   ┆ Microsoft         ┆ 385.83 ┆ 389.03   ┆ 383.05  ┆ 429.68    ┆ 297.61   │
10│ GOOG   ┆ Alphabet (Google) ┆ 152.67 ┆ 153.78   ┆ 151.17  ┆ 177.35    ┆ 111.43   │
11│ AMZN   ┆ Amazon            ┆ 172.84 ┆ 174.16   ┆ 172.88  ┆ 184.59    ┆ 108.58   │
12└────────┴───────────────────┴────────┴──────────┴─────────┴───────────┴──────────┘

上面的语句等价与:

Python
1exprs = [
2    (pl.col("price") / eur_usd_rate).round(2),
3    (pl.col("day_high") / eur_usd_rate).round(2),
4    (pl.col("day_low") / eur_usd_rate).round(2),
5    (pl.col("year_high") / eur_usd_rate).round(2),
6    (pl.col("year_low") / eur_usd_rate).round(2),
7]
8
9result2 = df.with_columns(exprs)
10
11print(result.equals(result2))
1True

本质就是把col函数扩展成了多个列名, 然后对每个列名进行相同的操作, 生成新的列。

按数据类型选择

在上例中,必须输入五个列名,但该函数 col 也可以方便地接受一种多种数据类型。如果提供的是数据类型而不是列名,则表达式将扩展为与提供的数据类型之一匹配的所有列。上面的示例也可以写成:

Python
1eur_usd_rate = 1.09
2
3result = df.with_columns(
4    (
5            pl.col(
6                "price",
7                "day_high",
8                "day_low",
9                "year_high",
10                "year_low",
11            )
12            / eur_usd_rate
13    ).round(2)
14)
15
16result2 = df.with_columns((pl.col(pl.Float64) / eur_usd_rate).round(2))
17print(result.equals(result2))
18print(result2)
1True
2shape: (5, 7)
3┌────────┬───────────────────┬────────┬──────────┬─────────┬───────────┬──────────┐
4│ ticker ┆ company_name      ┆ price  ┆ day_high ┆ day_low ┆ year_high ┆ year_low │
5│ ---    ┆ ---               ┆ ---    ┆ ---      ┆ ---     ┆ ---       ┆ ---      │
6│ str    ┆ str               ┆ f64    ┆ f64      ┆ f64     ┆ f64       ┆ f64      │
7╞════════╪═══════════════════╪════════╪══════════╪═════════╪═══════════╪══════════╡
8│ AAPL   ┆ Apple             ┆ 210.92 ┆ 212.21   ┆ 209.72  ┆ 217.64    ┆ 150.53   │
9│ NVDA   ┆ NVIDIA            ┆ 127.46 ┆ 128.07   ┆ 125.05  ┆ 129.14    ┆ 35.99    │
10│ MSFT   ┆ Microsoft         ┆ 385.83 ┆ 389.03   ┆ 383.05  ┆ 429.68    ┆ 297.61   │
11│ GOOG   ┆ Alphabet (Google) ┆ 152.67 ┆ 153.78   ┆ 151.17  ┆ 177.35    ┆ 111.43   │
12│ AMZN   ┆ Amazon            ┆ 172.84 ┆ 174.16   ┆ 172.88  ┆ 184.59    ┆ 108.58   │
13└────────┴───────────────────┴────────┴──────────┴─────────┴───────────┴──────────┘

甚至如果无法确定列的类型是Float64还是Float32时,可以同时指定两种数据类型:

Python
1result2 = df.with_columns(
2    (
3        pl.col(
4            pl.Float32,
5            pl.Float64,
6        )
7        / eur_usd_rate
8    ).round(2)
9)
10print(result.equals(result2))
1True
警告

col 函数可以接受任意数量的字符串或者任意数量的数据类型,但同一个函数中不能 同时 出现两者。如下错误示例:

Python
1result = df.with_columns(
2    (
3            pl.col(
4                "price",
5                "day_high",
6                pl.Float64
7            )
8            / eur_usd_rate
9    ).round(2)
10)

通过正则匹配选择

可以使用正则表达式来指定用于匹配的列名。为了区分常规列名和通过模式匹配扩展的列名, 正则表达式以^开头, 以$结尾。

  • 与数据类型扩展不同的是,正则表达式可以与常规列名混合使用
Python
1eur_usd_rate = 1.09
2
3result = df.with_columns(
4    (
5            pl.col("^.*_high$", "^.*_low$")
6            / eur_usd_rate
7    ).round(2)
8)
9
10print(result)
1shape: (5, 7)
2┌────────┬───────────────────┬────────┬──────────┬─────────┬───────────┬──────────┐
3│ ticker ┆ company_name      ┆ price  ┆ day_high ┆ day_low ┆ year_high ┆ year_low │
4│ ---    ┆ ---               ┆ ---    ┆ ---      ┆ ---     ┆ ---       ┆ ---      │
5│ str    ┆ str               ┆ f64    ┆ f64      ┆ f64     ┆ f64       ┆ f64      │
6╞════════╪═══════════════════╪════════╪══════════╪═════════╪═══════════╪══════════╡
7│ AAPL   ┆ Apple             ┆ 229.9  ┆ 212.21   ┆ 209.72  ┆ 217.64    ┆ 150.53   │
8│ NVDA   ┆ NVIDIA            ┆ 138.93 ┆ 128.07   ┆ 125.05  ┆ 129.14    ┆ 35.99    │
9│ MSFT   ┆ Microsoft         ┆ 420.56 ┆ 389.03   ┆ 383.05  ┆ 429.68    ┆ 297.61   │
10│ GOOG   ┆ Alphabet (Google) ┆ 166.41 ┆ 153.78   ┆ 151.17  ┆ 177.35    ┆ 111.43   │
11│ AMZN   ┆ Amazon            ┆ 188.4  ┆ 174.16   ┆ 172.88  ┆ 184.59    ┆ 108.58   │
12└────────┴───────────────────┴────────┴──────────┴─────────┴───────────┴──────────┘

同时存在列名和正则表达式时:

Python
1result = df.select(pl.col("ticker", "^.*_high$", "^.*_low$"))
2print(result)
1shape: (5, 5)
2┌────────┬──────────┬─────────┬───────────┬──────────┐
3│ ticker ┆ day_high ┆ day_low ┆ year_high ┆ year_low │
4│ ---    ┆ ---      ┆ ---     ┆ ---       ┆ ---      │
5│ str    ┆ f64      ┆ f64     ┆ f64       ┆ f64      │
6╞════════╪══════════╪═════════╪═══════════╪══════════╡
7│ AAPL   ┆ 231.31   ┆ 228.6   ┆ 237.23    ┆ 164.08   │
8│ NVDA   ┆ 139.6    ┆ 136.3   ┆ 140.76    ┆ 39.23    │
9│ MSFT   ┆ 424.04   ┆ 417.52  ┆ 468.35    ┆ 324.39   │
10│ GOOG   ┆ 167.62   ┆ 164.78  ┆ 193.31    ┆ 121.46   │
11│ AMZN   ┆ 189.83   ┆ 188.44  ┆ 201.2     ┆ 118.35   │
12└────────┴──────────┴─────────┴───────────┴──────────┘

选择所有列

Polars 提供该函数 all 作为简写符号来引用所有列:

1result = df.select(pl.all())
2print(result.equals(df))
3print(result)
1True
2shape: (5, 7)
3┌────────┬───────────────────┬────────┬──────────┬─────────┬───────────┬──────────┐
4│ ticker ┆ company_name      ┆ price  ┆ day_high ┆ day_low ┆ year_high ┆ year_low │
5│ ---    ┆ ---               ┆ ---    ┆ ---      ┆ ---     ┆ ---       ┆ ---      │
6│ str    ┆ str               ┆ f64    ┆ f64      ┆ f64     ┆ f64       ┆ f64      │
7╞════════╪═══════════════════╪════════╪══════════╪═════════╪═══════════╪══════════╡
8│ AAPL   ┆ Apple             ┆ 229.9  ┆ 231.31   ┆ 228.6   ┆ 237.23    ┆ 164.08   │
9│ NVDA   ┆ NVIDIA            ┆ 138.93 ┆ 139.6    ┆ 136.3   ┆ 140.76    ┆ 39.23    │
10│ MSFT   ┆ Microsoft         ┆ 420.56 ┆ 424.04   ┆ 417.52  ┆ 468.35    ┆ 324.39   │
11│ GOOG   ┆ Alphabet (Google) ┆ 166.41 ┆ 167.62   ┆ 164.78  ┆ 193.31    ┆ 121.46   │
12│ AMZN   ┆ Amazon            ┆ 188.4  ┆ 189.83   ┆ 188.44  ┆ 201.2     ┆ 118.35   │
13└────────┴───────────────────┴────────┴──────────┴─────────┴───────────┴──────────┘
TIP

函数 pl.allpl.col("*") 的语法糖,但由于参数 "*" 是特例并且all更加见其知意,因此更推荐使用 pl.all

排除列

Polars 还提供了一种从表达式扩展中排除某些列的机制。可以使用函数exclude来指定排除的列,该函数和col函数接受一样的参数类型:指定列名、指定列数据类型、正则匹配。

按列名排除

Python
1result = df.select(
2    pl.all().exclude("day_high","year_high")
3)
4print(result)
1shape: (5, 5)
2┌────────┬───────────────────┬────────┬─────────┬──────────┐
3│ ticker ┆ company_name      ┆ price  ┆ day_low ┆ year_low │
4│ ---    ┆ ---               ┆ ---    ┆ ---     ┆ ---      │
5│ str    ┆ str               ┆ f64    ┆ f64     ┆ f64      │
6╞════════╪═══════════════════╪════════╪═════════╪══════════╡
7│ AAPL   ┆ Apple             ┆ 229.9  ┆ 228.6   ┆ 164.08   │
8│ NVDA   ┆ NVIDIA            ┆ 138.93 ┆ 136.3   ┆ 39.23    │
9│ MSFT   ┆ Microsoft         ┆ 420.56 ┆ 417.52  ┆ 324.39   │
10│ GOOG   ┆ Alphabet (Google) ┆ 166.41 ┆ 164.78  ┆ 121.46   │
11│ AMZN   ┆ Amazon            ┆ 188.4  ┆ 188.44  ┆ 118.35   │
12└────────┴───────────────────┴────────┴─────────┴──────────┘

按数据类型排除

Python
1result = df.select(
2    pl.all().exclude(pl.Float64)
3)
4print(result)
1shape: (5, 2)
2┌────────┬───────────────────┐
3│ ticker ┆ company_name      │
4│ ---    ┆ ---               │
5│ str    ┆ str               │
6╞════════╪═══════════════════╡
7│ AAPL   ┆ Apple             │
8│ NVDA   ┆ NVIDIA            │
9│ MSFT   ┆ Microsoft         │
10│ GOOG   ┆ Alphabet (Google) │
11│ AMZN   ┆ Amazon            │
12└────────┴───────────────────┘

按正则匹配排除

Python
1result = df.select(
2    pl.all().exclude("^.*high$","^.*low$")
3)
4print(result)
1shape: (5, 3)
2┌────────┬───────────────────┬────────┐
3│ ticker ┆ company_name      ┆ price  │
4│ ---    ┆ ---               ┆ ---    │
5│ str    ┆ str               ┆ f64    │
6╞════════╪═══════════════════╪════════╡
7│ AAPL   ┆ Apple             ┆ 229.9  │
8│ NVDA   ┆ NVIDIA            ┆ 138.93 │
9│ MSFT   ┆ Microsoft         ┆ 420.56 │
10│ GOOG   ┆ Alphabet (Google) ┆ 166.41 │
11│ AMZN   ┆ Amazon            ┆ 188.4  │
12└────────┴───────────────────┴────────┘

其他示例

Python
1result = df.select(pl.col(pl.Float64).exclude("^day_.*$"))
2print(result)
1shape: (5, 3)
2┌────────┬───────────┬──────────┐
3│ price  ┆ year_high ┆ year_low │
4│ ---    ┆ ---       ┆ ---      │
5│ f64    ┆ f64       ┆ f64      │
6╞════════╪═══════════╪══════════╡
7│ 229.9  ┆ 237.23    ┆ 164.08   │
8│ 138.93 ┆ 140.76    ┆ 39.23    │
9│ 420.56 ┆ 468.35    ┆ 324.39   │
10│ 166.41 ┆ 193.31    ┆ 121.46   │
11│ 188.4  ┆ 201.2     ┆ 118.35   │
12└────────┴───────────┴──────────┘

列重命名

默认情况下,将表达式应用于列时,结果将与原始列保持相同的名称。保留列名在语义上可能是错误的,在某些情况下,如果出现重复的名称,Polars 甚至可能会引发错误,如下示例:

Python
1eur_usd_rate = 1.09
2gbp_usd_rate = 1.31
3
4result = df.select(
5    pl.col("price") / gbp_usd_rate,  # This would be named "price"...
6    pl.col("price") / eur_usd_rate,  # And so would this.
7)
8
9print(result)

运行上述代码将会报错。为了防止此类错误,允许用户在适当的时候重命名新生成的列,Polars 提供了一系列功能,可更改列或一组列的名称。

重命名单列

函数alias已在之前的文档中出现过很多次,它允许重命名新生成的单列,但alias不可用于表达式扩展,因为它是专门设计用于重命名单个列的。示例:

Python
1eur_usd_rate = 1.09
2gbp_usd_rate = 1.31
3
4result = df.select(
5    pl.col("price"),
6    (pl.col("price") / gbp_usd_rate).alias("price (GBP)"),
7    (pl.col("price") / eur_usd_rate).alias("price (EUR)"),
8)
9
10print(result)
1shape: (5, 3)
2┌────────┬─────────────┬─────────────┐
3│ price  ┆ price (GBP) ┆ price (EUR) │
4│ ---    ┆ ---         ┆ ---         │
5│ f64    ┆ f64         ┆ f64         │
6╞════════╪═════════════╪═════════════╡
7│ 229.9  ┆ 175.496183  ┆ 210.917431  │
8│ 138.93 ┆ 106.053435  ┆ 127.458716  │
9│ 420.56 ┆ 321.038168  ┆ 385.834862  │
10│ 166.41 ┆ 127.030534  ┆ 152.669725  │
11│ 188.4  ┆ 143.816794  ┆ 172.844037  │
12└────────┴─────────────┴─────────────┘

为列名添加前后缀

当只需为现有名称添加静态前缀或静态后缀时,可以使用命名空间name中的函数prefixsuffix来为列名添加前缀和后缀。

Python
1eur_usd_rate = 1.09
2gbp_usd_rate = 1.31
3
4result = df.select(
5    (pl.col("^year_.*$") / eur_usd_rate).name.prefix("in_eur_"),
6    (pl.col("day_high", "day_low") / gbp_usd_rate).name.suffix("_gbp"),
7)
8
9print(result)
1shape: (5, 4)
2┌──────────────────┬─────────────────┬──────────────┬─────────────┐
3│ in_eur_year_high ┆ in_eur_year_low ┆ day_high_gbp ┆ day_low_gbp │
4│ ---              ┆ ---             ┆ ---          ┆ ---         │
5│ f64              ┆ f64             ┆ f64          ┆ f64         │
6╞══════════════════╪═════════════════╪══════════════╪═════════════╡
7│ 217.642202       ┆ 150.53211       ┆ 176.572519   ┆ 174.503817  │
8│ 129.137615       ┆ 35.990826       ┆ 106.564885   ┆ 104.045802  │
9│ 429.678899       ┆ 297.605505      ┆ 323.694656   ┆ 318.717557  │
10│ 177.348624       ┆ 111.431193      ┆ 127.954198   ┆ 125.78626   │
11│ 184.587156       ┆ 108.577982      ┆ 144.908397   ┆ 143.847328  │
12└──────────────────┴─────────────────┴──────────────┴─────────────┘

列名转换

如果以上都不能满足需求, 命名空间name还提供了map方法, 该方法接收一个函数, 该函数接受旧的列名返回新的列名。

Python
1result = df.select(pl.all().name.map(str.upper))
2result2 = df.select(pl.all().name.map(lambda it: it.upper()))
3print(result.equals(result2))
4print(result)
1True
2shape: (5, 7)
3┌────────┬───────────────────┬────────┬──────────┬─────────┬───────────┬──────────┐
4│ TICKER ┆ COMPANY_NAME      ┆ PRICE  ┆ DAY_HIGH ┆ DAY_LOW ┆ YEAR_HIGH ┆ YEAR_LOW │
5│ ---    ┆ ---               ┆ ---    ┆ ---      ┆ ---     ┆ ---       ┆ ---      │
6│ str    ┆ str               ┆ f64    ┆ f64      ┆ f64     ┆ f64       ┆ f64      │
7╞════════╪═══════════════════╪════════╪══════════╪═════════╪═══════════╪══════════╡
8│ AAPL   ┆ Apple             ┆ 229.9  ┆ 231.31   ┆ 228.6   ┆ 237.23    ┆ 164.08   │
9│ NVDA   ┆ NVIDIA            ┆ 138.93 ┆ 139.6    ┆ 136.3   ┆ 140.76    ┆ 39.23    │
10│ MSFT   ┆ Microsoft         ┆ 420.56 ┆ 424.04   ┆ 417.52  ┆ 468.35    ┆ 324.39   │
11│ GOOG   ┆ Alphabet (Google) ┆ 166.41 ┆ 167.62   ┆ 164.78  ┆ 193.31    ┆ 121.46   │
12│ AMZN   ┆ Amazon            ┆ 188.4  ┆ 189.83   ┆ 188.44  ┆ 201.2     ┆ 118.35   │
13└────────┴───────────────────┴────────┴──────────┴─────────┴───────────┴──────────┘

还可以使用自定义函数进行列名的转换, 例如:

Python
1def new_name_example(old_name: str) -> str:
2    return old_name + "_" + old_name
3
4result = df.select(pl.all().name.map(new_name_example))
5print(result)
1shape: (5, 7)
2┌──────────────┬─────────────┬─────────────┬─────────────┬─────────────┬─────────────┬─────────────┐
3│ ticker_ticke ┆ company_nam ┆ price_price ┆ day_high_da ┆ day_low_day ┆ year_high_y ┆ year_low_ye │
4│ r            ┆ e_company_n ┆ ---         ┆ y_high      ┆ _low        ┆ ear_high    ┆ ar_low      │
5│ ---          ┆ ame         ┆ f64         ┆ ---         ┆ ---         ┆ ---         ┆ ---         │
6│ str          ┆ ---         ┆             ┆ f64         ┆ f64         ┆ f64         ┆ f64         │
7│              ┆ str         ┆             ┆             ┆             ┆             ┆             │
8╞══════════════╪═════════════╪═════════════╪═════════════╪═════════════╪═════════════╪═════════════╡
9│ AAPL         ┆ Apple       ┆ 229.9       ┆ 231.31      ┆ 228.6       ┆ 237.23      ┆ 164.08      │
10│ NVDA         ┆ NVIDIA      ┆ 138.93      ┆ 139.6       ┆ 136.3       ┆ 140.76      ┆ 39.23       │
11│ MSFT         ┆ Microsoft   ┆ 420.56      ┆ 424.04      ┆ 417.52      ┆ 468.35      ┆ 324.39      │
12│ GOOG         ┆ Alphabet    ┆ 166.41      ┆ 167.62      ┆ 164.78      ┆ 193.31      ┆ 121.46      │
13│              ┆ (Google)    ┆             ┆             ┆             ┆             ┆             │
14│ AMZN         ┆ Amazon      ┆ 188.4       ┆ 189.83      ┆ 188.44      ┆ 201.2       ┆ 118.35      │
15└──────────────┴─────────────┴─────────────┴─────────────┴─────────────┴─────────────┴─────────────┘

动态生成表达式

表达式扩展是一个非常有用的功能,但它并不能解决所有问题。如果需求稍微有些复杂,表达式扩展就显得有些力不从心了。例如,如果想计算每日和每年价格的高低差值,表达式扩展就没有办法,但可以使用动态生成表达式的方式来实现自己的需求。基于该需求,首先看看使用正常的方式如何实现:

Python
1result1 = df.with_columns(
2    (pl.col("day_high") - pl.col("day_low")).alias("day_amplitude"),
3    (pl.col("year_high") - pl.col("year_low")).alias("year_amplitude")
4)
5
6print(result1)
1shape: (5, 9)
2┌────────┬───────────────────┬────────┬──────────┬───┬──────────┬───────────────┬────────────────┐
3│ ticker ┆ company_name      ┆ price  ┆ day_high ┆ … ┆ year_low ┆ day_amplitude ┆ year_amplitude │
4│ ---    ┆ ---               ┆ ---    ┆ ---      ┆   ┆ ---      ┆ ---           ┆ ---            │
5│ str    ┆ str               ┆ f64    ┆ f64      ┆   ┆ f64      ┆ f64           ┆ f64            │
6╞════════╪═══════════════════╪════════╪══════════╪═══╪══════════╪═══════════════╪════════════════╡
7│ AAPL   ┆ Apple             ┆ 229.9  ┆ 231.31   ┆ … ┆ 164.08   ┆ 2.71          ┆ 73.15          │
8│ NVDA   ┆ NVIDIA            ┆ 138.93 ┆ 139.6    ┆ … ┆ 39.23    ┆ 3.3           ┆ 101.53         │
9│ MSFT   ┆ Microsoft         ┆ 420.56 ┆ 424.04   ┆ … ┆ 324.39   ┆ 6.52          ┆ 143.96         │
10│ GOOG   ┆ Alphabet (Google) ┆ 166.41 ┆ 167.62   ┆ … ┆ 121.46   ┆ 2.84          ┆ 71.85          │
11│ AMZN   ┆ Amazon            ┆ 188.4  ┆ 189.83   ┆ … ┆ 118.35   ┆ 1.39          ┆ 82.85          │
12└────────┴───────────────────┴────────┴──────────┴───┴──────────┴───────────────┴────────────────┘

然后可能还会想到使用for循环的方式来实现:

Python
1result2 = df
2for tp in ["day", "year"]:
3    result2 = result2.with_columns(
4        (pl.col(f"{tp}_high") - pl.col(f"{tp}_low")).alias(f"{tp}_amplitude")
5    )
6print(result2)
TIP

注意高亮的两行代码的作用!弄懂为什么要这么写!

1shape: (5, 9)
2┌────────┬───────────────────┬────────┬──────────┬───┬──────────┬───────────────┬────────────────┐
3│ ticker ┆ company_name      ┆ price  ┆ day_high ┆ … ┆ year_low ┆ day_amplitude ┆ year_amplitude │
4│ ---    ┆ ---               ┆ ---    ┆ ---      ┆   ┆ ---      ┆ ---           ┆ ---            │
5│ str    ┆ str               ┆ f64    ┆ f64      ┆   ┆ f64      ┆ f64           ┆ f64            │
6╞════════╪═══════════════════╪════════╪══════════╪═══╪══════════╪═══════════════╪════════════════╡
7│ AAPL   ┆ Apple             ┆ 229.9  ┆ 231.31   ┆ … ┆ 164.08   ┆ 2.71          ┆ 73.15          │
8│ NVDA   ┆ NVIDIA            ┆ 138.93 ┆ 139.6    ┆ … ┆ 39.23    ┆ 3.3           ┆ 101.53         │
9│ MSFT   ┆ Microsoft         ┆ 420.56 ┆ 424.04   ┆ … ┆ 324.39   ┆ 6.52          ┆ 143.96         │
10│ GOOG   ┆ Alphabet (Google) ┆ 166.41 ┆ 167.62   ┆ … ┆ 121.46   ┆ 2.84          ┆ 71.85          │
11│ AMZN   ┆ Amazon            ┆ 188.4  ┆ 189.83   ┆ … ┆ 118.35   ┆ 1.39          ┆ 82.85          │
12└────────┴───────────────────┴────────┴──────────┴───┴──────────┴───────────────┴────────────────┘

两者最后产生的结果是一样的:

Python
1print(result1.equals(result2))
1True
高亮代码含义

在用for循环实现效果时,在循环外面一个赋值语句,也许有人没明白为什么为要这么写,可能想的是下面这样:

Python
1# result2 = df
2for tp in ["day", "year"]:
3    result2 = df.with_columns(
4        (pl.col(f"{tp}_high") - pl.col(f"{tp}_low")).alias(f"{tp}_amplitude")
5    )

但是当你把结果 result2打印出来时,你会发现结果并不是和之前的一样,而是少了一列。对比后会发现少的是day_amplitude列,这是因为在循环中,result2被重新赋值了,所以day_amplitude列就消失了,只有最后一个year的计算值。而为什么上面示例中可以实现,核心就在于高亮的两行代码。

上面的示例中,使用for循环的方式,虽然实现了需求,但是不推荐这么使用。而更推荐使用下面这种方式:

Python
1def amplitude_expressions(time_periods):
2    for tp in time_periods:
3        yield (pl.col(f"{tp}_high") - pl.col(f"{tp}_low")).alias(f"{tp}_amplitude")
4
5result3 = df.with_columns(amplitude_expressions(["day", "year"]))
6print(result3)
1shape: (5, 9)
2┌────────┬───────────────────┬────────┬──────────┬───┬──────────┬───────────────┬────────────────┐
3│ ticker ┆ company_name      ┆ price  ┆ day_high ┆ … ┆ year_low ┆ day_amplitude ┆ year_amplitude │
4│ ---    ┆ ---               ┆ ---    ┆ ---      ┆   ┆ ---      ┆ ---           ┆ ---            │
5│ str    ┆ str               ┆ f64    ┆ f64      ┆   ┆ f64      ┆ f64           ┆ f64            │
6╞════════╪═══════════════════╪════════╪══════════╪═══╪══════════╪═══════════════╪════════════════╡
7│ AAPL   ┆ Apple             ┆ 229.9  ┆ 231.31   ┆ … ┆ 164.08   ┆ 2.71          ┆ 73.15          │
8│ NVDA   ┆ NVIDIA            ┆ 138.93 ┆ 139.6    ┆ … ┆ 39.23    ┆ 3.3           ┆ 101.53         │
9│ MSFT   ┆ Microsoft         ┆ 420.56 ┆ 424.04   ┆ … ┆ 324.39   ┆ 6.52          ┆ 143.96         │
10│ GOOG   ┆ Alphabet (Google) ┆ 166.41 ┆ 167.62   ┆ … ┆ 121.46   ┆ 2.84          ┆ 71.85          │
11│ AMZN   ┆ Amazon            ┆ 188.4  ┆ 189.83   ┆ … ┆ 118.35   ┆ 1.39          ┆ 82.85          │
12└────────┴───────────────────┴────────┴──────────┴───┴──────────┴───────────────┴────────────────┘

这样做的好处是,通过循环生成计算表达式,然后在上下文中使用一次即可;而不是在循环中生成计算表达式然后再在上下文中进行使用。而且这样写,Polars能够更好的优化查询以及并行进行计算。

遐想

循环中一个一个查询基金信息;循环外一次性查询出所有基金信息。

更多列选择方式

Polars附带子模块, selectors子模块提供了许多函数, 允许编写更灵活的列选择表达式。

通过selectors来实现使用string()函数和ends_with()函数选择所有字符串列或者以"_high"结尾的列。

Python
1import polars.selectors as cs
2
3result = df.select(cs.string() | cs.ends_with("_high"))
4print(result)
1shape: (5, 4)
2┌────────┬───────────────────┬──────────┬───────────┐
3│ ticker ┆ company_name      ┆ day_high ┆ year_high │
4│ ---    ┆ ---               ┆ ---      ┆ ---       │
5│ str    ┆ str               ┆ f64      ┆ f64       │
6╞════════╪═══════════════════╪══════════╪═══════════╡
7│ AAPL   ┆ Apple             ┆ 231.31   ┆ 237.23    │
8│ NVDA   ┆ NVIDIA            ┆ 139.6    ┆ 140.76    │
9│ MSFT   ┆ Microsoft         ┆ 424.04   ┆ 468.35    │
10│ GOOG   ┆ Alphabet (Google) ┆ 167.62   ┆ 193.31    │
11│ AMZN   ┆ Amazon            ┆ 189.83   ┆ 201.2     │
12└────────┴───────────────────┴──────────┴───────────┘

使用numeric()函数选择所有的数值列, 无论是整数还是浮点数:

Python
1import polars.selectors as cs
2
3result = df.select(cs.numeric())
4print(result)
1shape: (5, 5)
2┌────────┬──────────┬─────────┬───────────┬──────────┐
3│ price  ┆ day_high ┆ day_low ┆ year_high ┆ year_low │
4│ ---    ┆ ---      ┆ ---     ┆ ---       ┆ ---      │
5│ f64    ┆ f64      ┆ f64     ┆ f64       ┆ f64      │
6╞════════╪══════════╪═════════╪═══════════╪══════════╡
7│ 229.9  ┆ 231.31   ┆ 228.6   ┆ 237.23    ┆ 164.08   │
8│ 138.93 ┆ 139.6    ┆ 136.3   ┆ 140.76    ┆ 39.23    │
9│ 420.56 ┆ 424.04   ┆ 417.52  ┆ 468.35    ┆ 324.39   │
10│ 166.41 ┆ 167.62   ┆ 164.78  ┆ 193.31    ┆ 121.46   │
11│ 188.4  ┆ 189.83   ┆ 188.44  ┆ 201.2     ┆ 118.35   │
12└────────┴──────────┴─────────┴───────────┴──────────┘

选择器与数学集合操作

操作 解释
A | B 并集
A & B 交集
A - B 差集, 在A中不在B中
~A 补集, 不在A中

匹配名称中包含下划线的所有非字符串列:

Python
1result = df.select(cs.contains("_") - cs.string())
2print(result)
1shape: (5, 4)
2┌──────────┬─────────┬───────────┬──────────┐
3│ day_high ┆ day_low ┆ year_high ┆ year_low │
4│ ---      ┆ ---     ┆ ---       ┆ ---      │
5│ f64      ┆ f64     ┆ f64       ┆ f64      │
6╞══════════╪═════════╪═══════════╪══════════╡
7│ 231.31   ┆ 228.6   ┆ 237.23    ┆ 164.08   │
8│ 139.6    ┆ 136.3   ┆ 140.76    ┆ 39.23    │
9│ 424.04   ┆ 417.52  ┆ 468.35    ┆ 324.39   │
10│ 167.62   ┆ 164.78  ┆ 193.31    ┆ 121.46   │
11│ 189.83   ┆ 188.44  ┆ 201.2     ┆ 118.35   │
12└──────────┴─────────┴───────────┴──────────┘

匹配非数值列:

Python
1result = df.select(~cs.numeric())
2print(result)
1shape: (5, 2)
2┌────────┬───────────────────┐
3│ ticker ┆ company_name      │
4│ ---    ┆ ---               │
5│ str    ┆ str               │
6╞════════╪═══════════════════╡
7│ AAPL   ┆ Apple             │
8│ NVDA   ┆ NVIDIA            │
9│ MSFT   ┆ Microsoft         │
10│ GOOG   ┆ Alphabet (Google) │
11│ AMZN   ┆ Amazon            │
12└────────┴───────────────────┘

解决运算符歧义

~符号针对布尔表达式含义是取反, 针对选择器是取补集。当在使用选择器,然后又想要在表达式上下文中使用充当选择器的集合运算符的运算符之一时,可以使用函数as_expr解决歧义。例如,想把所有以"has_"开头的列的值取反。

1people = pl.DataFrame(
2    {
3        "name": ["Anna", "Bob"],
4        "has_partner": [True, False],
5        "has_kids": [False, False],
6        "has_tattoos": [True, False],
7        "is_alive": [True, True],
8    }
9)
10
11wrong_result = people.select((~cs.starts_with("has_")).name.prefix("not_"))
12print(wrong_result)
1shape: (2, 2)
2┌──────────┬──────────────┐
3│ not_name ┆ not_is_alive │
4│ ---      ┆ ---          │
5│ str      ┆ bool         │
6╞══════════╪══════════════╡
7│ Anna     ┆ true         │
8│ Bob      ┆ true         │
9└──────────┴──────────────┘

结果并不是我们想要的, 而是取出了所有不以"has_"开头的列,这就存在歧义。为了实现取反, 我们需要使用as_expr()方法:

Python
1result = people.select((~cs.starts_with("has_").as_expr()).name.prefix("not_"))
2print(result)
1shape: (2, 3)
2┌─────────────────┬──────────────┬─────────────────┐
3│ not_has_partner ┆ not_has_kids ┆ not_has_tattoos │
4│ ---             ┆ ---          ┆ ---             │
5│ bool            ┆ bool         ┆ bool            │
6╞═════════════════╪══════════════╪═════════════════╡
7│ false           ┆ true         ┆ false           │
8│ true            ┆ true         ┆ true            │
9└─────────────────┴──────────────┴─────────────────┘

正确的写法是使用as_expr()方法, as_expr方法将选择器转换为表达式, 然后取反, 我们看下面代码

1res = people.select((~cs.starts_with("has_").as_expr()))
2print(res)

可以细致对比一下结果, 是否按照我们的预期, 针对布尔值进行取反了

1shape: (2, 3)
2┌─────────────┬──────────┬─────────────┐
3│ has_partner ┆ has_kids ┆ has_tattoos │
4│ ---         ┆ ---      ┆ ---         │
5│ bool        ┆ bool     ┆ bool        │
6╞═════════════╪══════════╪═════════════╡
7│ false       ┆ true     ┆ false       │
8│ true        ┆ true     ┆ true        │
9└─────────────┴──────────┴─────────────┘

调试选择器

当不确定是否有一个Polars选择器时, 可以使用函数cs.is_selector来判断:

Python
1print(cs.is_selector(~cs.starts_with("has_")))
2print(cs.is_selector(~cs.starts_with("has_").as_expr()))
1True
2False

这可以帮助避免任何模棱两可的情况,即认为正在使用表达式进行操作,但实际上是使用选择器进行操作。另一个有用的调试器函数是 expand_selector, 对于给定的目标DataFrame或者选择器, 可以检查给定选择器将扩展为哪些列:

Python
1import polars.selectors as cs
2
3people = pl.DataFrame(
4    {
5        "name": ["Anna", "Bob"],
6        "has_partner": [True, False],
7        "has_kids": [False, False],
8        "has_tattoos": [True, False],
9        "is_alive": [True, True],
10    }
11)
12
13print(
14    cs.expand_selector(
15        people,
16        cs.starts_with("has_"),
17    )
18)
1('has_partner', 'has_kids', 'has_tattoos')

更多

子模块selectors的常见用法:

Python
1people = pl.DataFrame(
2    {
3        "name": ["Anna", "Bob"],
4        "age": [12, 21],
5        "weight": [42.3, 63.3],
6        "height": [161.3, 170.9],
7        "6month": [True, False],
8        "1year": [True, False],
9        "135": [True, False],
10        "has_partner": [True, False],
11        "has_kids": [False, False],
12        "has_tattoos": [True, False],
13        "is_alive": [True, True],
14    }
15)
16
17print("======= boolean")
18print(people.select(cs.boolean()))
19print("======= float")
20print(people.select(cs.float()))
21print("======= integer")
22print(people.select(cs.integer()))
23print("======= numeric")
24print(people.select(cs.numeric()))
25print("======= string")
26print(people.select(cs.string()))
27print("======= alphanumeric")
28print(people.select(cs.alphanumeric()))
29print("======= by_name")
30print(people.select(cs.by_name("age")))
31print("======= contains")
32print(people.select(cs.contains("h")))
33print("======= digit")
34print(people.select(cs.digit()))
1======= boolean
2shape: (2, 7)
3┌────────┬───────┬───────┬─────────────┬──────────┬─────────────┬──────────┐
4│ 6month ┆ 1year ┆ 135   ┆ has_partner ┆ has_kids ┆ has_tattoos ┆ is_alive │
5│ ---    ┆ ---   ┆ ---   ┆ ---         ┆ ---      ┆ ---         ┆ ---      │
6│ bool   ┆ bool  ┆ bool  ┆ bool        ┆ bool     ┆ bool        ┆ bool     │
7╞════════╪═══════╪═══════╪═════════════╪══════════╪═════════════╪══════════╡
8│ true   ┆ true  ┆ true  ┆ true        ┆ false    ┆ true        ┆ true     │
9│ false  ┆ false ┆ false ┆ false       ┆ false    ┆ false       ┆ true     │
10└────────┴───────┴───────┴─────────────┴──────────┴─────────────┴──────────┘
11======= float
12shape: (2, 2)
13┌────────┬────────┐
14│ weight ┆ height │
15│ ---    ┆ ---    │
16│ f64    ┆ f64    │
17╞════════╪════════╡
18│ 42.3   ┆ 161.3  │
19│ 63.3   ┆ 170.9  │
20└────────┴────────┘
21======= integer
22shape: (2, 1)
23┌─────┐
24│ age │
25│ --- │
26│ i64 │
27╞═════╡
28│ 12  │
29│ 21  │
30└─────┘
31======= numeric
32shape: (2, 3)
33┌─────┬────────┬────────┐
34│ age ┆ weight ┆ height │
35│ --- ┆ ---    ┆ ---    │
36│ i64 ┆ f64    ┆ f64    │
37╞═════╪════════╪════════╡
38│ 12  ┆ 42.3   ┆ 161.3  │
39│ 21  ┆ 63.3   ┆ 170.9  │
40└─────┴────────┴────────┘
41======= string
42shape: (2, 1)
43┌──────┐
44│ name │
45│ ---  │
46│ str  │
47╞══════╡
48│ Anna │
49│ Bob  │
50└──────┘
51======= alphanumeric
52shape: (2, 7)
53┌──────┬─────┬────────┬────────┬────────┬───────┬───────┐
54│ name ┆ age ┆ weight ┆ height ┆ 6month ┆ 1year ┆ 135   │
55│ ---  ┆ --- ┆ ---    ┆ ---    ┆ ---    ┆ ---   ┆ ---   │
56│ str  ┆ i64 ┆ f64    ┆ f64    ┆ bool   ┆ bool  ┆ bool  │
57╞══════╪═════╪════════╪════════╪════════╪═══════╪═══════╡
58│ Anna ┆ 12  ┆ 42.3   ┆ 161.3  ┆ true   ┆ true  ┆ true  │
59│ Bob  ┆ 21  ┆ 63.3   ┆ 170.9  ┆ false  ┆ false ┆ false │
60└──────┴─────┴────────┴────────┴────────┴───────┴───────┘
61======= by_name
62shape: (2, 1)
63┌─────┐
64│ age │
65│ --- │
66│ i64 │
67╞═════╡
68│ 12  │
69│ 21  │
70└─────┘
71======= contains
72shape: (2, 6)
73┌────────┬────────┬────────┬─────────────┬──────────┬─────────────┐
74│ weight ┆ height ┆ 6month ┆ has_partner ┆ has_kids ┆ has_tattoos │
75│ ---    ┆ ---    ┆ ---    ┆ ---         ┆ ---      ┆ ---         │
76│ f64    ┆ f64    ┆ bool   ┆ bool        ┆ bool     ┆ bool        │
77╞════════╪════════╪════════╪═════════════╪══════════╪═════════════╡
78│ 42.3   ┆ 161.3  ┆ true   ┆ true        ┆ false    ┆ true        │
79│ 63.3   ┆ 170.9  ┆ false  ┆ false       ┆ false    ┆ false       │
80└────────┴────────┴────────┴─────────────┴──────────┴─────────────┘
81======= digit
82shape: (2, 1)
83┌───────┐
84│ 135   │
85│ ---   │
86│ bool  │
87╞═══════╡
88│ true  │
89│ false │
90└───────┘
TIP

关于子模块selectors的更多用法, 请参阅 详见官方文档