Python实现 多进程导入CSV数据到 MySQL


Posted in Python onFebruary 26, 2017

前段时间帮同事处理了一个把 CSV 数据导入到 MySQL 的需求。两个很大的 CSV 文件, 分别有 3GB、2100 万条记录和 7GB、3500 万条记录。对于这个量级的数据,用简单的单进程/单线程导入 会耗时很久,最终用了多进程的方式来实现。具体过程不赘述,记录一下几个要点:

  1. 批量插入而不是逐条插入
  2. 为了加快插入速度,先不要建索引
  3. 生产者和消费者模型,主进程读文件,多个 worker 进程执行插入
  4. 注意控制 worker 的数量,避免对 MySQL 造成太大的压力
  5. 注意处理脏数据导致的异常
  6. 原始数据是 GBK 编码,所以还要注意转换成 UTF-8
  7. 用 click 封装命令行工具

具体的代码实现如下:

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import codecs
import csv
import logging
import multiprocessing
import os
import warnings

import click
import MySQLdb
import sqlalchemy

warnings.filterwarnings('ignore', category=MySQLdb.Warning)

# 批量插入的记录数量
BATCH = 5000

DB_URI = 'mysql://root@localhost:3306/example?charset=utf8'

engine = sqlalchemy.create_engine(DB_URI)


def get_table_cols(table):
  sql = 'SELECT * FROM `{table}` LIMIT 0'.format(table=table)
  res = engine.execute(sql)
  return res.keys()


def insert_many(table, cols, rows, cursor):
  sql = 'INSERT INTO `{table}` ({cols}) VALUES ({marks})'.format(
      table=table,
      cols=', '.join(cols),
      marks=', '.join(['%s'] * len(cols)))
  cursor.execute(sql, *rows)
  logging.info('process %s inserted %s rows into table %s', os.getpid(), len(rows), table)


def insert_worker(table, cols, queue):
  rows = []
  # 每个子进程创建自己的 engine 对象
  cursor = sqlalchemy.create_engine(DB_URI)
  while True:
    row = queue.get()
    if row is None:
      if rows:
        insert_many(table, cols, rows, cursor)
      break

    rows.append(row)
    if len(rows) == BATCH:
      insert_many(table, cols, rows, cursor)
      rows = []


def insert_parallel(table, reader, w=10):
  cols = get_table_cols(table)

  # 数据队列,主进程读文件并往里写数据,worker 进程从队列读数据
  # 注意一下控制队列的大小,避免消费太慢导致堆积太多数据,占用过多内存
  queue = multiprocessing.Queue(maxsize=w*BATCH*2)
  workers = []
  for i in range(w):
    p = multiprocessing.Process(target=insert_worker, args=(table, cols, queue))
    p.start()
    workers.append(p)
    logging.info('starting # %s worker process, pid: %s...', i + 1, p.pid)

  dirty_data_file = './{}_dirty_rows.csv'.format(table)
  xf = open(dirty_data_file, 'w')
  writer = csv.writer(xf, delimiter=reader.dialect.delimiter)

  for line in reader:
    # 记录并跳过脏数据: 键值数量不一致
    if len(line) != len(cols):
      writer.writerow(line)
      continue

    # 把 None 值替换为 'NULL'
    clean_line = [None if x == 'NULL' else x for x in line]

    # 往队列里写数据
    queue.put(tuple(clean_line))
    if reader.line_num % 500000 == 0:
      logging.info('put %s tasks into queue.', reader.line_num)

  xf.close()

  # 给每个 worker 发送任务结束的信号
  logging.info('send close signal to worker processes')
  for i in range(w):
    queue.put(None)

  for p in workers:
    p.join()


def convert_file_to_utf8(f, rv_file=None):
  if not rv_file:
    name, ext = os.path.splitext(f)
    if isinstance(name, unicode):
      name = name.encode('utf8')
    rv_file = '{}_utf8{}'.format(name, ext)
  logging.info('start to process file %s', f)
  with open(f) as infd:
    with open(rv_file, 'w') as outfd:
      lines = []
      loop = 0
      chunck = 200000
      first_line = infd.readline().strip(codecs.BOM_UTF8).strip() + '\n'
      lines.append(first_line)
      for line in infd:
        clean_line = line.decode('gb18030').encode('utf8')
        clean_line = clean_line.rstrip() + '\n'
        lines.append(clean_line)
        if len(lines) == chunck:
          outfd.writelines(lines)
          lines = []
          loop += 1
          logging.info('processed %s lines.', loop * chunck)

      outfd.writelines(lines)
      logging.info('processed %s lines.', loop * chunck + len(lines))


@click.group()
def cli():
  logging.basicConfig(level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(name)s - %(message)s')


@cli.command('gbk_to_utf8')
@click.argument('f')
def convert_gbk_to_utf8(f):
  convert_file_to_utf8(f)


@cli.command('load')
@click.option('-t', '--table', required=True, help='表名')
@click.option('-i', '--filename', required=True, help='输入文件')
@click.option('-w', '--workers', default=10, help='worker 数量,默认 10')
def load_fac_day_pro_nos_sal_table(table, filename, workers):
  with open(filename) as fd:
    fd.readline()  # skip header
    reader = csv.reader(fd)
    insert_parallel(table, reader, w=workers)


if __name__ == '__main__':
  cli()

以上就是本文给大家分享的全部没人了,希望大家能够喜欢

Python 相关文章推荐
Python中list列表的一些进阶使用方法介绍
Aug 15 Python
Python数据类型详解(三)元祖:tuple
May 08 Python
使用Python绘制图表大全总结
Feb 11 Python
python爬取网页内容转换为PDF文件
Jul 28 Python
Python中*args和**kwargs的区别详解
Sep 17 Python
利用Python小工具实现3秒钟将视频转换为音频
Oct 29 Python
浅析python redis的连接及相关操作
Nov 07 Python
解决django xadmin主题不显示和只显示bootstrap2的问题
Mar 30 Python
浅谈python锁与死锁问题
Aug 14 Python
基于python实现坦克大战游戏
Oct 27 Python
python opencv旋转图片的使用方法
Jun 04 Python
Python卷积神经网络图片分类框架详解分析
Nov 07 Python
python检查URL是否正常访问的小技巧
Feb 25 #Python
python解析基于xml格式的日志文件
Feb 25 #Python
Python中防止sql注入的方法详解
Feb 25 #Python
Python 数据结构之旋转链表
Feb 25 #Python
Python数据结构之翻转链表
Feb 25 #Python
浅析python中SQLAlchemy排序的一个坑
Feb 24 #Python
python函数的5种参数详解
Feb 24 #Python
You might like
使用PHP生成二维码的两种方法(带logo图像)
2014/03/14 PHP
PHP版QQ互联OAuth示例代码分享
2015/07/05 PHP
thinkphp框架下404页面设置 仅三步
2016/05/14 PHP
用客户端js实现带省略号的分页
2013/04/27 Javascript
使用js简单实现了tree树菜单
2013/11/20 Javascript
jQuery获取选中内容及设置元素属性的方法
2014/07/09 Javascript
JavaScript实现的浮动层框架用法实例分析
2015/10/10 Javascript
JS实现可自定义大小,可双击关闭的弹出层效果
2015/10/16 Javascript
jquery.cookie.js实现用户登录保存密码功能的方法
2016/04/15 Javascript
BootStrap Validator对于隐藏域验证和程序赋值即时验证的问题浅析
2016/12/01 Javascript
JavaScript实现星星等级评价功能
2017/03/22 Javascript
利用Node.js了解与测量HTTP所花费的时间详解
2017/09/22 Javascript
vue微信分享 vue实现当前页面分享其他页面
2017/12/02 Javascript
jQuery中的$是什么意思及 $. 和 $().的区别
2018/04/20 jQuery
vscode下vue项目中eslint的使用方法
2019/01/13 Javascript
微信小程序实现页面分享onShareAppMessage
2019/08/12 Javascript
Vue.js的模板语法详解
2020/02/16 Javascript
Vue两种组件类型:递归组件和动态组件的用法
2020/08/06 Javascript
在Python中操作字符串之replace()方法的使用
2015/05/19 Python
python使用PyCharm进行远程开发和调试
2017/11/02 Python
python 使用 requests 模块发送http请求 的方法
2018/12/09 Python
Python基于mysql实现学生管理系统
2019/02/21 Python
如何在Django配置文件里配置session链接
2019/08/06 Python
python实现修改固定模式的字符串内容操作示例
2019/12/30 Python
python numpy库linspace相同间隔采样的实现
2020/02/25 Python
利用python制作拼图小游戏的全过程
2020/12/04 Python
英国女士和男士时尚服装网上购物:Top Labels Online
2018/03/25 全球购物
个人简历中自我评价
2014/02/11 职场文书
临床专业自荐信
2014/06/22 职场文书
学校运动会报道稿
2014/09/23 职场文书
中学生逃课检讨书
2015/02/17 职场文书
论语读书笔记
2015/06/26 职场文书
公司宣传语大全
2015/07/13 职场文书
pytorch查看网络参数显存占用量等操作
2021/05/12 Python
PC版《死亡搁浅导剪版》现已发售 展开全新的探险
2022/04/03 其他游戏
Mysql 文件配置解析介绍
2022/05/06 MySQL