Airflow运行PySpark SQL管道时业务指标采集方案问询

嘿，我刚好在项目里做过类似的Airflow+PySpark监控需求，给你分享下两种方案的具体实现思路和伪代码，两种各有适用场景～

这种方案适合做常态化的监控告警，比如用Grafana做仪表盘展示，或者设置阈值告警（比如错误记录数超过某个值就触发告警）。核心思路是在PySpark任务里计算完指标后，推送到Prometheus的Pushgateway，Prometheus会定期拉取这些指标。

伪代码实现 #

1. PySpark处理脚本（带指标推送）

from pyspark.sql import SparkSession
from prometheus_client import CollectorRegistry, Gauge, push_to_gateway

def run_spark_pipeline():
    # 初始化Spark会话
    spark = SparkSession.builder.appName("BusinessMetricsPipeline").getOrCreate()
    
    # 1. 统计源数据读取量
    source_df = spark.read.parquet("s3://your-source-bucket/path/")
    source_record_count = source_df.count()
    
    # 2. 数据处理&错误记录统计
    # 这里模拟业务逻辑：过滤出有效数据和错误数据
    valid_df = source_df.filter("is_valid = true")
    error_df = source_df.filter("is_valid = false")
    error_record_count = error_df.count()
    
    # 3. 统计处理中DataFrame大小（两种方式可选）
    # 方式一：估算（大数据量推荐，性能好）
    processed_df_estimated_size = valid_df.count() * 150  # 假设单条记录约150字节
    # 方式二：精确计算（小数据量用，性能开销大）
    # processed_df_exact_size = valid_df.rdd.map(lambda row: len(str(row))).sum()
    
    # 4. 写入目标数据&统计写入量
    valid_df.write.parquet("s3://your-target-bucket/path/", mode="overwrite")
    target_record_count = valid_df.count()  # 或者重新读取目标数据校验，按需选择
    
    # 推送指标到Prometheus Pushgateway
    registry = CollectorRegistry()
    # 定义指标，注意命名符合Prometheus规范
    g_source = Gauge(
        'spark_pipeline_source_records_total', 
        'Total number of records read from source', 
        registry=registry
    )
    g_target = Gauge(
        'spark_pipeline_target_records_total', 
        'Total number of records written to target', 
        registry=registry
    )
    g_processed_size = Gauge(
        'spark_pipeline_processed_df_size_bytes', 
        'Size of processed valid DataFrame', 
        registry=registry
    )
    g_error = Gauge(
        'spark_pipeline_error_records_total', 
        'Total number of error records during processing', 
        registry=registry
    )
    
    # 设置指标值
    g_source.set(source_record_count)
    g_target.set(target_record_count)
    g_processed_size.set(processed_df_estimated_size)
    g_error.set(error_record_count)
    
    # 推送到Pushgateway（替换成你的Pushgateway地址）
    push_to_gateway('prometheus-pushgateway:9091', job='airflow_spark_pipeline', registry=registry)
    
    spark.stop()

if __name__ == "__main__":
    run_spark_pipeline()

2. Airflow DAG配置

from airflow import DAG
from airflow.providers.apache.spark.operators.spark_submit import SparkSubmitOperator
from datetime import datetime, timedelta

default_args = {
    'owner': 'data_engineering',
    'depends_on_past': False,
    'start_date': datetime(2024, 1, 1),
    'retries': 1,
    'retry_delay': timedelta(minutes=5)
}

with DAG(
    'spark_business_metrics_dag',
    default_args=default_args,
    schedule_interval='@daily',
    catchup=False
) as dag:
    spark_processing_task = SparkSubmitOperator(
        task_id='run_spark_pipeline',
        application='/path/to/your/spark_script.py',
        conn_id='spark_default',  # 提前在Airflow配置Spark连接
        executor_cores=3,
        executor_memory='4g',
        driver_memory='2g'
    )

注意事项 #

• 要先部署Prometheus Pushgateway，确保Spark任务所在节点能访问到它
• 指标命名尽量规范，方便后续用Grafana做可视化或者设置告警规则

如果只是想在Airflow里查看单次DAG运行的指标，或者在后续任务中复用这些指标（比如发送告警邮件），用XCom最方便。XCom是Airflow自带的任务间数据传递机制，适合存储小体积的数值型指标。

伪代码实现 #

这里用PythonOperator封装Spark任务，直接在Airflow任务里计算并推送XCom：

from airflow import DAG
from airflow.operators.python import PythonOperator
from pyspark.sql import SparkSession
from datetime import datetime, timedelta

def run_spark_with_xcom(**context):
    spark = SparkSession.builder.appName("XComMetricsPipeline").getOrCreate()
    
    # 1. 源数据读取量
    source_df = spark.read.parquet("s3://your-source-bucket/path/")
    source_count = source_df.count()
    
    # 2. 错误记录数
    valid_df = source_df.filter("is_valid = true")
    error_count = source_df.filter("is_valid = false").count()
    
    # 3. 处理中DF大小
    processed_size = valid_df.count() * 150  # 估算方式
    
    # 4. 目标数据写入量
    valid_df.write.parquet("s3://your-target-bucket/path/", mode="overwrite")
    target_count = valid_df.count()
    
    spark.stop()
    
    # 封装指标并推送到XCom
    metrics = {
        "source_records": source_count,
        "target_records": target_count,
        "processed_df_size_bytes": processed_size,
        "error_records": error_count
    }
    # 用task instance推送XCom
    context["ti"].xcom_push(key="spark_business_metrics", value=metrics)
    return metrics

def print_metrics(**context):
    # 从XCom拉取指标并打印（可以改成发送邮件、写入数据库等）
    metrics = context["ti"].xcom_pull(
        key="spark_business_metrics",
        task_ids="run_spark_pipeline"
    )
    print(f"本次Spark任务指标：{metrics}")

default_args = {
    'owner': 'data_engineering',
    'start_date': datetime(2024, 1, 1),
    'retries': 1
}

with DAG(
    'spark_xcom_metrics_dag',
    default_args=default_args,
    schedule_interval='@daily',
    catchup=False
) as dag:
    spark_task = PythonOperator(
        task_id="run_spark_pipeline",
        python_callable=run_spark_with_xcom,
        provide_context=True  # 必须开启，才能获取task instance
    )
    
    # 可选：后续任务处理指标
    print_task = PythonOperator(
        task_id="print_metrics",
        python_callable=print_metrics,
        provide_context=True
    )
    
    spark_task >> print_task

注意事项 #

• XCom适合存储小数据，你的指标都是数值，完全没问题；如果是大体积数据不建议用
• 可以在Airflow UI的「XCom」页面直接查看这些指标，不用额外部署工具

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