Python 다중 프로세스로 CSV를 데이터베이스로 가져오기

이 글에서는 Python을 사용하여 CSV 파일 데이터를 MySQL로 다중 프로세스 가져오기 및 특정 코드 공유를 구현하는 아이디어와 방법을 공유합니다. 동일한 요구 사항이 있는 친구가 참고할 수 있습니다

얼마 전에 동료가 이 문제를 처리하는 데 도움을 주었습니다. CSV 데이터를 MySQL로 가져오기 위한 요구 사항입니다. 두 개의 큰 CSV 파일(각각 2,100만 개의 레코드가 포함된 3GB 및 3,500만 개의 레코드가 포함된 7GB)입니다. 이 정도 규모의 데이터의 경우 간단한 단일 프로세스/단일 스레드 가져오기에는 시간이 오래 걸리므로 이를 구현하기 위해 마지막으로 다중 프로세스 접근 방식이 사용되었습니다. 구체적인 과정을 자세히 설명하지는 않지만 몇 가지 핵심 사항을 기록해 두겠습니다.

1. 한 번에 하나씩 넣지 말고 일괄적으로 넣으세요

2. 삽입 속도를 높이려면 Index

3. Producer 및 ConsumerModel를 빌드하지 마세요. 메인 프로세스에서 파일을 읽습니다. , 여러 작업자 프로세스가 삽입을 수행합니다

4. MySQL에 너무 많은 부담을 주지 않도록 작업자 수를 제어하는 ​​데 주의하세요

5. 주의하세요 더티 데이터 처리로 인한 예외

6. 원본 데이터가 GBK로 인코딩되어 있으므로 UTF-8로 변환할 때도 주의가 필요합니다

7. 클릭을 사용하여 명령줄 도구
   

특정 코드를 캡슐화합니다. 구현은 다음과 같습니다.

#!/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()

[관련 권장 사항]

1 . Python 무료 동영상 튜토리얼

2. Python 학습 매뉴얼

Geek Academy Python 동영상 튜토리얼

위 내용은 Python 다중 프로세스로 CSV를 데이터베이스로 가져오기의 상세 내용입니다. 자세한 내용은 PHP 중국어 웹사이트의 기타 관련 기사를 참조하세요!