Databricks流作业各目标表写入行数等批量指标收集方案咨询

针对Databricks流作业单源多Sink的行数统计需求，提供三种可行方案：

1. 手动统计+指标上报（最直接方案） #

在forEachBatch处理函数内，直接统计输入批次和各Sink数据集的行数，将指标上报到Databricks监控系统或存储到专门的指标表。

代码示例： #

from datetime import datetime

def process_batch(input_df, batch_id):
    # 缓存输入数据集，避免重复计算
    input_df.cache()
    
    # 统计输入行数
    input_row_count = input_df.count()
    # 上报输入行数到Databricks metrics
    dbutils.metrics.gauge("stream_input_total_rows", input_row_count, batch_id=batch_id)
    
    # 生成各目标Sink数据集
    sink_table1_df = input_df.filter("category = 'A'")
    sink_table2_df = input_df.filter("category = 'B'")
    
    # 统计各Sink数据集行数
    table1_row_count = sink_table1_df.count()
    table2_row_count = sink_table2_df.count()
    
    # 上报各Sink写入行数
    dbutils.metrics.gauge("stream_sink_table1_rows", table1_row_count, batch_id=batch_id)
    dbutils.metrics.gauge("stream_sink_table2_rows", table2_row_count, batch_id=batch_id)
    
    # 写入Delta表
    sink_table1_df.write.format("delta").mode("append").save("/dbfs/path/to/table1")
    sink_table2_df.write.format("delta").mode("append").save("/dbfs/path/to/table2")
    
    # 可选：将指标持久化到监控表，方便后续分析
    metrics_schema = "batch_id long, input_rows long, table1_rows long, table2_rows long, event_timestamp timestamp"
    metrics_df = spark.createDataFrame(
        [(batch_id, input_row_count, table1_row_count, table2_row_count, datetime.now())],
        schema=metrics_schema
    )
    metrics_df.write.format("delta").mode("append").save("/dbfs/path/to/stream_metrics_table")
    
    # 释放缓存
    input_df.unpersist()

优点：实现简单，实时性强，可直接在Databricks的Metrics页面查看各批次指标；
注意：使用cache()避免重复计算输入数据集，减少性能开销。

2. 利用Spark JobGroup+自定义Listener（自动化统计） #

通过为每个Sink的写入任务设置唯一JobGroup，结合自定义StreamingQueryListener捕获每个Job的输出行数，实现自动化统计。

步骤： #

1. 在forEachBatch中为每个Sink写入任务设置JobGroup：

def process_batch(input_df, batch_id):
    # 处理生成各Sink数据集...
    
    # 为table1写入任务设置JobGroup
    spark.sparkContext.setJobGroup(
        groupId=f"batch_{batch_id}_sink_table1",
        description=f"Batch {batch_id}: Write to table1"
    )
    sink_table1_df.write.format("delta").mode("append").save("/dbfs/path/to/table1")
    spark.sparkContext.clearJobGroup()
    
    # 为table2写入任务设置JobGroup
    spark.sparkContext.setJobGroup(
        groupId=f"batch_{batch_id}_sink_table2",
        description=f"Batch {batch_id}: Write to table2"
    )
    sink_table2_df.write.format("delta").mode("append").save("/dbfs/path/to/table2")
    spark.sparkContext.clearJobGroup()

2. 注册自定义StreamingQueryListener捕获Job指标：

from pyspark.sql.streaming import StreamingQueryListener

class SinkMetricsListener(StreamingQueryListener):
    def onQueryStarted(self, event):
        pass
    
    def onQueryProgress(self, event):
        batch_id = event.batchId
        # 遍历当前批次的所有Job
        for job in event.progress.jobs:
            job_group = job.jobGroup
            if job_group and "sink_table" in job_group:
                sink_name = job_group.split("_")[-1]
                output_rows = job.numOutputRows
                # 上报指标
                dbutils.metrics.gauge(f"stream_sink_{sink_name}_rows", output_rows, batch_id=batch_id)
    
    def onQueryTerminated(self, event):
        pass

# 注册Listener
spark.streams.addListener(SinkMetricsListener())

优点：无需手动统计行数，利用Spark内置的Job指标实现自动化；
注意：需要确保JobGroup命名规则唯一，方便Listener识别对应Sink。

3. 基于Delta Lake事务日志追溯（事后分析） #

利用Delta Lake的事务历史记录，查询每个批次写入各Sink表的行数，适合事后追溯或补全历史指标。

代码示例： #

def get_sink_batch_write_count(table_path, target_batch_id):
    # 查询Delta表的事务历史
    history_df = spark.sql(f"DESCRIBE HISTORY delta.`{table_path}`")
    # 过滤对应批次的写入操作（需在写入时添加batch_id元数据）
    batch_history = history_df.filter(
        f"operation = 'WRITE' AND userMetadata LIKE '%batch_id={target_batch_id}%'"
    ).select("numOutputRows")
    
    return batch_history.first()[0] if batch_history.count() > 0 else 0

# 写入时添加batch_id到userMetadata
sink_table1_df.write.format("delta").mode("append") \
    .option("userMetadata", f"batch_id={batch_id}") \
    .save("/dbfs/path/to/table1")

优点：无需修改处理逻辑，通过Delta内置功能获取准确的写入行数；
注意：必须在写入时通过userMetadata标记批次ID，否则无法关联到对应批次。

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