如何用DAG组织Spark数据处理？Dagster/Airflow适配难题求解

核心思路：避免直接传递Spark DataFrame #

Spark DataFrame是分布式内存中的数据结构，本身不适合序列化后跨任务传递（序列化成本极高，甚至直接失败）。正确的做法是传递数据的元信息（如存储路径、表名），而非DataFrame实例，让每个任务自行读取和处理数据。

针对Dagster的具体方案 #

1. 临时存储中转数据（手动管理）

   • 每个@op处理完DataFrame后，将其写入分布式存储（如HDFS、本地临时目录）的临时路径，或者写入Spark临时表/MySQL中间表
   • 输出该存储路径、表名等元信息给下一个@op，由下一个任务自行读取
   • 示例代码：

     from dagster import op, job
     from pyspark.sql import SparkSession
     
     @op
     def extract_from_mysql(spark: SparkSession):
         df = spark.read.format("jdbc").options(
             url="jdbc:mysql://host:port/db",
             dbtable="source_table",
             user="user",
             password="pass"
         ).load()
         # 写入临时Parquet文件
         temp_path = "/tmp/spark_temp/extract_output"
         df.write.mode("overwrite").parquet(temp_path)
         return temp_path
     
     @op
     def transform_data(spark: SparkSession, input_path: str):
         df = spark.read.parquet(input_path)
         transformed_df = df.filter(df["value"] > 0).withColumn("new_col", df["old_col"] * 2)
         temp_path = "/tmp/spark_temp/transform_output"
         transformed_df.write.mode("overwrite").parquet(temp_path)
         return temp_path
     
     @op
     def apply_ml_model(spark: SparkSession, input_path: str):
         df = spark.read.parquet(input_path)
         # 结合PyTorch执行机器学习逻辑
         result_df = df.withColumn("prediction", df["new_col"] + 1)
         result_df.write.mode("overwrite").jdbc(
             url="jdbc:mysql://host:port/db",
             dbtable="result_table",
             properties={"user": "user", "password": "pass"}
         )
     
     @job
     def spark_data_pipeline():
         extracted_path = extract_from_mysql()
         transformed_path = transform_data(extracted_path)
         apply_ml_model(transformed_path)
     

2. 利用Dagster Spark IO管理器（自动中转）

   • Dagster的dagster-spark库提供了SparkIOManager，可以自动处理Spark DataFrame的存储和读取，无需手动管理路径：

     from dagster import job, op
     from dagster_spark import spark_io_manager, SparkResource
     
     @op(required_resource_keys={"spark"})
     def extract_op(context):
         df = context.resources.spark.read.format("jdbc").options(
             url="jdbc:mysql://host:port/db",
             dbtable="source_table",
             user="user",
             password="pass"
         ).load()
         return df
     
     @op
     def transform_op(df):
         return df.filter(df["value"] > 0).withColumn("new_col", df["old_col"] * 2)
     
     @op(required_resource_keys={"spark"})
     def apply_ml_op(context, df):
         result_df = df.withColumn("prediction", df["new_col"] + 1)
         result_df.write.mode("overwrite").jdbc(
             url="jdbc:mysql://host:port/db",
             dbtable="result_table",
             properties={"user": "user", "password": "pass"}
         )
     
     @job(resource_defs={"spark": SparkResource(), "io_manager": spark_io_manager})
     def spark_pipeline():
         apply_ml_op(transform_op(extract_op()))
     

   • SparkIOManager会自动将每个op的输出DataFrame写入临时存储，下一个op读取时自动反序列化为Spark DataFrame，无需手动处理路径。

通用原则（适用于Airflow） #

• Airflow中同样避免直接传递DataFrame，用XCom传递存储路径/表名这类元信息
• 每个Task实例化自己的SparkSession，读取元信息对应的数据源进行处理
• 示例：

  from airflow import DAG
  from airflow.operators.python import PythonOperator
  from pyspark.sql import SparkSession
  import datetime
  
  def extract(**context):
      spark = SparkSession.builder.appName("AirflowSpark").getOrCreate()
      df = spark.read.jdbc(
          url="jdbc:mysql://host:port/db",
          table="source_table",
          properties={"user": "user", "password": "pass"}
      )
      temp_path = "/tmp/airflow_spark/extract_output"
      df.write.mode("overwrite").parquet(temp_path)
      context["task_instance"].xcom_push(key="input_path", value=temp_path)
  
  def transform(**context):
      spark = SparkSession.builder.appName("AirflowSpark").getOrCreate()
      input_path = context["task_instance"].xcom_pull(key="input_path", task_ids="extract_task")
      df = spark.read.parquet(input_path)
      transformed_df = df.filter(df["value"] > 0).withColumn("new_col", df["old_col"] * 2)
      temp_path = "/tmp/airflow_spark/transform_output"
      transformed_df.write.mode("overwrite").parquet(temp_path)
      context["task_instance"].xcom_push(key="transformed_path", value=temp_path)
  
  with DAG(
      "spark_airflow_pipeline",
      start_date=datetime.datetime(2024, 1, 1),
      schedule_interval="@daily"
  ) as dag:
      extract_task = PythonOperator(task_id="extract_task", python_callable=extract)
      transform_task = PythonOperator(task_id="transform_task", python_callable=transform)
      extract_task >> transform_task
  

关键注意事项 #

• 复用SparkSession：在Dagster中通过SparkResource共享Session，避免每个op重复创建；Airflow中可通过连接池或全局Session管理优化性能
• 临时存储清理：添加清理任务，在管道执行完成后删除临时文件，避免存储资源浪费
• 数据格式选择：优先使用Parquet、ORC等列存格式，序列化/反序列化效率远高于CSV，且保留完整Schema信息

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