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 相关文章推荐
sqlalchemy对象转dict的示例
Apr 22 Python
python中bisect模块用法实例
Sep 25 Python
零基础写python爬虫之抓取百度贴吧代码分享
Nov 06 Python
Python cookbook(数据结构与算法)通过公共键对字典列表排序算法示例
Mar 15 Python
深入浅析Python获取对象信息的函数type()、isinstance()、dir()
Sep 17 Python
Python数据预处理之数据规范化(归一化)示例
Jan 08 Python
Python判断对象是否为文件对象(file object)的三种方法示例
Apr 26 Python
如何在Django中使用聚合的实现示例
Mar 23 Python
python 识别登录验证码图片功能的实现代码(完整代码)
Jul 03 Python
django rest framework 自定义返回方式
Jul 12 Python
Python 抓取数据存储到Redis中的操作
Jul 16 Python
分位数回归模型quantile regeression应用详解及示例教程
Nov 02 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
Redis在Laravel项目中的应用实例详解
2017/08/11 PHP
Laravel5.1 框架响应基本用法实例分析
2020/01/04 PHP
JS模拟多线程
2007/02/07 Javascript
深入理解JavaScript高级之词法作用域和作用域链
2013/12/10 Javascript
Jquery validation remote 验证的缓存问题解决方法
2014/03/25 Javascript
jQuery源码解读之removeClass()方法分析
2015/02/20 Javascript
JavaScript中isPrototypeOf函数作用和使用实例
2015/06/01 Javascript
JS根据浏览器窗口大小实时动态改变网页文字大小的方法
2016/02/25 Javascript
基于JS实现翻书效果的页面切换样式
2017/02/16 Javascript
Angular2使用Augury来调试Angular2程序
2017/05/21 Javascript
微信小程序实现Session功能及无法获取session问题的解决方法
2019/05/07 Javascript
JavaScript如何实现监听键盘输入和鼠标监点击
2020/07/20 Javascript
详解vue中使用transition和animation的实例代码
2020/12/12 Vue.js
[57:12]完美世界DOTA2联赛循环赛 Inki vs Matador BO2第一场 10.31
2020/11/02 DOTA
python中plot实现即时数据动态显示方法
2018/06/22 Python
对Python3.6 IDLE常用快捷键介绍
2018/07/16 Python
django DRF图片路径问题的解决方法
2018/09/10 Python
Python3模拟登录操作实例分析
2019/03/12 Python
详解django实现自定义manage命令的扩展
2019/08/13 Python
Python如何实现强制数据类型转换
2019/11/22 Python
Python使用循环神经网络解决文本分类问题的方法详解
2020/01/16 Python
使用jupyter notebook将文件保存为Markdown,HTML等文件格式
2020/04/14 Python
keras 多gpu并行运行案例
2020/06/10 Python
Python JSON常用编解码方法代码实例
2020/09/05 Python
柒牌官方商城:中国男装优秀品牌
2017/06/30 全球购物
Fossil美国官网:化石手表、手袋、首饰及配饰
2019/02/17 全球购物
DataList 能否分页,请问如何实现?
2015/05/03 面试题
办公室文书岗位职责
2013/12/16 职场文书
快递业务员岗位职责
2014/01/06 职场文书
五型班组建设方案
2014/02/10 职场文书
学校领导班子对照检查材料
2014/09/24 职场文书
2014年药剂科工作总结
2014/11/26 职场文书
2016学习依法治国心得体会
2016/01/15 职场文书
爱国之歌(8首)
2019/09/29 职场文书
CSS实现切角+边框+投影+内容背景色渐变效果
2021/11/01 HTML / CSS
python区块链持久化和命令行接口实现简版
2022/05/25 Python