从零开始的分组实例

156 阅读8分钟

我发现在我的科学编码中成长的最好方法之一是花时间比较各种实现我认为有用的特定算法的方法的效率,以建立对科学Python生态系统的构建块的性能的直觉。

在这种情况下,今天我想看一下在许多方面都是数据驱动的探索的基本操作:group-by,也就是所谓的分割-应用-合并模式。图中显示了一个总和分组的典型例子,借用了《Python数据科学手册》中的聚合和分组部分。

split-apply-combine diagram

其基本思想是根据某些值将数据分成几组,对每组中的数据子集应用特定的操作(通常是聚合),然后将结果合并到一个输出数据框中。Python用户通常会求助于Pandas库来进行这种操作,它通过一个简洁的面向对象的API有效地实现。

在[1]中:

keys   = ['A', 'B', 'C', 'A', 'B', 'C']
vals = [ 1,   2,   3,   4,   5,   6 ]

In [2]:

import pandas as pd
pd.Series(vals).groupby(keys).sum()

输出[2]:

A    5
B    7
C    9
dtype: int64

在Python世界中经常重复的一个说法是,进行这种操作需要Pandas--NumPy和SciPy中提供的工具对元素操作是有用的,但对群组操作则不然。这是我在自己的教学和写作中经常重复的说法,但我最近发现自己在想这是否真的是真的。

下面我将从多个角度来探讨这个问题,利用内置的Python操作以及NumPy和SciPy提供的操作,来进行逐组求和。我希望通过这些,你能对PyData生态系统中各种算法构件的性能获得一些直觉,以及如何有效地应用它们来解决自己的问题。

计时功能

为了方便起见,让我们从定义一些功能开始,这些功能可以让我们快速剖析和可视化分组功能。我们假设要剖析的操作将以函数的形式出现,该函数接收一个键和值的列表,并返回一个组的总和字典。比如说。

In [3]:

def pandas_groupby(keys, vals):
    return pd.Series(vals).groupby(keys).sum().to_dict()

pandas_groupby(keys, vals)

Out[3]:

{'A': 5, 'B': 7, 'C': 9}

为了避免在笔记本中重复运行相同的基准,我将利用 Python 的functools.lru_cache 装饰器,它可以缓存长期运行的函数的结果,如果用相同的输入再次调用该函数,则返回缓存的值。

在 [4]:

%matplotlib inline

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
plt.style.use('seaborn-whitegrid')

from functools import lru_cache

# timeit magic requires global variables, so define them here
timeitfunc = None
timeitkeys = None
timeitvals = None


@lru_cache(maxsize=256)
def _bench_one(func, n_group, size, rseed=754389):
    """Compute a timing and cache it for later use"""
    global timeitfunc, timeitkeys, timeitvals
    rand = np.random.RandomState(rseed)
    timeitkeys = rand.randint(0, n_group, size)
    timeitvals = rand.rand(size)
    timeitfunc = func
    # Following is a magic function; only works in IPython/Jupyter
    t = %timeit -oq timeitfunc(timeitkeys, timeitvals)
    return t.best


def bench(funcs, n_groups, sizes):
    """Run a set of benchmarks and return as a dataframe"""
    n_groups, sizes = np.broadcast_arrays(n_groups, sizes)
    names = [func.__name__.split('_')[0] for func in funcs]
    timings = {name: [_bench_one(func, n_group, size)
                      for n_group, size in zip(n_groups, sizes)]
               for name, func in zip(names, funcs)}
    ind = pd.MultiIndex.from_arrays([n_groups, sizes],
                                    names=['num_groups', 'size'])
    return pd.DataFrame(timings, index=ind)[names]


def plot_benchmarks(funcs, n_groups=None, sizes=None):
    """Plot the benchmarks as a function of data size and number of groups"""
    if n_groups is None:
        n_groups = (10 ** np.linspace(0, 4, 10)).astype(int)
    if sizes is None:
        sizes = (10 ** np.linspace(2, 6, 10)).astype(int)
    timings_sizes = bench(funcs, 10, sizes).reset_index('num_groups', drop=True)
    timings_groups = bench(funcs, n_groups, 10000).reset_index('size', drop=True)
    
    fig, ax = plt.subplots(1, 2, figsize=(10, 3.5), sharey=True)
    timings_sizes.plot(ax=ax[0], logx=True, logy=True)
    ax[0].set(xlabel='size of dataset', ylabel='seconds',
              title='Varying dataset size with 10 groups')
    timings_groups.plot(ax=ax[1], logx=True, logy=True)
    ax[1].set(xlabel='number of groups', ylabel='seconds',
              title='Varying number of groups with 10000 elements')

    return fig

有了这个设置,我们就可以用一个简洁的命令来绘制潘达斯方法的一些基准。

在[5]中:

plot_benchmarks([pandas_groupby]);

对于Pandas Groupby操作,对于小的数据集有一些非实质性的扩展,随着数据的增长,它的执行时间与数据点的数量大约是线性的。pandas groupby是用高度优化的cython代码实现的,它为我们的探索提供了一个很好的比较基准线。

现在让我们来看看一些方法,我们可以从头开始实现这个操作。

用Python dicts进行分组

一个简单的Python实现可以在Python字典上使用一些简单的循环。

在[6]中:

result = {}
for key, val in zip(keys, vals):
    if key not in result:
        result[key] = 0
    result[key] += val
print(result)
{'A': 5, 'B': 7, 'C': 9}

这很简单,但是我们可以通过使用内置的defaultdict 集合来移除条件,当它们被引用时,它将自动初始化值。将其包裹在一个函数中,可以得到。

In [7]中:

from collections import defaultdict

def dict_groupby(keys, vals):
    count = defaultdict(int)
    for key, val in zip(keys, vals):
        count[key] += val
    return dict(count)

dict_groupby(keys, vals)

Out[7]:

{'A': 5, 'B': 7, 'C': 9}

有了这个函数的定义,我们可以将其性能与Pandas函数进行比较。

In [8]:

plot_benchmarks([pandas_groupby, dict_groupby]);

字典的实现随着数据的大小有非常可预测的线性扩展,因为它依赖于对数据集的一次传递。显然,在几千个元素以下,这种简单的Python方法甚至可以超过Pandas操作,但随着点数的增加,基于dict的实现变得比优化的Pandas代码慢得多。

Itertools方法¶

作为纯Python分组的另一种方法,你可能会想求助于标准库中的itertools.groupby 函数。经过进一步检查,这有点复杂,因为这个操作要求分组在输入中按顺序出现,这就需要对数据进行预排序,以便正确地对所有键进行分组。

这里有一个可能的实现。

In [9]:

from itertools import groupby
from operator import itemgetter

def itertools_groupby(keys, vals):
    get_key, get_val = itemgetter(0), itemgetter(1)
    sorted_pairs = sorted(zip(keys, vals), key=get_key)
    return {key: sum(map(get_val, values))
            for key, values in groupby(sorted_pairs, get_key)}

itertools_groupby(keys, vals)

Out[9]:

{'A': 5, 'B': 7, 'C': 9}

让我们来看看这样的比较。

在[10]中:

plot_benchmarks([pandas_groupby, dict_groupby, itertools_groupby]);

我们看到itertools的方法与基于dict的方法有相似的扩展性,但由于需要对数据进行多次处理,其中一次是相对昂贵的排序操作,所以不出意料地要慢一些。

基于NumPy的方法

当处理大数组时,加速的一个方法是使用NumPy的元素明智操作,将循环从解释的Python层推到编译的C层。例如,Numpy中的选择通常是通过掩码操作完成的,这可以与Python的dict理解相结合,从而实现一个很好的简洁的group-by。

In [11]:

import numpy as np

def masking_groupby(keys, vals):
    keys = np.asarray(keys)
    vals = np.asarray(vals)
    return {key: vals[keys == key].sum()
            for key in np.unique(keys)}
    
masking_groupby(keys, vals)

Out[11]:

{'A': 5, 'B': 7, 'C': 9}

每个屏蔽操作都是对数据的完整传递,并构建一个相关的布尔数组。这些传递发生在NumPy的编译层中,因此在实践中,最终比我们之前看到的纯Python循环要快。

在[12]中:

plot_benchmarks([pandas_groupby, dict_groupby,
                 itertools_groupby, masking_groupby]);

有趣的是,当组的数量较少时,这种方法甚至可以击败Pandas中的专门代码。然而,当组的数量很大时,每组一次的要求否定了这种好处。除非Numpy能够一次性识别所有组的数据,否则它将无法与Pandas的扩展性相媲美。

单次识别所有组实际上可以通过低级例程np.bincount ,它对整数数组中的数值出现的次数进行加权计算。为了使用它,我们必须首先将所有的值映射成唯一的整数,这可以通过np.unique 函数来完成。其结果是一个简洁而高效的两遍逐组实现。

In [13]:

import numpy as np

def bincount_groupby(keys, vals):
    unique_keys, group = np.unique(keys, return_inverse=True)
    counts = np.bincount(group, weights=vals)
    return dict(zip(unique_keys, counts))

bincount_groupby(keys, vals)

Out[13]:

{'A': 5.0, 'B': 7.0, 'C': 9.0}

让我们把这个加入到基准图中。

在[14]中:

plot_benchmarks([pandas_groupby, dict_groupby, itertools_groupby,
                 masking_groupby, bincount_groupby]);

与我们的基于掩码的方法相比,该结果与数据集的大小有相似的扩展性,同时避免了大量分组的不良行为。

虽然有趣的是,基于NumPy中的通用构件的实现可以接近Pandas中的目的优化代码的性能,但仍然存在一个问题,即bincount方法不是特别灵活:例如,minmax 聚合不能被实现为加权和。

稀疏矩阵方法

对于一个仍然基于标准堆栈中可用的通用构件的更灵活的方法,我将考虑一些需要有点发散性思维的东西。根据我的经验,PyData堆栈中一些最有用的底层例程可以在SciPy稀疏矩阵模块中找到:虽然它们是为快速操作稀疏矩阵而设计的,但它们往往可以被重新利用,用于解决概念上不同的问题。

在这种情况下,我们可以利用这些工具,构建一个稀疏矩阵,使每个组占据矩阵的一行,然后对该矩阵进行标准聚合,找到我们想要的结果。我们将再次利用np.unique ,为每个值确定适当的行/组,并以此来构建矩阵。

In[15]:

from scipy import sparse

def sparse_groupby(keys, vals):
    unique_keys, row = np.unique(keys, return_inverse=True)
    col = np.arange(len(keys))
    mat = sparse.coo_matrix((vals, (row, col)))
    return dict(zip(unique_keys, mat.sum(1).flat))

sparse_groupby(keys, vals)

Out[15]:

{'A': 5, 'B': 7, 'C': 9}

让我们把这个加入到我们的基准中。

在[16]中:

plot_benchmarks([pandas_groupby, dict_groupby, itertools_groupby,
                 masking_groupby, bincount_groupby, sparse_groupby]);

即使是较大的数组,这种稀疏方法也出乎意料地接近(在几个系数之内)Pandas中专门建立的group-by实现,并且还为稀疏矩阵提供了广泛的有效聚合选项。在这个特定的案例中,这种重新利用稀疏矩阵内部结构的做法可能看起来有点轻率,但在我自己的研究工作中,我经常发现类似的方法在高效实现算法方面非常有用,否则可能需要写一个自定义的C语言扩展。