前段時間幫同事處理了一個把 CSV 數據導入到 MySQL 的需求。兩個很大的 CSV 文件， 分別有 3GB、2100 萬條記錄和  7GB、3500 萬條記錄。對于這個量級的數據，用簡單的單進程／單線程導入  會耗時很久，最終用了多進程的方式來實現。具體過程不贅述，記錄一下幾個要點：

批量插入而不是逐條插入

為了加快插入速度，先不要建索引

生產者和消費者模型，主進程讀文件，多個 worker 進程執行插入

注意控制 worker 的數量，避免對 MySQL 造成太大的壓力

注意處理臟數據導致的異常

原始數據是 GBK 編碼，所以還要注意轉換成 UTF-8

用 click 封裝命令行工具

具體的代碼實現如下：

1. #!/usr/bin/env python 
2. # -*- coding: utf-8 -*- 
3.   
4. importcodecs 
5. importcsv 
6. importlogging 
7. importmultiprocessing 
8. importos 
9. importwarnings 
10.   
11. importclick 
12. importMySQLdb 
13. importsqlalchemy 
14.   
15. warnings.filterwarnings('ignore', category=MySQLdb.Warning) 
16.   
17. # 批量插入的記錄數量 
18. BATCH=5000 
19.   
20. DB_URI='mysql://root@localhost:3306/example?charset=utf8' 
21.   
22. engine=sqlalchemy.create_engine(DB_URI) 
23.   
24.   
25. defget_table_cols(table): 
26.   sql='SELECT * FROM `{table}` LIMIT 0'.format(table=table) 
27.   res=engine.execute(sql) 
28.   returnres.keys() 
29.   
30.   
31. definsert_many(table, cols, rows, cursor): 
32.   sql='INSERT INTO `{table}` ({cols}) VALUES ({marks})'.format( 
33.       table=table, 
34.       cols=', '.join(cols), 
35.       marks=', '.join(['%s']*len(cols))) 
36.   cursor.execute(sql,*rows) 
37.   logging.info('process %s inserted %s rows into table %s', os.getpid(),len(rows), table) 
38.   
39.   
40. definsert_worker(table, cols, queue): 
41.   rows=[] 
42.   # 每個子進程創建自己的 engine 對象 
43.   cursor=sqlalchemy.create_engine(DB_URI) 
44.   whileTrue: 
45.     row=queue.get() 
46.     ifrowisNone: 
47.       ifrows: 
48.         insert_many(table, cols, rows, cursor) 
49.       break 
50.   
51.     rows.append(row) 
52.     iflen(rows)==BATCH: 
53.       insert_many(table, cols, rows, cursor) 
54.       rows=[] 
55.   
56.   
57. definsert_parallel(table, reader, w=10): 
58.   cols=get_table_cols(table) 
59.   
60.   # 數據隊列，主進程讀文件并往里寫數據，worker 進程從隊列讀數據 
61.   # 注意一下控制隊列的大小，避免消費太慢導致堆積太多數據，占用過多內存 
62.   queue=multiprocessing.Queue(maxsize=w*BATCH*2) 
63.   workers=[] 
64.   foriinrange(w): 
65.     p=multiprocessing.Process(target=insert_worker, args=(table, cols, queue)) 
66.     p.start() 
67.     workers.append(p) 
68.     logging.info('starting # %s worker process, pid: %s...', i+1, p.pid) 
69.   
70.   dirty_data_file='./{}_dirty_rows.csv'.format(table) 
71.   xf=open(dirty_data_file,'w') 
72.   writer=csv.writer(xf, delimiter=reader.dialect.delimiter) 
73.   
74.   forlineinreader: 
75.     # 記錄并跳過臟數據: 鍵值數量不一致 
76.     iflen(line) !=len(cols): 
77.       writer.writerow(line) 
78.       continue 
79.   
80.     # 把 None 值替換為 'NULL' 
81.     clean_line=[Noneifx=='NULL'elsexforxinline] 
82.   
83.     # 往隊列里寫數據 
84.     queue.put(tuple(clean_line)) 
85.     ifreader.line_num%500000==0: 
86.       logging.info('put %s tasks into queue.', reader.line_num) 
87.   
88.   xf.close() 
89.   
90.   # 給每個 worker 發送任務結束的信號 
91.   logging.info('send close signal to worker processes') 
92.   foriinrange(w): 
93.     queue.put(None) 
94.   
95.   forpinworkers: 
96.     p.join() 
97.   
98.   
99. defconvert_file_to_utf8(f, rv_file=None): 
100.   ifnotrv_file: 
101.     name, ext=os.path.splitext(f) 
102.     ifisinstance(name,unicode): 
103.       name=name.encode('utf8') 
104.     rv_file='{}_utf8{}'.format(name, ext) 
105.   logging.info('start to process file %s', f) 
106.   withopen(f) as infd: 
107.     withopen(rv_file,'w') as outfd: 
108.       lines=[] 
109.       loop=0 
110.       chunck=200000 
111.       first_line=infd.readline().strip(codecs.BOM_UTF8).strip()+'/n' 
112.       lines.append(first_line) 
113.       forlineininfd: 
114.         clean_line=line.decode('gb18030').encode('utf8') 
115.         clean_line=clean_line.rstrip()+'/n' 
116.         lines.append(clean_line) 
117.         iflen(lines)==chunck: 
118.           outfd.writelines(lines) 
119.           lines=[] 
120.           loop+=1 
121.           logging.info('processed %s lines.', loop*chunck) 
122.   
123.       outfd.writelines(lines) 
124.       logging.info('processed %s lines.', loop*chunck+len(lines)) 
125.   
126.   
127. @click.group() 
128. defcli(): 
129.   logging.basicConfig(level=logging.INFO, 
130.             format='%(asctime)s - %(levelname)s - %(name)s - %(message)s') 
131.   
132.   
133. @cli.command('gbk_to_utf8') 
134. @click.argument('f') 
135. defconvert_gbk_to_utf8(f): 
136.   convert_file_to_utf8(f) 
137.   
138.   
139. @cli.command('load') 
140. @click.option('-t','--table', required=True,help='表名') 
141. @click.option('-i','--filename', required=True,help='輸入文件') 
142. @click.option('-w','--workers', default=10,help='worker 數量，默認 10') 
143. defload_fac_day_pro_nos_sal_table(table, filename, workers): 
144.   withopen(filename) as fd: 
145.     fd.readline() # skip header 
146.     reader=csv.reader(fd) 
147.     insert_parallel(table, reader, w=workers) 
148.  //Vevb.com 
149.   
150. if__name__=='__main__': 
151.   cli()