Dependencies & the TaskFlow API

Writing a DAG has two halves. You write each task, and you connect the tasks into a shape.

This chapter covers both:

1. How to connect tasks: >>, lists, chain() and cross_downstream().
2. How @task turns plain functions into tasks, and their calls into connections.
3. When a classic operator is still the better choice.

Every output here comes from a real run with Airflow 3.1.

a >> b means a must finish before b starts. Each such link is an edge in the graph. A list on one side means all of these tasks.

So one line can make a fan-out, where one task feeds many, and then a fan-in, where many feed one. Together they make a diamond:

start >> [load_orders, load_customers] >> join

A list cannot be on both sides of one >>. Python does not know how to join two lists:

[a, b] >> [c, d]

Two helpers from airflow.sdk cover this case, and they build different graphs:

chain([a, b], [c]) fails with a ValueError, because 2 and 1 cannot be paired. chain is most useful for tasks made in a loop: chain(*tasks) puts them in a row.

In the classic style, you create an operator object for each task. You set the order with >>, and you pull the value from XCom yourself:

def classic_load():
    return 6

def classic_build(ti):
    rows = ti.xcom_pull(task_ids="classic_load")
    print(f"classic: building revenue from {rows} orders")

load = PythonOperator(task_id="classic_load", python_callable=classic_load)
build = PythonOperator(task_id="classic_build", python_callable=classic_build)
load >> build

The TaskFlow style does the same with @task. You saw it in chapter 4:

@task
def load_orders() -> int:
    return 6

@task
def build_daily_revenue(rows: int):
    print(f"building revenue from {rows} orders")

build_daily_revenue(load_orders())

Underneath, each @task function still becomes an operator. So the scheduler runs both styles in the same way, and one DAG can mix them.

In TaskFlow, passing one task’s result to another task makes the edge. The order and the data are one statement, so they cannot disagree.

In the classic style they can disagree. You might write load >> build but pull a value from a third task. The graph looks right, and the task reads the wrong value.

When no value moves, use >> as before:

refresh_dashboard() >> send_email()

To give a classic operator’s result to a @task, use its .output:

count_files = BashOperator(task_id="count_files", bash_command="echo 3")

@task
def report_files(n):
    print(f"files: {n!r}")

report_files(count_files.output)

A task can return a dict. If its return type says dict, Airflow stores each key as its own XCom. This is called multiple_outputs:

@task
def order_stats() -> dict[str, int]:
    return {"rows": 6, "paid": 4}

@task
def report_paid(paid: int):
    print(f"paid orders: {paid}")

report_paid(order_stats()["paid"])

Now report_paid needs only the key paid. Other tasks can take rows, and each task depends only on the key it uses.

Without multiple_outputs, every task gets the whole dict, and must pick out its own key.

Some decorators change what a task does with its result, or where it runs:

short_circuit is good for days with nothing to do. In our run, it returned False, and the task after it was skipped:

@task.short_circuit
def has_orders():
    return False

has_orders() >> publish()

virtualenv and docker help when one task needs a library version that the rest cannot use. But both add start-up time to every run of that task.

Use @task for your own logic. Use a provider operator for work that a provider already does.

A provider operator for a system such as Databricks already handles waiting for the job, errors and logins. If you rewrite it as a @task, all of that becomes your code to maintain.

So most real DAGs mix both styles:

count = SQLExecuteQueryOperator(
    task_id="count_orders",
    conn_id="shop_db",
    sql="SELECT count(*) FROM orders WHERE order_date = '{{ ds }}'",
)

@task
def check_count(rows):
    print(f"rows: {rows!r}")
    if rows[0][0] == 0:
        raise ValueError("no orders for this day")

check_count(count.output) >> refresh_dashboard()

Choose the style for each task, not for the whole DAG.

The rules for shapes and @task:

1. a >> b means a finishes before b starts. A list means all of these.
2. Never put lists on both sides of >>. Use chain() to pair, or cross_downstream() to connect all.
3. In TaskFlow, passing a value makes the edge. Use >> only when no value moves.
4. Return small values. A return type of dict splits the keys into separate XComs.
5. Use @task for your logic, and provider operators for provider work.

Check yourself: you need a and b to both feed c and d. What do you write? (Answer: cross_downstream([a, b], [c, d]).)