Pandas性能优化:基础篇

Pandas 号称“数据挖掘瑞士军刀”，是数据处理最常用的库。在数据挖掘或者kaggle比赛中，我们经常使用pandas进行数据提取、分析、构造特征。而如果数据量很大，操作算法复杂，那么pandas的运行速度可能非常慢。本文根据实际工作中的经验，总结了一些pandas的使用技巧，帮助提高运行速度或减少内存占用。

1 按行迭代优化

很多时候，我们会按行对dataframe进行迭代，一般我们会用iterrows这个函数。在新版的pandas中，提供了一个更快的itertuples函数。

我们测试一下速度:

import pandas as pd
import numpy as np
import time
df = pd.DataFrame({'a': np.random.randn(1000),
                     'b': np.random.randn(1000),
                    'N': np.random.randint(100, 1000, (1000)),
                   'x':  np.random.randint(1, 10, (1000))})

%%timeit
a2=[]
for index,row in df.iterrows():
    temp=row['a']
    a2.append(temp*temp)
df['a2']=a2

67.6 ms ± 3.69 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%%timeit
a2=[]
for row in df.itertuples():
    temp=getattr(row, 'a')
    a2.append(temp*temp)
df['a2']=a2

1.54 ms ± 168 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

可以看到直接循环，itertuples速度是iterrows的很多倍。

所以如果在必须要对dataframe进行遍历的话，直接用itertuples替换iterrows。

2 apply 优化

一般情况下，如果要对dataframe里的数据逐行处理，而不需要上下文信息，可以使用apply函数。
对于上面的例子，我们使用apply看下：

%%timeit
df['a2']=df.a.apply(lambda x: x*x)

360 µs ± 355 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

可以看到，效率又有提升，apply的速度是itertuples的5倍左右。

注意一下，apply不光能对单个列值做处理，也能对多个列的值做处理。

%%timeit
df['a3']=df.apply( lambda row: row['a']*row['b'],axis=1)

15 ms ± 1.61 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

这里看出来，多值处理的时候几乎等于iterrows。因此比单列值apply慢了许多，所以这里不推荐对整行进行apply。

我们可以简单的这样改写：

%%timeit
df['a3']=df['a']*df['b']

204 µs ± 8.31 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

如果在计算的时候，使用series.values 提取numpy 数组并使用numpy原生函数计算，效率可能更高

%%timeit
df['a3']=np.multiply(df['a'].values,df['b'].values)

93.3 µs ± 1.45 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

3 聚合agg效率优化

有的时候，我们会对一列数据按值进行分组（groupby），再分组后，依次对每一组数据中的其他列进行聚合（agg）。

还是上面的那个dataframe，我们看下：
采用自定义函数的agg函数：

%%timeit
df.groupby("x")['a'].agg(lambda x:x.sum())

1.27 ms ± 45.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

采用agg内置函数：

%%timeit
df.groupby("x")['a'].agg(sum)
#等价df.groupby("x")['a'].sum()

415 µs ± 20.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

刚才是单列聚合，我们看下多列聚合：
自定义函数：

%%timeit
df.groupby("x").agg({"a":lambda x:x.sum(),"b":lambda x:x.count()})

2.6 ms ± 8.17 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

采用内置函数：

%%timeit
df.groupby("x").agg({"a":"sum","b":"count"})

1.33 ms ± 29.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

可以看出，对于多列聚合，内置函数仍然比自定义函数快1倍。所以再进行聚合操作时，尽量使用内置函数提高效率。

下面列出了一些内置函数：

当我们需要统计的方法没有对于内置函数的情况下，在自定义函数的时候，优先选用pandas或numpy内置的其他高效函数，如

df.groupby("x")['a'].agg(lambda x:np.sum(np.abs(x)))

在数据量很大的时候要比非numpy函数效率高（数据量小的时候差不多）。

4 数据读取优化

如果我们需要用pandas读取较多次文件，或者读取的文件较大，那么这方面占用的时间也会较长，我们也需要对其进行优化。pandas最常见读取函数是read_csv函数，可以对csv文件进行读取。但是pandas read_csv对读取较大的、数据结构复杂的文件，效率不是很理想。

这里有几种方法可以优化：

1. 分块读取csv文件
如果文件过大，内存不够或者时间上等不及，可以分块进行读取。

chunksize = 10 ** 6for chunk in pd.read_csv(filename,  chunksize=chunksize):
     process(chunk)

这里也可以使用多进程方式，具体我们在后面进阶篇介绍。

1. 过滤掉不需要的列

如果读取的文件列很多，可以使用usecols字段，只load需要的列，提高效率，节约内存。

df = pd.read_csv(filename,  usecols=["a","b","c"]):

1. 为列指定类型
   panda在read_csv的时候，会自动匹配列的数值类型，这样会导致速度很慢，并且占用内存较大。我们可以为每个列指定类型。

df = pd.read_csv("f500.csv", dtype = {"revenues" : "float64"，"name":str})

1. 保存为其他格式
   如果需要频繁的读取和写入，则可以将文件保存为其他格式的，如pickle或hdf。pickle和hdf读取速度是csv的数倍。这里注意一下，pickle比原csv略小，hdf比原csv略大。

import pandas as pd
 #读取csv
 df = pd.read_csv('filename.csv')
 
 #pickle格式
 df.to_pickle('filename.pkl') 
 df = pd.read_pickle('filename.pkl') 
 
 #hdf格式
 df.to_hdf('filename.hdf','df') 
 df = pd.read_hdf('filename.hdf','df')

1. 用第三方的包读取
   如可以使用Dask DataFrame读取大文件。第三方包放到最后细讲。

5 优化数据处理逻辑

这点算是业务角度优化。如果我们能直接数据处理的步骤，那么处理时间就少了很多。

这需要具体问题具体分析，举个例子，假设我们有一个通讯记录数据集：

call_seconds 是拨打的时长，单位是秒。
call_area=1 是拨打国内电话，费率是0.1/min，call_area=0 是拨打国外电话，费率是0.7/min
call_area=2 是接听电话，不收费。

我们随机生成一批数据：

import pandas as pd
import numpy as np
import time
import math
row_number=10000
df = pd.DataFrame({'call_area':  np.random.randint(0, 3, (row_number)),
                   'call_seconds':  np.random.randint(0, 10000, (row_number))})

如果我们想计算每个电话的费用：

def get_cost(call_area,call_seconds):
    if call_area==1:
        rate=0.1
    elif call_area==0:
        rate=0.7
    else:
        return 0
    return math.ceil(call_seconds/60)*rate

方法1，采用按行迭代循环

%%timeit
cost_list = []
for i,row in df.iterrows():
    call_area=row['call_area']
    call_seconds=row['call_seconds']
    cost_list.append(get_cost(call_area,call_seconds))
df['cost'] = cost_list

方法2，采用apply行的方式

%%timeit
df['cost']=df.apply(
         lambda row: get_cost(
             row['call_area'],
             row['call_seconds']),
         axis=1)

方法3，采用mask+loc，分组计算

%%timeit
mask1=df.call_area==1
mask2=df.call_area==0
mask3=df.call_area==2
df['call_mins']=np.ceil(df['call_seconds']/60)
df.loc[mask1,'cost']=df.loc[mask1,'call_mins']*0.1
df.loc[mask2,'cost']=df.loc[mask2,'call_mins']*0.7
df.loc[mask3,'cost']=0

方法4，使用numpy的方式，可以使用index的方式找到对应的费用。

%%timeit
prices = np.array([0.7, 0.1, 0])
mins=np.ceil(df['call_seconds'].values/60)
df['cost']=prices[df.call_area]*mins

测试结果：

方法4速度快是因为它采用了numpy向量化的数据处理方式。

总结

在优化尽量使用内置原生函数，写代码尽量避免循环，尽量写能够向量化计算的代码，最后别忘了按照自己业务需求进行算法优化。