Apache Airflow 系列教程 | 第12课：动态任务映射（Dynamic Task Mapping）

导读（Introduction）

欢迎来到 Apache Airflow 源码深度解析系列的第十二课。

在传统的工作流编排中，DAG 的拓扑结构在编写代码时就已经完全确定——每个任务的数量、名称、依赖关系都是静态的。然而在真实的数据工程场景中，我们经常遇到这样的需求：需要处理的文件数量每天不同、需要并行处理的分区数量取决于上游任务的输出、需要对动态列表中的每一项执行相同的操作。这就是动态任务映射（Dynamic Task Mapping）要解决的核心问题。

动态任务映射是 Airflow 2.3 引入的一个革命性特性（AIP-42），它允许一个任务在运行时根据数据动态生成多个并行的任务实例。这打破了"DAG 结构必须在解析时确定"的传统限制，使得 Airflow 能够优雅地处理可变规模的并行计算场景。

在 Airflow 3.x 的架构中，动态任务映射的实现被精心分布在task-sdk（用户侧 API 和定义模型）与airflow-core（调度器侧扩展逻辑和数据库模型）之间。本课将深入剖析这一机制的完整实现——从用户调用expand()的那一刻起，到调度器动态创建 TaskInstance 的全过程。

---

学习目标（Learning Objectives）

完成本课学习后，你将能够：

1. 理解动态任务映射的核心概念——掌握运行时决定任务数量的设计哲学
2. 深入expand()与partial()的工作原理——理解从 OperatorPartial 到 MappedOperator 的创建流程
3. 掌握 ExpandInput 模型体系——区分 DictOfListsExpandInput 与 ListOfDictsExpandInput 的使用场景
4. 理解映射任务的状态管理——TaskMap 模型如何跟踪映射长度和调度扩展
5. 分析调度器的扩展机制——如何从 MappedOperator 动态创建 N 个 TaskInstance
6. 运用实际例子——使用 TaskFlow API 和经典 Operator 构建动态并行管道

---

正文内容（Main Content）

1. 动态任务映射的核心概念

1.1 问题的本质：静态 DAG vs 动态需求

在 Airflow 的设计哲学中，DAG 是一个声明式的工作流描述——你在 Python 文件中定义任务和依赖关系，调度器在解析时生成 DAG 的拓扑结构。这种设计使得 DAG 可以被可视化、序列化和调度，但也带来了一个根本性限制：

# 静态方式：任务数量在编写代码时就固定了withDAG("static_etl"):process_file_1=process_file(file="data_1.csv")process_file_2=process_file(file="data_2.csv")process_file_3=process_file(file="data_3.csv")# 如果明天有 100 个文件呢？

这种方式的弊端是显而易见的：

问题	说明
硬编码数量	任务数量必须在代码中固定，无法适应变化
循环生成的局限	虽然可以用 for 循环生成任务，但循环的迭代对象也必须在解析时确定
上游依赖的数据	如果要处理的列表来自上游任务的输出，解析时根本无法获取
资源浪费	为了应对峰值，可能要预定义大量可能永远不会执行的任务

1.2 动态任务映射的解决方案

动态任务映射引入了一种新的范式：任务的数量可以在运行时由数据驱动地决定。

@taskdefget_files():"""运行时才知道有多少文件"""return["file_1.csv","file_2.csv",...,"file_N.csv"]@taskdefprocess_file(file_path:str):"""对每个文件执行相同的处理逻辑"""...files=get_files()process_file.expand(file_path=files)# 运行时生成 N 个并行任务实例

这里的关键洞察是：

• 解析时：DAG 中只有一个process_file节点（类型为MappedOperator），它的任务实例数量是未知的
• 运行时：当get_files完成并产生输出后，调度器才知道需要创建多少个process_file的任务实例
• 扩展时：调度器为每个映射值创建独立的 TaskInstance（map_index = 0, 1, 2, …）

1.3 Map-Reduce 模式在 Airflow 中的体现

动态任务映射天然地支持 Map-Reduce 模式：

[生成数据] → [Map: 并行处理每一项] → [Reduce: 聚合所有结果] get_files process_file.expand() aggregate_results

• Map 阶段：expand()将一个列表"展开"为多个并行任务实例
• Reduce 阶段：下游任务接收所有映射任务实例的输出列表

---

2.expand()与partial()的工作原理

2.1 API 设计：两步法创建映射任务

动态任务映射的 API 设计采用了"两步法"——先用partial()固定不变的参数，再用expand()指定需要映射的参数：

# TaskFlow API 风格@taskdefprocess(file_path:str,config:dict):...# file_path 是映射维度，config 对所有实例相同process.partial(config={"mode":"full"}).expand(file_path=["a.csv","b.csv","c.csv"])

# 经典 Operator 风格AddOneOperator.partial(task_id="add_one").expand(value=[1,2,3])

这种设计的精妙之处在于：partial()返回的不是最终的 Operator 实例，而是一个中间状态对象OperatorPartial，它"记住"了固定参数，等待expand()调用来完成创建。

2.2 OperatorPartial：中间状态对象

让我们深入OperatorPartial的实现（源码位于task-sdk/src/airflow/sdk/definitions/mappedoperator.py）：

@attrs.define(kw_only=True)classOperatorPartial:"""An "intermediate state" returned by BaseOperator.partial()."""operator_class:type[BaseOperator]kwargs:dict[str,Any]params:ParamsDict|dict_expand_called:bool=attrs.field(init=False,default=False)

OperatorPartial有几个关键设计点：

1. 防止遗忘 expand() 的安全守卫

def__del__(self):ifnotself._expand_called:try:task_id=repr(self.kwargs["task_id"])exceptKeyError:task_id=Nonewarnings.warn(f"Task{task_id}was created with `auto_partial` or `partial()` ""but never had `expand()` or `expand_kwargs()` called on it",UserWarning,stacklevel=1,)

这是一个非常贴心的用户体验设计：如果开发者调用了partial()但忘记调用expand()，对象被垃圾回收时会发出警告。这避免了一个常见的 bug——任务被"声明"了但从未真正添加到 DAG 中。

2. expand() 方法的验证流程

defexpand(self,**mapped_kwargs:OperatorExpandArgument)->MappedOperator:validate_mapping_kwargs(self.operator_class,"expand",mapped_kwargs)prevent_duplicates(self.kwargs,mapped_kwargs,fail_reason="unmappable or already specified")returnself._expand(DictOfListsExpandInput(mapped_kwargs),strict=False)

expand()在创建MappedOperator之前会执行严格的验证：

• validate_mapping_kwargs：遍历 Operator 的 MRO 链，确认所有参数名合法且值类型可映射
• prevent_duplicates：确保partial()和expand()之间没有参数重叠

3. expand_kwargs() 方法

defexpand_kwargs(self,kwargs:OperatorExpandKwargsArgument,*,strict:bool=True)->MappedOperator:ifisinstance(kwargs,Sequence):foriteminkwargs:ifnotisinstance(item,(XComArg,Mapping)):raiseTypeError(f"expected XComArg or list[dict], not{type(kwargs).__name__}")elifnotisinstance(kwargs,XComArg):raiseTypeError(f"expected XComArg or list[dict], not{type(kwargs).__name__}")returnself._expand(ListOfDictsExpandInput(kwargs),strict=strict)

expand_kwargs()是expand()的替代方案，当映射的参数组合来自一个字典列表时使用。两者的区别将在下一节详细讨论。

2.3 参数验证机制

validate_mapping_kwargs是映射参数验证的核心函数：

defvalidate_mapping_kwargs(op:type[BaseOperator],func:ValidationSource,value:dict[str,Any])->None:unknown_args=value.copy()forklassinop.mro():init=klass.__init__try:param_names=init._BaseOperatorMeta__param_namesexceptAttributeError:continuefornameinparam_names:value=unknown_args.pop(name,NOTSET)iffunc!="expand":continueifvalueisNOTSET:continueifis_mappable(value):continuetype_name=type(value).__name__ error=f"{op.__name__}.expand() got an unexpected type{type_name!r}for keyword argument{name}"raiseValueError(error)ifnotunknown_args:return

这个函数的工作方式非常精巧：

1. 遍历 Operator 类的 MRO（方法解析顺序）链
2. 通过_BaseOperatorMeta__param_names获取每个类的__init__参数名（这个属性由BaseOperatorMeta元类在类创建时注入）
3. 对于expand()调用，检查每个参数值是否是"可映射的"（is_mappable）

2.4 什么是"可映射的"值

is_mappable函数定义了什么类型的值可以作为映射输入：

defis_mappable(v:Any)->TypeGuard[OperatorExpandArgument]:fromairflow.sdk.definitions.xcom_argimportXComArgreturnisinstance(v,(MappedArgument,XComArg,Mapping,Sequence))andnotisinstance(v,str)

可映射的类型包括：

• XComArg：上游任务的输出引用（运行时解析）
• Sequence：列表、元组等有序集合（解析时已知长度）
• Mapping：字典类型（解析时已知长度）
• MappedArgument：任务组映射的中间参数

特别排除str：虽然str在 Python 中是Sequence，但将字符串映射到每个字符显然不是用户的意图，因此被明确排除。

2.5 防止重复参数

prevent_duplicates确保partial()和expand()之间没有参数冲突：

defprevent_duplicates(partial_kwargs,mapped_kwargs,fail_reason):duplicated=set(partial_kwargs).intersection