Polars分组 #

多线程 #

处理表状数据最高效的方式就是通过“分割-处理-组合”的方式并行地进行。这样的操作正是 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 有多少 Pro 和 Anti。我们可以不用 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 政治家的名字，并按照年龄排序。我们可以用 sort 和 groupby：

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!