Airflow分支逻辑不符合预期：如何实现二选一任务执行？

现有Airflow代码预期基于branch_decision()的结果，选择执行push_to_redshift_task或delete_from_redshift_task，但当前两个任务总是同时运行，需修改代码实现分支选择执行。原代码如下：

@task(task_id="list_files_task")
def list_files():
    return oldest_files

@task(task_id="transform_files_task")
def transform_file(source_s3_key: str):
    return ids

@task.branch(task_id="branch_decision_task")
def branch_decision(source_s3_key: str):
    if "delete" in source_s3_key:
        return 'delete_from_redshift_task'
    else:
        return 'push_to_redshift_task'

@task(task_id="push_to_redshift_task")
def push_to_redshift(source_s3_key: str):
    ...

@task(task_id="delete_from_redshift_task")
def delete_from_redshift(syncari_ids_to_delete):
    ...


list_files_task = list_files()
transform_files_task = transform_file.expand(source_s3_key=list_files_task)
branch_decision_task = branch_decision.expand(source_s3_key=list_files_task)
push_to_redshift_task = push_to_redshift.expand(source_s3_key=list_files_task)
delete_from_redshift_task = delete_from_redshift.expand(syncari_ids_to_delete=transform_files_task.output)
cleanup_files_task = cleanup_files.expand(source_s3_key=list_files_task)

list_files_task >> transform_files_task >> create_table_task >> branch_decision_task >> [push_to_redshift_task, delete_from_redshift_task] >> cleanup_files_task

核心问题是当前依赖将分支任务直接指向两个任务的列表，Airflow会默认触发所有任务。需通过分支映射的关联逻辑实现单任务选择执行，修改后的代码如下：

@task(task_id="list_files_task")
def list_files():
    return oldest_files

@task(task_id="transform_files_task")
def transform_file(source_s3_key: str):
    return ids

@task.branch(task_id="branch_decision_task")
def branch_decision(source_s3_key: str):
    if "delete" in source_s3_key:
        return 'delete_from_redshift_task'
    else:
        return 'push_to_redshift_task'

@task(task_id="push_to_redshift_task")
def push_to_redshift(source_s3_key: str):
    ...

@task(task_id="delete_from_redshift_task")
def delete_from_redshift(syncari_ids_to_delete):
    ...

@task(task_id="cleanup_files_task")
def cleanup_files(source_s3_key: str):
    ...


list_files_task = list_files()
transform_files_task = transform_file.expand(source_s3_key=list_files_task)
# 分支任务与文件一一映射
branch_decision_task = branch_decision.expand(source_s3_key=list_files_task)

# 生成映射任务实例
push_to_redshift_task = push_to_redshift.expand(source_s3_key=list_files_task)
delete_from_redshift_task = delete_from_redshift.expand(syncari_ids_to_delete=transform_files_task.output)
cleanup_files_task = cleanup_files.expand(source_s3_key=list_files_task)

# 调整依赖逻辑：分支任务分别关联两个目标任务
list_files_task >> transform_files_task >> create_table_task >> branch_decision_task
branch_decision_task >> push_to_redshift_task
branch_decision_task >> delete_from_redshift_task
# 清理任务等待所有分支触发的任务完成
[push_to_redshift_task, delete_from_redshift_task] >> cleanup_files_task

关键修改说明 #

1. 拆分依赖关系：不再将分支任务直接指向两个任务的列表，而是让分支任务分别与两个任务建立独立依赖。Airflow会根据分支任务的返回结果，自动跳过未被选中的任务实例
2. 映射任务的分支匹配：由于branch_decision_task是expand生成的映射任务，每个文件对应的分支实例会独立决策，对应的push或delete任务实例也会被精准触发/跳过
3. 利用Airflow分支机制：分支任务返回的任务ID会告诉Airflow仅执行该任务，未被指向的任务实例会被标记为skipped，不影响后续清理任务的执行

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