跳到内容

表达式和上下文

Polars 开发了自己的领域特定语言 (DSL) 用于数据转换。该语言易于使用,支持复杂查询,同时保持可读性。本文将介绍的表达式和上下文,对于实现这种可读性以及让 Polars 查询引擎尽可能快地优化和运行您的查询都非常重要。

表达式

在 Polars 中,一个表达式是数据转换的惰性表示。表达式是模块化且灵活的,这意味着您可以将它们用作构建更复杂表达式的积木。以下是一个 Polars 表达式的示例:

import polars as pl

pl.col("weight") / (pl.col("height") ** 2)

您可能已经猜到,这个表达式将名为“weight”的列的值除以名为“height”的列的值的平方,从而计算一个人的 BMI。

上述代码表达了一种抽象计算,我们可以将其保存到变量中,进行进一步操作,或者直接打印出来

bmi_expr = pl.col("weight") / (pl.col("height") ** 2)
print(bmi_expr)
[(col("weight")) / (col("height").pow([dyn int: 2]))]

因为表达式是惰性的,所以尚未进行任何计算。这就是我们需要上下文的原因。

上下文

Polars 表达式需要一个上下文来执行并产生结果。根据使用的上下文,相同的 Polars 表达式可以产生不同的结果。在本节中,我们将了解 Polars 提供的四种最常见的上下文1

  1. select
  2. with_columns
  3. filter
  4. group_by

我们使用下面的 DataFrame 来展示每个上下文如何工作。

from datetime import date

df = pl.DataFrame(
    {
        "name": ["Alice Archer", "Ben Brown", "Chloe Cooper", "Daniel Donovan"],
        "birthdate": [
            date(1997, 1, 10),
            date(1985, 2, 15),
            date(1983, 3, 22),
            date(1981, 4, 30),
        ],
        "weight": [57.9, 72.5, 53.6, 83.1],  # (kg)
        "height": [1.56, 1.77, 1.65, 1.75],  # (m)
    }
)

print(df)
use chrono::prelude::*;
use polars::prelude::*;

let df: DataFrame = df!(
    "name" => ["Alice Archer", "Ben Brown", "Chloe Cooper", "Daniel Donovan"],
    "birthdate" => [
        NaiveDate::from_ymd_opt(1997, 1, 10).unwrap(),
        NaiveDate::from_ymd_opt(1985, 2, 15).unwrap(),
        NaiveDate::from_ymd_opt(1983, 3, 22).unwrap(),
        NaiveDate::from_ymd_opt(1981, 4, 30).unwrap(),
    ],
    "weight" => [57.9, 72.5, 53.6, 83.1],  // (kg)
    "height" => [1.56, 1.77, 1.65, 1.75],  // (m)
)
.unwrap();
println!("{df}");
shape: (4, 4)
┌────────────────┬────────────┬────────┬────────┐
│ name           ┆ birthdate  ┆ weight ┆ height │
│ ---            ┆ ---        ┆ ---    ┆ ---    │
│ str            ┆ date       ┆ f64    ┆ f64    │
╞════════════════╪════════════╪════════╪════════╡
│ Alice Archer   ┆ 1997-01-10 ┆ 57.9   ┆ 1.56   │
│ Ben Brown      ┆ 1985-02-15 ┆ 72.5   ┆ 1.77   │
│ Chloe Cooper   ┆ 1983-03-22 ┆ 53.6   ┆ 1.65   │
│ Daniel Donovan ┆ 1981-04-30 ┆ 83.1   ┆ 1.75   │
└────────────────┴────────────┴────────┴────────┘

select

选择上下文 select 将表达式应用于列。select 上下文可以生成新的列,这些列可以是聚合结果、其他列的组合或字面值。

select

result = df.select(
    bmi=bmi_expr,
    avg_bmi=bmi_expr.mean(),
    ideal_max_bmi=25,
)
print(result)

select

let bmi = col("weight") / col("height").pow(2);
let result = df
    .clone()
    .lazy()
    .select([
        bmi.clone().alias("bmi"),
        bmi.clone().mean().alias("avg_bmi"),
        lit(25).alias("ideal_max_bmi"),
    ])
    .collect()?;
println!("{result}");

shape: (4, 3)
┌───────────┬───────────┬───────────────┐
│ bmi       ┆ avg_bmi   ┆ ideal_max_bmi │
│ ---       ┆ ---       ┆ ---           │
│ f64       ┆ f64       ┆ i32           │
╞═══════════╪═══════════╪═══════════════╡
│ 23.791913 ┆ 23.438973 ┆ 25            │
│ 23.141498 ┆ 23.438973 ┆ 25            │
│ 19.687787 ┆ 23.438973 ┆ 25            │
│ 27.134694 ┆ 23.438973 ┆ 25            │
└───────────┴───────────┴───────────────┘

select 上下文中的表达式必须产生长度全部相同的 Series,或者产生一个标量。标量将被广播以匹配其余 Series 的长度。字面值(如上面使用的数字)也会被广播。

请注意,广播也可以发生在表达式内部。例如,考虑下面的表达式

select

result = df.select(deviation=(bmi_expr - bmi_expr.mean()) / bmi_expr.std())
print(result)

select

let result = df
    .clone()
    .lazy()
    .select([((bmi.clone() - bmi.clone().mean()) / bmi.clone().std(1)).alias("deviation")])
    .collect()?;
println!("{result}");

shape: (4, 1)
┌───────────┐
│ deviation │
│ ---       │
│ f64       │
╞═══════════╡
│ 0.115645  │
│ -0.097471 │
│ -1.22912  │
│ 1.210946  │
└───────────┘

减法和除法都在表达式内部使用了广播,因为计算均值和标准差的子表达式都评估为单个值。

select 上下文非常灵活和强大,允许您独立且并行地评估任意表达式。对于接下来将看到的其他上下文也是如此。

with_columns

with_columns 上下文与 select 上下文非常相似。两者之间的主要区别在于,with_columns 上下文会创建一个新的 DataFrame,其中包含原始 DataFrame 中的列以及根据其输入表达式生成的新列,而 select 上下文只包含其输入表达式选择的列。

with_columns

result = df.with_columns(
    bmi=bmi_expr,
    avg_bmi=bmi_expr.mean(),
    ideal_max_bmi=25,
)
print(result)

with_columns

let result = df
    .clone()
    .lazy()
    .with_columns([
        bmi.clone().alias("bmi"),
        bmi.clone().mean().alias("avg_bmi"),
        lit(25).alias("ideal_max_bmi"),
    ])
    .collect()?;
println!("{result}");

shape: (4, 7)
┌────────────────┬────────────┬────────┬────────┬───────────┬───────────┬───────────────┐
│ name           ┆ birthdate  ┆ weight ┆ height ┆ bmi       ┆ avg_bmi   ┆ ideal_max_bmi │
│ ---            ┆ ---        ┆ ---    ┆ ---    ┆ ---       ┆ ---       ┆ ---           │
│ str            ┆ date       ┆ f64    ┆ f64    ┆ f64       ┆ f64       ┆ i32           │
╞════════════════╪════════════╪════════╪════════╪═══════════╪═══════════╪═══════════════╡
│ Alice Archer   ┆ 1997-01-10 ┆ 57.9   ┆ 1.56   ┆ 23.791913 ┆ 23.438973 ┆ 25            │
│ Ben Brown      ┆ 1985-02-15 ┆ 72.5   ┆ 1.77   ┆ 23.141498 ┆ 23.438973 ┆ 25            │
│ Chloe Cooper   ┆ 1983-03-22 ┆ 53.6   ┆ 1.65   ┆ 19.687787 ┆ 23.438973 ┆ 25            │
│ Daniel Donovan ┆ 1981-04-30 ┆ 83.1   ┆ 1.75   ┆ 27.134694 ┆ 23.438973 ┆ 25            │
└────────────────┴────────────┴────────┴────────┴───────────┴───────────┴───────────────┘

由于 selectwith_columns 之间的区别,with_columns 上下文使用的表达式必须生成与 DataFrame 中原始列长度相同的 Series,而 select 上下文中的表达式只需生成彼此长度相同的 Series 即可。

filter

filter 上下文根据一个或多个评估为布尔数据类型的表达式来筛选 DataFrame 的行。

filter

result = df.filter(
    pl.col("birthdate").is_between(date(1982, 12, 31), date(1996, 1, 1)),
    pl.col("height") > 1.7,
)
print(result)

filter

let result = df
    .clone()
    .lazy()
    .filter(
        col("birthdate")
            .is_between(
                lit(NaiveDate::from_ymd_opt(1982, 12, 31).unwrap()),
                lit(NaiveDate::from_ymd_opt(1996, 1, 1).unwrap()),
                ClosedInterval::Both,
            )
            .and(col("height").gt(lit(1.7))),
    )
    .collect()?;
println!("{result}");

shape: (1, 4)
┌───────────┬────────────┬────────┬────────┐
│ name      ┆ birthdate  ┆ weight ┆ height │
│ ---       ┆ ---        ┆ ---    ┆ ---    │
│ str       ┆ date       ┆ f64    ┆ f64    │
╞═══════════╪════════════╪════════╪════════╡
│ Ben Brown ┆ 1985-02-15 ┆ 72.5   ┆ 1.77   │
└───────────┴────────────┴────────┴────────┘

group_by 和聚合

group_by 上下文中,行根据分组表达式的唯一值进行分组。然后,您可以将表达式应用于结果组,这些组的长度可能是可变的。

在使用 group_by 上下文时,您可以使用表达式动态计算分组

group_by

result = df.group_by(
    (pl.col("birthdate").dt.year() // 10 * 10).alias("decade"),
).agg(pl.col("name"))
print(result)

group_by

let result = df
    .clone()
    .lazy()
    .group_by([(col("birthdate").dt().year() / lit(10) * lit(10)).alias("decade")])
    .agg([col("name")])
    .collect()?;
println!("{result}");

shape: (2, 2)
┌────────┬─────────────────────────────────┐
│ decade ┆ name                            │
│ ---    ┆ ---                             │
│ i32    ┆ list[str]                       │
╞════════╪═════════════════════════════════╡
│ 1980   ┆ ["Ben Brown", "Chloe Cooper", … │
│ 1990   ┆ ["Alice Archer"]                │
└────────┴─────────────────────────────────┘

使用 group_by 后,我们使用 agg 将聚合表达式应用于分组。由于在上述示例中我们只指定了列名,因此我们将该列的分组结果作为列表获取。

我们可以指定任意数量的分组表达式,group_by 上下文将根据指定表达式的 DISTINCT 值对行进行分组。在这里,我们根据出生年代和身高是否小于 1.7 米的组合进行分组

group_by

result = df.group_by(
    (pl.col("birthdate").dt.year() // 10 * 10).alias("decade"),
    (pl.col("height") < 1.7).alias("short?"),
).agg(pl.col("name"))
print(result)

group_by

let result = df
    .clone()
    .lazy()
    .group_by([
        (col("birthdate").dt().year() / lit(10) * lit(10)).alias("decade"),
        (col("height").lt(lit(1.7)).alias("short?")),
    ])
    .agg([col("name")])
    .collect()?;
println!("{result}");

shape: (3, 3)
┌────────┬────────┬─────────────────────────────────┐
│ decade ┆ short? ┆ name                            │
│ ---    ┆ ---    ┆ ---                             │
│ i32    ┆ bool   ┆ list[str]                       │
╞════════╪════════╪═════════════════════════════════╡
│ 1980   ┆ false  ┆ ["Ben Brown", "Daniel Donovan"… │
│ 1980   ┆ true   ┆ ["Chloe Cooper"]                │
│ 1990   ┆ true   ┆ ["Alice Archer"]                │
└────────┴────────┴─────────────────────────────────┘

应用聚合表达式后,结果 DataFrame 的左侧包含每个分组表达式对应的一列,然后是表示聚合表达式结果所需的任意数量的列。反过来,我们可以指定任意数量的聚合表达式

group_by

result = df.group_by(
    (pl.col("birthdate").dt.year() // 10 * 10).alias("decade"),
    (pl.col("height") < 1.7).alias("short?"),
).agg(
    pl.len(),
    pl.col("height").max().alias("tallest"),
    pl.col("weight", "height").mean().name.prefix("avg_"),
)
print(result)

group_by

let result = df
    .clone()
    .lazy()
    .group_by([
        (col("birthdate").dt().year() / lit(10) * lit(10)).alias("decade"),
        (col("height").lt(lit(1.7)).alias("short?")),
    ])
    .agg([
        len(),
        col("height").max().alias("tallest"),
        cols(["weight", "height"]).mean().name().prefix("avg_"),
    ])
    .collect()?;
println!("{result}");

shape: (3, 6)
┌────────┬────────┬─────┬─────────┬────────────┬────────────┐
│ decade ┆ short? ┆ len ┆ tallest ┆ avg_weight ┆ avg_height │
│ ---    ┆ ---    ┆ --- ┆ ---     ┆ ---        ┆ ---        │
│ i32    ┆ bool   ┆ u32 ┆ f64     ┆ f64        ┆ f64        │
╞════════╪════════╪═════╪═════════╪════════════╪════════════╡
│ 1980   ┆ true   ┆ 1   ┆ 1.65    ┆ 53.6       ┆ 1.65       │
│ 1980   ┆ false  ┆ 2   ┆ 1.77    ┆ 77.8       ┆ 1.76       │
│ 1990   ┆ true   ┆ 1   ┆ 1.56    ┆ 57.9       ┆ 1.56       │
└────────┴────────┴─────┴─────────┴────────────┴────────────┘

有关其他分组上下文,另请参阅 group_by_dynamicrolling

表达式扩展

最后一个示例包含两个分组表达式和三个聚合表达式,但结果 DataFrame 却包含了六列而非五列。如果仔细观察,最后一个聚合表达式提到了两个不同的列:“weight”和“height”。

Polars 表达式支持一项名为表达式展开的功能。表达式展开类似于一种速记法,用于当您希望将相同的转换应用于多个列时。正如我们所见,该表达式

pl.col("weight", "height").mean().name.prefix("avg_")

将计算“weight”和“height”列的均值,并分别将它们重命名为“avg_weight”和“avg_height”。实际上,上述表达式等同于使用以下两个表达式

[
    pl.col("weight").mean().alias("avg_weight"),
    pl.col("height").mean().alias("avg_height"),
]

在这种情况下,此表达式展开为两个独立的表达式,Polars 可以并行执行。在其他情况下,我们可能无法预先知道一个表达式将展开成多少个独立的表达式。

考虑这个简单但具有启发性的示例

(pl.col(pl.Float64) * 1.1).name.suffix("*1.1")

此表达式会将所有数据类型为 Float64 的列乘以 1.1。此表达式适用的列数取决于每个 DataFrame 的架构。在我们一直使用的 DataFrame 中,它适用于两列

group_by

expr = (pl.col(pl.Float64) * 1.1).name.suffix("*1.1")
result = df.select(expr)
print(result)

group_by

let expr = (dtype_col(&DataType::Float64) * lit(1.1))
    .name()
    .suffix("*1.1");
let result = df.clone().lazy().select([expr.clone()]).collect()?;
println!("{result}");

shape: (4, 2)
┌────────────┬────────────┐
│ weight*1.1 ┆ height*1.1 │
│ ---        ┆ ---        │
│ f64        ┆ f64        │
╞════════════╪════════════╡
│ 63.69      ┆ 1.716      │
│ 79.75      ┆ 1.947      │
│ 58.96      ┆ 1.815      │
│ 91.41      ┆ 1.925      │
└────────────┴────────────┘

在下面的 DataFrame df2 的情况下,同样的表达式展开为 0 列,因为没有列的数据类型是 Float64

group_by

df2 = pl.DataFrame(
    {
        "ints": [1, 2, 3, 4],
        "letters": ["A", "B", "C", "D"],
    }
)
result = df2.select(expr)
print(result)

group_by

let df2: DataFrame = df!(
    "ints" => [1, 2, 3, 4],
    "letters" => ["A", "B", "C", "D"],
)
.unwrap();
let result = df2.clone().lazy().select([expr.clone()]).collect()?;
println!("{result}");

shape: (0, 0)
┌┐
╞╡
└┘

同样容易想象,在某些场景下,相同的表达式会展开为数十列。

接下来,您将了解惰性 API 和 explain 函数,您可以使用它来预览给定架构下表达式将展开为何种形式。

总结

因为表达式是惰性的,当您在上下文中S使用表达式时,Polars 可以在运行其所表达的数据转换之前尝试简化您的表达式。上下文内的独立表达式是极其并行的,Polars 将利用这一点,同时在使用表达式展开时也会并行化表达式执行。使用接下来将介绍的Polars 的惰性 API 可以获得进一步的性能提升。

我们仅仅触及了表达式能力的冰山一角。还有大量的表达式,它们可以通过多种方式组合。有关可用不同类型表达式的更深入探讨,请参阅表达式章节


  1. 本指南后面还会介绍额外的列表 (List) 和 SQL 上下文。但为简单起见,我们暂时不讨论它们。