使用foreachBatch的Structured Stream Writer不遵循shuffle.partitions参数

一、foreachBatch 中 shuffle 分区数不生效的问题 #

1. 确保自适应执行开启 #

spark.sql.shuffle.partitions="auto" 依赖自适应查询执行（AQE），未开启AQE时该配置会失效，默认使用200分区。需在作业初始化或foreachBatch函数内显式开启：

spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.shuffle.partitions", "auto")

同时可配合调整AQE分区合并参数，优化后分区大小：

# 目标分区大小，可根据集群资源调整（例如设为128MB）
spark.conf.set("spark.sql.adaptive.shuffle.targetPostShuffleInputSize", "134217728")
# 最小合并分区数，避免分区过少
spark.conf.set("spark.sql.adaptive.coalescePartitions.minPartitionNum", "10")

2. 在foreachBatch 内部显式设置分区 #

foreachBatch的每个微批是独立Spark作业，全局配置可能未被正确继承。建议在处理每个batch的逻辑中，显式设置分区参数，或手动控制分区数量：

def process_batch(df, batch_id):
    # 在batch处理逻辑内设置分区
    spark.conf.set("spark.sql.shuffle.partitions", "auto")
    # 或者根据数据量手动 repartition/coalesce
    deduplicated_df = df.dropDuplicates(["key_col"]).repartition(50)
    deduplicated_df.write.mode("append").saveAsTable("target_table")

streaming_df.writeStream.foreachBatch(process_batch).start()

3. 检查数据源写入参数 #

如果写入JDBC等特定数据源，需同时设置数据源自身的分区控制参数（如numPartitions），否则shuffle分区数不会影响最终写入的并行度：

deduplicated_df.write \
    .mode("append") \
    .format("jdbc") \
    .option("url", "jdbc:mysql://host:port/db") \
    .option("dbtable", "target_table") \
    .option("numPartitions", "50")  # 与shuffle分区数匹配
    .save()

二、升级PySpark 3.5.0 后作业运行缓慢的问题 #

1. 排查AQE 参数变化 #

PySpark 3.5.0 对AQE默认配置有调整，例如spark.sql.adaptive.coalescePartitions.enabled默认开启，但可能因数据分布不均导致分区合并未生效。可检查并调整以下参数：

# 强制开启分区合并
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
# 调整合并阈值，允许更小的分区被合并
spark.conf.set("spark.sql.adaptive.coalescePartitions.initialPartitionNum", "200")

2. 对比执行计划，优化去重逻辑 #

升级后Spark查询优化器可能生成不同执行计划，可通过explain()查看去重步骤的shuffle情况：

deduplicated_df.explain(mode="extended")

如果用窗口函数（如row_number()）去重，可尝试调整分区键，减少shuffle数据量：

from pyspark.sql.window import Window
from pyspark.sql.functions import row_number

window_spec = Window.partitionBy("key_col").orderBy("event_time")
# 提前按分区键 repartition，减少 shuffle 开销
deduplicated_df = df.repartition("key_col") \
                    .withColumn("rn", row_number().over(window_spec)) \
                    .filter("rn = 1") \
                    .drop("rn")

3. 优化Python 代码执行效率 #

PySpark 3.5.0 对Python UDF执行机制有更新，若作业使用自定义UDF，可尝试：

• 替换为Spark SQL内置函数，避免Python-JVM交互开销
• 改用Pandas UDF（pandas_udf）替代普通UDF，提升批量处理效率
• 检查UDF内部逻辑，删除冗余计算或IO操作

4. 调整集群资源配置 #

200个微批作业运行缓慢可能是资源不足导致任务排队，可调整以下参数：

# 增加executor数量与核数
spark.conf.set("spark.executor.instances", "20")
spark.conf.set("spark.executor.cores", "4")
# 调整executor内存，避免GC开销
spark.conf.set("spark.executor.memory", "8g")
spark.conf.set("spark.executor.memoryOverhead", "2g")

5. 检查状态存储健康度 #

Structured Streaming状态存储在版本升级后可能存在兼容性问题，或状态数据过大导致加载缓慢：

• 若业务允许，可尝试重置状态（通过trigger(availableNow=True)或删除状态存储目录）
• 调整状态清理参数，定期清理过期状态：

spark.conf.set("spark.sql.streaming.stateStore.cleanupDelay", "3600")  # 1小时后清理过期状态
spark.conf.set("spark.sql.streaming.stateStore.minDeltasToRetain", "10")  # 保留最近10个版本的状态

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