PySpark（AWS Glue作业）插入PostgreSQL后获取主键并更新DataFrame

下面是几种比collect()循环更高效的方案，适合AWS Glue作业场景：

方案1：利用PostgreSQL RETURNING子句批量获取主键 #

PostgreSQL支持INSERT ... RETURNING语法，可以在插入数据的同时返回生成的主键，结合Spark的分区处理，避免单条插入的低效：

步骤： #

1. 给原DataFrame添加一个全局唯一的临时标识列，用于后续关联返回的主键与原数据；
2. 用mapPartitions按分区批量插入数据，每个分区建立一次数据库连接，执行带RETURNING的插入语句，收集返回的主键和临时标识；
3. 将返回的主键数据转为DataFrame，与原DataFrame通过临时标识关联，得到带主键的最终结果。

代码示例： #

from pyspark.sql import Row
from pyspark.sql.functions import monotonically_increasing_id
import psycopg2

# 1. 添加临时唯一标识列（全局唯一）
df_with_temp_id = df.withColumn("temp_unique_id", monotonically_increasing_id())

# 2. 定义分区插入函数
def insert_and_return_keys(partition):
    # 每个分区建立一次数据库连接，减少连接开销
    conn = psycopg2.connect(
        host="your-postgres-host",
        database="your-db-name",
        user="your-db-user",
        password="your-db-pass"
    )
    cur = conn.cursor()
    returned_rows = []

    for row in partition:
        # 替换为你的表名、字段列表，RETURNING主键和临时标识
        insert_sql = """
            INSERT INTO target_table (col1, col2, col3)
            VALUES (%s, %s, %s)
            RETURNING id, %s AS temp_unique_id
        """
        # 参数顺序对应字段+临时标识值
        cur.execute(insert_sql, (row.col1, row.col2, row.col3, row.temp_unique_id))
        # 获取插入后返回的主键和临时标识
        result = cur.fetchone()
        returned_rows.append(Row(id=result[0], temp_unique_id=result[1]))

    conn.commit()
    cur.close()
    conn.close()
    return returned_rows

# 执行分区插入，获取主键数据
returned_rdd = df_with_temp_id.rdd.mapPartitions(insert_and_return_keys)
returned_df = returned_rdd.toDF()

# 3. 关联原数据与主键，得到最终DataFrame
final_df = df_with_temp_id.join(returned_df, on="temp_unique_id").drop("temp_unique_id")

方案2：临时表批量插入+关联查询 #

适合数据量较大的场景，通过将原数据写入PostgreSQL临时表，再批量插入目标表并关联获取主键：

步骤： #

1. 将带临时标识的DataFrame写入PostgreSQL临时表；
2. 执行批量插入SQL，利用RETURNING或关联查询获取主键与临时标识的映射；
3. 读取映射结果，与原DataFrame关联得到最终数据。

代码示例： #

from pyspark.sql.functions import monotonically_increasing_id

# 1. 添加临时唯一标识并写入临时表
df_with_temp_id = df.withColumn("temp_unique_id", monotonically_increasing_id())

df_with_temp_id.write \
    .format("jdbc") \
    .option("url", "jdbc:postgresql://your-host:5432/your-db") \
    .option("dbtable", "temp_import_table") \
    .option("user", "your-db-user") \
    .option("password", "your-db-pass") \
    .option("driver", "org.postgresql.Driver") \
    .mode("overwrite") \
    .save()

# 2. 执行批量插入并读取主键映射
spark = sparkSession.getActiveSession()

# 构造插入并返回主键的SQL
insert_query = """
    INSERT INTO target_table (col1, col2, col3)
    SELECT col1, col2, col3 FROM temp_import_table
    RETURNING target_table.id, temp_import_table.temp_unique_id
"""

returned_df = spark.read \
    .format("jdbc") \
    .option("url", "jdbc:postgresql://your-host:5432/your-db") \
    .option("user", "your-db-user") \
    .option("password", "your-db-pass") \
    .option("driver", "org.postgresql.Driver") \
    .option("query", insert_query) \
    .load()

# 3. 关联得到最终DataFrame
final_df = df_with_temp_id.join(returned_df, on="temp_unique_id").drop("temp_unique_id")

# 清理临时表
spark.sql("DROP TABLE temp_import_table")

方案3：预先生成主键（适合UUID类型主键） #

如果目标表主键是UUID类型，可以直接在Spark端生成UUID，插入后无需回填：

代码示例： #

from pyspark.sql.functions import udf, StringType
import uuid

# 生成UUID的UDF
generate_uuid = udf(lambda: str(uuid.uuid4()), StringType())

# 直接添加主键列
df_with_id = df.withColumn("id", generate_uuid())

# 插入数据库，此时DataFrame已包含主键
df_with_id.write \
    .format("jdbc") \
    .option("url", "jdbc:postgresql://your-host:5432/your-db") \
    .option("dbtable", "target_table") \
    .option("user", "your-db-user") \
    .option("password", "your-db-pass") \
    .mode("append") \
    .save()

# df_with_id 即为带主键的DataFrame，可直接后续使用

注意事项： #

• 临时标识列推荐用monotonically_increasing_id()（全局唯一）或UUID，避免关联时出现冲突；
• 分区处理（mapPartitions）比单条插入（foreach）效率高很多，因为每个分区只建立一次数据库连接；
• AWS Glue作业中需确保PostgreSQL驱动已正确配置（可通过Glue作业依赖添加postgresql-jdbc包）。

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