如何优化Databricks中Spark对PostgreSQL大表的增删改操作

云原生阿岚

2026-6-13

优化PostgreSQL百万级数据批量删改插的Databricks Spark方案

问题背景

将ETL流程从Pentaho迁移至Databricks Spark平台，需针对PostgreSQL生产环境中百万级数据量的advdotp.MART表，基于Spark DataFrame生成的mart_deleteDF（86万+条）、mart_insertDF（56万+条）、mart_updateDF（45万+条）执行删改插操作。当前逐行删除的实现性能极差，亟需优化。

核心优化方向及实现方案

1. 替换逐行删除为批量删除（临时表关联法）

逐行执行DELETE是性能瓶颈的核心原因，改用临时表关联删除能将单条删除的O(N)复杂度降到O(1)级别的批量操作：

步骤：

将mart_deleteDF写入PostgreSQL的临时表（利用Spark JDBC）
执行批量DELETE语句，通过临时表关联原表进行删除

代码示例：

# 1. 配置PostgreSQL连接参数
pg_url = f"jdbc:postgresql://{ScardPg.host}:{ScardPg.port}/{ScardPg.database}"
pg_properties = {
    "user": ScardPg.user,
    "password": ScardPg.password,
    "driver": "org.postgresql.Driver"
}

# 2. 将deleteDF写入PostgreSQL临时表（临时表会话结束自动销毁）
mart_deleteDF.write \
    .mode("overwrite") \
    .jdbc(pg_url, "temp_mart_delete", properties=pg_properties)

# 3. 执行批量删除SQL
delete_sql = """
DELETE FROM advdotp.MART m
USING temp_mart_delete t
WHERE m.ORG_ID = t.ORG_ID 
  AND m.METRC_ID = t.METRC_ID 
  AND m.METRC_MO_NUM = t.METRC_MO_NUM 
  AND m.RPT_MO_NUM = t.RPT_MO_NUM 
  AND m.ACTV_DLR_ONLY_IND = t.ACTV_DLR_ONLY_IND 
  AND m.PRDCT_CTGRY = t.PRDCT_CTGRY;
"""

# 通过Spark执行SQL（或者用psycopg2执行一次即可）
spark.read.jdbc(pg_url, f"({delete_sql}) as del_result", properties=pg_properties)

2. 利用Spark JDBC的批量操作处理删改插

避免将DataFrame数据collect()到Driver节点（易OOM且慢），直接用Spark的分布式JDBC能力批量处理：

批量插入示例：

mart_insertDF.write \
    .mode("append") \
    .option("batchsize", 10000)  # 调整批量大小，根据数据量和数据库性能优化
    .jdbc(pg_url, "advdotp.MART", properties=pg_properties)

批量更新示例（结合临时表+UPDATE JOIN）：

# 1. 将updateDF写入临时表
mart_updateDF.write \
    .mode("overwrite") \
    .jdbc(pg_url, "temp_mart_update", properties=pg_properties)

# 2. 执行批量更新SQL
update_sql = """
UPDATE advdotp.MART m
SET 
  -- 列出需要更新的字段，例如：
  col1 = t.col1,
  col2 = t.col2
FROM temp_mart_update t
WHERE m.ORG_ID = t.ORG_ID 
  AND m.METRC_ID = t.METRC_ID 
  AND m.METRC_MO_NUM = t.METRC_MO_NUM 
  AND m.RPT_MO_NUM = t.RPT_MO_NUM 
  AND m.ACTV_DLR_ONLY_IND = t.ACTV_DLR_ONLY_IND 
  AND m.PRDCT_CTGRY = t.PRDCT_CTGRY;
"""

spark.read.jdbc(pg_url, f"({update_sql}) as update_result", properties=pg_properties)

3. 合并操作：用PostgreSQL UPSERT简化插改逻辑

如果业务允许，可将mart_insertDF和mart_updateDF合并为一个DataFrame，利用PostgreSQL的INSERT ... ON CONFLICT实现UPSERT（插入或更新），减少两次操作的开销：

# 假设主键是ORG_ID,METRC_ID,METRC_MO_NUM,RPT_MO_NUM,ACTV_DLR_ONLY_IND,PRDCT_CTGRY
upsert_sql = """
INSERT INTO advdotp.MART (ORG_ID, METRC_ID, METRC_MO_NUM, RPT_MO_NUM, ACTV_DLR_ONLY_IND, PRDCT_CTGRY, col1, col2)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT (ORG_ID, METRC_ID, METRC_MO_NUM, RPT_MO_NUM, ACTV_DLR_ONLY_IND, PRDCT_CTGRY)
DO UPDATE SET 
  col1 = EXCLUDED.col1,
  col2 = EXCLUDED.col2;
"""

# 用Spark的foreachBatch执行批量UPSERT
def upsert_batch(df, batch_id):
    df.write \
        .mode("append") \
        .option("truncate", False) \
        .option("query", upsert_sql) \
        .jdbc(pg_url, "advdotp.MART", properties=pg_properties)

# 合并insert和update DF（需确保字段一致）
combined_df = mart_insertDF.unionByName(mart_updateDF)
combined_df.writeStream.foreachBatch(upsert_batch).start().awaitTermination()
# 如果是批处理，直接用foreachBatch的批处理方式或者自定义批量写入

4. 数据库端性能优化

索引优化：确保advdotp.MART表中用于删除/更新条件的字段（ORG_ID、METRC_ID等）存在联合索引，避免全表扫描：
```
CREATE INDEX idx_mart_delete_keys ON advdotp.MART(ORG_ID, METRC_ID, METRC_MO_NUM, RPT_MO_NUM, ACTV_DLR_ONLY_IND, PRDCT_CTGRY);
```
调整PostgreSQL参数：临时增大work_mem（用于排序、哈希操作的内存）、maintenance_work_mem，提升批量操作的效率：
```
SET work_mem = '64MB';
SET maintenance_work_mem = '256MB';
```
关闭自动提交：批量操作前关闭自动提交，完成后一次性提交（Spark JDBC默认会处理，但自定义操作时需注意）

5. 避免Driver节点瓶颈

绝对不要用collect()将大数据量DataFrame拉到Driver节点，这会导致内存溢出且完全丧失Spark的分布式优势

所有操作尽量通过Spark的分布式JDBC调用完成，或者用foreachPartition代替collect()+循环，在Executor节点上批量处理数据：

def batch_delete(partition):
    conn = psycopg2.connect(
        dbname=ScardPg.database,
        user=ScardPg.user,
        password=ScardPg.password,
        host=ScardPg.host,
        port=ScardPg.port,
    )
    cursor = conn.cursor()
    # 预编译删除语句
    delete_query = """
    DELETE FROM advdotp.MART 
    WHERE ORG_ID = %s AND METRC_ID = %s AND METRC_MO_NUM = %s AND
          RPT_MO_NUM = %s AND ACTV_DLR_ONLY_IND = %s AND PRDCT_CTGRY = %s;
    """
    # 批量执行，每1000条提交一次
    batch_size = 1000
    batch = []
    for row in partition:
        batch.append((row["ORG_ID"], row["METRC_ID"], row["METRC_MO_NUM"], row["RPT_MO_NUM"], row["ACTV_DLR_ONLY_IND"], row["PRDCT_CTGRY"]))
        if len(batch) >= batch_size:
            cursor.executemany(delete_query, batch)
            conn.commit()
            batch = []
    if batch:
        cursor.executemany(delete_query, batch)
        conn.commit()
    cursor.close()
    conn.close()

# 用foreachPartition分布式执行
mart_deleteDF.foreachPartition(batch_delete)

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

如何优化Databricks中Spark对PostgreSQL大表的增删改操作

问题背景 #

核心优化方向及实现方案 #

1. 替换逐行删除为批量删除（临时表关联法） #

步骤：