Skip to content
Codeloom
Airflow

Airflow TaskFlow API: The Modern Way to Write DAGs

Master the TaskFlow API with @dag and @task decorators, implicit XCom passing, dynamic task mapping, and dataset-aware scheduling in Airflow 2.0+.

·11 min read · By Codeloom
Intermediate 13 min read

What you'll learn

  • Write DAGs using @dag and @task decorators
  • Pass data between tasks implicitly via return values
  • Use multiple_outputs for structured returns
  • Mix TaskFlow with traditional operators
  • Map tasks dynamically with .expand()
  • Schedule DAGs based on dataset updates

Prerequisites

  • Airflow fundamentals (DAGs, tasks, operators)
  • Python decorators
  • Basic understanding of XComs

The Problem TaskFlow Solves

Before diving into syntax, it is worth understanding the frustration that led to the TaskFlow API. In traditional Airflow, writing a simple three-step pipeline required a surprising amount of boilerplate. You had to instantiate operator objects, write separate functions for your logic, manually push and pull XCom values to pass data between tasks, and wire up dependencies with the >> operator. Even a straightforward ETL pipeline could feel verbose and scattered.

Here is what a traditional ETL pipeline looks like — notice how much ceremony is involved just to extract some data, transform it, and load it.

# Traditional approach -- verbose and manual
from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime

def extract(**kwargs):
    data = {"users": 150, "orders": 3200, "revenue": 48500.00}
    kwargs["ti"].xcom_push(key="raw_data", value=data)

def transform(**kwargs):
    raw = kwargs["ti"].xcom_pull(task_ids="extract", key="raw_data")
    raw["avg_order_value"] = raw["revenue"] / raw["orders"]
    kwargs["ti"].xcom_push(key="transformed", value=raw)

def load(**kwargs):
    data = kwargs["ti"].xcom_pull(task_ids="transform", key="transformed")
    print(f"Loading {data['orders']} orders, avg value: ${data['avg_order_value']:.2f}")

with DAG("traditional_etl", start_date=datetime(2026, 1, 1), schedule="@daily", catchup=False):
    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

Every task has to dig into kwargs["ti"] to push and pull XCom values. The data flow is hidden inside function bodies rather than being visible in the pipeline structure. And you need to match task_ids and key strings exactly — a typo means a silent None at runtime. The TaskFlow API, introduced in Airflow 2.0, eliminates all of this friction.


The @dag and @task Decorators

The TaskFlow API lets you write DAGs that look and feel like normal Python programs. Instead of instantiating operator objects, you decorate Python functions with @task. Instead of building a DAG context manager, you decorate a function with @dag. The result is dramatically cleaner code.

Here is the same ETL pipeline rewritten with TaskFlow. Compare it to the traditional version above.

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

@dag(
    dag_id="taskflow_etl",
    start_date=datetime(2026, 1, 1),
    schedule="@daily",
    catchup=False,
)
def taskflow_etl():

    @task()
    def extract():
        return {"users": 150, "orders": 3200, "revenue": 48500.00}

    @task()
    def transform(raw_data: dict):
        return {**raw_data, "avg_order_value": raw_data["revenue"] / raw_data["orders"]}

    @task()
    def load(transformed_data: dict):
        print(f"Loading {transformed_data['orders']} orders")
        print(f"Average order value: ${transformed_data['avg_order_value']:.2f}")

    raw = extract()
    transformed = transform(raw)
    load(transformed)

taskflow_etl()

Look at how much simpler this is. The functions return values directly — no xcom_push calls. The return values flow into downstream function parameters — no xcom_pull calls. The dependencies are inferred automatically from the function call chain — no >> operators needed. And the whole thing reads like a normal Python script.

The @dag decorator turns the outer function into a DAG factory. When you call taskflow_etl() at the bottom, Airflow registers the DAG. The @task decorator turns each inner function into a task, automatically wrapping it in a PythonOperator behind the scenes.


Implicit XCom Passing: The Magic Behind the Scenes

The most powerful feature of the TaskFlow API is implicit XCom passing, and it is worth understanding what is happening under the hood.

When a @task function returns a value, Airflow does not pass that value directly to the next function. Instead, it serializes the return value and stores it in the XCom table of the metadata database — just like the traditional approach. The difference is that Airflow handles this automatically. You never see the serialization, the storage, or the retrieval. You just write raw = extract() and pass raw into transform(raw), and Airflow figures out that transform depends on extract and should receive its XCom output.

This means the “magic” is really just syntactic sugar over the same XCom system. The data still goes through the metadata database, the same size limits apply, and the same serialization rules hold. But your code is dramatically cleaner because all the plumbing is hidden.

A single return value can even feed multiple downstream tasks. Airflow creates the correct dependency edges for each consumer.

@task()
def get_config():
    return {"batch_size": 1000, "retry_count": 3}

@task()
def process_batch(config: dict):
    print(f"Processing with batch size: {config['batch_size']}")
    return {"processed": 5000}

@task()
def report(config: dict, result: dict):
    print(f"Processed {result['processed']} items with config: {config}")

config = get_config()
result = process_batch(config)
report(config, result)  # config feeds both process_batch AND report

In the traditional approach, you would need xcom_pull in both process_batch and report to access config. With TaskFlow, you just pass it as an argument to both functions.


The multiple_outputs Parameter

By default, when a @task function returns a dictionary, the entire dictionary is stored as one XCom entry. But sometimes you want each key to become its own separate XCom value, so downstream tasks can depend on individual pieces of data rather than the whole dictionary.

This is what multiple_outputs=True enables. It splits the returned dictionary into separate XCom entries, one per key. This creates more granular dependencies — a task that only needs the user count does not need to wait for (or depend on) the order count.

Think of it this way: without multiple_outputs, you are handing someone a sealed envelope containing three pieces of information. With multiple_outputs, you are handing them three separate sticky notes. The second approach lets different people grab the specific note they need without waiting for someone to open the envelope.

@task(multiple_outputs=True)
def extract_metrics():
    return {
        "user_count": 1500,
        "order_count": 8200,
        "revenue": 125000.50,
    }

@task()
def process_users(count: int):
    print(f"Processing {count} users")

@task()
def process_orders(count: int):
    print(f"Processing {count} orders")

metrics = extract_metrics()
process_users(metrics["user_count"])
process_orders(metrics["order_count"])

The result is a cleaner dependency graph in the Airflow UI. The process_users task only depends on the user_count output, not the entire extract_metrics result.


Mixing TaskFlow with Traditional Operators

TaskFlow does not require an all-or-nothing commitment. You can mix @task functions with traditional operators in the same DAG, and this is actually the recommended approach. Use TaskFlow for Python logic and traditional operators for infrastructure tasks like running SQL, executing shell commands, or calling cloud services.

The key is understanding how to wire them together. You use the >> operator to set dependencies between TaskFlow tasks and traditional operators, just like you would between two traditional operators.

from airflow.decorators import dag, task
from airflow.operators.bash import BashOperator
from datetime import datetime

@dag(dag_id="mixed_pipeline", start_date=datetime(2026, 1, 1), schedule="@daily", catchup=False)
def mixed_pipeline():

    setup = BashOperator(task_id="setup_workspace", bash_command="mkdir -p /tmp/staging ")

    @task()
    def generate_data():
        return [{"name": "revenue", "value": 48500.00}, {"name": "orders", "value": 3200}]

    @task()
    def process_data(records: list):
        print(f"Processing {len(records)} records")
        return len(records)

    cleanup = BashOperator(task_id="cleanup", bash_command="rm -rf /tmp/staging ")

    data = generate_data()
    result = process_data(data)

    setup >> data       # BashOperator before TaskFlow task
    result >> cleanup   # TaskFlow task before BashOperator

mixed_pipeline()

This pattern is powerful because it lets you use the best tool for each job. SQL-heavy tasks are often cleaner with PostgresOperator or BigQueryInsertJobOperator. File operations are simpler with BashOperator. But Python transformation logic is much cleaner with @task than with PythonOperator.


Dynamic Task Mapping (Airflow 2.3+)

One of the most powerful features built on top of TaskFlow is dynamic task mapping. It solves a common problem: what do you do when you do not know at DAG parse time how many tasks you need?

For example, imagine your pipeline processes files from a directory, but the number of files varies each day. In traditional Airflow, you would have to either hardcode a maximum number of tasks or write complex DAG-generation logic. Dynamic task mapping lets you create task instances at runtime based on the output of an upstream task.

The .expand() method is the key. When you call process_file.expand(file_path=files), Airflow creates one instance of process_file for each element in the files list — and that list is determined at runtime, not parse time.

@dag(dag_id="dynamic_mapping", start_date=datetime(2026, 1, 1), schedule="@daily", catchup=False)
def dynamic_mapping():

    @task()
    def get_file_list():
        import os
        files = os.listdir("/data/incoming/")
        return [f"/data/incoming/{f}" for f in files if f.endswith(".csv")]

    @task()
    def process_file(file_path: str):
        print(f"Processing: {file_path}")
        return {"file": file_path, "rows": 1000}

    @task()
    def summarize(results: list):
        total = sum(r["rows"] for r in results)
        print(f"Processed {total} rows across {len(results)} files")

    files = get_file_list()
    results = process_file.expand(file_path=files)
    summarize(results)

dynamic_mapping()

If get_file_list returns three files, Airflow creates three instances of process_file. If it returns ten files, you get ten instances. The summarize task automatically receives a list of all the individual results.

When you need to map over one parameter while keeping others fixed, use .partial() in combination with .expand(). This is the equivalent of Python’s functools.partial — it pre-fills some arguments so you only need to expand the one that varies.

@task()
def transform_record(record: dict, schema: str, validate: bool):
    return {**record, "schema": schema}

records = get_records()
transform_record.partial(schema="analytics", validate=True).expand(record=records)

Datasets and Data-Aware Scheduling (Airflow 2.4+)

Traditional Airflow scheduling is time-based: “run this DAG every day at midnight.” But sometimes what you really want is event-based: “run this DAG whenever new data is available.” Datasets, introduced in Airflow 2.4, make this possible.

A Dataset is simply a URI that identifies a piece of data — an S3 path, a database table, a file location. A producer DAG declares that one of its tasks updates a dataset. A consumer DAG declares that it should run whenever that dataset is updated. Airflow connects the two automatically.

This is a fundamental shift in how you think about pipeline scheduling. Instead of coordinating DAGs through time-based schedules and ExternalTaskSensors, you express the actual data dependency: “this DAG produces orders data, and that DAG needs orders data.” Airflow handles the triggering.

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

orders_dataset = Dataset("s3://warehouse/orders/daily/")

# Producer: declares it updates the orders dataset
@dag(dag_id="produce_orders", start_date=datetime(2026, 1, 1), schedule="@daily", catchup=False)
def produce_orders():
    @task(outlets=[orders_dataset])
    def export_orders():
        print("Orders exported to S3")
    export_orders()

produce_orders()

# Consumer: triggered when the dataset is updated
@dag(dag_id="consume_orders", start_date=datetime(2026, 1, 1), schedule=[orders_dataset], catchup=False)
def consume_orders():
    @task()
    def build_report():
        print("Building report from fresh orders data")
    build_report()

consume_orders()

Notice that the consumer DAG has no cron schedule. Its schedule parameter is a list of datasets. Airflow triggers it automatically whenever all listed datasets have been updated since the last run. If you list multiple datasets, the consumer waits until all of them have been refreshed — a logical AND condition.


TaskFlow Best Practices

Now that you understand the features, here are the principles that separate clean TaskFlow DAGs from messy ones.

Keep tasks focused. Each @task function should do one thing. If a function extracts data, transforms it, and loads it, break it into three tasks. This maximizes retry granularity (you only re-run the step that failed) and parallelism (independent tasks can run simultaneously).

Watch XCom size. Remember that return values are stored in the metadata database. Keep them small — under 48 KB is the recommended limit. For large data, return a file path or S3 key instead of the data itself. The implicit XCom passing makes this easy to forget because you never see the serialization happening.

Use type hints. They do not affect runtime behavior, but they dramatically improve readability. When you see def transform(raw_data: dict), you immediately know what to expect. They also help IDEs provide autocompletion and catch errors.

Combine with traditional operators. Do not force everything into @task functions. SQL operations are cleaner with dedicated operators. Shell commands are simpler with BashOperator. Use TaskFlow where it shines: Python logic with data passing between tasks.

Prefer .expand() over loops. If you find yourself generating tasks in a Python loop at parse time, consider whether dynamic task mapping could replace it. Mapped tasks give you runtime flexibility and show up as expandable groups in the Airflow UI.


Next Steps

The TaskFlow API is the recommended way to write DAGs in modern Airflow. It reduces boilerplate, makes data flow explicit through function arguments, and integrates cleanly with traditional operators and features like dynamic task mapping and datasets. Here is where to go from here: