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 相关文章推荐
linux系统使用python监测系统负载脚本分享
Jan 15 Python
详解Python中的元组与逻辑运算符
Oct 13 Python
Python3使用requests登录人人影视网站的方法
May 11 Python
python3中str(字符串)的使用教程
Mar 23 Python
python获取时间及时间格式转换问题实例代码详解
Dec 06 Python
python矩阵/字典实现最短路径算法
Jan 17 Python
Python实现定时自动关闭的tkinter窗口方法
Feb 16 Python
Python调用C语言的实现
Jul 26 Python
详解python 中in 的 用法
Dec 12 Python
简单了解Python变量作用域正确使用方法
Jun 12 Python
Pytorch实现将模型的所有参数的梯度清0
Jun 24 Python
python安装mysql的依赖包mysql-python操作
Jan 01 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 excel类 phpExcel使用方法介绍
2010/08/21 PHP
php 伪静态之IIS篇
2014/06/02 PHP
Z-Blog中用到的js代码
2007/03/15 Javascript
jQuery 页面 Mask实现代码
2010/01/09 Javascript
js判断输入是否为正整数、浮点数等数字的函数代码
2010/11/17 Javascript
JavaScript flash复制库类 Zero Clipboard
2011/01/17 Javascript
解析Jquery中如何把一段html代码动态写入到DIV中(实例说明)
2013/07/09 Javascript
jQuery教程 $()包装函数来实现数组元素分页效果
2013/08/13 Javascript
javascript中负数算术右移、逻辑右移的奥秘探索
2013/10/17 Javascript
node中socket.io的事件使用详解
2014/12/15 Javascript
jquery简单实现网页层的展开与收缩效果
2015/08/07 Javascript
js数组如何添加json数据及js数组与json的区别
2015/10/27 Javascript
js中json处理总结之JSON.parse
2016/10/14 Javascript
jQuery实现联动下拉列表查询框
2017/01/04 Javascript
Angular实现响应式表单
2017/08/04 Javascript
仿淘宝JSsearch搜索下拉深度用法
2018/01/15 Javascript
Vue-router的使用和出现空白页,路由对象属性详解
2018/09/03 Javascript
element el-tree组件的动态加载、新增、更新节点的实现
2020/02/27 Javascript
vue项目打包之开发环境和部署环境的实现
2020/04/23 Javascript
Python yield 使用浅析
2015/05/28 Python
Python实现批量检测HTTP服务的状态
2016/10/27 Python
django用户注册、登录、注销和用户扩展的示例
2018/03/19 Python
python3解析库BeautifulSoup4的安装配置与基本用法
2018/06/26 Python
django中账号密码验证登陆功能的实现方法
2019/07/15 Python
Python range、enumerate和zip函数用法详解
2019/09/11 Python
Python实现计算长方形面积(带参数函数demo)
2020/01/18 Python
python利用faker库批量生成测试数据
2020/10/15 Python
HTML5 和小程序实现拍照图片旋转、压缩和上传功能
2018/10/08 HTML / CSS
时尚孕妇装:Ingrid & Isabel
2019/05/08 全球购物
软件测试工程师笔试题带答案
2015/03/27 面试题
家居设计专业个人自荐信范文
2013/11/26 职场文书
幼儿园保教管理制度
2014/02/03 职场文书
运动会开幕式主持词
2014/03/28 职场文书
体育系毕业生求职自荐信
2014/04/16 职场文书
效能监察建议书
2014/05/19 职场文书
银行柜员工作心得体会
2016/01/23 职场文书