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

前段时间帮同事处理了一个把 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  # csv  # 导入mysql  # mysql  # 多进程  # 使用python将excel数据导入数据库过程详解  # 用Python将Excel数据导入到SQL Server的例子  # Python 中导入csv数据的三种方法  # Python之csv文件从MySQL数据库导入导出的方法  # python批量导入数据进Elasticsearch的实例  # python Django批量导入数据  # python Django批量导入不重复数据  # Python制作数据导入导出工具  # Python导入oracle数据的方法  # python 导入数据及作图的实现  # 自己的  # 几个  # 太多  # 多个  # 很久  # 用了  # 太大  # 给大家  # 希望大家  # 人了  # 转换成  # 来实现  # 别有  # 前段时间  # 太慢  # 命令行  # 跳过  # 往里  # 键值  # 单线程 

相关文章： 如何通过西部数码建站助手快速创建专业网站？  东莞市网站制作公司有哪些,东莞找工作用什么网站好？  岳西云建站教程与模板下载_一站式快速建站系统操作指南  如何高效搭建专业期货交易平台网站？  东莞专业网站制作公司有哪些,东莞招聘网站哪个好？  深圳网站制作设计招聘,关于服装设计的流行趋势，哪里的资料比较全面？  小自动建站系统：AI智能生成+拖拽模板，多端适配一键搭建  如何配置支付宝与微信支付功能？  官网自助建站系统：SEO优化+多语言支持，快速搭建专业网站  如何在宝塔面板创建新站点？  如何在自有机房高效搭建专业网站？  道歉网站制作流程,世纪佳缘致歉小吴事件，相亲网站身份信息伪造该如何稽查？  如何通过IIS搭建网站并配置访问权限？  利用JavaScript实现拖拽改变元素大小  高端网站建设与定制开发一站式解决方案 中企动力  Python lxml的etree和ElementTree有什么区别  宁波自助建站系统如何快速打造专业企业网站？  如何在Windows虚拟主机上快速搭建网站？  单页制作网站有哪些,朋友给我发了一个单页网站，我应该怎么修改才能把他变成自己的呢，请求高手指点迷津？  高性价比服务器租赁——企业级配置与24小时运维服务  C++ static_cast和dynamic_cast区别_C++静态转换与动态类型安全转换  移动端手机网站制作软件,掌上时代，移动端网站的谷歌SEO该如何做？  如何通过万网虚拟主机快速搭建网站？  如何快速配置高效服务器建站软件？  家庭建站与云服务器建站，如何选择更优？  网站制作网站,深圳做网站哪家比较好？  建站之星×万网：智能建站系统+自助建站平台一键生成  如何快速搭建高效服务器建站系统？  弹幕视频网站制作教程下载,弹幕视频网站是什么意思？  广德云建站网站建设方案与建站流程优化指南  制作无缝贴图网站有哪些,3dmax无缝贴图怎么调？  如何确保西部建站助手FTP传输的安全性？  成都响应式网站开发,dw怎么把手机适应页面变成网页？  外贸公司网站制作,外贸网站建设一般有哪些步骤？  定制建站方案优化指南：企业官网开发与建站费用解析  如何在服务器上配置二级域名建站？  5种Android数据存储方式汇总  如何在Golang中处理模块冲突_解决依赖版本不兼容问题  天津个人网站制作公司,天津网约车驾驶员从业资格证官网？  如何快速搭建高效WAP手机网站吸引移动用户？  实例解析Array和String方法  如何在云指建站中生成FTP站点？  网站好制作吗知乎,网站开发好学吗？有什么技巧？  网站制作哪家好,cc、.co、.cm哪个域名更适合做网站？  深圳 网站制作,深圳招聘网站哪个比较好一点啊？  公司网站建设制作费用,想建设一个属于自己的企业网站，该如何去做？  宝塔面板如何快速创建新站点？  建站主机选购指南：核心配置优化与品牌推荐方案  如何在Windows服务器上快速搭建网站？  XML的“混合内容”是什么 怎么用DTD或XSD定义