如何高效向 PostgreSQL 插入 30 万条数据？

井九

发布于 2024-12-21 20:10:06

18800

代码可运行

文章被收录于专栏：四楼没电梯四楼没电梯

运行总次数：0

代码可运行

现代应用中，海量数据的存储和处理是一个常见需求。假如你需要将 30 万条数据插入到 PostgreSQL 数据库，你会想到什么？直接写一个循环一条一条地插？还是使用批量插入？

在这篇文章中，我们将以不同的场景出发，结合代码实例，分析完成这一任务所需的时间和优化策略。

数据库准备

首先，我们需要准备一个 PostgreSQL 数据库实例。以下是简单的数据库初始化步骤：

CREATE DATABASE test_db;

\c test_db

CREATE TABLE test_table (
    id SERIAL PRIMARY KEY,
    name VARCHAR(50),
    value INTEGER
);

表结构非常简单，仅包含一个自增主键和两个字段。接下来，我们用 Python 代码插入数据。

场景一：单条插入

最简单的方法是逐条插入。我们使用 Python 和 psycopg2 库来实现。

代码实现

import psycopg2
from time import time

# 数据库连接
conn = psycopg2.connect(
    dbname="test_db", user="postgres", password="yourpassword", host="localhost"
)
cursor = conn.cursor()

start_time = time()

# 单条插入
for i in range(300000):
    cursor.execute(
        "INSERT INTO test_table (name, value) VALUES (%s, %s)",
        (f"name_{i}", i),
    )

conn.commit()
end_time = time()

print(f"单条插入完成，耗时：{end_time - start_time:.2f} 秒")

cursor.close()
conn.close()

测试结果

实际测试发现，这种方式在插入 30 万条数据时，耗时约为 800~1200 秒，主要瓶颈在于每次执行 execute 都会触发一次数据库的网络通信和日志记录。

场景二：批量插入

为了减少网络通信的开销，我们可以一次性插入多条记录。

代码实现

# 批量插入
batch_size = 1000
start_time = time()

data = [(f"name_{i}", i) for i in range(300000)]
for i in range(0, len(data), batch_size):
    batch = data[i : i + batch_size]
    psycopg2.extras.execute_values(
        cursor,
        "INSERT INTO test_table (name, value) VALUES %s",
        batch
    )

conn.commit()
end_time = time()

print(f"批量插入完成，耗时：{end_time - start_time:.2f} 秒")

测试结果

使用批量插入，每批次插入 1000 条数据，耗时缩短到 30~50 秒，性能提升显著。

场景三：COPY 命令

PostgreSQL 的 COPY 命令专为批量数据导入设计，效率极高。我们可以先生成一个数据文件，再使用 COPY 导入。

代码实现

# 生成 CSV 文件
with open("data.csv", "w") as f:
    for i in range(300000):
        f.write(f"name_{i},{i}\n")

# 使用 COPY 导入
start_time = time()
with open("data.csv", "r") as f:
    cursor.copy_expert("COPY test_table (name, value) FROM STDIN WITH CSV", f)

conn.commit()
end_time = time()

print(f"COPY 导入完成，耗时：{end_time - start_time:.2f} 秒")

测试结果

在测试中，COPY 命令耗时约为 5~10 秒，是最优的插入方法。

不同方法对比

方法	耗时（秒）	优缺点
单条插入	800~1200	实现简单，但性能差，适合少量数据。
批量插入	30~50	性能提升显著，适合大部分批量插入场景。
COPY 命令	5~10	性能最佳，但需额外生成文件。