表达式和上下文
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
select
with_columns
filter
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
上下文可以生成新的列,这些列可以是聚合结果、其他列的组合或字面值。
result = df.select(
bmi=bmi_expr,
avg_bmi=bmi_expr.mean(),
ideal_max_bmi=25,
)
print(result)
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 的长度。字面值(如上面使用的数字)也会被广播。
请注意,广播也可以发生在表达式内部。例如,考虑下面的表达式
shape: (4, 1)
┌───────────┐
│ deviation │
│ --- │
│ f64 │
╞═══════════╡
│ 0.115645 │
│ -0.097471 │
│ -1.22912 │
│ 1.210946 │
└───────────┘
减法和除法都在表达式内部使用了广播,因为计算均值和标准差的子表达式都评估为单个值。
select
上下文非常灵活和强大,允许您独立且并行地评估任意表达式。对于接下来将看到的其他上下文也是如此。
with_columns
with_columns
上下文与 select
上下文非常相似。两者之间的主要区别在于,with_columns
上下文会创建一个新的 DataFrame,其中包含原始 DataFrame 中的列以及根据其输入表达式生成的新列,而 select
上下文只包含其输入表达式选择的列。
result = df.with_columns(
bmi=bmi_expr,
avg_bmi=bmi_expr.mean(),
ideal_max_bmi=25,
)
print(result)
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
上下文使用的表达式必须生成与 DataFrame 中原始列长度相同的 Series,而 select
上下文中的表达式只需生成彼此长度相同的 Series 即可。
filter
filter
上下文根据一个或多个评估为布尔数据类型的表达式来筛选 DataFrame 的行。
result = df.filter(
pl.col("birthdate").is_between(date(1982, 12, 31), date(1996, 1, 1)),
pl.col("height") > 1.7,
)
print(result)
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
上下文时,您可以使用表达式动态计算分组
result = df.group_by(
(pl.col("birthdate").dt.year() // 10 * 10).alias("decade"),
).agg(pl.col("name"))
print(result)
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 米的组合进行分组
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)
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 的左侧包含每个分组表达式对应的一列,然后是表示聚合表达式结果所需的任意数量的列。反过来,我们可以指定任意数量的聚合表达式
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)
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_dynamic
和 rolling
。
表达式扩展
最后一个示例包含两个分组表达式和三个聚合表达式,但结果 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 中,它适用于两列
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
df2 = pl.DataFrame(
{
"ints": [1, 2, 3, 4],
"letters": ["A", "B", "C", "D"],
}
)
result = df2.select(expr)
print(result)
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 可以获得进一步的性能提升。
我们仅仅触及了表达式能力的冰山一角。还有大量的表达式,它们可以通过多种方式组合。有关可用不同类型表达式的更深入探讨,请参阅表达式章节。
-
本指南后面还会介绍额外的列表 (List) 和 SQL 上下文。但为简单起见,我们暂时不讨论它们。 ↩