python大规模数据处理：pandas DataFrame代替字典dict实现快速查询，避免循环和长耗时

结论

先上结论：对于规模比较大的数据，不要创建以tuple等为键的字典，可以创建pandas DataFrame，利用DataFrame的多层Index来实现查询。这样做的原因有二：

• 创建大规模的字典时，尤其键是tuple这类数据是，非常慢。但是创建同规模的DataFrame时，效率很高。
• 字典一次只可以查询一个键，避免不了循环。DataFrame可以返回一组键对应的值，避免循环，耗时大大减少。

下面具体阐述一下我是在解决什么问题时运用了这个技巧以及怎么使用。

问题

先简化一下我要处理的问题。目前有四百万的有向点对，它们之间是不重复的。点对中第一个点和第二个点的id分别放在edge_index0和edge_index1的array中。每个点对都有一个值，放在edge_attr的array中。所以这三个array的长度是相同的，都是四百万。
现在给你三百万的有向点对，只知道这三百万点对在四百万点对中，如何获取这三百万点对的值？
第一反应是用点对作为键创建如下一个字典

{(edge_index0[i], edge_index1[i]):attr for i, attr in enumerate(edge_attr)}

数据量小的时候，确实是没有问题的。但是面对百万级数据，我这么创建了，就2000 years later。。。

解决方案

后来就创建了一个DataFrame，并将点对设为Index。非常快，顿时觉得之前建字典的我太傻了。另外在获取三百万点对值的时候，也千万不要忘记DataFrame是可以查询一个数组的，可别傻乎乎写一个循环。下面我用小数据说明一下。
先创建一下数据

import numpy as np
import pandas as pd

edge_index0 = np.arange(10)
edge_index1 = np.arange(10,0,-1)
edge_attr = np.random.randn(10)
edge_index0, edge_index1, edge_attr
>>> (array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
 array([10,  9,  8,  7,  6,  5,  4,  3,  2,  1]),
 array([ 0.3190391 , -0.24937038,  1.46210794, -2.06014071, -0.3224172 ,
        -0.38405435,  1.13376944, -1.09989127, -0.17242821, -0.87785842]))

接下来建DataFrame并将edge_index0和edge_index1设为index。

df = pd.DataFrame({'edge_index0':edge_index0, 'edge_index1':edge_index1, 'edge_attr':edge_attr}).set_index(['edge_index0', 'edge_index1'])

在给定需要查询的点对

edge_index0 = np.arange(5)
edge_index1 = np.arange(10, 5, -1)

查询

df.loc[list(zip(edge_index0, edge_index1))]

Done!