为什么Pandas很快：groupby

2019-12-12 11:06:46

为什么Pandas很快：groupby

In [1]:

%load_extcythonmagic

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]:

%timeitvs.groupby(ks, sort=False).sum()

10 loops, best of 3: 53 ms per loop

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

In [5]:

fromcollectionsimportdefaultdict

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]:

%timeitgroupby_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]:

%timeitgroupby_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)foriinxrange(N)]

ks2 = pd.Series(keys2)

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

In [11]:

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

1 loops, best of 3: 231 ms per loop

对应的Python代码如下：

In [12]:

defgroupby2_python(keys1, keys2, values):

d = defaultdict(float)

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

d[k1, k2] += v

return d

In [13]:

%timeitgroupby2_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]:

printks[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]:

%timeitnp.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字典，而是使用了更为高效的klibC语言库，并切操作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]: