如何使用foreachBatch统计微批次大小并保留原有Spark作业逻辑

核心实现逻辑 #

原有作业的Delta写入逻辑完全保留不变，通过foreachBatch钩子拿到每个微批次的数据集和批次ID，先执行原写入逻辑，再追加批次数据量统计、统计结果落盘的逻辑，既不改动核心数据传输链路，又能实现进度追踪。

该方案完全保留原作业的Exactly-Once语义、checkpoint偏移量管理逻辑，原有checkpoint可以直接复用，不需要重置消费进度。

改造后完整代码 #

from pyspark.sql import functions as F

def process_batch(batch_df, batch_id):
    # 原有核心写入逻辑，所有参数和改造前完全一致，无任何修改
    batch_df.write.format("delta").mode("append").save(s3_delta_path)

    # 新增进度追踪逻辑
    # 计算当前批次数据行数
    batch_row_num = batch_df.count()
    # 组装统计结果，包含批次ID、行数、处理时间
    track_df = batch_df.sparkSession.createDataFrame(
        [(batch_id, batch_row_num, F.current_timestamp())],
        schema="batch_id LONG, row_count LONG, process_time TIMESTAMP"
    )
    # 统计结果写入独立的存储路径，可根据实际需求替换为MySQL/Redis/普通Parquet路径等
    track_df.write.format("delta").mode("append").save(s3_progress_store_path)

# 流式作业启动逻辑，除替换为foreachBatch外，其余配置和原代码完全一致
stream_query = df.writeStream \
    .outputMode("append") \
    .option("checkpointLocation", s3_checkpoint_path) \
    .trigger(availableNow=True) \
    .foreachBatch(process_batch) \
    .start()

stream_query.awaitTermination()

关键说明 #

• 流启动代码中不需要再指定format("delta")，因为Delta写入逻辑已经下沉到每个批次的处理函数内，和原作业的写入行为、语义完全等价
• batch_df.count()开销极低，Spark会基于当前批次已加载到内存/本地磁盘的数据计算行数，不会重复拉取S3源数据；如果批次数据量极大，也可以替换为batch_df.rdd.countApprox(timeout=1000)做毫秒级近似统计，误差可控制在1%以内
• 统计结果的写入逻辑可以按需调整，不需要和原数据保持一致的格式/存储位置，只要保证写入幂等即可
• 不要调整process_batch中原有Delta写入的参数顺序、配置项，避免改变核心链路的写入行为

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