如何结合BranchPythonOperator在Airflow中按需并行执行任务

你遇到的问题核心在于原代码中的result函数逻辑会提前返回单个任务ID，而BranchPythonOperator在Airflow 1.10及以上版本其实支持返回任务ID列表，以此触发多个任务并行执行。咱们一步步来解决这个问题：

问题根源分析 #

你的result函数在遍历文件时，遇到第一个符合条件的文件就直接return了——比如先找到MEM开头的文件就返回mem_script，后面的FMS文件根本不会被处理。这就是为什么只能执行其中一个任务的核心原因。

解决方案 #

我们需要修改result函数，先收集所有符合条件的任务ID，再一次性返回；同时确保BranchPythonOperator能正确触发多个并行任务。

步骤1：修正result函数的逻辑与语法 #

先修复原函数的语法错误（比如elif的位置错误），再调整逻辑为收集所有需要执行的任务ID：

def result():
    # 初始化需要执行的任务列表
    tasks_to_run = []
    receipt_files = os.listdir(receiptPath)
    
    if not receipt_files:
        print('No script to launch')
        return "no_script"
    
    for file in receipt_files:
        if file.startswith('MEM') and file.endswith('.csv'):
            tasks_to_run.append('mem_script')
            print('Launching script for: '+file)
        elif file.startswith('FMS') and file.endswith('.csv'):
            tasks_to_run.append('fms_script')
            print('Launching script for: '+file)
    
    # 去重（避免多个同类型文件重复触发同一任务，按需取舍）
    tasks_to_run = list(set(tasks_to_run))
    
    # 如果有任务需要执行则返回列表，否则返回no_script
    return tasks_to_run if tasks_to_run else "no_script"

步骤2：确认BranchPythonOperator的并行触发能力 #

Airflow的BranchPythonOperator支持返回单个任务ID字符串，或者任务ID的列表。当返回列表时，所有对应的任务都会被并行触发，完全匹配你的需求。

步骤3：简化任务依赖写法（可选但更直观） #

可以用箭头语法替代原有的set_upstream写法，让依赖关系更清晰：

# 替代原有的set_upstream链式调用
file_sensor >> onlyCsvFiles
onlyCsvFiles >> [move_good_file, move_bad_file]
[move_good_file, move_bad_file] >> result_mv
result_mv >> [run_Mem_Script, run_Fms_Script, skip_script]
[run_Mem_Script, run_Fms_Script, skip_script] >> rerun_dag

最终效果 #

替换修改后的result函数并更新依赖后，你的DAG会实现：

• 仅存在MEM文件时，单独执行mem_script
• 仅存在FMS文件时，单独执行fms_script
• 两者同时存在时，并行执行mem_script和fms_script
• 无符合条件文件时，执行no_script

额外注意事项 #

• 确保你的Airflow版本在1.10及以上，低版本的BranchPythonOperator不支持返回任务列表
• 如果业务需要每个文件单独触发一次对应任务，可以去掉代码中的去重逻辑

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