Airflow同一DAG内如何基于前序任务结果创建动态数量的任务

核心问题说明 #

Airflow的DAG拓扑结构是DAG文件解析阶段就固定完成的，而XCom是任务运行时才产生的数据，你原来的写法中，解析阶段无法读取后续运行时才生成的分片数XCom值，自然无法动态生成对应数量的任务。

方案1（推荐）：使用Airflow 2.3+ 原生动态任务映射 #

这是Airflow官方针对运行时动态生成任务场景推出的原生能力，无需提前固定任务数量，完全适配你的需求。

修改后代码示例 #

1. 修正读取分片函数

import pandas as pd # 原代码聚合逻辑用到pandas，需要提前导入

def read_csv(s3_dir,client_s3,**kwargs):
    ti = kwargs['ti'] # 原代码漏了获取ti实例，会报错
    df_chunks = wr.s3.read_csv(s3_dir, sep='|', chunksize = 100,boto3_session = client_s3)
    intermediate_dir = Variable.get("s3_intermediate_dir")
    chunk_indexes = []
    for idx, df in enumerate(df_chunks):
        wr.s3.to_csv(df, f"{intermediate_dir}intermediate_{idx}")
        chunk_indexes.append(idx)
    # 推送分片索引列表到XCom
    ti.xcom_push(key='chunk_indexes', value=chunk_indexes)

2. 改造转换函数，接收分片索引参数

def transform_dataframes(chunk_idx, s3_dir, client_s3,**kwargs):
    intermediate_dir = Variable.get("s3_intermediate_dir")
    current_chunk_path = f"{intermediate_dir}intermediate_{chunk_idx}"
    # 此处写入对应分片的处理逻辑

3. 修正DAG任务定义

# 修正原DAG定义中拼写错误的concurrency参数
dag = DAG(
    dag_id='test_parallelism_dag',
    default_args=default_args,
    schedule_interval='@hourly',
    catchup=False,
    max_active_runs=3,
    concurrency = 10,
)

read_csv_task = PythonOperator(
    task_id='read_csv_func',
    python_callable=read_csv,
    provide_context=True,
    # 修正原代码中op_kwargs的括号错误
    op_kwargs={'s3_dir':Variable.get("init_s3_dir"),'client_s3':client_s3},
    dag=dag
)

# 使用partial固定公共参数，expand动态映射分片索引参数
transform_dataframes_task = PythonOperator.partial(
    task_id='transform_dataframe_chunk',
    python_callable=transform_dataframes,
    provide_context=True,
    op_kwargs={'s3_dir':Variable.get("init_s3_dir"),'client_s3':client_s3},
    dag=dag
).expand(chunk_idx=read_csv_task.output['chunk_indexes'])

# 聚合任务仅修正参数传递，逻辑不变
agg_df_one_task = PythonOperator(
    task_id='agg_df_one_func',
    python_callable=agg_df_one,
    provide_context=True,
    trigger_rule = 'none_failed',
    op_kwargs={'s3_dir':Variable.get("s3_intermediate_dir"),'client_s3':client_s3},
    dag=dag
)

read_csv_task >> transform_dataframes_task >> agg_df_one_task

动态映射会自动根据前序任务返回的索引列表长度生成对应数量的转换任务，所有任务自动并行执行，全部成功后才会触发下游聚合任务，完全匹配你的业务逻辑。

方案2：Airflow 2.3以下版本兼容方案 #

如果无法升级Airflow版本，可以通过预定义最大分片数，运行时判断是否执行的方式实现：

# 预估值，设置为业务场景中可能出现的最大分片数
MAX_ALLOWED_CHUNKS = 20

transform_tasks = []
for i in range(MAX_ALLOWED_CHUNKS):
    def run_transform(chunk_idx=i, **kwargs):
        ti = kwargs['ti']
        real_chunk_count = ti.xcom_pull(key='number_of_chunks', task_ids='read_csv_func')
        # 索引超过实际分片数直接跳过
        if chunk_idx >= real_chunk_count:
            return
        # 执行正常转换逻辑
        transform_dataframes(chunk_idx, s3_dir=Variable.get("init_s3_dir"), client_s3=client_s3)
    task = PythonOperator(
        task_id=f'transform_dataframe_chunk_{i}',
        python_callable=run_transform,
        provide_context=True,
        dag=dag
    )
    transform_tasks.append(task)

# 依赖关系设置不变
read_csv_task >> transform_tasks >> agg_df_one_task

该方案会固定生成MAX_ALLOWED_CHUNKS个任务，超出实际分片数的任务会空跑，仅作为低版本兼容方案使用。

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