Airflow TaskFlow API

Airflow TaskFlow API

介绍

Airflow TaskFlow API 是 Apache Airflow 2.0 引入的一种高级编程接口，旨在简化 DAG（有向无环图）的定义和任务之间的数据传递。它基于 Python 装饰器（`@task`），允许用户以更直观、函数式的方式编写任务逻辑，并自动处理 XComs（跨任务通信）的推送和拉取。TaskFlow API 特别适合数据管道开发，减少了传统 Operator 的样板代码，同时保持了 Airflow 的灵活性和可扩展性。

TaskFlow API 的核心优势包括：

• 减少显式 XComs 操作
• 更自然的函数式编程风格
• 自动依赖管理
• 与 Python 类型提示（Type Hints）集成

基本用法

以下是一个简单的 TaskFlow API 示例，展示如何定义任务并传递数据：

from airflow.decorators import dag, task
from datetime import datetime

@dag(schedule_interval=None, start_date=datetime(2023, 1, 1), catchup=False)
def taskflow_example():

    @task
    def extract() -> dict:
        return {"data": [1, 2, 3]}

    @task
    def transform(data: dict) -> dict:
        data["data"] = [x * 2 for x in data["data"]]
        return data

    @task
    def load(data: dict):
        print(f"Processed data: {data}")

    load(transform(extract()))

dag = taskflow_example()

代码解析： 1. 使用 `@dag` 装饰器定义 DAG 2. 每个任务函数用 `@task` 装饰器标记 3. 函数的返回值自动通过 XComs 传递 4. 任务依赖关系通过函数调用链 `load(transform(extract()))` 隐式定义

XComs 的自动处理

在传统 Airflow 中，XComs 需要显式调用 `xcom_push` 和 `xcom_pull`。TaskFlow API 自动完成这些操作：

• 函数返回值自动推送到 XComs
• 下游任务通过函数参数自动拉取 XComs 值
• 类型提示（如 `data: dict`）帮助 Airflow 验证数据类型

多返回值处理

TaskFlow API 支持返回多个值（通过元组或字典），下游任务可以按需接收：

@task
def extract_multiple() -> tuple:
    return [1, 2, 3], {"status": "ok"}

@task
def process(numbers: list, metadata: dict):
    print(f"Numbers: {numbers}, Metadata: {metadata}")

process(*extract_multiple())

高级特性

任务依赖的显式控制

除函数调用外，可以使用 `>>` 运算符定义依赖：

extract_task = extract()
transform_task = transform(extract_task)
load_task = load(transform_task)

extract_task >> transform_task >> load_task

与传统 Operator 混合使用

TaskFlow 任务可以与标准 Operator 共存：

from airflow.operators.python import PythonOperator

@task
def process_data(data):
    return data * 2

def traditional_task(**context):
    print("Traditional operator running")

with DAG(...) as dag:
    taskflow_result = process_data(42)
    PythonOperator(
        task_id="traditional",
        python_callable=traditional_task,
        op_kwargs={"input": taskflow_result}
    )

动态任务映射

Airflow 2.3+ 支持动态生成并行任务：

@task
def process_item(item):
    return item * 2

@task
def extract_items() -> list:
    return [1, 2, 3, 4]

process_item.expand(item=extract_items())

实际案例：数据ETL管道

以下是一个完整的数据处理管道示例：

from airflow.decorators import dag, task
from datetime import datetime
import pandas as pd

@dag(schedule="@daily", start_date=datetime(2023, 1, 1))
def sales_pipeline():

    @task(retries=3)
    def fetch_sales(date: datetime) -> pd.DataFrame:
        # 模拟从数据库获取数据
        return pd.DataFrame({
            "date": [date.date()],
            "revenue": [1000 * (date.day % 7 + 1)]
        })

    @task
    def calculate_stats(df: pd.DataFrame) -> dict:
        return {
            "total": df["revenue"].sum(),
            "average": df["revenue"].mean()
        }

    @task
    def store_report(stats: dict):
        print(f"Storing report: {stats}")

    sales_data = fetch_sales("{{ ds }}")
    stats = calculate_stats(sales_data)
    store_report(stats)

dag = sales_pipeline()

关键点：

• 使用 Pandas DataFrame 作为数据类型
• 任务自动重试机制（`retries=3`）
• 模板变量 `模板:Ds` 的用法
• 完整的类型安全数据流

性能考量

使用 TaskFlow API 时需注意： 1. XComs 大小限制：默认序列化为 JSON，大型数据应使用外部存储（如 S3） 2. 序列化开销：复杂对象可能需要自定义序列化 3. 任务映射内存消耗：动态生成大量任务时需监控资源

可视化

以下 Mermaid 图展示示例 DAG 的任务流：

数学表达

对于数据转换任务，可以形式化描述为：

$f_{transform}(x)=g\circ f_{extract}(x)$

其中：

• $f_{extract}$ 是提取任务
• $g$ 是转换逻辑
• $\circ$ 表示函数组合

总结

TaskFlow API 显著提升了 Airflow 的开发体验：

• 减少约 40% 的样板代码（基于典型 ETL 管道）
• 更符合 Python 开发者的直觉
• 保持与传统 Operator 的兼容性
• 通过类型提示提高代码可靠性

建议新项目优先采用 TaskFlow API，现有项目可以逐步迁移。对于复杂场景，可结合传统 Operator 实现最佳灵活性。