分组
多线程
处理表状数据最高效的方式就是通过“分割-处理-组合”的方式并行地进行。这样的操作正是 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!