从 Pandas 到 Polars 九：排序数据优化2-按组聚合

在之前的文章中，我们看到Polars具有针对排序数据计算某些统计信息的快速算法。在这篇文章中，我们了解到Polars还有一个针对在排序列上获取groupby键的快速算法。

性能比较

在这个例子中，我们创建了一个包含ID列的DataFrame，并且我们知道这个列是已排序的。在生成下面的DataFrame时，我们将N设置为10_000_000，以创建一个包含1000万行的DataFrame。

我们还将基数（cardinality）设置为10——这意味着我们有10个不同的ID进行分组。基数是一个关键变量，如下文所示。

import polars as pl
import numpy as np
N = 10_000_000
cardinality = 10
# 创建一个按整数id排序的数组
sorted_array = np.sort(np.random.randint(0,cardinality,N))
# 创建一个带有已排序ID列和随机值列的DataFrame
df = (
    pl.DataFrame(
        {
            "id":[i for i in sorted_array],
            "values":np.random.standard_normal(N)
        }
    )
)

对于这个比较，我们根据id列进行分组，并计算values列的平均值。

(
    df
    .groupby("id")
    .agg(
        pl.col("values").mean()
    )
)

目前，Polars不知道id列已经排序，因此不会使用快速算法。如果我们使用标准的非排序算法进行上述的分组聚合，它需要70毫秒。

正如我们在上一篇文章中看到的，我们使用set_sorted表达式告诉Polarsid列已经排序。

df = (
    df
    .with_column(
        pl.col("id").set_sorted()
    )
)

如果我们再次在Polars知道已排序的列上运行分组聚合，它只需要14毫秒——这是标准算法的5倍速度。

基数为王

所以，我们看到当Polars知道列已排序时，在排序列上进行分组聚合要快得多，对吧？不过，情况并不完全如此简单……

在上面的例子中，我们设置了基数为10，因此在1000万行中只有10个不同的组ID。

如果我们创建一个新的DataFrame，基数设为100_000，我们发现标准算法需要90毫秒，而快速排序算法需要60毫秒。在这种情况下，快速排序算法只快了1.5倍。

如果我们设置基数为1_000_000，我们发现这两种方法所需的时间差不多（或者标准方法稍微快一些）。

这种收敛的原因是快速排序算法不会逐个查看id列的元素，而是试图在id列的开头和结尾找到具有相同值的块。当算法找到这些块时，它可以假设整个块具有相同的值，因为它知道列已排序。

当基数较低时，存在许多这样的统一块，快速排序算法可以快速通过id列。随着基数的增加，统一块的长度减少，直到最终快速排序算法实际上在进行逐个元素的搜索。在这种情况下，快速排序算法在逐元素搜索之前仍在尝试跳过列中的元素，因此当基数较高时，排序方法甚至可能更慢。

考虑到这一点，这仍然是一个强大的优化方法，如果你知道你的数据已排序且基数不太高，就可以应用它。我还认为，当基数较高时，可以应用进一步的优化来加快速度。