pandas 性能问题 - 需要优化帮助

Question

我写了一些使用 pandas 库的 Python 代码。代码运行得有点慢，所以我用 cProfile 工具检查了一下，想看看哪里是瓶颈。

根据 cProfile 的结果，其中一个瓶颈是调用了 pandas.lib_scalar_compare：

1604  262.301    0.164  262.301    0.164 {pandas.lib.scalar_compare}

我想问的是，这个函数在什么情况下会被调用？我猜是当我选择 DataFrame 的一部分时。以下是我的代码：

if (var=='9999'):
    dataTable=resultTable.ix[(resultTable['col1'] == var1)  
                                             & (resultTable['col2']==var2)].copy() 
else:
    dataTable=resultTable.ix[(resultTable['col1'] == var1)  
                                           & (resultTable['col2']==var2)
                                           & (resultTable['col3']==int(val3))].copy() 

我有以下几个问题：

1. 这个代码片段是否最终调用了导致瓶颈的代码？
2. 如果是的话，有没有办法优化这个问题？我现在使用的 pandas 版本是 pandas-0.8。

任何帮助都将非常感激。

Answer 1

你可以使用列 col1 到 col3 来设置索引。这里有一个简单的例子：

In [1]: df = DataFrame(np.arange(20).reshape(5,4))

In [2]: df
Out[2]:
    0   1   2   3
0   0   1   2   3
1   4   5   6   7
2   8   9  10  11
3  12  13  14  15
4  16  17  18  19

In [3]: df2 = df.set_index(keys=[0,1,2])

In [4]: df2
Out[4]:
           3
0  1  2
0  1  2    3
4  5  6    7
8  9  10  11
12 13 14  15
16 17 18  19

多重索引元组：

In [5]: %timeit df2.ix[(4,5,6)]
10000 loops, best of 3: 99.5 us per loop

原始数据框：

In [6]: %timeit df.ix[(df[0]==4) & (df[1]==5) & (df[2]==6)][3]
1000 loops, best of 3: 515 us per loop

更新： 处理重复的索引

In [1]: df = DataFrame(np.arange(20).reshape(5,4))

In [2]: df = concat([df, df])

In [3]: df
Out[3]:
    0   1   2   3
0   0   1   2   3
1   4   5   6   7
2   8   9  10  11
3  12  13  14  15
4  16  17  18  19
0   0   1   2   3
1   4   5   6   7
2   8   9  10  11
3  12  13  14  15
4  16  17  18  19

这个方法失败了：

In [4]: df2 = df.set_index(keys=[0,1,2])

In [5]: df2.ix[(0,1,2)]

KeyError: u'no item named 1'

这个方法成功了：

In [6]: df2 = df.set_index(keys=[0,1,2]).sort()

In [7]: df2.ix[(0,1,2)]
Out[7]:
       3
0 1 2
0 1 2  3
    2  3

Answer 2

我的代码在处理 pandas.lib.scalar_compare 这个部分花了很多时间，通过把字符串类型的列转换成 'category' 类型，我的速度提高了10倍。

举个例子：

   $ df['ResourceName'] = df['ResourceName'].astype('category')

想了解更多信息，可以查看 https://pandas.pydata.org/pandas-docs/stable/user_guide/categorical.html