聚合
Polars 的 上下文 group_by 允许你对列的子集应用表达式，这些子集由数据分组所依据的列的唯一值定义。这是我们将在本用户指南部分探讨的一项非常强大的功能。
我们首先读取一个 美国国会 dataset
import polars as pl
url = "hf://datasets/nameexhaustion/polars-docs/legislators-historical.csv"
schema_overrides = {
"first_name": pl.Categorical,
"gender": pl.Categorical,
"type": pl.Categorical,
"state": pl.Categorical,
"party": pl.Categorical,
}
dataset = (
pl.read_csv(url, schema_overrides=schema_overrides)
.with_columns(pl.col("first", "middle", "last").name.suffix("_name"))
.with_columns(pl.col("birthday").str.to_date(strict=False))
)
DataFrame · Categorical · 通过特性 dtype-categorical 可用
use std::io::Cursor;
use polars::prelude::*;
use reqwest::blocking::Client;
let url = "https://hugging-face.cn/datasets/nameexhaustion/polars-docs/resolve/main/legislators-historical.csv";
let mut schema = Schema::default();
schema.with_column(
"first_name".into(),
DataType::Categorical(None, Default::default()),
);
schema.with_column(
"gender".into(),
DataType::Categorical(None, Default::default()),
);
schema.with_column(
"type".into(),
DataType::Categorical(None, Default::default()),
);
schema.with_column(
"state".into(),
DataType::Categorical(None, Default::default()),
);
schema.with_column(
"party".into(),
DataType::Categorical(None, Default::default()),
);
schema.with_column("birthday".into(), DataType::Date);
let data = Client::new().get(url).send()?.bytes()?;
let dataset = CsvReadOptions::default()
.with_has_header(true)
.with_schema_overwrite(Some(Arc::new(schema)))
.map_parse_options(|parse_options| parse_options.with_try_parse_dates(true))
.into_reader_with_file_handle(Cursor::new(data))
.finish()?
.lazy()
.with_columns([
col("first").name().suffix("_name"),
col("middle").name().suffix("_name"),
col("last").name().suffix("_name"),
])
.collect()?;
println!("{}", &dataset);
基本聚合
你可以轻松地将多个表达式应用于聚合值。只需在函数 agg 内部列出你想要的所有表达式即可。你可以进行的聚合数量没有上限，你可以进行任意组合。在下面的代码片段中，我们将根据“first_name”列对数据进行分组，然后应用以下聚合操作：
- 计算组中的行数（这意味着我们计算数据集中有多少人拥有每个唯一的名）；
- 通过引用“gender”列但省略聚合函数，将其值合并为一个列表；以及
- 获取组中“last_name”列的第一个值。
在计算聚合后，我们立即对结果进行排序，并将其限制为前五行，以便获得清晰的汇总概览。
q = (
dataset.lazy()
.group_by("first_name")
.agg(
pl.len(),
pl.col("gender"),
pl.first("last_name"), # Short for `pl.col("last_name").first()`
)
.sort("len", descending=True)
.limit(5)
)
df = q.collect()
print(df)
let df = dataset
.clone()
.lazy()
.group_by(["first_name"])
.agg([len(), col("gender"), col("last_name").first()])
.sort(
["len"],
SortMultipleOptions::default()
.with_order_descending(true)
.with_nulls_last(true),
)
.limit(5)
.collect()?;
println!("{df}");
shape: (5, 4)
┌────────────┬──────┬───────────────────┬───────────┐
│ first_name ┆ len ┆ gender ┆ last_name │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ u32 ┆ list[cat] ┆ str │
╞════════════╪══════╪═══════════════════╪═══════════╡
│ John ┆ 4227 ┆ ["M", "M", … "M"] ┆ Walker │
│ William ┆ 3309 ┆ ["M", "M", … "M"] ┆ Few │
│ James ┆ 2414 ┆ ["M", "M", … "M"] ┆ Armstrong │
│ Charles ┆ 1514 ┆ ["M", "M", … "M"] ┆ Carroll │
│ Thomas ┆ 1502 ┆ ["M", "M", … "M"] ┆ Tucker │
└────────────┴──────┴───────────────────┴───────────┘
就这么简单！让我们更进一步。
条件表达式
假设我们想知道一个州有多少代表是支持或反对政府的。我们可以在聚合中直接查询，而无需使用 lambda 或整理 DataFrame。
q = (
dataset.lazy()
.group_by("state")
.agg(
(pl.col("party") == "Anti-Administration").sum().alias("anti"),
(pl.col("party") == "Pro-Administration").sum().alias("pro"),
)
.sort("pro", descending=True)
.limit(5)
)
df = q.collect()
print(df)
let df = dataset
.clone()
.lazy()
.group_by(["state"])
.agg([
(col("party").eq(lit("Anti-Administration")))
.sum()
.alias("anti"),
(col("party").eq(lit("Pro-Administration")))
.sum()
.alias("pro"),
])
.sort(
["pro"],
SortMultipleOptions::default().with_order_descending(true),
)
.limit(5)
.collect()?;
println!("{df}");
shape: (5, 3)
┌───────┬──────┬─────┐
│ state ┆ anti ┆ pro │
│ --- ┆ --- ┆ --- │
│ cat ┆ u32 ┆ u32 │
╞═══════╪══════╪═════╡
│ CT ┆ 0 ┆ 5 │
│ NJ ┆ 0 ┆ 3 │
│ DE ┆ 1 ┆ 3 │
│ MD ┆ 0 ┆ 2 │
│ MA ┆ 0 ┆ 2 │
└───────┴──────┴─────┘
筛选
我们还可以过滤分组。假设我们想计算每个组的平均值，但我们不想包含该组中的所有值，也不想实际从 DataFrame 中过滤行，因为我们还需要这些行用于另一个聚合。
在下面的例子中，我们将展示如何实现这一点。
注意
请注意，我们可以定义 Python 函数以提高清晰度。这些函数不会增加开销，因为它们返回 Polars 表达式，我们不会在查询运行时将自定义函数应用于 Series。当然，你也可以在 Rust 中编写返回表达式的函数。
from datetime import date
def compute_age():
return date.today().year - pl.col("birthday").dt.year()
def avg_birthday(gender: str) -> pl.Expr:
return (
compute_age()
.filter(pl.col("gender") == gender)
.mean()
.alias(f"avg {gender} birthday")
)
q = (
dataset.lazy()
.group_by("state")
.agg(
avg_birthday("M"),
avg_birthday("F"),
(pl.col("gender") == "M").sum().alias("# male"),
(pl.col("gender") == "F").sum().alias("# female"),
)
.limit(5)
)
df = q.collect()
print(df)
fn compute_age() -> Expr {
lit(2024) - col("birthday").dt().year()
}
fn avg_birthday(gender: &str) -> Expr {
compute_age()
.filter(col("gender").eq(lit(gender)))
.mean()
.alias(format!("avg {gender} birthday"))
}
let df = dataset
.clone()
.lazy()
.group_by(["state"])
.agg([
avg_birthday("M"),
avg_birthday("F"),
(col("gender").eq(lit("M"))).sum().alias("# male"),
(col("gender").eq(lit("F"))).sum().alias("# female"),
])
.limit(5)
.collect()?;
println!("{df}");
shape: (5, 5)
┌───────┬────────────────┬────────────────┬────────┬──────────┐
│ state ┆ avg M birthday ┆ avg F birthday ┆ # male ┆ # female │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ cat ┆ f64 ┆ f64 ┆ u32 ┆ u32 │
╞═══════╪════════════════╪════════════════╪════════╪══════════╡
│ VT ┆ 204.59887 ┆ null ┆ 361 ┆ 0 │
│ WY ┆ 138.552632 ┆ 73.0 ┆ 114 ┆ 10 │
│ WV ┆ 138.98301 ┆ 130.0 ┆ 415 ┆ 7 │
│ MO ┆ 152.969302 ┆ 88.52381 ┆ 1086 ┆ 42 │
│ NH ┆ 192.462791 ┆ 69.363636 ┆ 438 ┆ 11 │
└───────┴────────────────┴────────────────┴────────┴──────────┘
平均年龄值看起来很荒谬吗？那是因为我们正在处理可追溯到 19 世纪的历史数据，并且我们在进行计算时假设数据集中所有人都还健在。
嵌套分组
前两个查询都可以使用嵌套的 group_by 来完成，但这无法让我们展示这些特性。😉 要进行嵌套 group_by，只需列出用于分组的列即可。
首先，我们使用嵌套的 group_by 来计算一个州有多少代表是支持或反对政府的
q = (
dataset.lazy()
.group_by("state", "party")
.agg(pl.len().alias("count"))
.filter(
(pl.col("party") == "Anti-Administration")
| (pl.col("party") == "Pro-Administration")
)
.sort("count", descending=True)
.limit(5)
)
df = q.collect()
print(df)
let df = dataset
.clone()
.lazy()
.group_by(["state", "party"])
.agg([len().alias("count")])
.filter(
col("party")
.eq(lit("Anti-Administration"))
.or(col("party").eq(lit("Pro-Administration"))),
)
.sort(
["count"],
SortMultipleOptions::default()
.with_order_descending(true)
.with_nulls_last(true),
)
.limit(5)
.collect()?;
println!("{df}");
shape: (5, 3)
┌───────┬─────────────────────┬───────┐
│ state ┆ party ┆ count │
│ --- ┆ --- ┆ --- │
│ cat ┆ cat ┆ u32 │
╞═══════╪═════════════════════╪═══════╡
│ CT ┆ Pro-Administration ┆ 5 │
│ VA ┆ Anti-Administration ┆ 5 │
│ DE ┆ Pro-Administration ┆ 3 │
│ NJ ┆ Pro-Administration ┆ 3 │
│ PA ┆ Anti-Administration ┆ 3 │
└───────┴─────────────────────┴───────┘
接下来，我们使用嵌套的 group_by 来计算每个州和每种性别的代表平均年龄
q = (
dataset.lazy()
.group_by("state", "gender")
.agg(
# The function `avg_birthday` is not needed:
compute_age().mean().alias("avg birthday"),
pl.len().alias("#"),
)
.sort("#", descending=True)
.limit(5)
)
df = q.collect()
print(df)
let df = dataset
.clone()
.lazy()
.group_by(["state", "gender"])
.agg([compute_age().mean().alias("avg birthday"), len().alias("#")])
.sort(
["#"],
SortMultipleOptions::default()
.with_order_descending(true)
.with_nulls_last(true),
)
.limit(5)
.collect()?;
println!("{df}");
shape: (5, 4)
┌───────┬────────┬──────────────┬──────┐
│ state ┆ gender ┆ avg birthday ┆ # │
│ --- ┆ --- ┆ --- ┆ --- │
│ cat ┆ cat ┆ f64 ┆ u32 │
╞═══════╪════════╪══════════════╪══════╡
│ NY ┆ M ┆ 164.204634 ┆ 3965 │
│ PA ┆ M ┆ 166.008592 ┆ 3205 │
│ OH ┆ M ┆ 156.579961 ┆ 2142 │
│ IL ┆ M ┆ 145.069482 ┆ 1895 │
│ CA ┆ M ┆ 114.400464 ┆ 1725 │
└───────┴────────┴──────────────┴──────┘
请注意，我们得到的结果相同，但数据格式不同。根据具体情况，一种格式可能比另一种更适用。
排序
常见的情况是，DataFrame 会被排序，仅仅是为了在分组操作期间管理其顺序。假设我们想获取每个州最年长和最年轻的政治家的姓名。我们可以先排序，然后分组
def get_name() -> pl.Expr:
return pl.col("first_name") + pl.lit(" ") + pl.col("last_name")
q = (
dataset.lazy()
.sort("birthday", descending=True)
.group_by("state")
.agg(
get_name().first().alias("youngest"),
get_name().last().alias("oldest"),
)
.limit(5)
)
df = q.collect()
print(df)
fn get_name() -> Expr {
col("first_name") + lit(" ") + col("last_name")
}
let df = dataset
.clone()
.lazy()
.sort(
["birthday"],
SortMultipleOptions::default()
.with_order_descending(true)
.with_nulls_last(true),
)
.group_by(["state"])
.agg([
get_name().first().alias("youngest"),
get_name().last().alias("oldest"),
])
.limit(5)
.collect()?;
println!("{df}");
shape: (5, 3)
┌───────┬────────────────┬──────────────────┐
│ state ┆ youngest ┆ oldest │
│ --- ┆ --- ┆ --- │
│ cat ┆ str ┆ str │
╞═══════╪════════════════╪══════════════════╡
│ CA ┆ Edward Gilbert ┆ William Gwin │
│ MI ┆ Edward Bradley ┆ Gabriel Richard │
│ OK ┆ Kendra Horn ┆ David Harvey │
│ NC ┆ John Ashe ┆ Samuel Johnston │
│ TX ┆ John Cranford ┆ Timothy Pilsbury │
└───────┴────────────────┴──────────────────┘
然而，如果我们还想按字母顺序对姓名进行排序，则需要执行额外的排序操作。幸运的是，我们可以在 group_by 上下文中进行排序，而无需更改底层 DataFrame 的排序。
q = (
dataset.lazy()
.sort("birthday", descending=True)
.group_by("state")
.agg(
get_name().first().alias("youngest"),
get_name().last().alias("oldest"),
get_name().sort().first().alias("alphabetical_first"),
)
.limit(5)
)
df = q.collect()
print(df)
let df = dataset
.clone()
.lazy()
.sort(
["birthday"],
SortMultipleOptions::default()
.with_order_descending(true)
.with_nulls_last(true),
)
.group_by(["state"])
.agg([
get_name().first().alias("youngest"),
get_name().last().alias("oldest"),
get_name()
.sort(Default::default())
.first()
.alias("alphabetical_first"),
])
.limit(5)
.collect()?;
println!("{df}");
shape: (5, 4)
┌───────┬──────────────────┬──────────────────┬────────────────────┐
│ state ┆ youngest ┆ oldest ┆ alphabetical_first │
│ --- ┆ --- ┆ --- ┆ --- │
│ cat ┆ str ┆ str ┆ str │
╞═══════╪══════════════════╪══════════════════╪════════════════════╡
│ NH ┆ John Sherburne ┆ Samuel Livermore ┆ Aaron Cragin │
│ MO ┆ Spencer Pettis ┆ Rufus Easton ┆ Abram Comingo │
│ WV ┆ Edward Rohrbough ┆ Daniel Polsley ┆ Adam Littlepage │
│ MA ┆ William Widgery ┆ Artemas Ward ┆ Aaron Hobart │
│ CA ┆ Edward Gilbert ┆ William Gwin ┆ Aaron Sargent │
└───────┴──────────────────┴──────────────────┴────────────────────┘
我们甚至可以根据另一列的顺序对某一列进行排序，这在 group_by 上下文中也同样适用。对前一个查询的此修改允许我们检查具有该名字的代表是男性还是女性
q = (
dataset.lazy()
.sort("birthday", descending=True)
.group_by("state")
.agg(
get_name().first().alias("youngest"),
get_name().last().alias("oldest"),
get_name().sort().first().alias("alphabetical_first"),
pl.col("gender").sort_by(get_name()).first(),
)
.sort("state")
.limit(5)
)
df = q.collect()
print(df)
let df = dataset
.clone()
.lazy()
.sort(
["birthday"],
SortMultipleOptions::default()
.with_order_descending(true)
.with_nulls_last(true),
)
.group_by(["state"])
.agg([
get_name().first().alias("youngest"),
get_name().last().alias("oldest"),
get_name()
.sort(Default::default())
.first()
.alias("alphabetical_first"),
col("gender")
.sort_by(["first_name"], SortMultipleOptions::default())
.first(),
])
.sort(["state"], SortMultipleOptions::default())
.limit(5)
.collect()?;
println!("{df}");
shape: (5, 5)
┌───────┬────────────────────┬──────────────────┬────────────────────┬────────┐
│ state ┆ youngest ┆ oldest ┆ alphabetical_first ┆ gender │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ cat ┆ str ┆ str ┆ str ┆ cat │
╞═══════╪════════════════════╪══════════════════╪════════════════════╪════════╡
│ TN ┆ William Cocke ┆ John Sevier ┆ Aaron Brown ┆ M │
│ CT ┆ Henry Edwards ┆ Roger Sherman ┆ Abner Sibal ┆ M │
│ GA ┆ Thomas Carnes ┆ George Mathews ┆ A. Ferguson ┆ M │
│ NH ┆ John Sherburne ┆ Samuel Livermore ┆ Aaron Cragin ┆ M │
│ NJ ┆ Lambert Cadwalader ┆ Abraham Clark ┆ Aaron Kitchell ┆ M │
└───────┴────────────────────┴──────────────────┴────────────────────┴────────┘
不要破坏并行化
仅限 Python 用户
以下部分仅适用于 Python，不适用于 Rust。在 Rust 中，代码块和闭包（lambda 表达式）可以并且将会并发执行。
Python 通常比 Rust 慢。除了运行“慢速”字节码的开销外，Python 还必须受限于全局解释器锁 (GIL)。这意味着如果你在并行化阶段使用 lambda 或自定义 Python 函数，Polars 的速度将受限于运行 Python 代码，从而阻止任何多线程执行该函数。
Polars 会尝试并行化聚合函数在组上的计算，因此建议你尽可能避免使用 lambda 表达式和自定义 Python 函数。相反，请尽量使用 Polars 表达式 API。然而，这并非总是可行，因此如果你想了解更多关于使用 lambda 表达式的信息，可以查阅用户指南中关于使用用户定义函数的部分。