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转换摩斯密码示例
Feb 16 Python
python快速查找算法应用实例
Sep 26 Python
Python 3.x读写csv文件中数字的方法示例
Aug 29 Python
完美解决Python 2.7不能正常使用pip install的问题
Jun 12 Python
Python 删除连续出现的指定字符的实例
Jun 29 Python
django query模块
Apr 20 Python
Python动态声明变量赋值代码实例
Dec 30 Python
Python requests获取网页常用方法解析
Feb 20 Python
pytorch 使用加载训练好的模型做inference
Feb 20 Python
使用Keras实现简单线性回归模型操作
Jun 12 Python
python实现启动一个外部程序,并且不阻塞当前进程
Dec 05 Python
python简单实现插入排序实例代码
Dec 16 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 switch语句多个值匹配同一代码块应用示例
2014/07/29 PHP
php实现的网页版剪刀石头布游戏示例
2016/11/25 PHP
PHP使用preg_split和explode分割textarea存放内容的方法分析
2017/07/03 PHP
IE浏览器兼容Firefox的JS脚本的代码
2008/10/23 Javascript
利用毫秒减值计算时长的js代码
2013/09/22 Javascript
js打造数组转json函数
2015/01/14 Javascript
12行javascript代码绘制一个八卦图
2015/04/02 Javascript
微信小程序 loading 详解及实例代码
2016/11/09 Javascript
微信小程序之获取当前位置经纬度以及地图显示详解
2017/05/09 Javascript
详解js静态资源文件请求的处理
2017/08/01 Javascript
总结js中的一些兼容性易错的问题
2017/12/18 Javascript
babel之配置文件.babelrc入门详解
2018/02/22 Javascript
javascript数据结构之多叉树经典操作示例【创建、添加、遍历、移除等】
2018/08/01 Javascript
layer弹出层 iframe层去掉滚动条的实例代码
2018/08/17 Javascript
layui 富文本编辑器和textarea值的相互传递方法
2019/09/18 Javascript
Element PageHeader页头的使用方法
2020/07/26 Javascript
仿照Element-ui实现一个简易的$message方法
2020/09/14 Javascript
[01:02:45]完美世界DOTA2联赛 LBZS vs Forest 第三场 11.07
2020/11/09 DOTA
Python命名空间详解
2014/08/18 Python
离线安装Pyecharts的步骤以及依赖包流程
2020/04/23 Python
Python利用Beautiful Soup模块创建对象详解
2017/03/27 Python
Python实现获取照片拍摄日期并重命名的方法
2017/09/30 Python
13个最常用的Python深度学习库介绍
2017/10/28 Python
Python使用dict.fromkeys()快速生成一个字典示例
2019/04/24 Python
python基于itchat模块实现微信防撤回
2019/04/29 Python
python Gunicorn服务器使用方法详解
2019/07/22 Python
Django框架 信号调度原理解析
2019/09/04 Python
Django微信小程序后台开发教程的实现
2020/06/03 Python
解决Keras自带数据集与预训练model下载太慢问题
2020/06/12 Python
工程质量承诺书范文
2014/03/27 职场文书
交通事故协议书
2014/04/15 职场文书
五水共治捐款倡议书
2014/05/14 职场文书
村主任群众路线个人对照检查材料
2014/09/26 职场文书
网络工程专业大学生求职信
2014/10/01 职场文书
幼儿园教师教学反思
2016/03/02 职场文书
python基础之//、/与%的区别详解
2022/06/10 Python