Airflow TaskFlow API: The Kwargs Typo That Broke Data Flow

TaskFlow is the modern way to write Airflow DAGs, and it's easy to love: decorators, plain functions, dependencies that wire themselves. It's also easy to over-trust. The same machinery that makes TaskFlow convenient — implicit XCom passing — is exactly where silent failures are born.

Our pipeline had a task that returned a dict of order metrics, and the next task read it with a wrong key. TaskFlow returned None. No error, no warning, DAG green. The warehouse filled with nulls for a day before anyone noticed.

Here's the TaskFlow model from decorator style to multiple outputs, and the contract discipline — types, explicit outputs, and unit tests — that keeps silent None out of your data flow.

1. Decorator Style vs Classic Operators

Classic Airflow is PythonOperator plus explicit dependency wiring: build the operator, then chain with >> or set_upstream. TaskFlow flips the model — you write plain functions, decorate them with @task, and the DAG wires itself.

The difference isn't aesthetics; it's testability. A @task function is a regular Python function. You can call it in a unit test, feed it fixtures, assert on the return value, all without a scheduler running. Classic PythonOperator callables are buried inside operator instantiation and hard to test in isolation.

That testability is the strongest argument for TaskFlow, and it's exactly what our incident lacked. The callable was never tested, so the wrong key survived review, CI, and two environments.

Version note: in Airflow 3 the stable authoring imports moved to airflow.sdk — from airflow.sdk import dag, task — with the legacy airflow.decorators paths deprecated. And decorated tasks are reusable: set retries directly on the decorator (@task(retries=3)) or reuse one task across DAGs with .override(task_id=..., retries=...).

# Classic style: operator + explicit dependency
from airflow.operators.python import PythonOperator

aggregate_task = PythonOperator(
    task_id="aggregate_revenue",
    python_callable=aggregate_revenue,
)
load_task = PythonOperator(
    task_id="load_revenue",
    python_callable=load_revenue,
)
aggregate_task >> load_task

# TaskFlow style: plain functions, wiring inferred
from airflow.decorators import task

@task
def aggregate_revenue():
    return 4320

@task
def load_revenue(order_count: int):
    print(f"Loading {order_count} orders")

load_revenue(aggregate_revenue())

2. How TaskFlow Wires Dependencies for You

TaskFlow reads the callable's signature. If load_revenue(order_count) takes the return value of aggregate_revenue(), TaskFlow creates the dependency for you and pushes the return value through XCom. That's the entire magic: dependencies are inferred from arguments.

The corollary is the trap. If the downstream function doesn't take the upstream result as an argument, there's no dependency and no data flow — and nothing tells you. A function that "doesn't need" the upstream value simply runs whenever its actual arguments allow.

Read your signatures like dependency graphs. Every parameter that comes from another task is an edge in the graph, and every missing parameter is a missing edge. The function signature is the DAG, in miniature.

Mechanically, calling a TaskFlow function returns an XComArg — a deferred reference to the eventual XCom, not the value itself — and any value passed as an argument must be serializable. Airflow supports built-in types out of the box plus objects decorated with @dataclass or @attr.define, so keep custom classes out of the signature unless they're serializable.

from airflow.decorators import dag, task
from airflow.utils.dates import days_ago

@dag(
    schedule="@daily",
    start_date=days_ago(2),
    catchup=False,
    default_args={"retries": 2},
)
def nightly_orders():
    @task
    def count_orders() -> int:
        return 4_320

    @task
    def group_by_region(order_count: int) -> dict:
        # dependency on count_orders is inferred from this argument
        return {"EMEA": order_count // 2, "AMER": order_count - order_count // 2}

    @task
    def publish_report(regional: dict):
        print(f"Regions: {regional}")

    publish_report(group_by_region(count_orders()))

nightly_orders()

3. Multiple Outputs: Unpacking and outputs Dicts

One task can produce several values. TaskFlow gives you two shapes: return a tuple and unpack it positionally, or return a dict and pass outputs={"key": "key"} so each key becomes its own XCom with a proper name.

The dict shape is the production shape. Positional unpacking depends on order — reorder the tuple and every consumer breaks in a way tests should catch but rarely do. The dict shape gives keys names, and named keys are debuggable: you can see them in the XCom tab and pull them by name anywhere.

Our incident's fix used exactly this: outputs={"result": {...}} turned an invisible return_value into a named, documented contract that the next task consumed explicitly.

The docs spell out the other side of the coin: if you return a dict without multiple_outputs=True, the whole dictionary is stored as a single XCom under return_value and must be pulled as a whole — no named keys, no per-key pulls.

from airflow.decorators import task

@task(multiple_outputs=True)
def compute_metrics() -> dict:
    return {
        "order_count": 4_320,
        "revenue": 210_500,
        "refund_rate": 0.012,
    }

@task
def publish_summary(order_count: int, revenue: float):
    print(f"{order_count} orders, {revenue} in revenue")

metrics = compute_metrics()
# keys of the returned dict become named XComs
publish_summary(metrics["order_count"], metrics["revenue"])

4. Partial Return Semantics

Here's the subtle one: a @task with multiple_outputs=True that returns only part of the expected dict. TaskFlow pushes only the keys that exist, and a consumer reading a missing key gets None — exactly the incident's failure mode, wearing a different coat.

Partial returns are a contract violation wearing the costume of a feature. The producer says "I'll give you four metrics," delivers three, and the consumer's fourth value silently becomes None. No exception, no warning, just a null in the warehouse.

The defense is the same one that fixed our incident: types and tests. Type the function's return as the full dict, and unit-test that every key the consumer reads exists in the producer's output. Contracts that can't be violated aren't enforced by discipline alone — they're enforced by tests.

5. Type Hints as Contracts (Incident Deep Dive)

Back to the incident: a task read output key "order_metrrics" instead of "order_metrics". The typo returned None, the DAG stayed green, and nulls flowed into the warehouse for a day. Why didn't anything catch it? Because nothing in the pipeline had a contract — no type hints, no outputs dict, no test.

The fix had three layers. First, explicit outputs: outputs={"result": {...}} made the shape visible and pullable. Second, type hints: the function declared its return type, so a mismatch became visible to any linter or test that cared. Third, a unit test of the callable itself: call the function, assert the keys, walk away.

Type hints aren't runtime enforcement in plain Python — Airflow won't fail a task for a wrong type. They're contracts for humans and CI. That's enough. A contract that's checked at review time beats one that's discovered at 2 AM.

from typing import TypedDict
from airflow.decorators import task

class OrderMetrics(TypedDict):
    order_count: int
    revenue: float
    refund_rate: float

@task(multiple_outputs=True)
def compute_metrics() -> OrderMetrics:
    return {"order_count": 4_320, "revenue": 210_500.0, "refund_rate": 0.012}

@task
def publish_summary(order_metrics: OrderMetrics):
    # typo here would fail type-checking, not production
    print(f"{order_metrics['order_count']} orders, {order_metrics['revenue']} revenue")

metrics = compute_metrics()
publish_summary(metrics)

# unit test, no scheduler required:
# assert compute_metrics()["order_count"] == 4_320
# assert "order_metrics" in OrderMetrics.__annotations__

6. Mixing TaskFlow and Classic Tasks in One DAG

You don't have to choose one style for the whole DAG. Legacy DAGs can adopt TaskFlow one task at a time, and new DAGs can call classic operators when the operator already exists — SnowflakeOperator, PostgresOperator, a sensor — and TaskFlow for everything custom.

The bridge between the styles is XCom. A classic operator pulls with ti.xcom_pull(task_ids=...) what a @task pushed, and vice versa. The mechanics are identical; only the ergonomics differ.

Migration advice: keep the classic operators you have, wrap new logic as @task functions, and let the mix prove itself before you rewrite anything. The goal is testable contracts, not stylistic purity.