如何以Airflow标准方式向SparkSubmitOperator传递参数？

在Airflow中，推荐通过以下几种符合规范的方式向SparkSubmitOperator传递输入输出等参数，实现配置与业务代码的分离，提升可维护性：

1. 使用application_args传递命令行参数 #

这是最直接的方式，把输入输出路径等参数作为命令行参数传给Spark脚本，脚本内部通过sys.argv接收。

DAG配置示例： #

spark_clean_store_data = SparkSubmitOperator(
    task_id="my_task_id",
    application="/path/to/my/dags/scripts/clean_store_data.py",
    conn_id="spark_conn",
    # 按顺序传递参数
    application_args=[
        "/path/to/input/store_data",
        "/path/to/output/cleaned_store_data"
    ],
    dag=dag,
)

对应Spark脚本（clean_store_data.py）： #

import sys
from pyspark.sql import SparkSession

if __name__ == "__main__":
    spark = SparkSession.builder.appName("CleanStoreData").getOrCreate()
    
    # 接收命令行参数
    input_path = sys.argv[1]
    output_path = sys.argv[2]
    
    # 业务处理逻辑
    df = spark.read.parquet(input_path)
    cleaned_df = df.filter(df["status"] == "valid")
    cleaned_df.write.mode("overwrite").parquet(output_path)
    
    spark.stop()

2. 结合Airflow Variables/Connections管理配置 #

如果参数需要频繁修改或统一管理，推荐将参数存储在Airflow的Variables或Connections中，避免硬编码在DAG里。

步骤： #

1. 在Airflow UI的「Admin」→「Variables」中创建变量，比如store_input_path和store_output_path
2. DAG中读取变量并传递：

from airflow.models import Variable

# 从Airflow Variables获取配置
input_path = Variable.get("store_input_path")
output_path = Variable.get("store_output_path")

spark_clean_store_data = SparkSubmitOperator(
    task_id="my_task_id",
    application="/path/to/my/dags/scripts/clean_store_data.py",
    conn_id="spark_conn",
    application_args=[input_path, output_path],
    dag=dag,
)

3. 通过XCom传递动态生成的参数 #

如果参数是上游任务动态生成的（比如临时路径、计算结果），可以用Airflow的XCom机制传递。

示例： #

from airflow.operators.python import PythonOperator

def generate_output_path(**context):
    # 模拟上游任务生成动态输出路径
    output_path = f"/path/to/dynamic/output/{context['execution_date'].strftime('%Y%m%d')}"
    # 将结果推送到XCom
    context["ti"].xcom_push(key="dynamic_output_path", value=output_path)
    return output_path

# 上游任务：生成动态路径
generate_path_task = PythonOperator(
    task_id="generate_output_path",
    python_callable=generate_output_path,
    provide_context=True,
    dag=dag,
)

# Spark任务：通过模板语法拉取XCom中的参数
spark_clean_store_data = SparkSubmitOperator(
    task_id="my_task_id",
    application="/path/to/my/dags/scripts/clean_store_data.py",
    conn_id="spark_conn",
    application_args=[
        "/path/to/input/store_data",
        "{{ ti.xcom_pull(task_ids='generate_output_path', key='dynamic_output_path') }}"
    ],
    dag=dag,
)

# 设置任务依赖
generate_path_task >> spark_clean_store_data

以上几种方式都遵循Airflow的最佳实践，将配置与业务代码解耦，便于后续的维护、修改和扩展。

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