Python多线程/多进程优化PostgreSQL批量CSV数据插入咨询

核心最优方案：PostgreSQL原生COPY命令 + 进程池处理 #

Python处理大CSV导入的效率远不如数据库原生工具，COPY是PostgreSQL专为批量导入设计的功能，速度比Python逐行/批量插入快10-100倍。结合进程池处理拆分好的8个CSV文件，既能利用多核优势，又能避免内存过载。

你之前方案的问题根源 #

• multiprocessing触发SIGKILL：大概率是用pandas把1000万条CSV全加载到内存，多进程同时运行导致系统内存耗尽，被内核强制杀掉。
• 多线程效率下降：线程数超过8后，数据库连接竞争、线程上下文切换开销剧增，反而拖慢速度；单条数据开线程完全是资源浪费，线程创建销毁的成本远高于实际操作收益。

具体实现代码 #

方案1：进程池+psycopg2.copy_from（推荐） #

用psycopg2的copy_from调用原生COPY能力，配合进程池并行处理8个文件，每个进程对应一个文件，内存占用可控：

import os
import multiprocessing
import psycopg2
from psycopg2 import sql

# 数据库配置
DB_CONFIG = {
    "dbname": "your_db",
    "user": "your_user",
    "password": "your_pwd",
    "host": "localhost",
    "port": "5432"
}

def import_csv(csv_path, table_name):
    conn = None
    try:
        conn = psycopg2.connect(**DB_CONFIG)
        cur = conn.cursor()
        with open(csv_path, 'r', encoding='utf-8') as f:
            next(f)  # 跳过表头（如果CSV有表头）
            cur.copy_from(
                file=f,
                table=sql.Identifier(table_name),
                sep=',',
                null='',
                columns=('col1', 'col2', 'col3')  # 替换为你的表字段顺序
            )
        conn.commit()
        print(f"完成导入: {csv_path}")
    except Exception as e:
        if conn:
            conn.rollback()
        print(f"导入失败 {csv_path}: {str(e)}")
    finally:
        if conn:
            conn.close()

if __name__ == '__main__':
    csv_files = [
        "/data/file1.csv",
        "/data/file2.csv",
        "/data/file3.csv",
        "/data/file4.csv",
        "/data/file5.csv",
        "/data/file6.csv",
        "/data/file7.csv",
        "/data/file8.csv"
    ]
    target_table = "your_target_table"
    
    # 进程数设为文件数（8个），刚好匹配拆分的文件
    with multiprocessing.Pool(processes=8) as pool:
        pool.starmap(import_csv, [(file, target_table) for file in csv_files])

方案2：优化后的多线程方案（若必须用线程） #

如果因限制无法用进程，控制线程数在8个（你的最优值），同时用分块读取+批量插入，避免加载全量数据：

import threading
import pandas as pd
import psycopg2
from psycopg2 import sql
from queue import Queue

DB_CONFIG = {
    "dbname": "your_db",
    "user": "your_user",
    "password": "your_pwd",
    "host": "localhost",
    "port": "5432"
}
BATCH_SIZE = 10000  # 批量插入大小，根据内存调整

def worker(queue, table_name):
    conn = psycopg2.connect(**DB_CONFIG)
    cur = conn.cursor()
    # 预编译插入语句，减少SQL解析开销
    insert_sql = sql.SQL("INSERT INTO {} (col1, col2, col3) VALUES (%s, %s, %s)").format(
        sql.Identifier(table_name)
    )
    while True:
        batch = queue.get()
        if batch is None:
            break
        try:
            cur.executemany(insert_sql, batch)
            conn.commit()
        except Exception as e:
            conn.rollback()
            print(f"批量插入失败: {str(e)}")
        queue.task_done()
    conn.close()

def batch_import(csv_files, table_name, thread_count=8):
    queue = Queue(maxsize=thread_count*2)
    # 启动线程
    threads = []
    for _ in range(thread_count):
        t = threading.Thread(target=worker, args=(queue, table_name))
        t.start()
        threads.append(t)
    
    # 分块读取CSV，分批入队
    for csv_path in csv_files:
        for chunk in pd.read_csv(csv_path, chunksize=BATCH_SIZE):
            batch = [tuple(row) for row in chunk.values]
            queue.put(batch)
    
    # 发送结束信号
    for _ in range(thread_count):
        queue.put(None)
    
    # 等待所有线程完成
    for t in threads:
        t.join()

if __name__ == '__main__':
    csv_files = [
        "/data/file1.csv",
        "/data/file2.csv",
        # ... 剩余6个文件
    ]
    batch_import(csv_files, "your_target_table", thread_count=8)

额外提速建议 #

1. 数据库层面：
   • 导入前关闭目标表的索引和外键约束，导入完成后再重建，索引会大幅拖慢插入速度。
   • 临时调整PostgreSQL配置：增大work_mem、maintenance_work_mem，调高max_wal_size减少日志写入开销。
   • 用UNLOGGED表导入，完成后转为普通表（跳过WAL日志，速度更快，适合非核心数据导入）。
2. 文件层面：
   • 确保CSV文件和数据库在同一机器，避免网络传输耗时。
   • 用gzip压缩CSV，PostgreSQL的COPY支持直接读取压缩文件（COPY ... FROM PROGRAM 'gzip -dc file.csv.gz'）。

内容的提问来源于stack exchange，提问作者Tatia -