泄漏 memory 解析 TSV 並將 CSV 寫入 Python

Question

我正在 Python 中編寫一個簡單的腳本作為學習練習。 我有一個從俄亥俄州選舉委員會下載的 TSV 文件，我想處理一些數據並寫出一個 CSV 文件以導入另一個系統。

我的問題是它像篩子一樣漏水 memory。 在我停止它之前，單次運行 154MB TSV 文件會消耗 2GB 的 memory。

代碼如下，有人可以幫我確定 Python 缺少什么嗎？

import csv
import datetime
import re

def formatAddress(row):
    address = ''
    if str(row['RES_HOUSE']).strip():
        address += str(row['RES_HOUSE']).strip()
    if str(row['RES_FRAC']).strip():
        address += '-' + str(row['RES_FRAC']).strip()
    if str(row['RES STREET']).strip():
        address += ' ' + str(row['RES STREET']).strip()
    if str(row['RES_APT']).strip():
        address += ' APT ' + str(row['RES_APT']).strip()
    return address

vote_type_map = {
    'G': 'General',
    'P': 'Primary',
    'L': 'Special'
}

def formatRow(row, fieldnames):
    basic_dict = {
        'Voter ID': str(row['VOTER ID']).strip(),
        'Date Registered': str(row['REGISTERED']).strip(),
        'First Name': str(row['FIRSTNAME']).strip(),
        'Last Name': str(row['LASTNAME']).strip(),
        'Middle Initial': str(row['MIDDLE']).strip(),
        'Name Suffix': str(row['SUFFIX']).strip(),
        'Voter Status': str(row['STATUS']).strip(),
        'Current Party Affiliation': str(row['PARTY']).strip(),
        'Year Born': str(row['DATE OF BIRTH']).strip(),
        #'Voter Address': formatAddress(row),
        'Voter Address': formatAddress({'RES_HOUSE': row['RES_HOUSE'], 'RES_FRAC': row['RES_FRAC'], 'RES STREET': row['RES STREET'], 'RES_APT': row['RES_APT']}),
        'City': str(row['RES_CITY']).strip(),
        'State': str(row['RES_STATE']).strip(),
        'Zip Code': str(row['RES_ZIP']).strip(),
        'Precinct': str(row['PRECINCT']).strip(),
        'Precinct Split': str(row['PRECINCT SPLIT']).strip(),
        'State House District': str(row['HOUSE']).strip(),
        'State Senate District': str(row['SENATE']).strip(),
        'Federal Congressional District': str(row['CONGRESSIONAL']).strip(),
        'City or Village Code': str(row['CITY OR VILLAGE']).strip(),
        'Township': str(row['TOWNSHIP']).strip(),
        'School District': str(row['SCHOOL']).strip(),
        'Fire': str(row['FIRE']).strip(),
        'Police': str(row['POLICE']).strip(),
        'Park': str(row['PARK']).strip(),
        'Road': str(row['ROAD']).strip()
    }

    for field in fieldnames:
        m = re.search('(\d{2})(\d{4})-([GPL])', field)
        if m:
            vote_type = vote_type_map[m.group(3)] or 'Other'
            #print { 'k1': m.group(1), 'k2': m.group(2), 'k3': m.group(3)}
            d = datetime.date(year=int(m.group(2)), month=int(m.group(1)), day=1)
            csv_label = d.strftime('%B %Y') + ' ' + vote_type + ' Ballot Requested'
            d = None
            basic_dict[csv_label] = row[field]
        m = None

    return basic_dict

output_rows = []
output_fields = []
with open('data.tsv', 'r') as f:
    r = csv.DictReader(f, delimiter='\t')
    #f.seek(0)
    fieldnames = r.fieldnames
    for row in r:
        output_rows.append(formatRow(row, fieldnames))
f.close()

if output_rows:
    output_fields = sorted(output_rows[0].keys())
    with open('data_out.csv', 'wb') as f:
        w = csv.DictWriter(f, output_fields, quotechar='"')
        w.writeheader()
        for row in output_rows:
            w.writerow(row)
    f.close()

Answer 1

您正在將所有數據累積到一個巨大的列表output_rows 。 您需要在讀取每一行時對其進行處理，而不是將所有行都保存到內存昂貴的Python列表中。

with open('data.tsv', 'rb') as fin, with open('data_out.csv', 'wb') as fout:
    reader = csv.DictReader(fin, delimiter='\t')
    firstrow = next(r)
    fieldnames = reader.fieldnames
    basic_dict = formatRow(firstrow, fieldnames)
    output_fields = sorted(basic_dict.keys())
    writer = csv.DictWriter(fout, output_fields, quotechar='"')
    writer.writeheader()
    writer.writerow(basic_dict)
    for row in reader:
        basic_dict = formatRow(row, fieldnames)        
        writer.writerow(basic_dict)

Answer 2

您沒有泄漏任何內存，只是在使用大量內存。

您要將每一行文本轉換成Python字符串的字典，這比單個字符串占用更多的內存。 有關詳細信息，請參閱為什么我的100MB文件占用1GB內存？

解決方案是迭代地執行此操作。 實際上，您不需要整個列表，因為您永遠不會引用任何先前的值。 所以：

with open('data.tsv', 'r') as fin, open('data_out.csv', 'w') as fout:
    r = csv.DictReader(fin, delimiter='\t')
    output_fields = sorted(r.fieldnames)
    w = csv.DictWriter(fout, output_fields, quotechar='"')
    w.writeheader()
    for row in r:
        w.writerow(formatRow(row, fieldnames))

或者，更簡單地說：

    w.writerows(formatRow(row, fieldnames) for row in r)

---

當然，這與您的原始代碼稍有不同，即使輸入文件為空，它也會創建輸出文件。 如果重要的話，您可以輕松地解決該問題：

with open('data.tsv', 'r') as fin:
    r = csv.DictReader(fin, delimiter='\t')
    first_row = next(r)
    if row:
        with open('data_out.csv', 'wb') as fout:
            output_fields = sorted(r.fieldnames)
            w = csv.DictWriter(fout, output_fields, quotechar='"')
            w.writeheader()
            w.writerow(formatRow(row, fieldnames))
        for row in r:
            w.writerow(formatRow(row, fieldnames))

Answer 3

也許它可以幫助一些有類似問題的人..

在逐行讀取一個普通的 CSV 文件並通過一個字段決定是否將其保存在文件 A 或文件 B 中時，發生了 memory 溢出，我的 kernel 死了。 因此，我分析了我的 memory 用法和這個小改動 1. 通過削減將迭代次數增加了三倍 2. 修復了 memory 泄漏的問題

那是我的代碼，有 memory 泄漏和長時間運行

with open('input_file.csv', 'r') as input_file, open('file_A.csv', 'w') as file_A, open('file_B.csv', 'w') as file_B):
   input_csv = csv.reader(input_file)
   file_A_csv = csv.writer(file_A)
   file_B_csv = csv.writer(file_B)
   for row in input_file:
       condition_row = row[1]
       if condition_row == 'condition':
           file_A.writerow(row)
       else: 
           file_B.write(row)

但是如果你沒有像這樣聲明變量（或你閱讀文件的更多變量）：

with open('input_file.csv', 'r') as input_file, open('file_A.csv', 'w') as file_A, open('file_B.csv', 'w') as file_B):
input_csv = csv.reader(input_file)
file_A_csv = csv.writer(file_A)
file_B_csv = csv.writer(file_B)
for row in input_file:
    if row[1] == 'condition':
        file_A.writerow(row)
    else: 
        file_B.write(row)

我無法解釋為什么會這樣，但經過一些測試后我可以確定我的平均速度是原來的 3 倍，而我的 RAM 接近於零。