如何为PySpark DataFrame每行生成MySQL更新查询并批量执行

方案1：本地循环执行（适合小数据量场景） #

先把PySpark DataFrame的全量数据收集到Driver节点，转换为可遍历的行对象，再循环拼接SQL执行即可：

# 1. 拉取DataFrame全量数据到Driver端
row_list = df.collect()

# 2. 定义SQL更新模板，预留参数位
update_template = "UPDATE DB.TABLE SET DATE = '{}', VALUE = {} WHERE NAME = '{}'"

# 3. 逐行生成SQL并执行
for row in row_list:
    run_date = row["RUN_DATE"]
    name = row["NAME"]
    value = row["VALUE"]
    current_sql = update_template.format(run_date, value, name)
    # 调用自定义MySQL执行函数
    mysql_conn(host, user_name, password, current_sql)

注意事项 #

• 该方案仅适用于数据量较小的场景，collect()操作会把全量DataFrame数据加载到Driver内存，数据量过大时会触发OOM报错
• 上述SQL模板已经匹配你示例中的字段类型规则：DATE、NAME为字符串类型自动加单引号，VALUE为数值类型不加单引号
• 如果需要避免SQL注入风险，可以修改你的mysql_conn函数支持参数化查询，不用手动拼接SQL，参考逻辑如下：

# 参数化查询模板，不需要手动加引号
update_template = "UPDATE DB.TABLE SET DATE = %s, VALUE = %s WHERE NAME = %s"
for row in row_list:
    params = (row["RUN_DATE"], row["VALUE"], row["NAME"])
    # mysql_conn内部使用游标execute(sql, params)的方式执行即可
    mysql_conn(host, user_name, password, update_template, params)

大数据量优化方案 #

如果待更新的数据量较大，逐行执行性能极低，建议用批量更新方案：

1. 先把df通过PySpark JDBC接口写入MySQL的临时中间表
2. 执行一条关联更新SQL，用临时表数据批量更新目标表
3. 清理临时中间表

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