Python实现多进程导入CSV大文件到数据库，pandas分块读取

本文链接：https://blog.csdn.net/qq_38723814/article/details/129559735

Python实现多进程导入CSV大文件到数据库

对于比较大的CSV文件，直接读取所有数据到内存肯定是万万不得行滴，文件稍稍大一点可能读一万行需要两分钟或者直接卡死，所以需要使用 pandas 分块读取

一、数据读取：Pandas 的 read_csv 函数

先生成一个测试文件

import pandas as pd
import numpy as np

# filename_ = r'D:\Projects\test\pandas_study\bigexcel/data.csv'

np.random.seed = 2023
df_size = 10000000  # 10000000

df = pd.DataFrame({
    'a': np.random.randint(1000, size=df_size),
    'b': np.random.randint(1000, size=df_size),
    'c': np.random.randint(1000, size=df_size),
    'd': np.random.randint(1000, size=df_size),
    'e': np.random.randint(1000, size=df_size)
})
df.head()
# print(df)
df.to_csv('data.csv')

Pandas 的 read_csv 函数提供2个参数：chunksize、iterator ，可实现按行多次读取文件，避免内存不足情况

* iterator : boolean, default False
返回一个TextFileReader 对象，以便逐块处理文件。

* chunksize : int, default None
文件块的大小， See IO Tools docs for more informationon iterator and chunksize.

1、指定 chunksize 分块读取文件

pandas.read_csv 参数 chunksize 通过指定一个分块大小（每次读取多少行）来读取大数据文件，可避免一次性读取内存不足，返回的是一个可迭代对象 TextFileReader

import pandas as pd
reader = pd.read_csv('data.csv', sep=',', chunksize=10)
# <pandas.io.parsers.TextFileReader at 0x1fc81f905e0>

for chunk in reader:
    # df = chunk
    # 对 chunk 进行数据处理
    print(type(chunk), chunk.shape)
'''
<class 'pandas.core.frame.DataFrame'> (10, 6)
<class 'pandas.core.frame.DataFrame'> (10, 6)
<class 'pandas.core.frame.DataFrame'> (10, 6)
<class 'pandas.core.frame.DataFrame'> (10, 6)
<class 'pandas.core.frame.DataFrame'> (10, 6)
'''

for chunk in reader:
    # df = chunk
    # 对 chunk 进行数据处理
    chunk.rename(columns={'Unnamed: 0':'index2'}, inplace=True) # 修改列名
    print(chunk.columns)

2、指定 iterator=True

指定 iterator=True 也可以返回一个可迭代对象 TextFileReader 。
get_chunk(size) – 返回一个N行的数据块
每次执行获取N行数据，再次执行，获取下一个数据块

iterator=True 和 chunksize 可以同时指定使用

reader = pd.read_csv('data.csv', sep=',', iterator=True)
data = reader.get_chunk(5) # 返回N行数据块
data
'''
   Unnamed: 0         a         b         c         d         e
0           0  0.289972  0.717806  0.886283  0.522148  0.976798
1           1  0.254952  0.048073  0.464765  0.138978  0.983041
2           2  0.634708  0.533182  0.855981  0.456156  0.620018
3           3  0.812648  0.024870  0.536520  0.894937  0.102704
4           4  0.699629  0.038305  0.379534  0.876242  0.906875
'''

3、获取文件行数

count = 0
file = open('data_csv.csv', 'r', encoding='utf-8')
while 1:
    buffer = file.read(8*1024*1024) # 可大概设置
    if not buffer:
        break
    count += buffer.count('\n')
print(count)
file.close()

4、分块拆分文件 and 合并数据

import pandas as pd
reader = pd.read_csv('data_csv.csv', sep=',', chunksize=2000000)
for i, chunk in enumerate(reader):
    print(i, '  ', len(chunk))
    chunk.to_csv('./data/data_' + str(i) + '.csv', index=False)


import pandas as pd
df = [pd.read_csv('./data/data_' + str(i) + '.csv') for i in range(5)] # 列表推导式
data = pd.concat(df, axis=0).reset_index(drop=True) # 合并
data.head()
data.tail()

5、分块读取文件

import feather
import pandas as pd

filePath = r'data_csv.csv'

def read_csv_feature(filePath):
    # 读取文件
    f = open(filePath, encoding='utf-8')
    reader = pd.read_csv(f, sep=',', iterator=True)
    loop = True
    chunkSize = 1000000
    chunks = []
    while loop:
        try:
            chunk = reader.get_chunk(chunkSize)
            chunks.append(chunk)
        except StopIteration:
            loop = False
            print('Iteration is END!!!')
    df = pd.concat(chunks, axis=0, ignore_index=True)
    f.close()
    return df 

data = read_csv_feature(filePath)

二、数据写入数据库

1、批量插入而不是逐条插入
2、为了加快插入速度，先不要建索引
3、生产者和消费者模型，主进程读文件，多个 worker 进程执行插入
4、注意控制 worker 的数量，避免对 MySQL 造成太大的压力
5、注意处理脏数据导致的异常

三、具体代码实现

import pandas as pd
import multiprocessing
import pymysql

# 进程数
w = 4
# 批量插入的记录数量
BATCH = 5000

# 数据库连接信息
DB_URI = 'mysql://*******'
connection = pymysql.connect(DB_URI)
# 文件路径
filename_ = r'D:\Projects\test\pandas_study\bigexcel/data.csv'


def insert_worker(queue):
    rows = []
    # 每个子进程创建自己的 游标对象
    cursor = connection.cursor()  # 生成游标对象
    while True:
        row = queue.get()
        if row is None:
            if rows:
                # 批量插入方法 或者 sql
                print(f"插入{len(rows)}条")
                insert_many(rows, cursor)
            break

        rows.append(row)
        if len(rows) == BATCH:
            # 批量插入方法 或者 sql
            print("插入5000条")
            insert_many(rows, cursor)
            rows = []


def csv():
    # 创建数据队列，主进程读文件并往里写数据，worker 进程从队列读数据
    # 注意一下控制队列的大小，避免消费太慢导致堆积太多数据，占用过多内存
    queue_ = multiprocessing.Queue(maxsize=w * BATCH * 2)
    workers = []
    for i in range(w):
        # 增加进程任务
        p = multiprocessing.Process(target=insert_worker, args=(queue_,))
        p.start()
        workers.append(p)
        print('starting # %s worker process, pid: %s...', i + 1, p.pid)

    dirty_data_file = r'D:\Projects\test\pandas_study\bigexcel/error.csv'
    # 表字段
    cols = ["a", "b", "c", "d", "e"]
    # 错误数据列表
    dirty_data_list = []

    reader = pd.read_csv(filename_, sep=',', iterator=True)
    loop = True
    while loop:
        try:
            data = reader.get_chunk(BATCH)  # 返回N行数据块  DataFrame类型
            data_list = data.values.tolist()
            for line in data_list:
                # 记录并跳过脏数据: 键值数量不一致 或者 其他检验数据错误逻辑
                # cols 表字段
                if len(line) != (len(cols) + 1):
                    dirty_data_list.append(line[1:])
                    continue
                # 把 None 值替换为 'NULL'
                clean_line = [None if x == 'NULL' else x for x in line]

                # 往队列里写数据
                queue_.put(tuple(clean_line))
        except StopIteration:
            # 错误数据写入文件
            dirty_data_frame = pd.DataFrame(dirty_data_list, columns=cols)
            dirty_data_frame.to_csv(dirty_data_file)
            loop = False

    # 给每个 worker 发送任务结束的信号
    print('send close signal to worker processes')
    for i in range(w):
        queue_.put(None)

    for p in workers:
        p.join()


if __name__ == '__main__':
    csv()