折叠
Polars 提供了许多表达式来跨列执行计算,例如 sum_horizontal、mean_horizontal 和 min_horizontal。然而,这些只是称为折叠的通用算法的特殊情况,当 Polars 的专用版本不够用时,Polars 提供了一种通用机制,供您计算自定义折叠。
使用函数fold计算的折叠操作在完整列上以获取最大速度。它们非常高效地利用数据布局,并且通常具有向量化执行。
基本示例
作为第一个例子,我们将使用函数fold重新实现sum_horizontal:
import operator
import polars as pl
df = pl.DataFrame(
{
"label": ["foo", "bar", "spam"],
"a": [1, 2, 3],
"b": [10, 20, 30],
}
)
result = df.select(
pl.fold(
acc=pl.lit(0),
function=operator.add,
exprs=pl.col("a", "b"),
).alias("sum_fold"),
pl.sum_horizontal(pl.col("a", "b")).alias("sum_horz"),
)
print(result)
use polars::lazy::dsl::sum_horizontal;
use polars::prelude::*;
let df = df!(
"label" => ["foo", "bar", "spam"],
"a" => [1, 2, 3],
"b" => [10, 20, 30],
)?;
let result = df
.clone()
.lazy()
.select([
fold_exprs(
lit(0),
|acc, val| (&acc + &val).map(Some),
[col("a"), col("b")],
)
.alias("sum_fold"),
sum_horizontal([col("a"), col("b")], true)?.alias("sum_horz"),
])
.collect()?;
println!("{:?}", result);
shape: (3, 2)
┌──────────┬──────────┐
│ sum_fold ┆ sum_horz │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞══════════╪══════════╡
│ 11 ┆ 11 │
│ 22 ┆ 22 │
│ 33 ┆ 33 │
└──────────┴──────────┘
函数 fold 期望一个函数 f 作为参数 function,并且 f 应该接受两个参数。第一个参数是累积结果,我们将其初始化为零,第二个参数接受参数 exprs 中列出的表达式的连续值。在我们的例子中,它们是两列“a”和“b”。
下面的代码片段包括第三个显式表达式,它表示函数 fold 在上面所做的操作:
acc = pl.lit(0)
f = operator.add
result = df.select(
f(f(acc, pl.col("a")), pl.col("b")),
pl.fold(acc=acc, function=f, exprs=pl.col("a", "b")).alias("sum_fold"),
)
print(result)
let acc = lit(0);
let f = |acc: Expr, val: Expr| acc + val;
let result = df
.clone()
.lazy()
.select([
f(f(acc, col("a")), col("b")),
fold_exprs(
lit(0),
|acc, val| (&acc + &val).map(Some),
[col("a"), col("b")],
)
.alias("sum_fold"),
])
.collect()?;
println!("{:?}", result);
shape: (3, 2)
┌─────────┬──────────┐
│ literal ┆ sum_fold │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════════╪══════════╡
│ 11 ┆ 11 │
│ 22 ┆ 22 │
│ 33 ┆ 33 │
└─────────┴──────────┘
fold in Python
大多数编程语言都包含一个高阶函数,该函数实现了Polars中fold函数所实现的算法。
Polars的fold与Python的functools.reduce非常相似。
你可以在这篇文章中了解更多关于functools.reduce的强大功能。
初始值 acc
为累加器acc选择的初始值通常是,但不总是,你想要应用的操作的单位元。例如,如果我们想要跨列相乘,如果我们的累加器设置为零,我们将不会得到正确的结果:
result = df.select(
pl.fold(
acc=pl.lit(0),
function=operator.mul,
exprs=pl.col("a", "b"),
).alias("prod"),
)
print(result)
let result = df
.clone()
.lazy()
.select([fold_exprs(
lit(0),
|acc, val| (&acc * &val).map(Some),
[col("a"), col("b")],
)
.alias("prod")])
.collect()?;
println!("{:?}", result);
shape: (3, 1)
┌──────┐
│ prod │
│ --- │
│ i64 │
╞══════╡
│ 0 │
│ 0 │
│ 0 │
└──────┘
要解决这个问题,累加器 acc 应设置为 1:
result = df.select(
pl.fold(
acc=pl.lit(1),
function=operator.mul,
exprs=pl.col("a", "b"),
).alias("prod"),
)
print(result)
let result = df
.clone()
.lazy()
.select([fold_exprs(
lit(1),
|acc, val| (&acc * &val).map(Some),
[col("a"), col("b")],
)
.alias("prod")])
.collect()?;
println!("{:?}", result);
shape: (3, 1)
┌──────┐
│ prod │
│ --- │
│ i64 │
╞══════╡
│ 10 │
│ 40 │
│ 90 │
└──────┘
条件
在您希望跨数据帧的所有列应用条件/谓词的情况下,fold 可以是一种非常简洁的表达方式。
df = pl.DataFrame(
{
"a": [1, 2, 3],
"b": [0, 1, 2],
}
)
result = df.filter(
pl.fold(
acc=pl.lit(True),
function=lambda acc, x: acc & x,
exprs=pl.all() > 1,
)
)
print(result)
let df = df!(
"a" => [1, 2, 3],
"b" => [0, 1, 2],
)?;
let result = df
.clone()
.lazy()
.filter(fold_exprs(
lit(true),
|acc, val| (&acc & &val).map(Some),
[col("*").gt(1)],
))
.collect()?;
println!("{:?}", result);
shape: (1, 2)
┌─────┬─────┐
│ a ┆ b │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 3 ┆ 2 │
└─────┴─────┘
上面的代码片段过滤了所有列都大于1的行。
折叠和字符串数据
折叠可用于连接字符串数据。然而,由于中间列的具体化,此操作将具有平方复杂度。
因此,我们建议使用函数 concat_str 来实现这一点:
df = pl.DataFrame(
{
"a": ["a", "b", "c"],
"b": [1, 2, 3],
}
)
result = df.select(pl.concat_str(["a", "b"]))
print(result)
concat_str · 在功能 concat_str 上可用
let df = df!(
"a" => ["a", "b", "c"],
"b" => [1, 2, 3],
)?;
let result = df
.lazy()
.select([concat_str([col("a"), col("b")], "", false)])
.collect()?;
println!("{:?}", result);
shape: (3, 1)
┌─────┐
│ a │
│ --- │
│ str │
╞═════╡
│ a1 │
│ b2 │
│ c3 │
└─────┘