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中使用PIL模块处理图像的教程
Apr 29 Python
详解Python中的__getitem__方法与slice对象的切片操作
Jun 27 Python
Window 64位下python3.6.2环境搭建图文教程
Sep 19 Python
PyQt 图解Qt Designer工具的使用方法
Aug 06 Python
Python在cmd上打印彩色文字实现过程详解
Aug 07 Python
python关于变量名的基础知识点
Mar 03 Python
Django 404、500页面全局配置知识点详解
Mar 10 Python
Python3实现飞机大战游戏
Apr 24 Python
python读取yaml文件后修改写入本地实例
Apr 27 Python
Python参数传递对象的引用原理解析
May 22 Python
解决c++调用python中文乱码问题
Jul 29 Python
Python实现迪杰斯特拉算法过程解析
Sep 18 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 curl批处理实现可控并发异步操作示例
2018/05/09 PHP
PHP获取HTTP body内容的方法
2018/12/31 PHP
php框架知识点的整理和补充
2021/03/01 PHP
firefox和IE系列的相关区别整理 以备后用
2009/12/28 Javascript
$.ajax返回的JSON无法执行success的解决方法
2011/09/09 Javascript
Raphael一个用于在网页中绘制矢量图形的Javascript库
2013/01/08 Javascript
使用js实现关闭js弹出层的窗口
2014/02/10 Javascript
用Jquery实现滚动新闻
2014/02/12 Javascript
让alert不出现弹窗的两种方法
2014/05/18 Javascript
webapp框架AngularUI的demo改造之路
2014/12/21 Javascript
JS & JQuery 动态添加 select option
2016/06/08 Javascript
浅谈MVC+EF easyui dataGrid 动态加载分页表格
2016/11/10 Javascript
BootStrap Validator对于隐藏域验证和程序赋值即时验证的问题浅析
2016/12/01 Javascript
ES6新特性三: Generator(生成器)函数详解
2017/04/21 Javascript
小程序登录/注册页面设计的实现代码
2019/05/24 Javascript
使用vue引入maptalks地图及聚合效果的实现
2020/08/10 Javascript
Python正则表达式介绍
2012/08/06 Python
Python urllib、urllib2、httplib抓取网页代码实例
2015/05/09 Python
python图像处理之反色实现方法
2015/05/30 Python
Python爬取京东的商品分类与链接
2016/08/26 Python
浅谈pandas中shift和diff函数关系
2018/04/08 Python
dataframe 按条件替换某一列中的值方法
2019/01/29 Python
python创建与遍历List二维列表的方法
2019/08/16 Python
keras的backend 设置 tensorflow,theano操作
2020/06/30 Python
纯css3无js实现的Android Logo(有简单动画)
2013/01/21 HTML / CSS
绩效工资分配方案
2014/01/18 职场文书
互联网创业计划书的书写步骤
2014/01/28 职场文书
董事长助理岗位职责
2014/02/18 职场文书
夫妻双方自愿离婚协议书怎么写
2014/12/01 职场文书
庆祝教师节主题班会
2015/08/17 职场文书
创业计划书之书店
2019/09/10 职场文书
python自动化之如何利用allure生成测试报告
2021/05/02 Python
为什么代码规范要求SQL语句不要过多的join
2021/06/23 MySQL
vue实现移动端div拖动效果
2022/03/03 Vue.js
Java 多态分析
2022/04/26 Java/Android
SpringBoot全局异常处理方案分享
2022/05/25 Java/Android