分组

多线程

处理表状数据最高效的方式就是通过“分割-处理-组合”的方式并行地进行。这样的操作正是 Polars 的 分组操作的核心,也是 Polars 如此高效的秘密。特别指出,分割和处理都是多线程执行的。

下面的例子展示了分组操作的流程:

对于分割阶段的哈希操作,Polars 使用了无锁多线程方式,如下图所示:

这样的并行操作可以让分组和联合操作非常非常高效。

更多解释参考 这篇博客

不要“杀死”并行

众所周知,Python 慢、水平拓展不好。除了因为是解释型语言,Python 还收到全局解释器锁,GIL。 这就意味着,如果你传入一个 lambda 或者 Python 自定义函数,Polars 速度会被限制,即 无法使用多核进行并行计算。

这是个很糟糕的情况,特别我们在做 .groupby 的时候会经常传入 lambda 函数。虽然 Polars 支持这种操作,但是请注意 Python 的限制,特别是解释器和GIL。

为了解决这个问题,Polars 实现了一种非常强大的语法,在其延迟执行API和即时执行API上都有定义。

Polars Expressions

刚才我们提到自定义 Python 函数会损伤并行能力,Polars 提供了惰性 API 来应对这种情况。接下来 我们看看这是什么意思。

我们可以从这个数据集开始:US congress dataset.

import polars as pl

from .dataset import dataset

q = (
    dataset.lazy()
    .groupby("first_name")
    .agg(
        [
            pl.count(),
            pl.col("gender"),
            pl.first("last_name"),
        ]
    )
    .sort("count", descending=True)
    .limit(5)
)

df = q.collect()

基本聚合操作

你可以轻松地把多个聚合表达式放在一个 list 里面,并没有数量限制,你可以任意组合你放入任何数量的表达式。 下面这段代码中我们做如下聚合操作:

对于每一个 first_name 分组:

  • 统计每组的行数:
    • 短版:pl.count("party")
    • 长版:pl.col("party").count()
  • 把每组的性别放入一个列表:
    • 长版: pl.col("gender").list()
  • 找到每组的第一个 last_name
    • 短版: pl.first("last_name")
    • 长版: pl.col("last_name").first()

除了聚合,我们还立即对结果进行排序,并取其中前5条记录,这样我们能更好地从宏观角度理解这组数据的特征。

import polars as pl

from .dataset import dataset

q = (
    dataset.lazy()
    .groupby("first_name")
    .agg(
        [
            pl.count(),
            pl.col("gender"),
            pl.first("last_name"),
        ]
    )
    .sort("count", descending=True)
    .limit(5)
)

df = q.collect()
shape: (5, 4)
┌────────────┬───────┬───────────────────┬───────────┐
│ first_name ┆ count ┆ 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     ┆ 454   ┆ ["M", "M", … "M"] ┆ Tucker    │
│ Charles    ┆ 439   ┆ ["M", "M", … "M"] ┆ Carroll   │
└────────────┴───────┴───────────────────┴───────────┘

条件

简单吧!我们加点料!假设我们想要知道对于每个 state 有多少 ProAnti。我们可以 不用 lambda 而直接查询。

import polars as pl

from .dataset import dataset

q = (
    dataset.lazy()
    .groupby("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()
shape: (5, 3)
┌───────┬──────┬─────┐
│ state ┆ anti ┆ pro │
│ ---   ┆ ---  ┆ --- │
│ cat   ┆ u32  ┆ u32 │
╞═══════╪══════╪═════╡
│ NJ    ┆ 0    ┆ 3   │
│ CT    ┆ 0    ┆ 3   │
│ NC    ┆ 1    ┆ 2   │
│ VA    ┆ 3    ┆ 1   │
│ MA    ┆ 0    ┆ 1   │
└───────┴──────┴─────┘

类似的,我们可以通过多层聚合实现,但是这不利于我显摆这些很酷的特征😉!

import polars as pl

from .dataset import dataset

q = (
    dataset.lazy()
    .groupby(["state", "party"])
    .agg([pl.count("party").alias("count")])
    .filter((pl.col("party") == "Anti-Administration") | (pl.col("party") == "Pro-Administration"))
    .sort("count", descending=True)
    .limit(5)
)

df = q.collect()
shape: (5, 3)
┌───────┬─────────────────────┬───────┐
│ state ┆ party               ┆ count │
│ ---   ┆ ---                 ┆ ---   │
│ cat   ┆ cat                 ┆ u32   │
╞═══════╪═════════════════════╪═══════╡
│ VA    ┆ Anti-Administration ┆ 3     │
│ CT    ┆ Pro-Administration  ┆ 3     │
│ NJ    ┆ Pro-Administration  ┆ 3     │
│ NC    ┆ Pro-Administration  ┆ 2     │
│ VA    ┆ Pro-Administration  ┆ 1     │
└───────┴─────────────────────┴───────┘

过滤

我们也可以过滤分组。假设我们想要计算每组的均值,但是我们不希望计算所有值的均值,我们也不希望直接 从 DataFrame 过滤,因为我们后需还需要那些行做其他操作。

下面的例子说明我们是如何做到的。注意,我们可以写明 Python 的自定义函数,这些函数没有什么 运行时开销。因为这些函数返回了 Polars 表达式,我们并没在运行时让 Series 调用自动函数。

from datetime import date

import polars as pl

from .dataset import dataset


def compute_age() -> pl.Expr:
    return date(2021, 1, 1).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()
    .groupby(["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()
shape: (5, 5)
┌───────┬────────────────┬────────────────┬────────┬──────────┐
│ state ┆ avg M birthday ┆ avg F birthday ┆ # male ┆ # female │
│ ---   ┆ ---            ┆ ---            ┆ ---    ┆ ---      │
│ cat   ┆ f64            ┆ f64            ┆ u32    ┆ u32      │
╞═══════╪════════════════╪════════════════╪════════╪══════════╡
│ WI    ┆ 152.939698     ┆ null           ┆ 199    ┆ 0        │
│ LA    ┆ 157.195531     ┆ 97.8           ┆ 194    ┆ 5        │
│ OH    ┆ 171.836735     ┆ 79.444444      ┆ 672    ┆ 9        │
│ MO    ┆ 163.741433     ┆ 81.625         ┆ 329    ┆ 8        │
│ PA    ┆ 179.724846     ┆ 91.857143      ┆ 1050   ┆ 7        │
└───────┴────────────────┴────────────────┴────────┴──────────┘

排序

我们经常把一个 DataFrame 排序为了在分组操作的时候保持某种顺序。假设我们我们希望知道 每个 state 政治家的名字,并按照年龄排序。我们可以用 sortgroupby

import polars as pl

from .dataset import dataset


def get_person() -> pl.Expr:
    return pl.col("first_name") + pl.lit(" ") + pl.col("last_name")


q = (
    dataset.lazy()
    .sort("birthday", descending=True)
    .groupby(["state"])
    .agg(
        [
            get_person().first().alias("youngest"),
            get_person().last().alias("oldest"),
        ]
    )
    .limit(5)
)

df = q.collect()
shape: (5, 3)
┌───────┬──────────────────┬─────────────────┐
│ state ┆ youngest         ┆ oldest          │
│ ---   ┆ ---              ┆ ---             │
│ cat   ┆ str              ┆ str             │
╞═══════╪══════════════════╪═════════════════╡
│ VT    ┆ Benjamin Deming  ┆ Moses Robinson  │
│ MT    ┆ Greg Gianforte   ┆ James Cavanaugh │
│ MN    ┆ Erik Paulsen     ┆ Cyrus Aldrich   │
│ AS    ┆ Eni Faleomavaega ┆ Fofó Sunia      │
│ NC    ┆ James McKay      ┆ Samuel Johnston │
└───────┴──────────────────┴─────────────────┘

但是,如果我们想把名字也按照字母排序,上面的代码就不行了。 幸运的是,我们可以在 groupby 上下文中进行排序,与 DataFrame 无关。

import polars as pl

from .dataset import dataset


def get_person() -> pl.Expr:
    return pl.col("first_name") + pl.lit(" ") + pl.col("last_name")


q = (
    dataset.lazy()
    .sort("birthday", descending=True)
    .groupby(["state"])
    .agg(
        [
            get_person().first().alias("youngest"),
            get_person().last().alias("oldest"),
            get_person().sort().first().alias("alphabetical_first"),
        ]
    )
    .limit(5)
)

df = q.collect()
shape: (5, 4)
┌───────┬─────────────────────┬─────────────────┬────────────────────┐
│ state ┆ youngest            ┆ oldest          ┆ alphabetical_first │
│ ---   ┆ ---                 ┆ ---             ┆ ---                │
│ cat   ┆ str                 ┆ str             ┆ str                │
╞═══════╪═════════════════════╪═════════════════╪════════════════════╡
│ OH    ┆ Amos Townsend       ┆ Paul Fearing    ┆ Aaron Harlan       │
│ KY    ┆ Benjamin Grey       ┆ Matthew Lyon    ┆ Aaron Harding      │
│ HI    ┆ Tulsi Gabbard       ┆ Robert Wilcox   ┆ Cecil Heftel       │
│ LA    ┆ John Slidell        ┆ Thomas Posey    ┆ Adolph Meyer       │
│ PR    ┆ Aníbal Acevedo-Vilá ┆ Tulio Larrinaga ┆ Antonio Colorado   │
└───────┴─────────────────────┴─────────────────┴────────────────────┘

我们甚至可以在 groupby 上下文中增加另一个列,并且按照男女排序: pl.col("gender").sort_by("first_name").first().alias("gender")

import polars as pl

from .dataset import dataset


def get_person() -> pl.Expr:
    return pl.col("first_name") + pl.lit(" ") + pl.col("last_name")


q = (
    dataset.lazy()
    .sort("birthday", descending=True)
    .groupby(["state"])
    .agg(
        [
            get_person().first().alias("youngest"),
            get_person().last().alias("oldest"),
            get_person().sort().first().alias("alphabetical_first"),
            pl.col("gender").sort_by("first_name").first().alias("gender"),
        ]
    )
    .sort("state")
    .limit(5)
)

df = q.collect()
shape: (5, 5)
┌───────┬────────────────┬─────────────────┬────────────────────┬────────┐
│ state ┆ youngest       ┆ oldest          ┆ alphabetical_first ┆ gender │
│ ---   ┆ ---            ┆ ---             ┆ ---                ┆ ---    │
│ cat   ┆ str            ┆ str             ┆ str                ┆ cat    │
╞═══════╪════════════════╪═════════════════╪════════════════════╪════════╡
│ CT    ┆ Samuel Simons  ┆ Roger Sherman   ┆ Abner Sibal        ┆ M      │
│ KY    ┆ Benjamin Grey  ┆ Matthew Lyon    ┆ Aaron Harding      ┆ M      │
│ FL    ┆ George Hawkins ┆ Joseph White    ┆ Abijah Gilbert     ┆ M      │
│ NY    ┆ Robert Baker   ┆ Philip Schuyler ┆ A. Foster          ┆ M      │
│ MI    ┆ Samuel Clark   ┆ Gabriel Richard ┆ Aaron Bliss        ┆ M      │
└───────┴────────────────┴─────────────────┴────────────────────┴────────┘

结论

上面的例子中我们知道通过组合表达式可以完成复杂的查询。而且,我们避免了使用自定义 Python 函数 带来的性能损失 (解释器和 GIL)。

如果这里少了哪类表达式,清在这里开一个 Issue: feature request!