TimescaleDB使用pandas批量插入速度慢的优化问题咨询

当前的写入操作功能逻辑正常，但未针对2亿条级别的批量写入场景做适配优化，直接按现有方案执行会遇到严重的性能瓶颈：

• 默认pandas.to_sql基于psycopg2的普通insert实现，未使用PostgreSQL最高效的COPY写入能力，写入效率仅为COPY方案的1/5~1/3
• 插入前提前创建了二级索引idx_symbol，插入过程中每一行数据都需要同步更新索引，2亿条数据场景下该部分开销会占总写入耗时的60%以上
• 未对写入数据按时间戳排序，TimescaleDB超表按时间分片的特性会导致乱序写入产生大量随机IO，额外拉高写入开销
• 未调整数据库写入相关配置，默认配置下WAL checkpoint频繁触发，会拖慢整体写入速度

分Python端、数据库端两个维度优化，优化后整体写入速度可提升5~15倍：

一、Python端写入逻辑优化 #

优先使用COPY FROM替代默认to_sql写入，这是PostgreSQL生态下批量写入性能最高的实现方式，示例代码如下：

import io
from sqlalchemy import create_engine

engine = create_engine(f'postgresql+psycopg2://{config.DB_USER}:{config.DB_PASS}@{config.DB_HOST}:{config.DB_PORT}/{config.DB_NAME}')

def batch_copy_write(df, table_name, chunk_size=200000):
    # 按时间戳排序，适配TimescaleDB分片写入逻辑降低IO开销
    df_sorted = df.sort_values('timestamp', ignore_index=True)
    conn = engine.raw_connection()
    cur = conn.cursor()
    # 分批写入避免内存溢出
    for i in range(0, len(df_sorted), chunk_size):
        chunk = df_sorted.iloc[i:i+chunk_size]
        output = io.StringIO()
        chunk.to_csv(output, sep='\t', header=False, index=False, na_rep='\\N')
        output.seek(0)
        cur.copy_from(output, table_name, null='\\N', columns=chunk.columns.tolist())
        conn.commit()
    cur.close()
    conn.close()

# 调用写入
batch_copy_write(df_downloaded_grand, target_table_name)

如果继续使用to_sql，需添加以下参数提升性能：

df_downloaded_grand.to_sql(
    target_table_name, 
    engine, 
    if_exists="append",
    index=False,
    chunksize=10000,
    method='multi'
)

二、数据库端配置调整 #

写入前操作 #

• 提前删除二级索引idx_symbol，等所有数据全部写入完成后再重建索引，可大幅降低插入阶段的额外开销
• 临时调整数据库配置（写入完成后可改回原有配置）：

  -- 仅非生产归档场景可调整WAL级别，生产环境可跳过该配置
  ALTER SYSTEM SET wal_level = minimal;
  ALTER SYSTEM SET wal_buffers = '1GB';
  ALTER SYSTEM SET maintenance_work_mem = '8GB';
  ALTER SYSTEM SET checkpoint_timeout = '60min';
  ALTER SYSTEM SET max_wal_size = '20GB';
  -- 配置生效
  SELECT pg_reload_conf();
  

• 可根据数据的时间跨度调整超表的chunk大小，比如时间跨度超过1年可将chunk大小调整为30天，减少分片数量：

  SELECT set_chunk_time_interval('db_009a005a_df_downloaded_grand', interval '30 days');
  

写入后操作 #

• 重建二级索引idx_symbol
• 恢复数据库原有配置，执行VACUUM ANALYZE更新表统计信息

三、额外可选优化 #

如果是一次性全量导入2亿条数据，可直接使用TimescaleDB官方提供的timescaledb-parallel-copy工具，将数据导出为CSV后用8并发进程写入，写入速度还可再提升2~3倍。

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