Skip to content

表达式和上下文

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

我们使用下面的数据框来展示每个上下文的工作原理。

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必须生成长度相同的序列，或者它们必须生成一个标量。标量将被广播以匹配剩余序列的长度。字面量，如上面使用的数字，也会被广播。

请注意，广播也可以在表达式中发生。例如，考虑以下表达式：

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 创建一个新的数据框，其中包含原始数据框中的列以及根据其输入表达式生成的新列，而上下文 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            │
└────────────────┴────────────┴────────┴────────┴───────────┴───────────┴───────────────┘

由于select和with_columns之间的这种差异，在with_columns上下文中使用的表达式必须生成与数据框中的原始列长度相同的系列，而在select上下文中，表达式生成的系列只需要在它们之间具有相同的长度即可。

filter

上下文 filter 根据一个或多个评估为布尔数据类型的表达式过滤数据框的行。

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将根据指定表达式中的不同值对行进行分组。在这里，我们根据出生年代和身高是否低于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   ┆ true   ┆ ["Chloe Cooper"]                │
│ 1980   ┆ false  ┆ ["Ben Brown", "Daniel Donovan"… │
│ 1990   ┆ true   ┆ ["Alice Archer"]                │
└────────┴────────┴─────────────────────────────────┘

应用聚合表达式后生成的数据框，左侧每个分组表达式对应一列，然后根据需要表示聚合表达式结果的列数。反过来，我们可以指定任意数量的聚合表达式：

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        │
╞════════╪════════╪═════╪═════════╪════════════╪════════════╡
│ 1990   ┆ true   ┆ 1   ┆ 1.56    ┆ 57.9       ┆ 1.56       │
│ 1980   ┆ true   ┆ 1   ┆ 1.65    ┆ 53.6       ┆ 1.65       │
│ 1980   ┆ false  ┆ 2   ┆ 1.77    ┆ 77.8       ┆ 1.76       │
└────────┴────────┴─────┴─────────┴────────────┴────────────┘

另请参阅 group_by_dynamic 和 group_by_rolling 以了解其他分组上下文。

表达式扩展

最后一个例子包含两个分组表达式和三个聚合表达式，但结果的数据框却包含六列而不是五列。如果我们仔细观察，最后一个聚合表达式提到了两个不同的列：“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。此操作适用的列数取决于每个数据框架的模式。在我们一直使用的数据框架的情况下，它适用于两列：

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      │
└────────────┴────────────┘

在下面的数据框 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，您可以使用它 来预览表达式在给定模式下的扩展情况。

结论

因为表达式是惰性的，当你在上下文中使用表达式时，Polars 可以在运行数据转换之前尝试简化你的表达式。上下文中的独立表达式是令人尴尬的并行，Polars 将利用这一点，同时在使用表达式扩展时并行化表达式执行。当使用 Polars 的惰性 API 时，可以进一步获得性能提升，这将在接下来介绍。

我们仅仅触及了表达式功能的表面。还有大量的表达式，它们可以以多种方式组合。有关可用表达式的不同类型的更深入探讨，请参见表达式部分。

  1. 本指南后面会介绍更多的列表和SQL上下文。但为了简单起见，我们现在暂时不讨论它们。↩