Airflow：如何获取BranchPythonOperator选中任务的Task ID以拉取XCom？

你可以通过以下几种方式获取分支选中的任务返回的文件名：

方法1：直接拉取两个任务的XCom，过滤有效结果 #

由于分支逻辑只会执行其中一个数据准备任务，未执行的任务不会生成XCom结果（或结果为None）。利用Jinja2过滤器可以直接筛选出有效结果：

修改ParametersConstructor的参数定义：

opr_parameters_constructor = ParametersConstructor(
    task_id='parameters_constructor',
    file='{{ [ti.xcom_pull(task_ids="data_preparation_normal"), ti.xcom_pull(task_ids="data_preparation_arrivals_change")] | select("defined") | list | first }}',
    initial_time='{{ dag_run.conf.get("initial_time") }}',
    final_time='{{ dag_run.conf.get("final_time") }}',
)

逻辑说明：

• 同时拉取两个数据准备任务的XCom结果
• select("defined")过滤掉未定义的结果（即未执行任务的返回值）
• 转成列表后取第一个元素，即为执行过的任务返回的文件名

方法2：让分支任务推送选中的Task ID到XCom #

修改分支函数，将选中的数据准备任务ID存入XCom，后续任务通过这个ID拉取对应结果：

第一步：修改分支函数 #

def choose_data_preparation_operator(**kwargs):
    ti = kwargs['ti']
    arrival_factor = float(kwargs.get("arrival_factor"))
    
    if arrival_factor != 1.0:
        selected_task = 'data_preparation_arrivals_change'
        ti.xcom_push(key='selected_data_prep_task', value=selected_task)
        return [selected_task, 'parameters_constructor']
    else:
        selected_task = 'data_preparation_normal'
        ti.xcom_push(key='selected_data_prep_task', value=selected_task)
        return [selected_task, 'parameters_constructor']

第二步：修改ParametersConstructor的参数 #

opr_parameters_constructor = ParametersConstructor(
    task_id='parameters_constructor',
    file='{% set task_id = ti.xcom_pull(task_ids="choose_data_preparation_path", key="selected_data_prep_task") %}{{ ti.xcom_pull(task_ids=task_id) }}',
    initial_time='{{ dag_run.conf.get("initial_time") }}',
    final_time='{{ dag_run.conf.get("final_time") }}',
)

逻辑说明：

• 先从分支任务的XCom中拉取选中的任务ID
• 再用这个ID拉取对应数据准备任务返回的文件名

方法3：在ParametersConstructor的执行逻辑中动态获取 #

如果允许修改ParametersConstructor的代码，可以在execute方法中直接处理：

class ParametersConstructor(BaseOperator):
    template_fields = ['initial_time', 'final_time']  # 移除file的模板字段，改为动态获取

    def __init__(self, initial_time, final_time, **kwargs):
        super().__init__(**kwargs)
        self.initial_time = initial_time
        self.final_time = final_time

    def execute(self, context):
        ti = context['ti']
        # 优先取arrivals任务的结果，为空则取normal任务的结果
        filename = ti.xcom_pull(task_ids='data_preparation_arrivals_change') or ti.xcom_pull(task_ids='data_preparation_normal')
        
        # 后续使用filename进行参数构造逻辑
        ...

这种方式无需在参数中写复杂的Jinja2模板，直接在Operator内部处理XCom拉取逻辑。

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