在Python中快速从文件中grep多个值的方法

Question

• 我有一个包含3亿行的文件（inputFile），每行有两列，用制表符分隔。
• 我还有一个包含1000个独特项目的列表（vals）。

我想创建一个字典，字典的键是第一列的值，值是第二列的值，条件是第一列的值必须在vals中出现。vals中有一些项目在文件中没有出现，这些值需要保存到一个新的列表里。我可以使用最多20个线程来加快这个过程。

有什么最快的方法可以做到这一点呢？

到目前为止，我的最佳尝试是：

newDict = {}
foundVals = []
cmd = "grep \"" + vals[0]
for val in vals:
     cmd = cmd + "\|^"+val+"[[:space:]]"
cmd = cmd + "\" " + self.inputFile
p = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
for line in iter(p.stdout.readline, ''):
    info = line.split()
    foundVals.append(info[0])
    newDict.update({info[0]:info[1]})
p.wait()
notFound = [x for x in vals if x not in set(foundVals)]

示例 inputFile:

2       9913
3       9913
4       9646
...
594592886       32630
594592888       32630
594592890       32630

vals:

[1,2,594592888]

想要的字典:

{2:9913,594592888:32630}

在notFound中:

[1]

Answer 1

如果我理解得没错的话，你是想要过滤掉那些不符合你指定的vals的文件行。因为你提到的是很大的文件，而你想要的值又比较少，我建议你可以试试下面的方法：

vals_set = set(vals)
found_vals = {}

with open(inputfile,"r") as in_file:
    for line in in_file:
        line = line.split() # Assuming tabs or whitespaces
        if line[0] in vals_set:
            found_vals[line[0]] = line[1]


not_found_vals = vals_set.difference(found_vals)

这样做会比较节省内存，你会在found_vals里得到找到的值，而在not_found_vals里得到没找到的值。实际上，内存的使用量只会跟你想要查找的vals的数量有关，而和文件的大小无关。

编辑：

我觉得最简单的并行处理这个任务的方法就是把文件分成几块，然后用不同的进程分别去查找每一块。这样你就不用担心线程之间的通信问题（我觉得这样更简单也更快）。

如果你在用BASH（因为你用了grep :P），一个不错的方法就是参考这个回答里提到的：

split -l 1000000 filename

这样会生成每个文件有1000000行。

你可以很容易地修改你的脚本，把找到的结果保存到每个进程的新文件里，然后再把不同的输出文件合并起来。

Answer 2

这样做在内存使用上不是特别高效（对于一个有3亿行的文件，这可能会成为一个问题）。我想不出有什么办法可以在列表推导式中保存那些找不到的值，除了保存所有的值（或者读取文件两次）。我觉得使用线程也没什么大用，因为文件的输入输出可能会成为性能的瓶颈。

我假设文件中的分隔符是制表符（你没有说明，但示例数据看起来是用制表符分隔的）。

vals = [1,2,594592888]

with open(self.inputfile,'r') as i_file:
    all_vals = {
        int(t[0]):int(t[1])
        for t in (
                line.strip().split('\t')
                for line in i_file
        )
    }

newDict = {
    t[0]:t[1] for t in filter(lambda t: t[0] in vals, all_vals.items())
}

notFound = list(set(all_vals.keys()).difference(newDict.keys()))

Answer 3

你在评论中说明了每个键在数据中最多只出现一次。根据这一点，以及只有1000个键的事实，可以得出在Python中所做的工作量是微不足道的；你几乎所有的时间都在等待grep的输出。这没问题；你把行提取的工作交给一个专门的工具的策略是合理的。但这也意味着性能提升必须在行提取的那一部分找到。

你可以通过优化你的正则表达式来加快速度。例如，代替

^266[[:space:]]\|^801[[:space:]]\|^810[[:space:]]

你可以使用：

^\(266\|801\|810\)[[:space:]]

这样就不需要为每个选项单独匹配锚点。我在测试数据（1000万行，25个键）上看到这个改动大约提升了15%的速度。

进一步的优化是将常见的前缀统一在选择中：266\|801\|810可以替换为等效的266\|8\(01\|10\)。以这种方式重写25个键的正则表达式在测试数据上可以接近50%的速度提升。

到这个时候，grep开始显示出它的局限性。看起来它是受CPU限制的：iostat显示每次改进正则表达式时，grep运行时每秒的IO请求数量在增加。而且在页面缓存已经预热的情况下重新运行grep并使用--mmap选项并没有加快速度（如果文件IO是瓶颈的话，它可能会加速）。因此，要想更快，可能需要一个具有更快正则引擎的工具。

其中一个就是ag（源代码可以在这里找到），它的正则实现还会自动优化，所以你不需要做太多手动调整。虽然我在我的机器上无法让grep处理测试数据少于大约12秒，但ag在处理上述所有正则变体时只需大约0.5秒。