DataFrame通过修改索引提升查找速度

下面以MovieLens-100K(可以在网上自己下载为例)来展示不同语法的查询效果。

import os
import pandas as pd
import time
import random

data=[]
with open(os.path.join('ml_data','u.data'),'r') as r_data:
    for line in r_data:
        data.append(line.split())
data = pd.DataFrame(data,columns=['User_ID','Item_ID','Rating','Timestamp'])
#要求取出每个user_id对每个item_id的timestamp字段
#现将user_id和item_id保存成tuple对
#原始数据共有100000行，这里只随机从中抽取10000个tuple对
user_item=[(user,item) for user,item in data[['User_ID','Item_ID']].values]
random.shuffle(user_item)
#第一种写法
start=time.time()
for i in user_item[:5000]:
    date1=data[(data['User_ID']==i[0])&(data['Item_ID']==i[1])]['Timestamp']
end=time.time()
print("Method 1-total time: ",end-start)
print("Method 2-average time: ",(end-start)/5000)
#第二种写法
user_item_data=data.set_index(['User_ID','Item_ID'])
start=time.time()
for i in user_item:
    date1=user_item_data.loc[i]['Timestamp']
end=time.time()
print("Method 2-total time: ",end-start)
print("Method 2-average time: ",(end-start)/5000)

其代码运行结果如下：

 

方法1和方法2的区别在于查找方法的不同，方法1没查找一次需要把所有数据都扫描一遍。而方法2通过将要筛选的字段做成索引而大大提升了查询速度。在大批量数据下，这种方法非常有效。