The TaskFlow API: DAGs That Look Like Python

Decorate a function with @task, call it, and Airflow builds the graph for you. The TaskFlow API turns a DAG into code you can almost read aloud.

The classic DAG from the operators post works, but read it back and you’ll notice something: you spent as many lines wiring tasks together as writing what they do. You named a callable, wrapped it in a PythonOperator, gave the operator a task_id that repeats the function’s name, then drew the arrows by hand. The TaskFlow API deletes almost all of that ceremony by making a plain Python function be a task and a plain function call be a dependency.

@dag and @task

Two decorators from the Task SDK carry the whole thing:

from airflow.sdk import dag, get_current_context, task

@dag(schedule="@daily", catchup=False)
def bookshop_daily():
    @task
    def extract():
        ...

    @task
    def load(rows):
        ...

    load(extract())

bookshop_daily()

@dag turns the function into a DAG definition; calling it at the bottom is what registers it. Inside, each @task function is a task, and the task_id defaults to the function name — no more repeating yourself. The name of the outer function becomes the dag_id.

The line doing the real work is load(extract()). You didn’t write extract >> load. You called extract, passed its result into load, and Airflow read that call graph to infer the dependency: load needs extract’s output, so extract must run first. Dependencies fall out of how the data flows, which is how you’d reason about the code anyway.

Return values become XComs

extract() doesn’t return a value at authoring time — it returns a placeholder that stands in for “whatever this task produces when it actually runs.” When extract executes, its return value is stored as an XCom (Airflow’s mechanism for passing small values between tasks), and when load runs, that value is pulled back out and handed to the rows argument. You wrote it like an ordinary function call; Airflow turned it into a cross-task handoff.

That handoff has rules — mainly, keep it small — and they get their own post next. What matters here is that the plumbing is invisible. You return a value, you receive a value, and the fact that the two tasks might run in different processes on different machines is Airflow’s problem, not yours.

Reading the run context

A task almost always needs to know which run it belongs to. What day is this? Was it triggered manually or on schedule? What did the operator override in params? Airflow keeps all of that in a per-run context dictionary, and TaskFlow gives you two ways to reach it.

The first is to declare a parameter whose name matches a context key. Airflow inspects your function’s signature and injects the matching values for you:

@task
def extract(data_interval_start=None, data_interval_end=None, ti=None):
    day = data_interval_start.date().isoformat()
    print(f"loading orders for {day}, try #{ti.try_number}")
    return f"/data/orders/{day}.csv"

You never call extract(data_interval_start=...) yourself — those arguments aren’t part of the data flow, so you leave them out of the call and Airflow fills them in at execution time. The keys you’ll reach for most:

  • data_interval_start / data_interval_end — the half-open time window this run covers. For an @daily DAG, start is midnight of the day being processed and end is midnight of the next. This is the pair you use to pick which slice of data to read.
  • logical_date — the run’s nominal timestamp. Useful, but note that on asset-triggered or manually kicked runs it can be absent or unintuitive; when you want the data window, prefer data_interval_start over logical_date.
  • ti — the task instance, your handle for this task’s own metadata: ti.try_number, ti.map_index, and the lower-level ti.xcom_push / ti.xcom_pull.
  • dag_run — the run object. dag_run.run_type tells you "scheduled" vs "manual" vs "backfill"; dag_run.conf holds the JSON payload passed to a triggered run.
  • params — the merged parameters (defaults plus per-run overrides), which we’ll come back to.

Two of these are worth dwelling on because newcomers reach for the wrong one. logical_date feels like “the date this run is about,” and for a simple daily schedule it roughly is — but it’s the run’s timestamp, not a window, and it wobbles once you introduce assets or manual triggers. data_interval_start and data_interval_end are the honest answer to “what slice of data does this run own”: a clean half-open [start, end) interval that lines up exactly with the schedule. When you’re filtering a source table with WHERE created_at >= :start AND created_at < :end, those two are the values you bind.

The second way is get_current_context(), which hands you the whole dictionary from anywhere the task is running — handy inside a helper function that isn’t itself the decorated task, so you don’t have to thread the context through every call:

from airflow.sdk import get_current_context, task

def _target_path() -> str:
    context = get_current_context()
    day = context["data_interval_start"].date().isoformat()
    return f"/data/orders/{day}.csv"

@task
def extract() -> str:
    return _target_path()

One habit to build now: read the date from the context, never from a rendered template baked into an f-string. It’s tempting to write f"/data/orders/{{ ds }}.csv" inside a @task — but {{ ds }} is Jinja, and Python string templating doesn’t render Jinja. Inside a Python function you’d just get the literal braces back. Templating like {{ ds }} renders in an operator’s templated fields (a bash_command, a SQL file path); inside your Python, reach into the context object and call .date().isoformat() on data_interval_start. Same value, but this one is real.

Tuning a task from the decorator

Everything you could configure on a PythonOperator you can pass to @task, because under the hood that’s still what you’re building. The decorator takes the operator’s arguments:

from datetime import timedelta

@task(
    retries=3,
    retry_delay=timedelta(minutes=2),
    execution_timeout=timedelta(minutes=15),
    pool="warehouse",
    task_id="extract_orders",
)
def extract() -> str:
    ...
  • retries / retry_delay — how many times Airflow re-runs the task on failure, and how long it waits between attempts. Because a task should be idempotent — running it twice leaves the world in the same state as running it once — retries are free safety, not a gamble. The ti.try_number you saw in the context is what makes this observable: on attempt three, the task can see it’s on attempt three.
  • execution_timeout — a wall-clock ceiling. If the task runs longer, Airflow kills it and marks it failed rather than letting a hung API call block the pool forever.
  • pool — a named concurrency budget. Point ten warehouse-touching tasks at a warehouse pool with four slots and only four run at once, no matter how much else is idle. This is how you stop a fan-out from stampeding a database.
  • task_id — override the auto-derived name. By default the task_id is the function name; pass task_id= when you want the UI to read differently from the code, or when the same function backs more than one task.

That last case is worth its own paragraph, because it’s a sharp edge.

Calling a task more than once

A decorated function is a factory, not a single task. Call it twice and you get two task instances — which means two task_ids, and they can’t collide. Airflow resolves the collision by auto-suffixing:

@task
def fetch(source: str) -> str:
    ...

a = fetch("orders")       # task_id: fetch
b = fetch("customers")    # task_id: fetch__1
c = fetch("inventory")    # task_id: fetch__2

That works, but fetch__1 is a terrible name to see in a log three weeks later. When you reuse a task, name each call explicitly with .override():

orders    = fetch.override(task_id="fetch_orders")("orders")
customers = fetch.override(task_id="fetch_customers")("customers")
inventory = fetch.override(task_id="fetch_inventory")("inventory")

.override() returns a fresh copy of the task with any decorator arguments you like changed — not just task_id but retries, pool, anything — so one function can spawn several tasks with tailored settings. Reach for it whenever you’d otherwise be staring at __1 and __2 in the grid.

The @task family: beyond plain Python

@task runs a Python callable in the worker’s own interpreter. But the decorator has a family of variants, @task.<flavor>, each of which changes where or how the body runs while keeping the same call-it-to-wire-it ergonomics. You reach for a flavor when plain Python isn’t the right execution environment.

@task.bash — the function returns a shell command string, and Airflow runs it as a BashOperator would. You get Python to compute the command and Bash to run it:

@task.bash
def archive(path: str) -> str:
    return f"gzip -f {path} && echo archived {path}"

@task.branch — the function returns the task_id (or a list of them) to run next; every other downstream task is skipped. This is how you put an if in a DAG:

@task.branch
def choose_path(dag_run=None) -> str:
    return "full_reload" if dag_run.run_type == "manual" else "incremental"

@task.short_circuit — the function returns a truthy/falsy value. Falsy skips everything downstream; truthy lets the flow continue. It’s the “only proceed if there’s actually work” gate:

@task.short_circuit
def has_new_orders(path: str) -> bool:
    import os
    return os.path.getsize(path) > 0

@task.virtualenv and @task.external_python — run the body in a different Python environment. @task.virtualenv builds a throwaway venv with the requirements you list, ideal when a task needs a library that would clash with Airflow’s own dependencies. @task.external_python runs against a pre-built interpreter you point at by path, so you pay the setup cost once:

@task.virtualenv(requirements=["pandas==2.2.3"], system_site_packages=False)
def summarize(path: str) -> dict:
    import pandas as pd
    df = pd.read_csv(path)
    return {"rows": len(df), "revenue": float(df["total"].sum())}

@task.docker and @task.kubernetes — run the body inside a container. @task.docker executes it in a Docker image on the same host; @task.kubernetes launches a fresh pod in your cluster. Both isolate the task completely — its own image, dependencies, and resource limits — which is how you run a heavyweight or untrusted step without polluting the worker:

@task.docker(image="python:3.12-slim", auto_remove="force")
def render_report(rows: int) -> str:
    return f"generated a report for {rows} rows"

@task.sensor — turns a function into a sensor, a task that waits for a condition by polling. The body returns whether the condition is met (a bool, or a PokeReturnValue carrying data forward), and Airflow re-invokes it on the interval you set until it’s satisfied or times out:

from airflow.sdk import PokeReturnValue, task

@task.sensor(poke_interval=60, timeout=3600, mode="reschedule")
def wait_for_file(path: str) -> PokeReturnValue:
    import os
    ready = os.path.exists(path)
    return PokeReturnValue(is_done=ready, xcom_value=path if ready else None)

Two of these deserve a footnote. On @task.sensor, the mode argument decides how the wait is paid for. The default, mode="poke", holds a worker slot for the entire wait — fine for a few seconds, ruinous if you’re waiting an hour for an upstream file, because that slot does nothing but sleep. mode="reschedule" releases the slot between checks and lets the scheduler wake the task up again on the next interval, so a hundred sensors waiting on slow inputs don’t starve your workers. For anything that might wait minutes, reach for reschedule.

And @task.branch only chooses — it doesn’t clean up after itself. The branches it doesn’t pick are marked skipped, and skips propagate downstream, so a task that joins the branches back together needs the right trigger rule or it’ll skip too. That interaction is a chapter of its own; for now, know that a branch returns the id(s) of what to run and Airflow skips the rest.

You don’t need every one of these on day one. The point is that “TaskFlow” isn’t only @task — it’s a family, and when a step wants Bash, a branch, a clean venv, or a container, there’s a decorator that keeps the syntax identical and just swaps the engine underneath.

multiple_outputs, and its edges

If a task needs to hand off more than one thing, tell the decorator:

@task(multiple_outputs=True)
def extract():
    return {"path": "/data/orders/2026-05-31.csv", "row_count": 4212}

With multiple_outputs=True, each key of the returned dict becomes its own XCom, so a downstream task can depend on extract()["path"] without dragging along the whole dict:

result = extract()
load(result["path"])                 # depends on just the path
audit(result["row_count"])           # depends on just the count

A few edges are worth knowing before this bites you:

  • It’s dict-only. multiple_outputs splits a returned dictionary into per-key XComs. Return a list or a tuple with the flag on and you’ll get an error, because there are no keys to split by.
  • Keys must be valid. Each key becomes an XCom key, so they need to be strings that are legal identifiers. A dict keyed by integers or spaces won’t map cleanly.
  • Typing can infer it for you. Annotate the return as -> dict and Airflow turns multiple_outputs on automatically — the type hint is the signal. That’s convenient, but it also means a stray -> dict annotation can silently change behaviour. When you want the split, the explicit multiple_outputs=True flag is the clearer way to say so; when you don’t, keep the annotation off or return a non-dict.

If you only ever consume the whole object downstream, you don’t need the flag at all — a plain return {...} hands the dict across as one XCom, and the receiving task indexes it in Python.

Fanning out: a first look at mapping

Sometimes you don’t know how many tasks you need until the run is underway — one load per file the extract discovered, say. TaskFlow expresses that with .expand():

@task
def list_files() -> list[str]:
    return ["/data/orders/a.csv", "/data/orders/b.csv", "/data/orders/c.csv"]

@task
def load_one(path: str):
    ...

load_one.expand(path=list_files())

load_one.expand(path=...) creates one mapped task instance per item in the list, at run time — the DAG file was parsed before any run existed, so it couldn’t have known the count in advance, and that’s exactly the gap mapping fills. In the grid you’ll see load_one render as a single node with a bracketed count, each instance addressable by its map_index. Its siblings .partial() (bind arguments that don’t vary across the fan-out) and .expand_kwargs() (map a list of dicts when each instance needs several fields) round out the toolkit. There’s real depth here — how to cap the fan-out so a bad source file doesn’t spawn forty thousand tasks, how to avoid accidental Cartesian products — and it gets its own chapter. For now, just register that the pattern exists and reads as naturally as the rest of TaskFlow. We cover it properly in Dynamic Task Mapping and TaskGroups.

Grouping and bracketing: @task_group, @setup, @teardown

Two more decorators shape the graph rather than the work.

@task_group bundles related tasks into a labelled box in the UI and namespaces their task_ids, without creating any runtime isolation — it’s purely for readability and reuse. Define one, call it, and wire it like any task:

from airflow.sdk import dag, task, task_group

@dag(schedule="@daily", catchup=False)
def bookshop_daily():
    @task_group
    def process(path: str):
        @task
        def clean(p: str) -> str:
            ...
        @task
        def validate(p: str) -> str:
            ...
        return validate(clean(path))

    @task
    def extract() -> str:
        ...

    @task
    def publish(p: str):
        ...

    publish(process(extract()))

bookshop_daily()

In the grid, clean and validate collapse into a single process group you can expand — and because the group is a function, you can call it more than once (with .override(group_id=...)) to reuse the same mini-pipeline over several inputs. TaskGroups replace the old SubDAGs, which were removed; if you learned Airflow years ago and remember SubDagOperator, this is what to use instead, and it’s strictly better.

@setup and @teardown mark tasks that bracket a piece of work: a setup builds a resource (spin up a temp warehouse, create a scratch schema), and its paired teardown tears it down (drop the schema) — critically, the teardown runs even if the work in between fails, and its failure doesn’t by default fail the whole run. That’s the difference from an ordinary downstream task, which would be skipped on upstream failure:

from airflow.sdk import dag, setup, task, teardown

@dag(schedule="@daily", catchup=False)
def bookshop_daily():
    @setup
    def make_scratch() -> str:
        # create a temp schema, return its name
        return "scratch_20260530"

    @task
    def build(schema: str):
        ...

    @teardown
    def drop_scratch(schema: str):
        # always runs, even if build() failed
        ...

    scratch = make_scratch()
    work = build(scratch)
    drop_scratch(scratch).as_teardown(setups=scratch)

bookshop_daily()

The as_teardown(setups=...) call pairs the teardown to its setup, and that pairing is what gives the teardown its special behaviour: it’s scheduled after the work regardless of how the work ends, and by default a failing teardown doesn’t cascade into a failed DAG run — because “I couldn’t drop the scratch schema” usually shouldn’t nullify a pipeline that already produced the data you wanted. If you clear and re-run the work later, the setup re-runs with it, so the resource is always freshly built for the retry rather than assumed to still exist.

You won’t need setup/teardown on a first pipeline, but the moment you’re renting an expensive resource per run — a cluster, a scratch database — this is the pattern that guarantees you turn it off.

The bookshop pipeline, cleaned up

Here’s the whole daily pipeline — extract, transform, load — written in TaskFlow. Set it beside the classic version from the operators post and count what’s gone: no PythonOperator wrappers, no python_callable=, no hand-drawn arrows, no task_id that echoes the function name.

from airflow.sdk import dag, get_current_context, task

@dag(schedule="@daily", catchup=False, tags=["bookshop"])
def bookshop_daily():
    @task
    def extract() -> str:
        context = get_current_context()
        day = context["data_interval_start"].date().isoformat()
        path = f"/data/orders/{day}.csv"
        # fetch the day's orders CSV, write to path
        return path

    @task
    def transform(path: str) -> str:
        # read raw CSV, clean, write a tidy parquet, return its path
        return path.replace(".csv", ".parquet")

    @task
    def load(path: str):
        # load the tidy file into DuckDB
        ...

    load(transform(extract()))

bookshop_daily()

The DAG reads top to bottom like a description of the work: extract a path, transform it into another path, load that. load(transform(extract())) is the entire dependency graph, and it’s the same expression you’d write if these were ordinary functions in a script. That’s the point — the orchestration disappears into the code.

And notice params still works exactly as it did in the classic style — declare it, override it per run:

@dag(schedule="@daily", catchup=False, params={"source": "prod"})
def bookshop_daily():
    @task
    def extract(params: dict) -> str:
        source = params["source"]
        ...

TaskFlow also composes with operators rather than replacing them. A @task function and a SQLExecuteQueryOperator share a DAG happily; wire a TaskFlow output into a classic operator, or set order between them with the familiar >>. Use @task for your Python, prebuilt operators for the systems you integrate with, and don’t feel you’ve mixed metaphors — that’s the intended shape.

Final thoughts

For a new DAG, TaskFlow is the default, and not because it’s newer. It collapses the gap between how you think about a pipeline and how you write it: a chain of functions passing data along. The classic operator style hasn’t gone anywhere — you’ll use it every time you touch a prebuilt integration — but the scaffolding around your own logic should be as thin as this. And the family fills in the rest: @task.bash when you want a shell, @task.branch for an if, @task.virtualenv or @task.docker for an isolated environment, @task_group to keep the graph legible, @setup/@teardown to bracket a resource. When the wiring vanishes, the only thing left to read is what the pipeline actually does, which is the only thing you wanted to read in the first place.

Next: XComs: How Tasks Pass Notes

Comments