如何在Apache Airflow中基于Operator实现TaskGroup级别的动态任务映射？

我明白你现在的需求：已经能在TaskGroup里给单个Operator做动态任务映射，但想把整个TaskGroup作为一个单元来做映射，而不是每个任务单独处理。我来给你梳理清楚怎么解决这个问题。

先回顾下你的现有场景 #

你已经有能正常工作的代码：用SQLExecuteQueryOperator动态获取id列表，然后在TaskGroup里给每个SQLCheckOperator和SQLExecuteQueryOperator单独调用expand做映射，比如：

# 动态获取id列表的前置任务
task_id = SQLExecuteQueryOperator(
    task_id='scope', 
    sql='''SELECT DISTINCT id FROM idtable''',
    handler=lambda x: [dict(zip(['id'], row)) for row in fetch_all_handler(x)]
)

# 带任务级别映射的TaskGroup
with TaskGroup('check_group') as check:
    t1 = SQLCheckOperator.partial(
        task_id='check1',
        sql='SELECT count(*)> 0 FROM data_table WHERE id = {{ task.parameters["id"] }}'
    ).expand(parameters=task_id.output)

    t2 = SQLExecuteQueryOperator.partial(
        task_id='update1',
        map_index_template="{{ task.parameters['id'] }}",
        sql='''UPDATE table SET date = NOW() WHERE id = {{ task.parameters["id"] }} '''
    ).expand(parameters=task_id.output)

这段代码没问题，但你想把整个TaskGroup作为映射单元，就像用@task_group装饰器配合普通@task的示例那样——但直接套Operator的时候报错了。

你尝试的错误方式分析 #

你之前写的带@task_group的代码之所以报错，是因为你把Operator当成了普通的@task来传参，这两者的工作逻辑不一样：普通@task是Python函数，可以直接接收参数，但Operator是Airflow的任务实例，不能直接把TaskGroup的映射参数当成函数参数传进去。

正确的实现方式 #

要在TaskGroup级别用Operator做动态映射，核心是把TaskGroup的映射参数传递给内部Operator的配置（比如params或者parameters），然后对整个TaskGroup调用expand。下面分两种常用方式给你示例：

方式一：用@task_group装饰器（推荐）

这种方式最清晰，把TaskGroup作为一个可复用的模板，每个映射实例对应一个id：

from airflow.utils.task_group import TaskGroup
from airflow.providers.common.sql.operators.sql import SQLCheckOperator, SQLExecuteQueryOperator

# 前置任务：获取需要映射的id列表
task_id = SQLExecuteQueryOperator(
    task_id='scope', 
    sql='''SELECT DISTINCT id FROM idtable''',
    handler=lambda x: [dict(zip(['id'], row)) for row in fetch_all_handler(x)]
)

@task_group(group_id="dynamic_check_update_group")
def dynamic_task_group(id_value):
    # 检查任务：直接使用TaskGroup传递的id_value作为参数
    check_task = SQLCheckOperator(
        task_id='verify_id_exists',
        sql='SELECT count(*) > 0 FROM data_table WHERE id = {{ params.id }}',
        params={"id": id_value}
    )

    # 更新任务：同样使用id_value
    update_task = SQLExecuteQueryOperator(
        task_id='mark_id_processed',
        sql='''UPDATE table SET date = NOW() WHERE id = {{ params.id }}''',
        params={"id": id_value},
        map_index_template="{{ params.id }}"  # 用来在UI里显示对应的id
    )

    # 设置任务依赖
    check_task >> update_task

# 关键：在TaskGroup级别调用expand，传入前置任务的输出
mapped_task_groups = dynamic_task_group.expand(id_value=task_id.output)

方式二：用类式TaskGroup的partial+expand

如果你更习惯用with TaskGroup(...)的写法，也可以这样做（Airflow 2.5+版本支持更好）：

# 先定义TaskGroup模板
def create_check_update_group():
    with TaskGroup(group_id="check_update_group") as tg:
        check = SQLCheckOperator(
            task_id='check_id',
            sql='SELECT count(*) > 0 FROM data_table WHERE id = {{ params.id }}'
        )

        update = SQLExecuteQueryOperator(
            task_id='update_id',
            sql='''UPDATE table SET date = NOW() WHERE id = {{ params.id }}''',
            map_index_template="{{ params.id }}"
        )

        check >> update
    return tg

# 前置任务还是之前的task_id
# 对整个TaskGroup做映射
mapped_tg = create_check_update_group().partial().expand(params=task_id.output)

关键注意点 #

1. Airflow版本要求：TaskGroup级别的动态映射在Airflow 2.3及以上版本才稳定支持，如果你用的是更早的版本，建议先升级。
2. 参数传递逻辑：当你对TaskGroup调用expand时，Airflow会为每个输入参数生成一个独立的TaskGroup实例，每个实例内部的Operator会拿到对应的参数值——不需要再给内部Operator单独调用expand（除非你在每个组内还要做更细粒度的映射）。
3. 模板变量的使用：在Operator的sql里，用{{ params.id }}来引用参数，而不是{{ task.parameters["id"] }}，因为这里的参数是直接传递给Operator的params字段，不是通过expand(parameters=...)传递的。

这样调整后，你就能实现把整个TaskGroup作为一个单元做动态映射，每个id对应一个包含检查和更新任务的独立组，和你想要的效果一致。

备注：内容来源于stack exchange，提问作者Petitbreton