对超大型文件的优化处理

Question

我的任务相对简单:对于输入文件中的每一行，测试该行是否满足给定的一组条件，如果满足，则将该行的特定列写入一个新文件。我已经写了一个python脚本来做这件事，但是我希望在以下方面得到一些帮助: 1)提高速度，2)根据列名工作的最佳方式(因为列号因文件而异)，以及3)指定过滤条件和期望输出列的最佳方式。

1)我使用的文件包含天文图像的光度测量。每个文件大约有1e6行乘以150列的浮点数，大小通常超过1 1GB。我有一个旧的AWK脚本，可以在大约1分钟内处理这样的文件；我的python脚本需要5到7分钟。我经常需要调整过滤条件并重新运行几次，直到输出文件是我想要的，所以速度绝对是我想要的。我发现for循环相当快；是我在循环中做事情的方式减慢了它的速度。使用itemgetter挑选出我想要的列比将整行读入内存有很大的改进，但我不确定我能做些什么来进一步提高速度。它能像AWK一样快吗？

2)我希望使用列名而不是列号，因为特定数量的列号(光子计数、背景、信噪比等)可以在不同文件之间改变。在我的AWK脚本中，我总是需要检查指定了条件和输出列的列号是否正确，即使过滤和输出应用于相同的数量。我在python中的解决方案是创建一个字典，为每个数量分配一个列号。当一个文件有不同的列时，我只需要指定一个新的字典。也许有更好的方法来做到这一点呢？

3)理想情况下，我只需要指定输入和输出文件的名称、过滤条件和要输出的列，它们将在我的脚本顶部找到，这样我就不需要通过搜索代码来调整一些东西。我的主要问题是未定义的变量。例如，一个典型的条件是'SNR > 4'，但是直到开始从光度文件中读取行，'SNR‘(信噪比)才被实际赋值。我的解决方案是使用字符串和eval/exec的组合。再说一次，也许有更好的方法？

我根本没有接受过计算机科学方面的培训(我是一个天文学研究生)--我通常只是把一些东西拼凑在一起，然后进行调试，直到它正常工作。然而，关于以上三点的优化对我的研究来说已经变得非常重要。我为这篇冗长的帖子道歉，但我觉得细节会有所帮助。你给我的任何和所有的建议，除了只是清理东西/编码风格，都将非常感谢。

非常感谢，杰克

#! /usr/bin/env python2.6

from operator import itemgetter


infile = 'ugc4305_1.phot'
outfile = 'ugc4305_1_filt.phot'

# names must belong to dicitonary
conditions = 'OBJ <= 2 and SNR1 > 4 and SNR2 > 4 and FLAG1 < 8 and FLAG2 < 8 and (SHARP1 + SHARP2)**2 < 0.075 and (CROWD1 + CROWD2) < 0.1'

input = 'OBJ, SNR1, SNR2, FLAG1, FLAG2, SHARP1, SHARP2, CROWD1, CROWD2'
    # should contain all quantities used in conditions

output = 'X, Y, OBJ, COUNTS1, BG1, ACS1, ERR1, CHI1, SNR1, SHARP1, ROUND1, CROWD1, FLAG1, COUNTS2, BG2, ACS2, ERR2, CHI2, SNR2, SHARP2, ROUND2, CROWD2, FLAG2'

# dictionary of col. numbers for the more important qunatities
columns = dict(EXT=0, CHIP=1, X=2, Y=3, CHI_GL=4, SNR_GL=5, SHARP_GL=6, ROUND_GL=7, MAJAX_GL=8, CROWD_GL=9, OBJ=10, COUNTS1=11, BG1=12, ACS1=13, STD1=14, ERR1=15, CHI1=16, SNR1=17, SHARP1=18, ROUND1=19, CROWD1=20, FWHM1=21, ELLIP1=22, PSFA1=23, PSFB1=24, PSFC1=25, FLAG1=26, COUNTS2=27, BG2=28, ACS2=29, STD2=30, ERR2=31, CHI2=32, SNR2=33, SHARP2=34, ROUND2=35, CROWD2=36, FWHM2=37, ELLIP2=38, PSFA2=39, PSFB2=40, PSFC2=41, FLAG2=42)



f = open(infile)
g = open(outfile, 'w')


# make string that extracts values for testing
input_items = []
for i in input.replace(',', ' ').split():
    input_items.append(columns[i])
input_items = ', '.join(str(i) for i in input_items)

var_assign = '%s = [eval(i) for i in itemgetter(%s)(line.split())]' % (input, input_items) 


# make string that specifies values for writing
output_items = []
for i in output.replace(',', ' ').split():
    output_items.append(columns[i])
output_items = ', '.join(str(i) for i in output_items)

output_values = 'itemgetter(%s)(line.split())' % output_items


# make string that specifies format for writing
string_format = []
for i in output.replace(',', ' ').split():
    string_format.append('%s')
string_format = ' '.join(string_format)+'\n'


# main loop
for line in f:
   exec(var_assign)
   if eval(conditions):
      g.write(string_format % tuple(eval(output_values)))
f.close()
g.close()

Answer 1

我想你没有提到它，但看起来你的数据是csv格式的。您可能会从使用csv.DictReader中获益良多。您可以一次迭代一行文件(避免将整个文件加载到内存中)，并通过它们的名称引用列。

如果您还没有看过cProfile，那么也应该看看Python的分析器。它会告诉你程序的哪些部分花费了最多的执行时间。

Answer 2

在这里，我的第一步是去掉exec()和eval()调用。每次对字符串求值时，都必须对其进行编译，然后执行，这增加了在文件的每一行上调用函数的开销。更不用说，eval往往会导致混乱、难以调试的代码，通常应该避免。

你可以通过将你的逻辑放到一个小的，易于理解的函数中来开始重构。例如，您可以将eval(conditions)替换为函数，例如：

def conditions(d):
    return (d[OBJ] <= 2 and
            d[SNRI] > 4 and
            d[SNR2] > 4 and
            d[FLAG1] < 8 and ...

提示:如果您的一些条件语句失败的概率较高，请将它们放在第一位，python将跳过其余条件语句的计算。

我会去掉列名字典，只需在文件的顶部设置一组变量，然后通过line[COLNAME]引用列。这可以帮助您简化某些部分，如conditions函数，并且您可以通过名称引用列，而不必为每个变量赋值。

Answer 3

这就是我如何处理这样的事情...

在我的机器上运行您的原始脚本只需大约35秒，而您的原始脚本只需要大约3分钟。可以添加更多的优化(例如，我们只需要将一些列转换为浮点数)，但这只会减少几秒钟的运行时间。

你也可以很容易地在这里使用csv.DictReader，正如一些人所建议的那样。我避免使用它，因为你必须定义一个自定义的方言，而且只需要额外的几行代码就可以做同样的事情。(各种csv模块类还会检查更复杂的行为(例如，带引号的字符串等)，在这种情况下，您无需担心这些行为。在许多情况下，它们非常非常方便，但在这种情况下，它们有点过头了。)

注意，当你调用脚本时，你也可以很容易地添加你的infile和outfile名称作为参数，而不是硬编码它们(例如infile = sys.argv[0]等)。这也将允许您通过管道轻松地输入或输出数据……(您可以检查sys.argv的长度，并相应地将infile或outfile设置为sys.stdin和/或sys.stdout )

def main():
    infile = 'ugc4305_1.phot'
    outfile = 'ugc4305_1_filt.phot'
    process_data(infile, outfile)

def filter_conditions(row):
    for key, value in row.iteritems():
        row[key] = float(value)

    cond = (row['OBJ'] <= 2 and row['SNR1'] > 4 
       and row['SNR2'] > 4 and row['FLAG1'] < 8 
       and row['FLAG2'] < 8 
       and (row['SHARP1'] + row['SHARP2'])**2 < 0.075 
       and (row['CROWD1'] + row['CROWD2']) < 0.1
       )
    return cond

def format_output(row):
    output_columns = ('X', 'Y', 'OBJ', 'COUNTS1', 'BG1', 'ACS1', 'ERR1', 'CHI1', 
                     'SNR1', 'SHARP1', 'ROUND1', 'CROWD1', 'FLAG1', 'COUNTS2', 
                     'BG2', 'ACS2', 'ERR2', 'CHI2', 'SNR2', 'SHARP2', 'ROUND2', 
                     'CROWD2', 'FLAG2')
    delimiter = '\t'
    return delimiter.join((row[name] for name in output_columns))

def process_data(infilename, outfilename):
    column_names = ('EXT', 'CHIP', 'X', 'Y', 'CHI_GL', 'SNR_GL', 'SHARP_GL', 
                    'ROUND_GL', 'MAJAX_GL', 'CROWD_GL', 'OBJ', 'COUNTS1', 
                    'BG1', 'ACS1', 'STD1', 'ERR1', 'CHI1', 'SNR1', 'SHARP1', 
                    'ROUND1', 'CROWD1', 'FWHM1', 'ELLIP1', 'PSFA1', 'PSFB1', 
                    'PSFC1', 'FLAG1', 'COUNTS2', 'BG2', 'ACS2', 'STD2', 
                    'ERR2', 'CHI2', 'SNR2', 'SHARP2', 'ROUND2', 'CROWD2', 
                    'FWHM2', 'ELLIP2', 'PSFA2', 'PSFB2', 'PSFC2', 'FLAG2')

    with open(infilename) as infile:
        with open(outfilename, 'w') as outfile:
            for line in infile:
                line = line.strip().split()
                row = dict(zip(column_names, line))
                if filter_conditions(row.copy()):
                    outfile.write(format_output(row) + '\n')

if __name__ == '__main__':
    main()