Airflow TaskFlow API

The TaskFlow API, introduced in Airflow 2.0, lets you write DAGs using Python decorators instead of manually creating Operator objects. It makes DAG code shorter, more readable, and handles XCom data exchange automatically.

Traditional DAG vs TaskFlow DAG

Traditional Style (Airflow 1.x / early 2.x)

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime

def extract(): return {"sales": 1500, "returns": 50}
def transform(**context):
    data = context["ti"].xcom_pull(task_ids="extract")
    return data["sales"] - data["returns"]
def load(**context):
    result = context["ti"].xcom_pull(task_ids="transform")
    print(f"Loading net sales: {result}")

with DAG("sales_dag", start_date=datetime(2024,1,1), schedule="@daily") as dag:
    t1 = PythonOperator(task_id="extract",   python_callable=extract)
    t2 = PythonOperator(task_id="transform", python_callable=transform)
    t3 = PythonOperator(task_id="load",      python_callable=load)
    t1 >> t2 >> t3

TaskFlow Style (Airflow 2.0+)

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

@dag(start_date=datetime(2024, 1, 1), schedule="@daily", catchup=False)
def sales_dag():

    @task
    def extract():
        return {"sales": 1500, "returns": 50}

    @task
    def transform(data: dict):
        return data["sales"] - data["returns"]

    @task
    def load(net_sales: int):
        print(f"Loading net sales: {net_sales}")

    raw = extract()
    net = transform(raw)     # XCom push/pull happens automatically
    load(net)                 # dependencies set automatically

sales_dag()

The TaskFlow version is shorter, reads like regular Python code, and handles XCom data as normal function arguments.

The Three Key Decorators

@dag

Turns a Python function into a DAG. Parameters go into the decorator — the same parameters as the DAG() constructor:

@dag(
    dag_id="my_pipeline",
    start_date=datetime(2024, 1, 1),
    schedule="@daily",
    catchup=False,
    tags=["etl", "sales"],
)
def my_pipeline():
    ...

my_pipeline()  # ← This line registers the DAG with Airflow

@task

Turns a Python function into a task. The function's return value becomes an XCom value automatically. When you pass the return value of one @task function to another, Airflow sets the task dependency for you.

@task
def get_price():
    return 99.99

@task
def apply_discount(price: float):
    return price * 0.9

@task
def display_final(final_price: float):
    print(f"Final price: ${final_price:.2f}")

# Inside the @dag function:
price    = get_price()
discount = apply_discount(price)
display_final(discount)

Dependency chain created automatically:
[get_price] → [apply_discount] → [display_final]

@task.branch

Turns a function into a branching task. Return the task_id (or list of task_ids) you want to run:

@task.branch
def check_stock(units: int):
    if units < 100:
        return "reorder_items"
    else:
        return "send_ok_email"

@task
def reorder_items():
    print("Placing reorder with supplier")

@task
def send_ok_email():
    print("Sending all-clear email")

# Inside @dag:
stock_level = 45
branch = check_stock(stock_level)
branch >> [reorder_items(), send_ok_email()]

Passing Multiple Values Between Tasks

@task
def get_report_info():
    return {
        "filename": "q1_report.csv",
        "rows": 2500,
        "status": "complete"
    }

@task
def email_summary(info: dict):
    print(f"File: {info['filename']}, Rows: {info['rows']}")

# Inside @dag:
info = get_report_info()
email_summary(info)

The entire dictionary travels via XCom automatically. No manual xcom_push or xcom_pull calls needed.

Multiple Outputs With @task(multiple_outputs=True)

@task(multiple_outputs=True)
def extract_stats():
    return {
        "total_sales": 50000,
        "num_orders": 320,
        "top_product": "Widget A"
    }

@task
def report_sales(total_sales: int):
    print(f"Total sales: ${total_sales}")

@task
def report_orders(num_orders: int):
    print(f"Orders placed: {num_orders}")

# Inside @dag:
stats = extract_stats()
report_sales(stats["total_sales"])     # pulls only total_sales
report_orders(stats["num_orders"])     # pulls only num_orders

Each key of the returned dictionary becomes a separate XCom entry. Downstream tasks pull only the specific values they need.

Mixing TaskFlow and Traditional Operators

TaskFlow tasks and traditional Operator tasks work together in the same DAG:

from airflow.decorators import dag, task
from airflow.providers.postgres.operators.postgres import PostgresOperator
from datetime import datetime

@dag(start_date=datetime(2024, 1, 1), schedule="@daily", catchup=False)
def mixed_dag():

    create_table = PostgresOperator(
        task_id="create_table",
        postgres_conn_id="my_db",
        sql="CREATE TABLE IF NOT EXISTS results (id SERIAL, value INT);",
    )

    @task
    def compute_value():
        return 42

    @task
    def log_result(value: int):
        print(f"Result is {value}")

    value = compute_value()
    create_table >> value >> log_result(value)

mixed_dag()

@task.virtualenv: Run in a Separate Python Environment

Use @task.virtualenv when a task needs a library version that conflicts with the main Airflow environment:

@task.virtualenv(
    task_id="run_in_venv",
    requirements=["pandas==1.5.3", "numpy==1.24.0"],
    system_site_packages=False,
)
def process_with_old_pandas():
    import pandas as pd
    df = pd.DataFrame({"a": [1, 2, 3]})
    return df["a"].sum()

Airflow creates a temporary virtual environment, installs the specified packages, runs the function, and tears down the environment when done.

When to Use TaskFlow vs Traditional Style