Skip to content

聚合

Polars 的 context group_by 允许您对列的子集应用表达式,这些子集由数据分组的列的唯一值定义。这是一个非常强大的功能,我们将在用户指南的这一部分中进行探讨。

我们首先读取一个 美国国会 dataset

DataFrame · Categorical

import polars as pl

url = "https://theunitedstates.io/congress-legislators/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("birthday").str.to_date(strict=False)
)

DataFrame · Categorical · 在功能 dtype-categorical 上可用

use std::io::Cursor;

use polars::prelude::*;
use reqwest::blocking::Client;

let url = "https://theunitedstates.io/congress-legislators/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: Vec<u8> = Client::new().get(url).send()?.text()?.bytes().collect();

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()?;

println!("{}", &dataset);

基本聚合

您可以轻松地将多个表达式应用于您的聚合值。只需在函数agg中列出所有您想要的表达式。您可以进行的聚合数量没有上限,并且您可以进行任何您想要的组合。在下面的代码片段中,我们将基于“first_name”列对数据进行分组,然后我们将应用以下聚合:

  • 计算组中的行数(这意味着我们计算数据集中有多少人具有每个唯一的名字);
  • 通过引用列但省略聚合函数,将列“gender”的值合并到一个列表中;以及
  • 获取组内列“last_name”的第一个值。

在计算聚合之后,我们立即对结果进行排序并将其限制在前五行,以便我们有一个很好的摘要概览:

group_by

q = (
    dataset.lazy()
    .group_by("first_name")
    .agg(
        pl.len(),
        pl.col("gender"),
        pl.first("last_name"),  # `pl.col("last_name").first()` 的简写
    )
    .sort("len", descending=True)
    .limit(5)
)

df = q.collect()
print(df)

group_by

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 │
│ ---        ┆ ---  ┆ ---               ┆ ---       │
│ cat        ┆ u32  ┆ list[cat]         ┆ str       │
╞════════════╪══════╪═══════════════════╪═══════════╡
│ John       ┆ 1256 ┆ ["M", "M", … "M"] ┆ Walker    │
│ William    ┆ 1022 ┆ ["M", "M", … "M"] ┆ Few       │
│ James      ┆ 714  ┆ ["M", "M", … "M"] ┆ Armstrong │
│ Thomas     ┆ 453  ┆ ["M", "M", … "M"] ┆ Tucker    │
│ Charles    ┆ 439  ┆ ["M", "M", … "M"] ┆ Carroll   │
└────────────┴──────┴───────────────────┴───────────┘

就是这么简单!让我们再提升一个档次。

条件语句

假设我们想知道一个州有多少代表是“支持”或“反对”政府。我们可以在聚合中直接查询,而不需要使用lambda或整理数据框:

group_by

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)

group_by

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    ┆ 3   │
│ NJ    ┆ 0    ┆ 3   │
│ NC    ┆ 1    ┆ 2   │
│ VA    ┆ 3    ┆ 1   │
│ SC    ┆ 0    ┆ 1   │
└───────┴──────┴─────┘

过滤

我们也可以对组进行过滤。假设我们想要计算每组的平均值,但我们不想包括该组中的所有值,同时我们也不想实际从数据框中过滤掉这些行,因为我们需要这些行进行另一个聚合。

在下面的示例中,我们展示了如何完成此操作。

注意

请注意,为了清晰起见,我们可以定义Python函数。 这些函数不会给我们带来任何成本,因为它们返回Polars表达式,我们在查询运行时不会在系列上应用自定义函数。 当然,你也可以用Rust编写返回表达式的函数。

group_by

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)

group_by

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 {} birthday", gender))
}

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      │
╞═══════╪════════════════╪════════════════╪════════╪══════════╡
│ MA    ┆ 202.757212     ┆ 106.5          ┆ 423    ┆ 4        │
│ SD    ┆ 145.085106     ┆ 95.0           ┆ 47     ┆ 4        │
│ AR    ┆ 158.073394     ┆ 125.4          ┆ 112    ┆ 5        │
│ SC    ┆ 187.018349     ┆ 125.8          ┆ 247    ┆ 5        │
│ VT    ┆ 218.309735     ┆ null           ┆ 116    ┆ 0        │
└───────┴────────────────┴────────────────┴────────┴──────────┘

平均年龄值看起来不合理吗?这是因为我们使用的是可以追溯到1800年代的历史数据,并且我们在进行计算时假设数据集中代表的每个人都还活着并且活跃。

嵌套分组

前面的两个查询可以使用嵌套的group_by来完成,但那样就无法展示这些功能了。😉 要进行嵌套的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)

group_by

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  ┆ 3     │
│ VA    ┆ Anti-Administration ┆ 3     │
│ NJ    ┆ Pro-Administration  ┆ 3     │
│ NC    ┆ Pro-Administration  ┆ 2     │
│ VT    ┆ Anti-Administration ┆ 1     │
└───────┴─────────────────────┴───────┘

接下来,我们使用嵌套的group_by来计算每个州和每种性别代表的平均年龄:

group_by

q = (
    dataset.lazy()
    .group_by("state", "gender")
    .agg(
        # 不需要函数 `avg_birthday`:
        compute_age().mean().alias("avg birthday"),
        pl.len().alias("#"),
    )
    .sort("#", descending=True)
    .limit(5)
)

df = q.collect()
print(df)

group_by

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      ┆ 189.330663   ┆ 1457 │
│ PA    ┆ M      ┆ 183.724846   ┆ 1050 │
│ OH    ┆ M      ┆ 175.672414   ┆ 673  │
│ IL    ┆ M      ┆ 157.710638   ┆ 478  │
│ VA    ┆ M      ┆ 195.542781   ┆ 430  │
└───────┴────────┴──────────────┴──────┘

请注意,我们得到了相同的结果,但数据的格式不同。根据具体情况,一种格式可能比另一种更合适。

排序

通常可以看到,为了在分组操作期间管理排序,会对数据框进行排序。假设我们想要获取每个州最年长和最年轻的政治家的名字。我们可以从排序开始,然后进行分组:

group_by

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)

group_by

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              │
╞═══════╪═══════════════════════╪══════════════════╡
│ WI    ┆ Mike Gallagher        ┆ Henry Dodge      │
│ KY    ┆ John Edwards          ┆ Matthew Lyon     │
│ NC    ┆ John Ashe             ┆ Samuel Johnston  │
│ TX    ┆ John Cranford         ┆ Timothy Pilsbury │
│ NY    ┆ Cornelius Schoonmaker ┆ Philip Schuyler  │
└───────┴───────────────────────┴──────────────────┘

然而,如果我们还想按字母顺序对名称进行排序,我们需要执行额外的排序操作。幸运的是,我们可以在group_by上下文中进行排序,而不会改变底层数据框的排序:

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"),
    )
    .limit(5)
)

df = q.collect()
print(df)

group_by

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                │
╞═══════╪═════════════════════╪══════════════════╪════════════════════╡
│ TX    ┆ John Cranford       ┆ Timothy Pilsbury ┆ Abraham Kazen      │
│ KS    ┆ Steven Watkins      ┆ James Lane       ┆ Abel Wilder        │
│ AK    ┆ Mark Begich         ┆ Thomas Cale      ┆ Anthony Dimond     │
│ DK    ┆ George Mathews      ┆ John Todd        ┆ George Mathews     │
│ IL    ┆ Benjamin Stephenson ┆ Shadrack Bond    ┆ Aaron Schock       │
└───────┴─────────────────────┴──────────────────┴────────────────────┘

我们甚至可以按照另一列的顺序对列进行排序,这同样适用于group_by上下文。对之前查询的修改让我们可以检查名字为第一个的代表是男性还是女性:

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)

group_by

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    │
╞═══════╪═══════════════════╪═══════════════════╪════════════════════╪════════╡
│ DE    ┆ Samuel White      ┆ George Read       ┆ Albert Polk        ┆ M      │
│ VA    ┆ William Grayson   ┆ Robert Rutherford ┆ A. McEachin        ┆ M      │
│ SC    ┆ Ralph Izard       ┆ Thomas Sumter     ┆ Abraham Nott       ┆ M      │
│ MD    ┆ Benjamin Contee   ┆ William Smith     ┆ Albert Blakeney    ┆ M      │
│ PA    ┆ Thomas Fitzsimons ┆ Israel Jacobs     ┆ Aaron Kreider      ┆ M      │
└───────┴───────────────────┴───────────────────┴────────────────────┴────────┘

不要杀死并行化

仅限Python用户

以下部分特定于Python,不适用于Rust。 在Rust中,块和闭包(lambda)可以并且将会并发执行。

Python 通常比 Rust 慢。除了运行“慢”字节码的开销外,Python 还必须保持在全局解释器锁(GIL)的限制内。这意味着如果你在并行化阶段使用 lambda 或自定义 Python 函数来应用,Polars 的速度将受到运行 Python 代码的限制,阻止任何多个线程执行该函数。

Polars 会尝试并行化聚合函数在组上的计算,因此建议您尽可能避免使用 lambda 和自定义 Python 函数。相反,尽量保持在 Polars 表达式 API 的范围内。然而,这并不总是可能的,所以如果您想了解更多关于使用 lambda 的信息,您可以访问 用户指南中关于使用用户定义函数的部分