原文地址:https://github.com/danielbeach/testingArrowADBCdrivers
测试数据库插入性能
本项目比较了将 CSV 文件读取并插入数据库的不同方法,包括:
- PyArrow + ADBC - 使用 PyArrow 读取 CSV,并使用 ADBC 驱动插入 Postgres
- Polars + ADBC - 使用 Polars 读取 CSV,并使用 ADBC 驱动插入 Postgres
- psycopg2 + COPY - 使用 Python 的 csv 模块和 psycopg2 的 COPY 命令
- DuckDB + COPY - 使用 DuckDB 的原生 COPY 命令
前提条件
- Python 3.12+
- Docker 和 Docker Compose(用于本地 Postgres)
uv包管理器(或手动安装依赖项)
设置
- 安装依赖项:
bash
uv sync
或使用 pip:
bash
pip install -e .
启动本地 Postgres 数据库
使用 Docker Compose 启动本地 PostgreSQL 数据库:
bash
docker compose up -d
这将启动一个 Postgres 16 容器,配置如下:
- 主机:localhost
- 端口:5432
- 数据库:postgres
- 用户名:postgres
- 密码:postgres
容器将被命名为 adbc-postgres,数据将持久保存在 Docker 卷中。
检查 Postgres 状态
检查 Postgres 是否正在运行:
bash
docker-compose ps
查看日志:
bash
docker-compose logs postgres
停止 Postgres
停止 Postgres 容器:
bash
docker-compose down
停止并删除所有数据:
bash
docker-compose down -v
可用脚本
所有脚本从 data/uncompressed/ 目录读取 CSV 文件并将其插入数据库。每个脚本创建不同的表名,以便您可以比较结果。
1. PyArrow + ADBC (pyarrow_adbc_driver.py)
使用 PyArrow 的 dataset API 读取 CSV,并使用 ADBC 驱动插入 Postgres。
bash
python pyarrow_adbc_driver.py
表名 :divvy_tripdata
特点:
- 使用 PyArrow 高效的 dataset API 读取多个 CSV 文件
- 转换为 Arrow 表格式
- 使用 ADBC 的
adbc_ingest()进行批量插入 - 创建具有适当类型映射的表
2. Polars + ADBC (polars_adbc_driver.py)
使用 Polars 读取 CSV,并使用 Polars 内置的 write_database() 方法与 ADBC 引擎。
bash
python polars_adbc_driver.py
表名 :divvy_tripdata_polars
特点:
- 使用 Polars 快速读取 CSV
- 利用 Polars 原生的
write_database()方法 - 自动模式推断和表创建
- 底层使用 ADBC 引擎
3. psycopg2 + COPY (psycopg2_driver.py)
使用 Python 内置的 csv 模块和 psycopg2 的 COPY 命令(最快的 psycopg2 方法)。
bash
python psycopg2_driver.py
表名 :divvy_tripdata_psycopg2_noarrow
特点:
- 纯 Python 读取 CSV(无需 Arrow 依赖)
- 通过
copy_expert()使用 PostgreSQL 的原生 COPY 命令 - 基于样本数据的类型推断
- 针对批量插入进行优化
4. DuckDB + COPY (duckdb_copy_driver.py)
使用 DuckDB 的原生 COPY 命令直接加载 CSV 文件。
bash
python duckdb_copy_driver.py
表名 :divvy_tripdata_duckdb(位于 DuckDB 数据库文件中)
特点:
- 创建本地 DuckDB 数据库文件(
divvy_data.duckdb) - 使用 DuckDB 的
COPY FROM命令 - 自动模式推断
- 以 DuckDB 的列式格式存储数据
连接配置
使用默认连接(Docker Postgres)
所有 Postgres 脚本(DuckDB 除外)默认使用 Docker Postgres 设置:
bash
# 启动 Postgres
docker-compose up -d
# 运行任何脚本
python pyarrow_adbc_driver.py
python polars_adbc_driver.py
python psycopg2_driver.py
使用自定义连接
您可以使用环境变量覆盖 Postgres 连接参数:
bash
export PGHOST=your-host
export PGPORT=5432
export PGDATABASE=your-database
export PGUSER=your-user
export PGPASSWORD=your-password
python pyarrow_adbc_driver.py
性能比较
所有脚本都包含计时信息,显示:
- 总插入时间
- 每秒行数吞吐量
- 行数验证
这使您可以比较不同方法的性能:
- ADBC 方法(PyArrow/Polars) - 利用 Arrow 格式实现零拷贝传输
- psycopg2 COPY - 使用 PostgreSQL 的原生批量加载
- DuckDB COPY - 优化的列式数据库,具有原生 CSV 加载功能
运行所有脚本并比较计时结果,以了解哪种方法最适合您的用例!