TaskFlow模式下如何向第三方Airflow Operator传递XCom？

解决方案 #

核心原因 #

传统Operator与TaskFlow API的参数传递机制不同：TaskFlow任务返回的是PlainXComArg对象，而EmrServerlessCreateApplicationOperator的config参数期望直接接收字典类型，因此直接传递会触发类型错误，且无法自动建立依赖。

方法一：用TaskFlow包装传统Operator #

将传统Operator的调用逻辑包装在@task装饰的函数中，这样可以直接接收上游TaskFlow任务返回的字典，同时自动建立依赖关系：

from airflow.providers.amazon.aws.operators.emr import EmrServerlessCreateApplicationOperator
from airflow.decorators import task, task_group

@task_group(group_id="create_emr_serverless_app")
def create_emr_serverless_app_group():
    @task(task_id="network_configuration")
    def network_configuration() -> dict:
        return {
            "subnetIds": ["subnet-ABC"],
            "securityGroupIds": ["sg-12345"],
        }

    @task(task_id="maximum_capacity")
    def maximum_capacity() -> dict:
        return {
            "cpu": "20vCPU",
            "memory": "64GB",
        }

    @task(task_id="auto_stop_configuration")
    def auto_stop_configuration() -> dict:
        return {
            "enabled": True,
            "idleTimeoutMinutes": 15,
        }

    @task(task_id="config")
    def config(network_conf: dict, max_cap: dict, auto_stop_conf: dict) -> dict:
        return {
            "name": "my-emr-serverless-app",
            "networkConfiguration": network_conf,
            "maximumCapacity": max_cap,
            "autoStopConfiguration": auto_stop_conf,
        }

    @task(task_id="create_spark_app_on_emr_serverless")
    def create_spark_app(config_dict: dict):
        operator = EmrServerlessCreateApplicationOperator(
            task_id="inner_create_app",
            job_type="SPARK",
            release_label="emr-6.9.0",
            config=config_dict,
            executor_config=executor_config,
            aws_conn_id="",
        )
        # 手动执行Operator逻辑
        return operator.execute(context=None)

    # 构建任务依赖链
    net_conf = network_configuration()
    max_cap = maximum_capacity()
    auto_stop_conf = auto_stop_configuration()
    config_dict = config(net_conf, max_cap, auto_stop_conf)
    create_spark_app(config_dict)

方法二：使用Airflow模板语法拉取XCom #

利用Airflow的模板系统，在传统Operator的config参数中直接指定XCom拉取逻辑，同时显式建立任务依赖：

from airflow.providers.amazon.aws.operators.emr import EmrServerlessCreateApplicationOperator
from airflow.decorators import task, task_group

@task_group(group_id="create_emr_serverless_app")
def create_emr_serverless_app_group():
    @task(task_id="network_configuration")
    def network_configuration() -> dict:
        return {
            "subnetIds": ["subnet-ABC"],
            "securityGroupIds": ["sg-12345"],
        }

    @task(task_id="maximum_capacity")
    def maximum_capacity() -> dict:
        return {
            "cpu": "20vCPU",
            "memory": "64GB",
        }

    @task(task_id="auto_stop_configuration")
    def auto_stop_configuration() -> dict:
        return {
            "enabled": True,
            "idleTimeoutMinutes": 15,
        }

    @task(task_id="config")
    def config(network_conf: dict, max_cap: dict, auto_stop_conf: dict) -> dict:
        return {
            "name": "my-emr-serverless-app",
            "networkConfiguration": network_conf,
            "maximumCapacity": max_cap,
            "autoStopConfiguration": auto_stop_conf,
        }

    # 调用上游任务并保存引用
    net_conf = network_configuration()
    max_cap = maximum_capacity()
    auto_stop_conf = auto_stop_configuration()
    config_task = config(net_conf, max_cap, auto_stop_conf)

    # 使用模板语法从XCom拉取字典
    create_slo_app = EmrServerlessCreateApplicationOperator(
        task_id="create_spark_app_on_emr_serverless",
        job_type="SPARK",
        release_label="emr-6.9.0",
        config="{{ ti.xcom_pull(task_ids='create_emr_serverless_app.config') }}",
        executor_config=executor_config,
        aws_conn_id="",
    )

    # 显式设置依赖关系
    config_task >> create_slo_app

方法三：使用XComArg.resolve()（Airflow 2.3+） #

如果使用Airflow 2.3及以上版本，可以直接调用TaskFlow任务返回值的.resolve()方法，在运行时获取实际的字典值，同时显式建立依赖：

# 其他上游任务代码同方法二

# 调用上游任务并保存引用
net_conf = network_configuration()
max_cap = maximum_capacity()
auto_stop_conf = auto_stop_configuration()
config_task = config(net_conf, max_cap, auto_stop_conf)

# 使用resolve()获取实际字典值
create_slo_app = EmrServerlessCreateApplicationOperator(
    task_id="create_spark_app_on_emr_serverless",
    job_type="SPARK",
    release_label="emr-6.9.0",
    config=config_task.resolve(),
    executor_config=executor_config,
    aws_conn_id="",
)

# 显式设置依赖关系
config_task >> create_slo_app

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