PostgreSQL多查询场景：异步编程与多线程方案选型

我有一个Python脚本，需要读取约一百万行CSV文件数据，每行数据需执行至少1次、最多2次PostgreSQL表查询。为加速处理过程，我希望引入并行机制，但对Python中多线程、多进程及异步编程的适用场景并不熟悉，因此想了解该如何选择合适的方案。以下是当前的同步代码：

def distribute_rows_to_files(file_path: str) -> None:
    exists_file = "exists_data.csv"
    c_not_exists_file = "c_not_exists_data.csv"
    i_not_exists_file = "i_not_exists_data.csv"
    exception_file = "exceptions.csv"

    # Open the files before the loop to reduce overhead
    exists_file_handle = open(exists_file, mode='a', newline='')
    c_not_exists_file_handle = open(c_not_exists_file, mode='a', newline='')
    i_not_exists_file_handle = open(i_not_exists_file, mode='a', newline='')
    exception_handle = open(exception_file, mode='a', newline='')

    with psycopg2.connect(**AUTHENTICATOR_PG_DB_CREDENTIALS) as conn:
        with open(file_path, mode='r', newline='') as file:
            reader = csv.reader(file)
            next(reader)  # Skip the header line

            count = 0
            for row in reader:
                count += 1
                if count == 100:
                    break

                # Process each row here
                i_code, t_id, __, ___, ____ = row
                try:
                    cur = conn.cursor()
                    query = """
                        SELECT customer_id
                        FROM buckets
                        WHERE i_code = %(i_code)s
                        LIMIT 1
                    """
                    cur.execute(query, {"i_code": i_code})
                    result = cur.fetchone()
                    cur.close()

                    if result:
                        try:
                            cur = conn.cursor()
                            second_query = """
                                SELECT EXISTS (
                                    SELECT 1
                                    FROM customers
                                    WHERE customer_id = %(customer_id)s
                                    AND t_id = %(t_id)s
                                )
                            """
                            cur.execute(second_query, {"customer_id": result[0], "t_id": t_id})
                            exists = cur.fetchone()[0]
                            cur.close()

                            file_handle = exists_file_handle if exists else c_not_exists_file_handle
                            writer = csv.writer(file_handle)
                            writer.writerow(row)
                        except Exception as e:
                            row_with_exception = row + (str(e),)
                            writer = csv.writer(exception_handle)
                            writer.writerow(row_with_exception)
                    else:
                        writer = csv.writer(i_not_exists_file_handle)
                        writer.writerow(row)

                except Exception as e:
                    row_with_exception = row + (str(e),)
                    writer = csv.writer(exception_handle)
                    writer.writerow(row_with_exception)

    exists_file_handle.close()
    c_not_exists_file_handle.close()
    i_not_exists_file_handle.close()
    exception_handle.close()

注：已知异步开发中推荐使用asyncpg而非psycopg2，但需先确认异步编程是否为合适方案。

各方案适用场景 #

• 多进程：适合CPU密集型任务（如大量数据计算）。但你的场景以数据库IO等待为主，多进程会带来额外的进程间通信开销，且每个进程需单独建立PostgreSQL连接，资源消耗大，性价比极低，不推荐。
• 多线程：Python的*GIL（全局解释器锁）*在IO阻塞（如等待数据库响应）时会自动释放，因此多线程可有效利用等待时间处理其他行，开销远低于多进程，适合IO密集场景。需注意：数据库连接不能跨线程共享，每个线程需单独创建连接；文件写入需加锁保证线程安全。
• 异步编程：基于asyncpg实现全异步IO，避免线程切换开销，在高并发IO场景下效率最高。百万行数据的批量IO操作，异步能最大化利用系统资源，减少等待时间，是最优选择。需注意：整个流程需适配异步模式（包括文件读写、CSV处理），且文件写入需用异步锁保证安全。

优先推荐方案 #

优先选择异步编程，其次是多线程。以下分别给出两种方案的代码示例。

异步方案（基于asyncpg） #

import asyncio
import csv
import asyncpg
from asyncio import Lock

# 定义文件锁，保证异步写入安全
file_locks = {
    "exists": Lock(),
    "c_not_exists": Lock(),
    "i_not_exists": Lock(),
    "exception": Lock()
}

async def process_row(row, conn):
    i_code, t_id, _, _, _ = row
    try:
        # 第一个查询
        result = await conn.fetchrow(
            """
            SELECT customer_id
            FROM buckets
            WHERE i_code = $1
            LIMIT 1
            """,
            i_code
        )

        if result:
            # 第二个查询
            exists = await conn.fetchval(
                """
                SELECT EXISTS (
                    SELECT 1
                    FROM customers
                    WHERE customer_id = $1
                    AND t_id = $2
                )
                """,
                result["customer_id"],
                t_id
            )

            # 写入对应文件
            if exists:
                async with file_locks["exists"]:
                    with open("exists_data.csv", "a", newline='') as f:
                        writer = csv.writer(f)
                        writer.writerow(row)
            else:
                async with file_locks["c_not_exists"]:
                    with open("c_not_exists_data.csv", "a", newline='') as f:
                        writer = csv.writer(f)
                        writer.writerow(row)
        else:
            async with file_locks["i_not_exists"]:
                with open("i_not_exists_data.csv", "a", newline='') as f:
                    writer = csv.writer(f)
                    writer.writerow(row)
    except Exception as e:
        async with file_locks["exception"]:
            with open("exceptions.csv", "a", newline='') as f:
                writer = csv.writer(f)
                writer.writerow(row + (str(e),))

async def main(file_path):
    # 初始化数据库连接池
    conn_pool = await asyncpg.create_pool(**AUTHENTICATOR_PG_DB_CREDENTIALS)

    # 读取CSV文件
    with open(file_path, "r", newline='') as f:
        reader = csv.reader(f)
        next(reader)  # 跳过表头
        rows = list(reader)
        
        # 注：百万行直接批量处理可能导致连接数过载，建议分批（如每1000行一批）
        # rows_batch = []
        # for idx, row in enumerate(reader):
        #     if idx == 0: continue
        #     rows_batch.append(row)
        #     if len(rows_batch) == 1000:
        #         await asyncio.gather(*[process_row(row, await conn_pool.acquire()) for row in rows_batch])
        #         rows_batch = []
        # if rows_batch:
        #     await asyncio.gather(*[process_row(row, await conn_pool.acquire()) for row in rows_batch])

    # 并发处理所有行
    tasks = [process_row(row, await conn_pool.acquire()) for row in rows]
    await asyncio.gather(*tasks)

    await conn_pool.close()

if __name__ == "__main__":
    asyncio.run(main("your_input.csv"))

多线程方案（基于threading） #

import csv
import threading
import psycopg2
from queue import Queue

# 定义队列和锁
row_queue = Queue(maxsize=1000)
file_locks = {
    "exists": threading.Lock(),
    "c_not_exists": threading.Lock(),
    "i_not_exists": threading.Lock(),
    "exception": threading.Lock()
}

def worker():
    # 每个线程单独建立数据库连接
    conn = psycopg2.connect(**AUTHENTICATOR_PG_DB_CREDENTIALS)
    while True:
        row = row_queue.get()
        if row is None:
            break
        process_row_thread(row, conn)
        row_queue.task_done()
    conn.close()

def process_row_thread(row, conn):
    i_code, t_id, _, _, _ = row
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT customer_id
                FROM buckets
                WHERE i_code = %(i_code)s
                LIMIT 1
                """,
                {"i_code": i_code}
            )
            result = cur.fetchone()

            if result:
                cur.execute(
                    """
                    SELECT EXISTS (
                        SELECT 1
                        FROM customers
                        WHERE customer_id = %(customer_id)s
                        AND t_id = %(t_id)s
                    )
                    """,
                    {"customer_id": result[0], "t_id": t_id}
                )
                exists = cur.fetchone()[0]

                if exists:
                    with file_locks["exists"]:
                        with open("exists_data.csv", "a", newline='') as f:
                            writer = csv.writer(f)
                            writer.writerow(row)
                else:
                    with file_locks["c_not_exists"]:
                        with open("c_not_exists_data.csv", "a", newline='') as f:
                            writer = csv.writer(f)
                            writer.writerow(row)
            else:
                with file_locks["i_not_exists"]:
                    with open("i_not_exists_data.csv", "a", newline='') as f:
                        writer = csv.writer(f)
                        writer.writerow(row)
    except Exception as e:
        with file_locks["exception"]:
            with open("exceptions.csv", "a", newline='') as f:
                writer = csv.writer(f)
                writer.writerow(row + (str(e),))

def main(file_path):
    # 启动线程池，建议线程数等于数据库连接池最大连接数（如10-20）
    thread_count = 10
    threads = []
    for _ in range(thread_count):
        t = threading.Thread(target=worker)
        t.start()
        threads.append(t)

    # 读取CSV并放入队列
    with open(file_path, "r", newline='') as f:
        reader = csv.reader(f)
        next(reader)
        for row in reader:
            row_queue.put(row)

    # 等待队列处理完成
    row_queue.join()

    # 发送终止信号给线程
    for _ in range(thread_count):
        row_queue.put(None)
    for t in threads:
        t.join()

if __name__ == "__main__":
    main("your_input.csv")

内容的提问来源于stack exchange，提问作者Duck Dodgers