AWS Glue中pg8000插入DataFrame获取主键的执行异常问题

Spark采用惰性求值机制：所有RDD/DataFrame的转换操作（比如map）不会立即执行，只有调用show()、count()、write()这类行动操作时，才会触发整个计算链的执行。

你的代码里，rdd.map(lambda x:(*x, insertProcess(x)))属于转换操作，没有行动操作触发的话，insertProcess里的插入逻辑根本不会运行；而两次调用show()会触发两次计算，导致insertProcess被执行两次，自然重复插入。

另外，现有代码还有个严重问题：每行数据都创建一次数据库连接，这会极大消耗数据库资源，甚至触发连接数限制。

改用Spark结构化API的foreachBatch方法处理插入，同时优化数据库连接复用，并且确保主键正确回填。

优化后的代码示例 #

from pyspark.sql import functions as sf
import pg8000
from pg8000.native import Connection

# 复用数据库连接，每个批次创建一次连接
def batch_insert(batch_df, batch_id):
    # 获取批次数据的列（除了要回填的idAAA）
    cols = [col for col in batch_df.columns if col != "idAAA"]
    # 构造批量插入语句，用占位符避免SQL注入
    insert_sql = """
        INSERT INTO tabAAA (colBBB) 
        VALUES {} 
        RETURNING idAAA
    """.format(','.join(['(%s)' for _ in range(batch_df.count())]))
    
    # 提取要插入的数值，注意顺序要和表字段对应
    data = [row.asDict()["colBBB"] for row in batch_df.collect()]
    
    # 创建连接（每个批次一次）
    conn = Connection(
        user="...",
        host="...",
        database="...",
        password="...",
        ssl_context=True
    )
    
    try:
        # 执行插入并获取返回的主键
        results = conn.run(insert_sql, *data)
        # 提取主键列表
        ids = [res[0] for res in results]
        # 将主键和原批次DataFrame合并
        batch_with_id = batch_df.withColumn("idAAA", sf.array(*[sf.lit(id) for id in ids]).getItem(sf.monotonically_increasing_id() % batch_df.count()))
        # 这里可以将带主键的批次数据写入临时表或者返回，根据需求调整
        batch_with_id.createOrReplaceTempView(f"batch_{batch_id}")
        conn.commit()
    except Exception as e:
        conn.rollback()
        raise e
    finally:
        conn.close()

# 主逻辑
def insertPerQuery(myDataframe):
    # 先移除原有的idAAA列（如果存在）
    if "idAAA" in myDataframe.columns:
        myDataframe = myDataframe.drop("idAAA")
    
    # 用foreachBatch处理每个批次
    myDataframe.foreachBatch(batch_insert)
    
    # 从临时表中合并所有带主键的数据
    final_df = sqlContext.sql("SELECT * FROM (SELECT *, row_number() OVER (ORDER BY colBBB) AS rn FROM batch_*) t ORDER BY rn")
    return final_df

df = insertPerQuery(df)
# 这里调用一次行动操作即可触发插入，比如write或者show
df.write.mode("overwrite").saveAsTable("your_target_table")
# 或者 df.show()

关键优化点 #

• 避免重复计算：foreachBatch是行动操作的一种，只会触发一次插入逻辑；如果需要后续使用带主键的DataFrame，通过临时表合并结果，避免重复执行插入。
• 连接复用：每个批次创建一次数据库连接，而不是每行都创建，大幅提升性能。
• 批量插入：用批量插入语句替代单条插入，减少数据库IO开销。
• 主键回填：利用RETURNING语句直接获取插入后的主键，再通过monotonically_increasing_id()关联回原数据（如果需要严格顺序，可根据业务字段排序）。

• 如果你的DataFrame数据量很大，collect()可能会导致Driver内存溢出，此时可以将批次拆小，或者用分区插入的方式。
• 建议开启Glue作业的日志，方便排查插入过程中的异常。
• 确保PostgreSQL用户有足够的权限执行INSERT和RETURNING操作。

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