SQLAlchemy多进程处理PostgreSQL数据时重复行问题解决

问题根源分析 #

1. 无排序的row_number()导致批次拆分不稳定
   查询pairs批次时使用row_number() over ()但未指定ORDER BY子句，PostgreSQL中这种情况下行的顺序是不确定的（依赖物理存储和查询优化策略）。多进程并行查询时，不同进程的子查询可能返回不同的行顺序，导致多个进程拿到重复的索引对，最终处理后写入重复行。

2. 代码变量引用缺失
   do_something函数中，执行批次查询后未将结果赋值给变量就直接调用pairs.source_table_1_ind，属于未定义变量的错误（你说无报错可能是实际代码中漏写了pairs = pd.read_sql(query, engine)），即使补上这行，若批次拆分逻辑不稳定，依然会出现重复问题。

3. 多进程任务执行的惰性风险
   使用pool.imap时，若不主动迭代结果（比如转换为list），进程可能还未完成任务就执行了pool.close()和pool.join()，虽你提到行数匹配，但该问题会导致任务执行不完整的潜在风险。

具体修复步骤 #

1. 固定row_number()的排序规则 #

修改pairs批次查询的SQL，添加明确的ORDER BY，确保每次查询的行顺序一致，拆分的批次不会重叠：

query = f"""select source_table_1_ind, source_table_2_ind from (
        select pairs.*, row_number()
        over (ORDER BY source_table_1_ind, source_table_2_ind) as rnum
        from pairs) df
        where df.rnum between {lower} and {upper};"""

优先选择pairs表的主键或唯一标识列作为排序依据，保证rnum的生成结果完全固定。

2. 修正do_something中的变量错误 #

补全批次数据读取步骤，并优化查询逻辑减少重复数据读取：

# 读取当前批次的pairs数据
pairs_df = pd.read_sql(query, engine)

# 查询source_table_1，用unique()减少重复索引查询
source_1_filter = tuple(pairs_df.source_table_1_ind.unique())
df1_query = f"select * from source_table_1 where index in {source_1_filter};"
df1 = pd.read_sql(df1_query, engine)

# 查询source_table_2，同理优化
source_2_filter = tuple(pairs_df.source_table_2_ind.unique())
df2_query = f"select * from source_table_2 where index in {source_2_filter};"  
df2 = pd.read_sql(df2_query, engine)

3. 确保多进程任务完全执行 #

将pool.imap改为pool.map，或者主动迭代imap的结果，避免惰性执行导致的任务中断：

# 替换原imap代码，保证所有任务执行完成
pool.map(do_something, arg_list[1:])

4. 可选：数据库层面的重复防护 #

若担心进程间写入仍有重复，可给target_table的唯一索引列添加UNIQUE约束，或使用PostgreSQL的ON CONFLICT语法避免重复插入：

# 先写入临时表，再合并到目标表
output.to_sql("tmp_target_table", engine, schema="...", if_exists="replace", index=False, chunksize=10000)
with engine.connect() as conn:
    conn.execute("""
        INSERT INTO ...target_table
        SELECT * FROM tmp_target_table
        ON CONFLICT (your_unique_column) DO NOTHING;
    """)
    conn.commit()

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