为什么Pandas很快：groupby

In [1]:

%load_ext cythonmagic 

import numpy as np 

import pandas as pd 

import random 

一个例子

让我们先通过一个例子演示Pandas的速度。首先创建100万个浮点数values，并创建由200个不同值的100万个字符串keys，并分别将它们转换成Series对象：vs和ks。

In [2]:

N = 1000000 

uniques_keys = [pd.core.common.rands(3) for i in xrange(200)] 

keys = [random.choice(uniques_keys) for i in xrange(N)] 

values = np.random.rand(N).tolist() 

vs = pd.Series(values) 

ks = pd.Series(keys) 

下面的程序对ks进行分类，并对vs进行求和：

In [3]:

vs.groupby(ks, sort=False).sum() 

Out[3]:

ZDI    2567.780079

yYV    2452.373746

bgi    2459.354122

IP7    2464.018642

Bc9    2542.754923

hJb    2538.633765

NOk    2521.267062

K4A    2505.704467

weQ    2472.289371

ikS    2480.189153

paS    2519.596946

wzg    2438.843985

Zco    2549.183511

51x    2581.934765

46g    2446.169356

...

LQI    2483.944911

2im    2604.985061

HPZ    2486.839048

0dv    2497.814490

EXG    2484.701354

Vdb    2506.095656

4sR    2565.334928

E7M    2430.046894

Mz2    2514.600613

wN5    2474.341671

VJT    2495.642607

yo8    2495.442933

2ru    2453.049701

c2a    2465.111358

nph    2505.492098

Length: 200

让我们看看它的运算速度：

In [4]:

%timeit vs.groupby(ks, sort=False).sum() 

10 loops, best of 3: 53 ms per loop

如果用Python标准库来实现这个功能的话，可以使用defaultdict。下面的程序对列表keys和values进行迭代，因为Python列表的存取速度比Series要快很多。

In [5]:

from collections import defaultdict 

from itertools import izip 

def groupby_python(keys, values): 

    d = defaultdict(float) 

    for k, v in izip(keys, values): 

        d[k] += v 

    return d 

In [6]:

%timeit groupby_python(keys, values) 

1 loops, best of 3: 183 ms per loop

Pandas的Series.groupby比用defaultdict实现的要快接近4倍。下面检查二者的结果是否相同，Series.groupby返回的是一个Series对象，调用to_dict()可以将其转换成字典：

In [7]:

r1 = vs.groupby(ks, sort=False).sum() 

r2 = groupby_python(keys, values) 

r1.to_dict() == r2 

Out[7]:

True

下面让我们用Cython实现相同的逻辑，Cython编译器能将list和dict的操作转换成相应的C API调用，因此在Cython中使用dict会更快一些。程序中尽可能增加类型信息，以实现最快速度：

In [88]:

%%cython -c=-O3 

cimport cython 

@cython.boundscheck(False) 

@cython.wraparound(False) 

def groupby_sum(list keys not None, list values not None): 

    cdef dict d = {} 

    cdef int i 

    cdef double v 

    cdef str k 

    for i in range(len(keys)): 

        k = keys[i] 

        v = values[i] 

        d[k] = <double>d.get(k, 0.0) + v 

    return d 

In [89]:

%timeit groupby_sum(keys, values) 

10 loops, best of 3: 67.6 ms per loop

由上面的结果可知，Pandas的groupby效率甚至超过了采用类型声明的Cython代码的速度了。

下面让我们看看多层groupby的运行速度。创建第二个随机字符串列表，并将其转换成Series对象：

In [10]:

keys2 = [random.choice(uniques_keys) for i in xrange(N)] 

ks2 = pd.Series(keys2) 

下面的代码计算ks和ks2中的组合进行分类，并对vs进行求和：

In [11]:

%timeit vs.groupby([ks, ks2], sort=False).sum() 

1 loops, best of 3: 231 ms per loop

对应的Python代码如下：

In [12]:

def groupby2_python(keys1, keys2, values): 

    d = defaultdict(float) 

    for k1, k2, v in izip(keys, keys2, values): 

        d[k1, k2] += v 

    return d 

In [13]:

%timeit groupby2_python(keys, keys2, values) 

1 loops, best of 3: 505 ms per loop

下面比较二者的结果：

In [39]:

r3 = groupby2_python(keys, keys2, values) 

vs.groupby([ks, ks2], sort=False).sum().to_dict() == r3 

Out[39]:

True

factorize

Factorize是Pandas的核心功能之一，它将一个数组转换成两个数组：labels和levels。其中levels包含原数组中的所有不重复的值，labels则是一个长度和原始数组相同，值为从0到len(levels)-1的整数数组，其中的每个值表示原始数组中对应值在levels中的下标。即：ks[i]与levels[labels[i]]相等。

In [14]:

labels, levels = pd.factorize(ks) 

print labels 

print levels 

[  0   1   2 ..., 133 180 192]

['ZDI' 'yYV' 'bgi' 'IP7' 'Bc9' 'hJb' 'NOk' 'K4A' 'weQ' 'ikS' 'paS' 'wzg'

 'Zco' '51x' '46g' 'dVY' 'stZ' 'S2A' 'wV1' 'fkE' 'RK9' 'lxJ' 'mxl' 'Zkl'

 'ghV' 'v4o' 'aeN' 'XRI' '1Gm' '5kD' 'ACc' 'liA' 'IRz' 'qGh' 'fFD' 'X7K'

 'JWc' 'ayB' 'xja' 'We8' '8PE' 'bAq' 'EHx' 'yhH' 'eig' 'K2o' '93U' 'QVG'

 'lok' '5r7' 'Z48' 'MUn' 'yNm' 'SUw' 'tYe' 'XYj' 'Dgh' 'Hb5' 'h1E' 'Sea'

 '50j' 'l31' '2Il' 'iH2' '2Zv' 'xQu' 'B5n' 'uI9' 'G69' 'lTs' 'x3R' 'BTH'

 '95L' 'Vrb' 'OxZ' '3QE' 'Tmp' 'rtN' 'SkB' 'VK3' 'P9n' 'AB5' 'pFP' 'thk'

 'GnS' 'NAv' 'YEB' 'lYM' 'oCn' 'CRr' 'scd' 'Bvg' 'yva' 'eyE' 'tbX' 'VjA'

 'niB' 'pUD' '26p' 'Bf8' '2In' '9Rs' 'czj' '0N9' 'iSq' 'Fnp' 'nS1' 'auL'

 'Lc6' 'vRb' 'cq4' 'abe' 'UlG' 'QHi' '54f' 'xL1' 'NgQ' 'ew7' '3RA' 'WYe'

 'jbO' 'dkg' 'u4c' 'wIG' '3Y9' 'pRE' 'rMg' 'USA' 'ypF' '91U' '19q' 'T5J'

 'vuH' 'bM9' 'Dy5' 'plV' 'nyg' 'hy4' '0kt' 'ae3' 'FRK' 'US1' 'EqR' 'lVW'

 'PUy' 'qpF' 'Vp9' 'YKw' 'HxZ' 'tTO' 'IwD' 'Feq' 'ugs' '7lS' 'Nmf' 'v16'

 'aCY' 'FTR' '9Sp' 'OHu' 'LUh' 'x6N' '9w8' 'COQ' 'THC' 'RVJ' '1Ub' 'EUt'

 'YM4' 'M6K' 'kmI' 'Is2' 'afO' 'cOV' 'k3j' 'gCo' 'lgP' '1hy' 'YS8' 'qfl'

 'k7Y' 'ZWd' 'Xr2' 'xRf' 'OHz' 'LQI' '2im' 'HPZ' '0dv' 'EXG' 'Vdb' '4sR'

 'E7M' 'Mz2' 'wN5' 'VJT' 'yo8' '2ru' 'c2a' 'nph']

In [15]:

print ks[100], levels[labels[100]] 

lYM lYM

通过levels[labels]可得到原始数组：

In [20]:

(levels[labels] == ks).all() 

Out[20]:

True

一旦对原始的字符串数组进行factorize运算之后，我们就很容易使用labels数组实现快速的groupby运算了。在NumPy中这个运算为bincount()。它的第一个参数为labels数组，第二个数组为需要求和的数组。下面的程序使用bincount()求和之后，将其转换成Series对象，然后再转换成字典，与groupby_python()的结果比较：

In [23]:

pd.Series(np.bincount(labels, vs.values), index=levels).to_dict() == r2 

Out[23]:

True

bincount(labels, values)的运算是非常快的，因为它只需创建一个长度为len(levels)的数组result，对两个参数数组进行一次循环，计算result[labels[i]] += values[i]即可：

In [24]:

%timeit np.bincount(labels, vs.values) 

100 loops, best of 3: 3.5 ms per loop

而两层groupby操作可以通过三次factorize()运算得到。请读者思考下面的程序是如何实现两层groupby运算的。

In [61]:

labels, levels = pd.factorize(ks) 

labels2, levels2 = pd.factorize(ks2) 

groups_id = labels * len(levels2) + labels2 

labels_g, levels_g = pd.factorize(groups_id) 

index_g = zip(levels[levels_g // len(levels2)], levels2[levels_g % len(levels2)]) 

value_g = np.bincount(labels_g, vs.values) 

pd.Series(value_g, index=index_g).to_dict() == r3 

Out[61]:

True

factorize的高效实现

一旦字符串数组被factorize之后，groupby操作就由字典操作简化成了数组操作。然而factorize本身仍然是需要字典的帮助才能实现的。Pandas的factorize操作之所以能如此迅速，是因为它并没有采用Python的dict字典，而是使用了更为高效的C语言库，并切操作klib字典的程序实在Cython中实现的。

klib字典较dict有两个优势：

• klib字典可以指定初始大小，对于factorize这样的操作，我们是事先就知道最终的字典的最大长度的。而dict则会在添加元素的过程中多次重新分配和复制内存。

• klib字典中可以不保存Python的对象，只保存单纯的数值。

Pandas使用klib实现的字典都可以在pandas.hashtable模块中找到。例如下面使用StringHashTable()对ks进行factorize运算，其结果和之前的结果相同。

In [76]:

sht = pd.hashtable.StringHashTable() 

levels_sht, labels_sht = sht.factorize(ks) 

np.all(labels == labels_sht) 

Out[76]:

True

调用StringHashTable.factorize()之后，sht是一个将字符串映射到其位置的字典，例如：

In [87]:

np.array([sht.get_item(x) for x in levels]) 

Out[87]:

array([  0,   1,   2,   3,   4,   5,   6,   7,   8,   9,  10,  11,  12,

        13,  14,  15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,

        26,  27,  28,  29,  30,  31,  32,  33,  34,  35,  36,  37,  38,

        39,  40,  41,  42,  43,  44,  45,  46,  47,  48,  49,  50,  51,

        52,  53,  54,  55,  56,  57,  58,  59,  60,  61,  62,  63,  64,

        65,  66,  67,  68,  69,  70,  71,  72,  73,  74,  75,  76,  77,

        78,  79,  80,  81,  82,  83,  84,  85,  86,  87,  88,  89,  90,

        91,  92,  93,  94,  95,  96,  97,  98,  99, 100, 101, 102, 103,

       104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116,

       117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129,

       130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142,

       143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155,

       156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168,

       169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181,

       182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194,

       195, 196, 197, 198, 199])