使用Pandas筛选属于某个列表中的数据

Jan. 15th, 2018

Pandas

最近一个小项目中，需要完成以下事情。有一个csv文件，里面的每一条记录是一个计算任务的编号，内容描述，在超算系统里的process ID以及计算结果的md5映射（列的名称分别为“task_id”，“task_desc”，“process_id”和“md5”）。另外，有一个文件存放着已完成的任务的process_ID。我的任务是查找出所有完成的任务的md5映射值。

在SQL数据库里，这是一个非常标准的查询：

SELECT md5 FROM task_table WHERE process_id IN process_list;

使用Pandas我们也可以非常容易地完成这个查询。

关于更多SQL语句的Pandas对应实现，可参见这一篇文章。

构建process_list：文件读取

所有已完成任务的process_id被存在文件"completed_task.txt"中。文件的每一行是一个id。因此得到process_list列表非常容易：

process_id = open("completed_task.txt", 'r').read().split('\n')

这里假设文件中没有空行，而且换行符都是'\n'。由于csv文件中"process_id"都是以字符串存放（会有补零位）， 因此我们保留其字符串格式而不转换为整数。

Pandas Dataframe预处理和筛选

接下来我们读取tasks.csv文件并进行预处理，筛选和结果存放。

import pandas as pd
from math import nan # 用于处理空缺位。
df = pd.read_csv('tasks.csv', sep='\t') # 此文件以tab符分割数据。
df['md5'].replace(nan, "", inplace=True) # 此问题并不需要这样做，不过有字符串处理的任务就需要
md5_selected = df[df['process_id'].isin(process_id)]['md5']
md5_list = md5_selected.tolist()
open('completed_md5.txt', 'w').write('\n'.join(md5_list)) # 将筛选的数据存储到另一个文件中。

这里用到了以下知识点（敲黑板）：

1. pandas读取csv文档，设定分隔符。（作业：自己回忆如何读取Excel文档。）
2. 空白数据清洗：当一个值缺失时，Pandas默认使用math.nan代替。如果不对其预处理，过后在进行字符串操作时有可能产生错误。这里将其替换成空行。
3. Pandas查询。md5_selected = df[df['process_id'].isin(process_id)]['md5']就是Pandas实现前面提到的那条SQL查询语句的方法。