使用Python读取和处理大型CSV

时间:2018-07-17 05:41:15

标签: python python-3.x

我有一个问题与此previously asked question相似。尽管如此,我似乎无法找到合适的解决方案。

输入:我有

的CSV数据
id,prescriber_last_name,prescriber_first_name,drug_name,drug_cost
1000000001,Smith,James,AMBIEN,100
1000000002,Garcia,Maria,AMBIEN,200
1000000003,Johnson,James,CHLORPROMAZINE,1000
1000000004,Rodriguez,Maria,CHLORPROMAZINE,2000
1000000005,Smith,David,BENZTROPINE MESYLATE,1500

输出:从这里,我只需要输出每种药物,所有处方上加起来的总费用,我就需要计算出唯一的处方药数量。

drug_name,num_prescriber,total_cost
AMBIEN,2,300.0
CHLORPROMAZINE,2,3000.0
BENZTROPINE MESYLATE,1,1500.0

我能够使用Python轻松完成此任务。但是,当我尝试使用更大的输入(1gb)输入运行代码时,我的代码不会在合理的时间内终止。

import sys, csv

def duplicate_id(id, id_list):
    if id in id_list:
        return True
    else:
        return False

def write_file(d, output):
    path = output
    # path = './output/top_cost_drug.txt'
    with open(path, 'w', newline='') as csvfile:
        fieldnames = ['drug_name', 'num_prescriber', 'total_cost']
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        writer.writeheader()
        for key, value in d.items():
            print(key, value)
            writer.writerow({'drug_name': key, 'num_prescriber': len(value[0]), 'total_cost': sum(value[1])})

def read_file(data):
    # TODO: https://codereview.stackexchange.com/questions/88885/efficiently-filter-a-large-100gb-csv-file-v3
    drug_info = {}
    with open(data) as csvfile:
        readCSV = csv.reader(csvfile, delimiter=',')
        next(readCSV)
        for row in readCSV:
            prescriber_id = row[0]
            prescribed_drug = row[3]
            prescribed_drug_cost = float(row[4])

            if prescribed_drug not in drug_info:
                drug_info[prescribed_drug] = ([prescriber_id], [prescribed_drug_cost])
            else:
                if not duplicate_id(prescriber_id, drug_info[prescribed_drug][0]):
                    drug_info[prescribed_drug][0].append(prescriber_id)
                    drug_info[prescribed_drug][1].append(prescribed_drug_cost)
                else:
                    drug_info[prescribed_drug][1].append(prescribed_drug_cost)
    return(drug_info)

def main():
    data = sys.argv[1]
    output = sys.argv[2]
    drug_info = read_file(data)
    write_file(drug_info, output)

if __name__ == "__main__":
    main()

我在弄清楚如何重构它以处理较大的输入时遇到了麻烦,希望有人可以看看并为我提供一些有关如何解决此问题的建议。

2 个答案:

答案 0 :(得分:1)

如果可以使用熊猫,请尝试以下操作。熊猫会读取您的文件并将其存储在数据框中。这比使用迭代器手动处理文件要快得多。

import pandas as pd


df = pd.read_csv('sample_data.txt')

columns = ['id','drug_name','drug_cost']


df1 = df[columns]
gd = df1.groupby('drug_name')
cnt= gd.count()
s=gd.sum()

out = s.join(cnt,lsuffix='x')
out['total_cost']=out['drug_costx']
out['num_prescriber']=out['drug_cost']
fout = out[['num_prescriber','total_cost']]

fout.to_csv('out_data.csv')

我得到以下输出。

drug_name,num_prescriber,total_cost
AMBIEN,2,300
BENZTROPINE MESYLATE,1,1500
CHLORPROMAZINE,2,3000

希望这会有所帮助。

答案 1 :(得分:0)

列表在测试成员资格方面效率不高,尤其是当长度为数千时,因为它花费O(n)。使用集来存储您的处方者ID,这会将成员资格测试的成本降低到O(1)。

        if prescribed_drug not in drug_info:
            drug_info[prescribed_drug] = ({prescriber_id}, [prescribed_drug_cost])
        else:
            if prescriber_id not in drug_info[prescribed_drug][0]:
                drug_info[prescribed_drug][0].add(prescriber_id)
                drug_info[prescribed_drug][1].append(prescribed_drug_cost)
            else:
                drug_info[prescribed_drug][1].append(prescribed_drug_cost)