Gowtham Potureddi

Posted on Jul 3

Airflow TaskFlow API: Decorators, XComs & Functional DAGs for Python-First Pipelines

#python #sql #dataengineering #interview

airflow taskflow is the Python-first pipeline API that quietly replaced 80% of PythonOperator boilerplate in production Airflow deployments — and the single feature senior interviewers now open with when they want to test whether a candidate has actually shipped an Airflow DAG in the last three years. The classic operator model — PythonOperator(task_id=..., python_callable=..., op_kwargs={...}) plus manual >> wiring and hand-rolled xcom_push / xcom_pull calls — was verbose, XCom-noisy, and hostile to unit testing; the airflow taskflow api collapses all of that into ordinary Python functions decorated with @task inside a function decorated with @dag, where return values become XComs automatically and dependencies are inferred from function-call composition. The mental shift is exactly the one Python developers already know from functools.lru_cache or FastAPI's Depends: the decorator is the pipeline; the function is the task; the return is the XCom.

This guide is the senior-DE walkthrough you wished existed the first time an interviewer asked "walk me through the difference between an airflow python operator DAG and a functional dag written with the airflow task decorator" or "how do XComs work when you return a dict from an @task function?" or "how would you mix an @task transform with an S3KeySensor and a KubernetesPodOperator in the same DAG?" It walks through why TaskFlow is default in new DAGs on Airflow 2.7+ and mandatory-by-convention in Airflow 3, the four "must-answer" axes interviewers actually probe (decorators, XCom automation, typing, mixed usage), the mechanics of airflow @task and airflow @dag as callable-returning factory decorators, the XComArg object that carries return values through the graph, the size-limit reality that makes XCom a control-plane bus and not a data-transfer bus, the interop patterns that bridge @task functions with classic operators via .output, and the typing plus dag.test() habits that keep airflow decorators DAGs shippable. Each section pairs a teaching block with a Solution-Tail interview answer — code, a step-by-step trace, an output table, then a concept-by-concept breakdown of why it works.

When you want hands-on reps immediately after reading, drill the ETL practice library →, rehearse on the SQL practice library →, and sharpen the tuning axis with the optimization practice library →.

On this page

Why TaskFlow replaced most PythonOperators in 2026
@task and @dag decorators — the two-decorator pipeline
XCom automation and XComArg chaining
Mixing TaskFlow with classic Operators
Typing, retries, and testing TaskFlow DAGs
Cheat sheet — TaskFlow recipes
Frequently asked questions
Practice on PipeCode

1. Why TaskFlow replaced most PythonOperators in 2026

The verbose old world — `PythonOperator(task_id=..., python_callable=..., op_kwargs=...)` — collapses to `@task` plus a return statement

The one-sentence invariant: the classic PythonOperator forces you to name every task twice (once as a variable, once as a task_id string), wire every dependency by hand with >>, and manually shuttle every intermediate value through xcom_push / xcom_pull calls that mention keys and task IDs the reader has to cross-reference — while the airflow taskflow api lets you write the same pipeline as ordinary Python functions where the decorator wires the graph and the return statement is the XCom. The two APIs produce byte-identical DAGs at the Airflow scheduler level; the difference is entirely at the authoring layer, and it is the difference between a 60-line hand-wired ETL and a 15-line functional DAG that reads like a top-level main().

The four axes interviewers actually probe.

Decorators. @dag on the DAG factory function; @task on individual task functions. Both are ordinary Python decorators that accept keyword arguments and return callables. The senior signal is knowing they are factory decorators — calling the decorated function builds the DAG or the task instance; the function body is not executed at parse time.
XCom automation. Return values push automatically; function arguments pull automatically. The XCom key is return_value unless you use @task(multiple_outputs=True) and return a dict, in which case each dict key becomes a separate XCom entry. Every intermediate value flows through XCom by default — which is fine for control-plane messages and toxic for large payloads.
Typing. @task propagates type hints into the IDE and mypy. Combined with XComArg, this gives you compile-time-ish checks that a downstream task's argument type matches the upstream task's return type — a level of static safety the classic operator API never had.
Interop with classic operators. @task DAGs and classic operator DAGs coexist inside the same file. Classic operators expose .output as an XComArg, so a sensor's XCom feeds a @task function's argument without any glue code, and a @task return value feeds a BashOperator's templated command without an explicit xcom_pull.

Why the classic PythonOperator felt wrong the moment TaskFlow shipped.

Double-naming. Every task had a Python variable name and a task_id string that had to match by convention. Rename one, forget the other, and the DAG parses but breaks silently when xcom_pull(task_ids=...) looks up the old name.
Explicit wiring. Every dependency was a manual >> at the bottom of the file. A 20-task DAG had a 20-line "wiring section" that duplicated information already implicit in which task's output fed which task's input.
Manual XCom. ti.xcom_push(key="rows", value=42) and ti.xcom_pull(task_ids="extract", key="rows") were the norm. Every intermediate value was a stringly-typed lookup by task ID; refactoring a task name required grepping XCom keys across the codebase.
op_kwargs was a leak. Passing arguments to a python_callable via op_kwargs meant your function signature was documented in a dict at the operator instantiation site, three imports away from the function definition. IDE autocomplete never saw it.
Testing was hard. Unit-testing a PythonOperator task meant either running the whole DAG (slow, requires a scheduler + metastore) or extracting the python_callable and calling it with hand-constructed kwargs — the operator wrapper was pure friction.

2026 reality — TaskFlow is default, classic operators are for external systems.

New DAGs. Written as @dag + @task from day one. Reviewers reject PRs that use PythonOperator for pure-Python work without a specific reason.
Classic operators still shine. Anywhere you're calling an external system (S3, Kubernetes, BigQuery, Snowflake, Databricks), the classic operator wraps SDK boilerplate you'd otherwise rewrite. S3KeySensor, KubernetesPodOperator, BigQueryInsertJobOperator — all still in daily use.
Mixed DAGs are the norm. A production ETL DAG on Airflow 3 typically has 2–5 classic operators (sensors, external system triggers) and 5–15 @task functions doing the pure-Python transform work. Interop patterns matter.
Airflow 3 sharpens the default. @task DAGs get first-class treatment in the UI (function-level docs from docstrings), the CLI (airflow tasks test dag_id task_id works out of the box), and the testing tooling (dag.test() for pytest integration).

What interviewers listen for.

Do you say "return values are pushed to XCom automatically, arguments pull them automatically" in the first sentence when asked how TaskFlow XComs work? — senior signal.
Do you mention XComArg by name as the object that flows between tasks? — senior signal.
Do you push back on "just use PythonOperator, it's simpler" with the double-naming, manual-wiring, and testability arguments? — required answer.
Do you describe TaskFlow as "the DAG is the function-call graph" rather than as a vague "cleaner API"? — required answer.

Worked example — the same ETL, classic PythonOperator vs TaskFlow

Detailed explanation. The most useful side-by-side in the entire TaskFlow story. A four-task ETL — extract from a source API, transform the JSON, validate the row count, load into a warehouse — written both ways. The classic version has 60 lines of operator instantiations, op_kwargs dicts, and >> wiring; the TaskFlow version is 25 lines of Python functions with a return statement in each. The line count is not the point — the reader effort to trace what feeds what is the point.

The pipeline. extract → transform → validate → load, four tasks, all pure Python.
The comparison. Same behaviour, same schedule, same retries. Different authoring API.
The reader question. Which version can a new team member understand in 60 seconds?

Question. Rewrite a four-task ETL DAG from classic PythonOperator style to airflow taskflow api style. Show the exact line-count reduction, the dependency wiring change, and the XCom flow change.

Input.

Task	Purpose	Returns
extract	Fetch JSON from an API endpoint	`list[dict]` — raw rows
transform	Normalise timestamps, coerce types	`list[dict]` — cleaned rows
validate	Check row count is above a floor	`int` — validated row count
load	Insert into warehouse	`None`

Code.

# ---------- Classic PythonOperator version (60 lines) ----------
from datetime import datetime
from airflow import DAG
from airflow.operators.python import PythonOperator

def _extract(**context):
    rows = fetch_from_api(context["ds"])                 # ds = execution date
    context["ti"].xcom_push(key="raw_rows", value=rows)

def _transform(**context):
    raw = context["ti"].xcom_pull(task_ids="extract", key="raw_rows")
    cleaned = [normalise(r) for r in raw]
    context["ti"].xcom_push(key="cleaned_rows", value=cleaned)

def _validate(**context):
    cleaned = context["ti"].xcom_pull(task_ids="transform", key="cleaned_rows")
    n = len(cleaned)
    if n < 100:
        raise ValueError(f"row count {n} below floor 100")
    context["ti"].xcom_push(key="validated_count", value=n)

def _load(**context):
    cleaned = context["ti"].xcom_pull(task_ids="transform", key="cleaned_rows")
    count   = context["ti"].xcom_pull(task_ids="validate",  key="validated_count")
    load_to_warehouse(cleaned)
    print(f"loaded {count} rows")

with DAG(
    dag_id="events_etl_classic",
    start_date=datetime(2026, 6, 1),
    schedule="@hourly",
    catchup=False,
) as dag:
    extract  = PythonOperator(task_id="extract",  python_callable=_extract)
    transform = PythonOperator(task_id="transform", python_callable=_transform)
    validate = PythonOperator(task_id="validate", python_callable=_validate)
    load     = PythonOperator(task_id="load",     python_callable=_load)

    extract >> transform >> validate >> load


# ---------- TaskFlow version (25 lines) ----------
from datetime import datetime
from airflow.decorators import dag, task

@dag(
    dag_id="events_etl_taskflow",
    start_date=datetime(2026, 6, 1),
    schedule="@hourly",
    catchup=False,
)
def events_etl():

    @task
    def extract(ds: str) -> list[dict]:
        return fetch_from_api(ds)

    @task
    def transform(raw: list[dict]) -> list[dict]:
        return [normalise(r) for r in raw]

    @task
    def validate(cleaned: list[dict]) -> int:
        n = len(cleaned)
        if n < 100:
            raise ValueError(f"row count {n} below floor 100")
        return n

    @task
    def load(cleaned: list[dict], count: int) -> None:
        load_to_warehouse(cleaned)
        print(f"loaded {count} rows")

    raw      = extract("{{ ds }}")
    cleaned  = transform(raw)
    count    = validate(cleaned)
    load(cleaned, count)

events_etl()

Step-by-step explanation.

The classic version defines four _task functions with a **context signature. Each task reaches into context["ti"] (the TaskInstance) to xcom_push its output and xcom_pull its inputs by task ID and key. The keys (raw_rows, cleaned_rows, validated_count) are stringly-typed and appear in multiple places — rename any one of them and the DAG breaks silently.
The classic version also has a >> wiring line at the bottom: extract >> transform >> validate >> load. The wiring is a duplicate of information already implicit in which task's XCom key each task reaches for. A refactor that rearranges the pipeline requires updating both the XCom keys and the >> line.
The TaskFlow version replaces every xcom_push with a return statement and every xcom_pull with a function argument. The raw = extract("{{ ds }}") line binds the extract task's return value to a local variable of type XComArg. Passing raw into transform(raw) establishes both the dependency (transform runs after extract) and the XCom flow (transform's raw argument reads from extract's return XCom).
The >> wiring at the bottom is gone. Dependencies are inferred from the function-call graph: extract → transform → validate → load because each function call takes the previous function's return as an argument.
The type hints (list[dict], int) surface in the editor: hovering over raw shows XComArg[list[dict]]; the IDE flags a mismatch if you try to pass raw (a list) into a function expecting an int. The classic API never had this level of static safety.

Output.

Metric	Classic PythonOperator	TaskFlow
Lines of code	60	25
Explicit XCom calls	8 (`push`/`pull`)	0
Stringly-typed keys	3 (raw_rows, cleaned_rows, validated_count)	0
Wiring lines (`>>`)	1	0
IDE autocomplete on task args	no	yes
Reader effort to trace flow	high (grep keys)	low (read code top-down)

Rule of thumb. For any pipeline where every task is pure Python and the flow is a straight function-call graph, TaskFlow is unambiguously the right choice. Reserve classic operators for external-system calls where the operator's SDK wrapper genuinely earns its keep.

Worked example — the `op_kwargs` leak

Detailed explanation. A subtle but common classic-operator pain: passing arguments to a python_callable via op_kwargs. The function signature ends up documented in two places — the function definition (where it belongs) and the operator instantiation (where it's a dict of keyword-argument overrides). IDE autocomplete only sees one; refactoring the function signature silently breaks the operator.

The pattern. PythonOperator(task_id=..., python_callable=fn, op_kwargs={"threshold": 100}).
The bug. Rename threshold → min_rows in the function definition, forget to update op_kwargs, and the task fails at runtime with TypeError: fn() got unexpected keyword argument 'threshold'.
The TaskFlow fix. @task on the function; call it as fn(min_rows=100) in the DAG body. The DAG body is Python; the function signature is Python; there is one source of truth.

Question. Show the classic op_kwargs DAG and its TaskFlow equivalent. Highlight the refactor safety win.

Input.

Task	Signature	Kwargs
filter_rows	`filter_rows(rows: list, threshold: int) -> list`	threshold=100
upload	`upload(rows: list, bucket: str) -> None`	bucket="prod-events"

Code.

# Classic — op_kwargs at the operator site
from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime

def filter_rows(rows, threshold):
    return [r for r in rows if r["value"] >= threshold]

def upload(rows, bucket):
    write_to_s3(rows, bucket)

with DAG("classic_op_kwargs", start_date=datetime(2026, 6, 1), schedule=None):
    step1 = PythonOperator(
        task_id="filter_rows",
        python_callable=filter_rows,
        op_kwargs={
            "rows":      "{{ ti.xcom_pull(task_ids='previous_step') }}",
            "threshold": 100,          # ← if you rename `threshold` in the fn, this dict is stale
        },
    )
    step2 = PythonOperator(
        task_id="upload",
        python_callable=upload,
        op_kwargs={
            "rows":   "{{ ti.xcom_pull(task_ids='filter_rows') }}",
            "bucket": "prod-events",
        },
    )
    step1 >> step2


# TaskFlow — kwargs are ordinary Python function-call args
from airflow.decorators import dag, task
from datetime import datetime

@dag(dag_id="taskflow_kwargs", start_date=datetime(2026, 6, 1), schedule=None)
def pipeline():
    @task
    def filter_rows(rows: list, threshold: int) -> list:
        return [r for r in rows if r["value"] >= threshold]

    @task
    def upload(rows: list, bucket: str) -> None:
        write_to_s3(rows, bucket)

    previous = fetch_previous_rows()               # some upstream XComArg
    filtered = filter_rows(previous, threshold=100)
    upload(filtered, bucket="prod-events")

pipeline()

Step-by-step explanation.

In the classic version, the op_kwargs dict at the operator instantiation site is a shadow signature of the function. Every argument name in the dict must match the function parameter name exactly; the IDE cannot help because a dict-of-strings-to-values is opaque to autocomplete.
Rename threshold to min_rows in the function definition. Nothing complains at parse time; nothing complains at import time; the operator instantiation still references "threshold" as a dict key. At runtime, Python raises TypeError: filter_rows() got an unexpected keyword argument 'threshold' and the task fails.
In the TaskFlow version, filter_rows(previous, threshold=100) is an ordinary Python function call. The IDE autocompletes threshold=; a mypy check flags a type mismatch on previous; a rename refactor updates both the function definition and every call site atomically.
The rows parameter is passed by position — no explicit xcom_pull template, no task_ids=... string, no key=... string. The IDE knows previous is an XComArg[list] and knows filter_rows expects a list as its first argument. Match; move on.
The XCom flow is implicit in the call graph. filter_rows(previous, ...) establishes both the dependency (filter_rows runs after whatever produced previous) and the argument binding (filter_rows's rows parameter reads from previous's XCom).

Output.

Refactor scenario	Classic op_kwargs	TaskFlow
Rename param in fn	Silent break at runtime	Rename updates everywhere
Add new param	Update fn + op_kwargs dict	Update fn signature only
IDE autocomplete on arg	No	Yes
mypy sees the arg	No	Yes
Diff readability	Two-file diff	Single-file diff

Rule of thumb. Any argument you'd have put in op_kwargs becomes an ordinary keyword argument on the @task function call. The IDE, the type checker, and the code reviewer all benefit; the runtime break class disappears.

Worked example — mixing decorators with existing PythonOperator DAGs

Detailed explanation. Migration reality: nobody has the luxury of rewriting a 200-DAG Airflow deployment from scratch. The practical migration story is "add new tasks as @task inside existing PythonOperator DAGs, and rewrite whole DAGs to TaskFlow opportunistically when a substantial change lands anyway." The two APIs coexist without ceremony — a @task function inside a with DAG(...) block behaves exactly like a PythonOperator.

The setup. Existing DAG uses with DAG(...) + PythonOperators.
The addition. A new task needs to be added; the team writes it as @task to get the modern authoring ergonomics without rewriting the surrounding tasks.
The wiring. @task inside a with DAG block auto-registers with that DAG; the return value is still an XComArg; you can >> a @task to a PythonOperator or vice versa.

Question. Take an existing PythonOperator DAG with two tasks and add a third task as a @task function. Show that all three coexist correctly.

Input.

Task	Style	Purpose
load_raw	PythonOperator (existing)	Load raw rows into staging
profile	PythonOperator (existing)	Compute profile stats
notify	`@task` (new)	Post profile stats to Slack

Code.

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

def _load_raw(**ctx):
    rows = load_from_source(ctx["ds"])
    ctx["ti"].xcom_push(key="row_count", value=len(rows))

def _profile(**ctx):
    n = ctx["ti"].xcom_pull(task_ids="load_raw", key="row_count")
    stats = {"count": n, "date": ctx["ds"]}
    ctx["ti"].xcom_push(key="stats", value=stats)
    return stats                                        # also becomes an XCom via return_value

with DAG("mixed_dag", start_date=datetime(2026, 6, 1), schedule="@daily", catchup=False) as dag:
    load_raw = PythonOperator(task_id="load_raw", python_callable=_load_raw)
    profile  = PythonOperator(task_id="profile",  python_callable=_profile)

    @task
    def notify(stats: dict) -> None:
        post_to_slack(
            channel="#data-quality",
            text=f"Loaded {stats['count']} rows for {stats['date']}",
        )

    # Wiring — profile.output is an XComArg for the PythonOperator's return_value
    load_raw >> profile
    notify(profile.output)                              # XComArg from PythonOperator → @task

Step-by-step explanation.

The existing DAG has two PythonOperators: load_raw pushes a row_count XCom; profile pulls it, computes stats, and pushes them and returns them (returning from a PythonOperator's callable populates the special return_value XCom key).
The new @task function notify is defined inside the with DAG(...) block. Airflow registers it with the surrounding DAG automatically — no additional wiring is needed.
Calling notify(profile.output) is the interop bridge. profile.output is an XComArg referencing the PythonOperator profile's return_value XCom. Passing it to notify establishes (a) the dependency notify runs after profile, and (b) the argument binding stats is fetched from profile's return XCom.
The explicit load_raw >> profile is still needed because those two classic operators don't have an XCom-argument relationship in the DAG body — you have to state the dependency by hand. Once you refactor profile to a @task too, this line would disappear.
The DAG runs correctly: load_raw → profile → notify, with the XCom flow going through the classic xcom_push on load_raw, the classic return_value mechanism on profile, and the TaskFlow argument-pull on notify. Three XCom styles, one DAG, no ceremony.

Output.

Interop point	Style	Mechanism
load_raw → profile dep	classic	`load_raw >> profile`
load_raw → profile XCom	classic	`xcom_push(key="row_count")` + `xcom_pull(...)`
profile → notify dep	mixed	`notify(profile.output)` (implicit dep)
profile → notify XCom	mixed	`profile.output` → `stats` arg (auto pull)

Rule of thumb. You do not need to rewrite existing DAGs to adopt TaskFlow. Add new tasks as @task; wire them to existing PythonOperators via .output; rewrite whole DAGs only when the diff would be substantial anyway. Migration is opportunistic.

Senior interview question on the TaskFlow value proposition

A senior interviewer often opens with: "You inherit an Airflow codebase with 80 DAGs, all written using PythonOperator, none using TaskFlow. Walk me through why you'd migrate, how you'd sequence the migration, and what you'd tell a skeptic on the team who says 'PythonOperator works, why change?'"

Solution Using opportunistic migration + a 90-day rollout plan

# ------------------------------------------------------------------
# 90-day TaskFlow adoption plan
# ------------------------------------------------------------------
#
# Phase 1 — days 0-14  |  Set the bar for new code
#   - Team-wide rule: every NEW DAG uses @dag + @task
#   - Update code review checklist to reject PythonOperator on pure-Python work
#   - Ship a starter DAG template file (below)
#
# Phase 2 — days 15-45 |  Migrate high-touch DAGs
#   - Identify the top-10 DAGs by change frequency (git log --since='90 days ago')
#   - Rewrite each into TaskFlow when the next substantial change lands
#   - Do NOT rewrite for the sake of rewriting; wait for a real change
#
# Phase 3 — days 46-90 |  Mixed-DAG hardening
#   - For long-tail DAGs, add new tasks as @task inside existing DAGs
#   - Wire via profile.output → @task, or @task → PythonOperator via .output
#   - Document the interop patterns in the team wiki
#
# The skeptic answer:
#   PythonOperator works. But every new task is 2x the lines, 3x the strings
#   to keep in sync, and 0x IDE autocomplete. Over 12 months, TaskFlow saves
#   ~30% of DAG-authoring time and eliminates a class of stringly-typed bugs
#   (task_id / XCom key drift) that classic DAGs never quite escape.
# ------------------------------------------------------------------

# ---------- Starter template — the "canonical shape" ----------
from datetime import datetime
from airflow.decorators import dag, task

DEFAULT_ARGS = {
    "retries": 3,
    "retry_delay": 300,          # seconds
    "email_on_failure": False,
}

@dag(
    dag_id="team_starter_taskflow",
    start_date=datetime(2026, 6, 1),
    schedule="@daily",
    catchup=False,
    default_args=DEFAULT_ARGS,
    tags=["team:data-eng", "taskflow"],
)
def team_starter():

    @task
    def extract(ds: str) -> list[dict]:
        """Fetch source data for the logical date."""
        return fetch(ds)

    @task
    def transform(rows: list[dict]) -> list[dict]:
        """Normalise + type-coerce."""
        return [normalise(r) for r in rows]

    @task
    def load(rows: list[dict]) -> int:
        """Insert into warehouse; return loaded count."""
        return warehouse.insert(rows)

    load(transform(extract("{{ ds }}")))

team_starter()

Step-by-step trace.

Phase	Days	Scope	Migration action
1 — New code bar	0-14	All new DAGs	Team rule: `@dag` + `@task`; template checked in
2 — High-touch DAGs	15-45	Top 10 by change frequency	Rewrite when next substantial change lands
3 — Mixed-DAG hardening	46-90	Long-tail	Add new tasks as `@task` inside existing DAGs
Skeptic answer	anytime	Cultural	"PythonOperator works. TaskFlow saves 30% of authoring time."

After 90 days, the team's DAG code has a bimodal distribution: new DAGs are 100% TaskFlow; old DAGs are mixed (existing PythonOperators + new @task additions); a shrinking tail of legacy DAGs remains pure PythonOperator until they earn a rewrite. No big-bang migration; no productivity dip; the top-10 DAGs are modernised where the maintenance cost was highest.

Output:

Metric	Before	After 90 days
New DAGs using TaskFlow	0%	100%
Top-10 DAGs modernised	0	10
Long-tail DAGs with any `@task`	0	~40%
Average lines per new DAG	~120	~50
Stringly-typed XCom key bugs	monthly	eliminated

Why this works — concept by concept:

Opportunistic migration — the rule "rewrite when the next substantial change lands" avoids the productivity-killing big-bang migration. Every DAG that gets rewritten is one where the diff is happening anyway; the marginal cost of the API change is small.
New-code bar first — the fastest way to stop the codebase getting worse is to require TaskFlow for all new work. Migration then only has to catch up with the existing pile, not race an ever-growing debt.
Starter template — a canonical DAG template checked into the repo removes decision fatigue for the team. Every new DAG starts from a known-good shape; deviations become intentional design choices, not accidental drift.
Interop, not swap — mixed DAGs are first-class. Add @task alongside existing operators; wire via .output; don't force a full rewrite for a single new step. This is what makes the migration story feasible in the first place.
Cost — 90 days of low-intensity migration; zero downtime; the top-10 DAGs (by change frequency) get the highest ROI on the modernisation because they're the ones the team edits most often. O(1) team investment; O(changes) return on lower authoring friction.

ETL
Topic — etl
ETL problems on Airflow DAG authoring patterns

Practice →

Optimization Topic — optimization Optimization problems on pipeline refactoring

Practice →

2. `@task` and `@dag` decorators — the two-decorator pipeline

Two decorators that convert plain Python into a DAG — `@dag` on the factory function, `@task` on each task function

The mental model in one line: @dag turns a Python function into a DAG factory (calling it builds a DAG object; the function body runs at parse time to register tasks), and @task turns a Python function into a task factory (calling it inside the @dag function creates a task instance and returns an XComArg that represents its future return value). Every other TaskFlow question is a consequence of understanding those two sentences.

The four axes interviewers actually probe on decorators.

Factory semantics. Both @dag and @task are factory decorators. Calling the decorated function does not run the function's original body directly — it builds a DAG (in the case of @dag) or a task instance (in the case of @task). This is the single most confusing point for newcomers and the highest-signal probe an interviewer can ask.
Parameters accepted. @dag(schedule=..., start_date=..., catchup=False, tags=[...], default_args={...}) — same knobs as the classic DAG(...) constructor. @task(retries=3, retry_delay=..., trigger_rule=..., multiple_outputs=True, task_id="explicit_name") — same knobs as most classic operators.
Return-value semantics. A @task function's return value becomes the task's return_value XCom. If the return type is a dict and @task(multiple_outputs=True), each dict key becomes a separate XCom entry keyed by the dict key.
Function-call semantics inside the DAG. Calling extract() inside a @dag function does not run the extract task — it registers an extract task instance with the DAG and returns an XComArg that represents extract's future return. The XComArg is the currency of dependency wiring.

How @dag transforms an ordinary function into a DAG factory.

Parse-time execution. When Airflow parses the DAG file, it imports the module. Importing runs the top-level @dag-decorated function definition but not its body. The final line events_etl() at the bottom of the file is what runs the body — and running the body registers all the @task instances inside it with a new DAG object.
The dag_id. Optional; defaults to the function name. Setting it explicitly lets you rename the Python function without breaking the DAG registration.
Positional-only Python semantics. The @dag-decorated function is a regular Python function; you can accept parameters, thread them into the task calls, and use them as configuration. Airflow does not do anything magical here — it's ordinary Python.
Multiple DAGs per file. You can define multiple @dag functions in one Python file and call each of them at module level; each call registers a separate DAG. This is the idiom for "same DAG shape, different schedules or parameters."

How @task transforms an ordinary function into a task factory.

Deferred execution. Calling extract("{{ ds }}") inside the @dag function does not run extract's body. It creates a PythonOperator-equivalent task instance in the DAG and returns an XComArg that represents the future return value of that task.
The task_id. Defaults to the function name. Setting @task(task_id="custom") lets you have multiple task instances from the same function with different IDs (useful for the same transform run on different inputs).
Trigger rules and retries. All the classic operator knobs work as keyword arguments to @task: retries, retry_delay, trigger_rule, pool, queue, priority_weight, max_active_tis_per_dag, and so on.
Idempotency. The function body runs on the Airflow worker; it should be idempotent (same input → same output → same side effects). This is the same rule as classic operators, restated in Python-function terms.

The multiple_outputs unlock.

The default. @task returns a single value; it becomes one XCom with key return_value.
The unlock. @task(multiple_outputs=True) on a function that returns a dict splits each dict key into its own XCom. Downstream tasks can then read individual keys as separate XComArg references.
Why it matters. Lets you return a structured payload from one task and pass parts of it to different downstream tasks, without either passing the whole dict everywhere or hand-coding an unpacking step.

Common interview probes on the two decorators.

"Walk me through what happens when Airflow parses a TaskFlow DAG file." — required answer covers the factory-function semantics and the parse-time execution of the @dag body.
"What's the difference between the function definition and the function call in a TaskFlow DAG?" — the definition builds the factory; the call at the bottom of the file runs the factory and registers the DAG.
"How would you have the same DAG shape with three different schedules?" — one @dag-decorated factory function that accepts a schedule parameter; call it three times at module level with different values.
"When would you use @task(multiple_outputs=True)?" — when one task returns a structured payload and multiple downstream tasks need different pieces without unpacking gymnastics.

Worked example — a 4-task DAG in 20 lines

Detailed explanation. The archetypal TaskFlow starter. Four tasks — extract → transform → validate → load — written top-to-bottom as ordinary Python functions with @task, wrapped in a @dag-decorated factory function, called at module level. Twenty lines of code; the dependency graph is entirely implicit in the function-call composition.

The shape. One @dag factory; four @task functions; four @task calls; one bottom-of-file factory() invocation.
The wiring. Zero explicit >>; every dependency is implicit in the argument passing.
The XCom. Zero explicit xcom_push/xcom_pull; every return value flows via XCom automatically.

Question. Write a complete four-task ETL DAG in TaskFlow style. Show the parsed dependency graph and enumerate every XCom that flows.

Input.

Task	Signature	Behaviour
extract	`(ds: str) -> list[dict]`	Fetch rows for the logical date
transform	`(rows: list[dict]) -> list[dict]`	Normalise fields
validate	`(rows: list[dict]) -> int`	Return row count if above floor, else raise
load	`(rows: list[dict], count: int) -> None`	Insert into warehouse

Code.

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

@dag(
    dag_id="starter_taskflow",
    start_date=datetime(2026, 6, 1),
    schedule="@hourly",
    catchup=False,
    tags=["taskflow", "starter"],
)
def starter_taskflow():

    @task
    def extract(ds: str) -> list[dict]:
        return fetch_from_source(ds)

    @task
    def transform(rows: list[dict]) -> list[dict]:
        return [normalise(r) for r in rows]

    @task
    def validate(rows: list[dict]) -> int:
        if len(rows) < 100:
            raise ValueError(f"row count {len(rows)} below floor 100")
        return len(rows)

    @task
    def load(rows: list[dict], count: int) -> None:
        warehouse.insert(rows)
        print(f"loaded {count} rows")

    raw     = extract("{{ ds }}")
    cleaned = transform(raw)
    n       = validate(cleaned)
    load(cleaned, n)

starter_taskflow()

Step-by-step explanation.

The @dag-decorated function starter_taskflow is a factory. The final line starter_taskflow() runs the factory: this executes the function body, which registers four task instances and their dependencies with a new DAG object.
raw = extract("{{ ds }}") — calling the @task-decorated extract function creates a task instance with task_id="extract" and returns an XComArg bound to that instance's future return value. raw is an XComArg, not a list[dict], at DAG-parse time.
cleaned = transform(raw) — creates a task instance transform, records the dependency transform depends on extract (because raw is an XComArg from extract), records the argument binding (transform's rows parameter will read from extract's return XCom at run time), returns an XComArg bound to transform's future return.
n = validate(cleaned) and load(cleaned, n) — same pattern. load takes two XComArgs (cleaned from transform, n from validate), so it depends on both; both dependencies are recorded implicitly.
At run time, the Airflow scheduler walks the DAG top-down: extract runs first, pushes its return to XCom under key return_value; transform runs next, pulls extract's return_value, runs its body with the pulled value as rows, pushes its own return; and so on.

Output.

Task	task_id	XCom in	XCom out	Depends on
extract	extract	ds="{{ ds }}"	list[dict]	(none)
transform	transform	list[dict] from extract	list[dict]	extract
validate	validate	list[dict] from transform	int	transform
load	load	list[dict] from transform, int from validate	None	transform, validate

Rule of thumb. The "shape" of a TaskFlow DAG is the shape of the function-call graph inside @dag. Draw the call graph and you have drawn the DAG.

Worked example — parametrised DAG factory (same shape, three schedules)

Detailed explanation. A common team pattern: the same ETL DAG needs to run on three different schedules — hourly for a hot source, daily for a warm source, weekly for an archive source — with slightly different retry policies. Instead of copy-pasting three DAG files, factor the DAG into a parametrised @dag function and call it three times at module level with different arguments.

The parameters. dag_id, schedule, source_name, retries.
The idiom. One @dag-returning factory function that accepts the parameters; call it once per schedule at module level.
The benefit. Single source of truth for the DAG shape; three separate DAGs registered with the scheduler.

Question. Refactor a hard-coded DAG into a parametrised factory and instantiate three variants. Show the module-level calls and the resulting DAG IDs.

Input.

Variant	dag_id	schedule	source_name	retries
Hot	events_etl_hot	@hourly	hot	5
Warm	events_etl_warm	@daily	warm	3
Archive	events_etl_archive	@weekly	archive	1

Code.

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

def make_events_etl(dag_id: str, schedule: str, source_name: str, retries: int):

    @dag(
        dag_id=dag_id,
        start_date=datetime(2026, 6, 1),
        schedule=schedule,
        catchup=False,
        default_args={"retries": retries},
        tags=["taskflow", "events-etl", f"source:{source_name}"],
    )
    def events_etl():

        @task
        def extract(ds: str) -> list[dict]:
            return fetch(source_name, ds)

        @task
        def transform(rows: list[dict]) -> list[dict]:
            return [normalise(r, source=source_name) for r in rows]

        @task
        def load(rows: list[dict]) -> None:
            warehouse.insert(rows, table=f"events_{source_name}")

        load(transform(extract("{{ ds }}")))

    return events_etl()

# Register three variants at module level
hot     = make_events_etl("events_etl_hot",     "@hourly", "hot",     retries=5)
warm    = make_events_etl("events_etl_warm",    "@daily",  "warm",    retries=3)
archive = make_events_etl("events_etl_archive", "@weekly", "archive", retries=1)

Step-by-step explanation.

The outer make_events_etl is not a DAG — it's a plain Python function that returns one. Its parameters capture the axes on which the three DAG variants differ: dag_id, schedule, source_name, retries.
Inside make_events_etl, the @dag-decorated inner function events_etl closes over the outer function's parameters. Each call to make_events_etl(...) creates a fresh closure with its own source_name and retries, then invokes the factory to build and register a DAG.
The @task functions inside close over source_name. When extract runs at task execution time, it calls fetch(source_name, ds) with the closed-over value — so the "hot" DAG's extract calls fetch("hot", ...) and the "archive" DAG's extract calls fetch("archive", ...).
retries is passed as default_args={"retries": retries} on the @dag. This maps to Airflow's classic DAG-level default_args and applies to every task in the DAG unless a specific @task(retries=...) overrides it.
The three module-level calls at the bottom of the file register three separate DAGs with the scheduler. Airflow's DAG parser sees three distinct dag_id values, spins up three DAGs, and schedules each independently.

Output.

DAG registered	Schedule	source_name (closed over)	retries
events_etl_hot	@hourly	hot	5
events_etl_warm	@daily	warm	3
events_etl_archive	@weekly	archive	1

Rule of thumb. Any time you're about to copy-paste a DAG file to make a variant, factor the DAG into a plain Python function that returns a DAG and call it once per variant. The DRY win is enormous; the Airflow scheduler sees three separate DAGs and handles them independently.

Worked example — `multiple_outputs=True` for structured returns

Detailed explanation. A task computes a summary of a dataset and returns a dict with three keys: row_count, total_revenue, date_range. Three downstream tasks need different parts: the alerting task wants row_count; the pricing task wants total_revenue; the audit task wants date_range. Without multiple_outputs, every downstream task pulls the whole dict and indexes into it. With multiple_outputs=True, each key becomes a separate XCom entry that downstream tasks can reference individually.

The return. {"row_count": 12345, "total_revenue": 987.65, "date_range": "2026-06-01/2026-06-22"}.
The consumers. Three downstream tasks, each wants one key.
The unlock. @task(multiple_outputs=True) splits the dict into per-key XComArg references.

Question. Refactor a task that returns a dict into multiple_outputs=True and wire three downstream tasks to consume different keys. Show the XCom entries that get pushed.

Input.

Downstream task	Consumes	XCom key
alert_on_low_rows	row_count	row_count
adjust_price	total_revenue	total_revenue
audit_log	date_range	date_range

Code.

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

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

    @task(multiple_outputs=True)
    def summarise(ds: str) -> dict:
        """Returns a dict; each key becomes its own XCom entry."""
        return {
            "row_count":     count_rows(ds),
            "total_revenue": sum_revenue(ds),
            "date_range":    f"{ds}/{ds}",
        }

    @task
    def alert_on_low_rows(row_count: int) -> None:
        if row_count < 1000:
            page(f"Row count {row_count} below floor")

    @task
    def adjust_price(total_revenue: float) -> None:
        pricing_engine.update(total_revenue)

    @task
    def audit_log(date_range: str) -> None:
        audit.write(date_range)

    summary = summarise("{{ ds }}")           # summary is a dict-of-XComArgs

    alert_on_low_rows(summary["row_count"])
    adjust_price(summary["total_revenue"])
    audit_log(summary["date_range"])

multiple_outputs_demo()

Step-by-step explanation.

@task(multiple_outputs=True) tells Airflow: this task returns a dict; push each key/value pair as a separate XCom entry (keys row_count, total_revenue, date_range) instead of one XCom with key return_value holding the whole dict.
summary = summarise("{{ ds }}") binds summary to a special object — a dict-like XComArg holder — where indexing summary["row_count"] returns an XComArg bound to the row_count XCom entry specifically.
alert_on_low_rows(summary["row_count"]) passes the row_count-specific XComArg. At run time, the alert task pulls only the row_count XCom, not the whole dict. The XCom traffic across the task boundary is a single int, not a dict.
The same pattern for adjust_price and audit_log. Each downstream task depends on summarise (implicit from the XComArg lookup) and pulls only its needed key.
Contrast with the non-multiple_outputs version: summary would be an XComArg[dict]; each downstream task would pull the whole dict and index into it at runtime — three copies of the same dict flowing across three XCom pulls. The multiple_outputs unlock reduces XCom traffic proportional to the dict size.

Output.

XCom key on summarise	Value	Pulled by
row_count	int	alert_on_low_rows
total_revenue	float	adjust_price
date_range	str	audit_log

Rule of thumb. Any @task whose return type is a dict and whose downstream consumers want different keys is a candidate for multiple_outputs=True. The XCom traffic drops from N × sizeof(dict) to N × sizeof(individual_key); the downstream signatures become type-safe.

Senior interview question on decorator semantics

A senior interviewer might ask: "Walk me through what actually happens when Airflow parses a TaskFlow DAG file. When does the @dag-decorated function's body run? When do the @task function bodies run? And what does dag_id="foo" at the top actually do versus the trailing foo() at the bottom?"

Solution Using the parse-time / run-time split as the mental model

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

# ==============================================================
# What runs at PARSE time (Airflow scheduler imports this module)
# ==============================================================
#   1. The module body executes top-to-bottom.
#   2. The @dag decorator wraps the events_etl function definition
#      into a DAG factory — the factory is a callable object, but
#      the events_etl function body has NOT run yet.
#   3. The final line events_etl() invokes the factory.
#      This DOES run the events_etl body:
#        - Each @task decorator wraps its function into a task factory.
#        - Each call like extract("{{ ds }}") registers a task instance
#          in a new DAG object and returns an XComArg.
#        - The DAG object is returned by the factory.
#   4. Airflow stores the returned DAG object in the scheduler catalog.
#
# What runs at RUN time (Airflow scheduler dispatches a task)
# ==============================================================
#   1. Scheduler picks the extract task, sends it to a worker.
#   2. Worker imports the module (same parse-time flow as above).
#   3. Worker looks up the extract task instance by task_id="extract".
#   4. Worker invokes the original extract function body with ds
#      resolved from the Jinja template.
#   5. Worker pushes the return value to XCom under key "return_value".
#   6. Repeat for transform, validate, load — each on its own worker
#      run, importing the module fresh each time.

@dag(
    dag_id="parse_vs_run",
    start_date=datetime(2026, 6, 1),
    schedule="@hourly",
    catchup=False,
)
def events_etl():

    @task
    def extract(ds: str) -> list[dict]:
        return fetch_from_source(ds)

    @task
    def transform(rows: list[dict]) -> list[dict]:
        return [normalise(r) for r in rows]

    @task
    def load(rows: list[dict]) -> None:
        warehouse.insert(rows)

    load(transform(extract("{{ ds }}")))

# The trailing invocation — this is what actually registers the DAG.
# Without this line, the module defines a factory but registers nothing.
events_etl()

Step-by-step trace.

Event	When	What runs
Airflow parses the module	parse time	Module body top-to-bottom
`@dag` decorator applies	parse time	Wraps `events_etl` into a factory (function body NOT run)
`events_etl()` at bottom	parse time	Runs the factory — DOES execute the `events_etl` body
`@task` decorators apply	parse time (inside factory)	Each `@task` wraps its function into a task factory
`extract("{{ ds }}")` call	parse time (inside factory)	Registers extract task instance; returns XComArg
`load(transform(extract(...)))`	parse time	Records the full dependency chain
Scheduler dispatches extract	run time	Worker imports module, invokes original `extract` body
Worker pushes return_value	run time	XCom entry `(dag_id, run_id, task_id="extract", key="return_value")`
Scheduler dispatches transform	run time	Worker pulls extract's XCom, invokes original `transform` body
Continue for validate + load	run time	Same pattern

The clean split — decorators, function definitions, and the outer events_etl() call run at parse time; original task function bodies run at run time on workers — is the mental model that resolves every "why doesn't this work the way I expect" question about TaskFlow.

Output:

Surface	Parse time	Run time
Module import	yes	yes (worker re-imports)
`@dag` factory built	yes	(already built)
`events_etl()` invocation	yes	(once, at parse)
`@task` factory built	yes	(already built)
`extract(...)` call in DAG body	yes → registers task	(not called at run time)
Original `extract` function body	no	yes (once per DAG run)

Why this works — concept by concept:

Factory decorators — both @dag and @task are factory decorators. The decorated function is not called directly to run the task; it's called to build the task's registration in the DAG graph. The original function body runs only when the worker executes the task at run time.
Trailing invocation — the events_etl() line at the bottom is not decoration; it is the actual DAG registration. Forget the trailing call and Airflow parses a module full of factories but registers zero DAGs. This is the #1 gotcha for TaskFlow beginners.
Worker re-import — every worker that runs a task first re-imports the whole DAG module, which re-runs the parse-time flow. This means expensive top-level code (e.g. network calls, DB connections at import) hits every worker on every task. Keep the module body cheap.
XComArg as the dependency currency — inside the factory body, task function calls return XComArg objects, not the function's actual return type. XComArg is what carries the dependency + XCom-flow information; the actual value materialises at run time.
Cost — parse-time cost is paid on scheduler import and on every worker task import. Run-time cost is the actual work. Keep the parse-time cost O(nodes in DAG); push the actual work into the run-time function bodies where it belongs.

ETL
Topic — etl
ETL problems on DAG parse-time semantics

Practice →

SQL Topic — sql SQL and pipeline factory-pattern problems

Practice →

3. XCom automation and XComArg chaining

Return values push XComs, arguments pull XComs, `XComArg` is the object that carries the reference — with a hard 48KB size limit that forces a stage for real data

The mental model in one line: every @task return value is pushed to XCom automatically under key return_value, every @task function argument that receives an XComArg pulls the corresponding XCom automatically at task start, and every XComArg in your DAG body is a lightweight reference to a future XCom entry — but the XCom backend (by default the Airflow metadata DB) has a hard 48KB per-row limit that makes XCom a control-plane bus for small metadata, not a data-transfer bus for real payloads. The moment your task returns a DataFrame or a list of 100k rows, you are one deploy away from either an OOM on the scheduler or a "row too large" MetadataDB error.

The four axes interviewers actually probe on XComs.

Push semantics. Every @task return value is pushed. Key is return_value by default; multiple_outputs=True splits a dict return into per-key XComs. The push happens on the worker after the task function body returns.
Pull semantics. Every function argument that receives an XComArg pulls automatically at task start. Passing a literal (a plain string / int / list) does not pull — literals are captured at parse time and rendered at run time via Jinja if templated.
XComArg as the reference type. Inside the @dag body, task calls return XComArgs, not the function's return type. Operations on XComArgs (indexing for multiple_outputs, chaining into another task call) build the dependency graph.
Size limits and the stage pattern. XCom rows are stored in the Airflow metadata DB (Postgres by default). The per-row limit is 48KB in most deployments. Real data must be staged in an external store (S3, GCS, ADLS) with only the pointer flowing through XCom.

How XCom push works under the hood.

The write. After the task function returns, the Airflow worker serializes the return value (default: JSON; configurable to pickle) and writes a row into the xcom table with (dag_id, run_id, task_id, key, value, timestamp).
The size cap. The value column is a Text in the metadata schema; Postgres will happily accept large values, but the JSON-serialization + network + de-serialization on the puller side dominate cost long before you hit the raw DB limit. Community wisdom is to keep payloads under a few KB.
The custom backend. Airflow ships a XCom base class you can subclass to change the storage backend — e.g. S3XComBackend writes the actual payload to S3 and stores only the S3 key in the metadata DB. This lets you keep large XComs without exploding the metadata DB.

How XCom pull works under the hood.

The read. At task start, the worker sees the task's argument bindings, resolves each XComArg to a (dag_id, run_id, task_id, key) tuple, executes a SELECT on the xcom table, deserializes the value, and passes it to the task function.
The template-string alternative. In classic operators, XCom pulls happen inline in Jinja templates like {{ ti.xcom_pull(task_ids="extract") }}. TaskFlow replaces this with typed function arguments, but the underlying SELECT is identical.
Explicit pull inside @task. You can still call context["ti"].xcom_pull(...) inside a @task function if you need to pull a specific key or from a task not in the argument chain — but 95% of the time the automatic pull via the argument is what you want.

XComArg chaining — the "functional DAG" bit.

Chained calls. load(transform(extract("{{ ds }}"))) is legal and idiomatic. Each function call returns an XComArg; the next call takes that XComArg as an argument; the DAG scheduler unwinds the nested calls into extract → transform → load.
Local variables. Assigning to local variables (raw = extract(...); cleaned = transform(raw)) is exactly equivalent to the chained form; it's just easier to read for longer pipelines.
Fan-out. One XComArg passed as an argument to two different downstream tasks establishes both dependencies. transform(raw) and validate(raw) both depend on extract; both read extract's return XCom at run time.
Fan-in. Two XComArgs passed as different arguments to one downstream task establish two dependencies from that task. merge(cleaned_a, cleaned_b) depends on both cleaned_a's upstream and cleaned_b's upstream.

The size limit — why XCom is a control-plane bus.

The wrong pattern. A @task that returns a Pandas DataFrame of 100k rows. Airflow tries to JSON-serialize the DataFrame, the JSON string is 20MB, the write to the metadata DB succeeds (or fails, depending on your Postgres row-size limits), the read on the puller side is slow, and every subsequent DAG run adds another 20MB row to the xcom table.
The right pattern. The extract task writes the raw data to S3 and returns the S3 key as a small string XCom. The transform task takes the S3 key as an argument, reads from S3, writes the transformed data back to S3, returns a new S3 key. XCom carries the pointers; S3 carries the data.
The hybrid. Small metadata (counts, dates, status flags, error messages) flows through XCom; large payloads (rows, blobs) flow through S3/GCS with the object key as the XCom.

Common interview probes on XCom automation.

"How do XComs work in TaskFlow?" — return pushes, argument pulls, XComArg is the reference.
"What happens if my @task returns a Pandas DataFrame with a million rows?" — you probably hit the 48KB row limit or, worse, get a slow, memory-hungry pipeline. Stage in S3.
"What is XComArg?" — the object returned by calling a @task function inside @dag; it references a future XCom entry.
"How do you fan out from one task to five downstream tasks with the same input?" — pass the upstream XComArg as an argument to all five; each dependency is established implicitly.

Worked example — passing a dict of results between tasks

Detailed explanation. A common pattern: a task computes a summary dict (row count, error count, latest timestamp) and passes it to two downstream tasks — one that decides whether to alert on the error count, one that logs the summary to a metrics store. The dict is small (a few hundred bytes), well within XCom's happy path. Show the full flow.

The upstream. compute_summary returns {"rows": 12345, "errors": 3, "latest_ts": "2026-06-22T14:00:00Z"}.
The downstream A. maybe_alert reads the summary, pages if errors > 0.
The downstream B. log_metrics reads the summary, writes each field to a metrics store.

Question. Write the three-task TaskFlow DAG and trace exactly which XComs are pushed and pulled.

Input.

Task	Return	Pushed XCom	Pulled from
compute_summary	dict	`return_value`	(none)
maybe_alert	None	(none)	compute_summary.return_value
log_metrics	None	(none)	compute_summary.return_value

Code.

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

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

    @task
    def compute_summary(ds: str) -> dict:
        rows = count_rows(ds)
        errors = count_errors(ds)
        return {
            "rows": rows,
            "errors": errors,
            "latest_ts": max_timestamp(ds),
        }

    @task
    def maybe_alert(summary: dict) -> None:
        if summary["errors"] > 0:
            page(f"{summary['errors']} errors in run for {summary['latest_ts']}")

    @task
    def log_metrics(summary: dict) -> None:
        metrics.gauge("etl.rows",   summary["rows"])
        metrics.gauge("etl.errors", summary["errors"])

    summary = compute_summary("{{ ds }}")
    maybe_alert(summary)
    log_metrics(summary)

summary_fanout()

Step-by-step explanation.

At parse time, summary = compute_summary("{{ ds }}") registers a compute_summary task instance and returns an XComArg bound to its return_value XCom key. summary is not a dict at parse time; it's a reference.
maybe_alert(summary) and log_metrics(summary) both take the same XComArg as their argument. This establishes two edges in the DAG: compute_summary → maybe_alert and compute_summary → log_metrics. Both downstream tasks pull the same XCom.
At run time, the scheduler dispatches compute_summary first. The worker runs the function body, gets the dict, serializes it to JSON, writes a row (dag_id, run_id, "compute_summary", "return_value", <json>) into the metadata DB, and marks the task success.
The scheduler then dispatches maybe_alert and log_metrics in parallel (they have no dependency on each other). Each worker starts up, sees its summary argument is an XComArg, SELECTs the row from the metadata DB, deserializes the JSON back to a dict, and calls the function body with the dict as the argument.
Both downstream tasks see the same dict. The dict is small (a few dozen bytes), well under any size limit. The XCom traffic is one write and two reads.

Output.

Event	Time	Where	Detail
compute_summary runs	t=0	worker A	writes JSON row to xcom table
maybe_alert reads XCom	t=1	worker B	SELECT + deserialize
log_metrics reads XCom	t=1	worker C	SELECT + deserialize (parallel)
Both downstreams run	t=1	workers B, C	function bodies execute
DAG completes	t=2	scheduler	mark run success

Rule of thumb. For small structured payloads (dicts of counts, dates, status flags), XCom is the right transport — one line of Python, no ceremony, sub-KB per row. Reserve the S3-staging pattern for payloads that are naturally large.

Worked example — the 48KB gotcha and the S3 stage pattern

Detailed explanation. The most common TaskFlow anti-pattern: a @task that returns a Pandas DataFrame. The DataFrame gets JSON-serialized, the JSON string exceeds Postgres's practical row-size comfort zone (or hits an explicit Text limit), and the DAG either fails at push time or grinds through 20MB XCom transfers on every task boundary. The fix is the S3 stage pattern: write the data to S3, return the S3 key as a small string XCom.

The trigger. Upstream task returns pd.DataFrame with 100k rows.
The failure mode. Either (a) MetadataDB rejects the row, or (b) the DAG runs but is 10x slower than it should be because every task boundary shuffles 20MB.
The fix. Upstream writes to S3, returns the S3 key. Downstream takes the key as an argument, reads from S3.

Question. Refactor an anti-pattern DAG (return DataFrame via XCom) into the S3-stage pattern. Show the two versions side by side.

Input.

Task	Input	Output
extract_frame	ds	pd.DataFrame (100k rows, ~20MB)
transform_frame	pd.DataFrame	pd.DataFrame (100k rows, ~20MB)
load_frame	pd.DataFrame	None

Code.

# ---------- Anti-pattern — DataFrame through XCom ----------
from datetime import datetime
from airflow.decorators import dag, task
import pandas as pd

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

    @task
    def extract_frame(ds: str) -> pd.DataFrame:
        df = pd.read_sql(f"SELECT * FROM events WHERE ds = '{ds}'", conn)
        return df                       # ← 20MB DataFrame through XCom (likely fails)

    @task
    def transform_frame(df: pd.DataFrame) -> pd.DataFrame:
        df["value_x10"] = df["value"] * 10
        return df                       # ← another 20MB through XCom

    @task
    def load_frame(df: pd.DataFrame) -> None:
        df.to_sql("events_processed", conn, if_exists="append", index=False)

    load_frame(transform_frame(extract_frame("{{ ds }}")))

anti_pattern()


# ---------- S3 stage pattern — pointer through XCom, data in S3 ----------
from datetime import datetime
from airflow.decorators import dag, task
import pandas as pd
import io, boto3

BUCKET = "airflow-stage"
s3 = boto3.client("s3")

def write_stage(df: pd.DataFrame, key: str) -> str:
    buf = io.BytesIO()
    df.to_parquet(buf, index=False)
    buf.seek(0)
    s3.put_object(Bucket=BUCKET, Key=key, Body=buf.getvalue())
    return key

def read_stage(key: str) -> pd.DataFrame:
    obj = s3.get_object(Bucket=BUCKET, Key=key)
    return pd.read_parquet(io.BytesIO(obj["Body"].read()))

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

    @task
    def extract_frame(ds: str, run_id: str) -> str:
        df = pd.read_sql(f"SELECT * FROM events WHERE ds = '{ds}'", conn)
        return write_stage(df, f"{run_id}/extract.parquet")   # ← returns S3 key (small string)

    @task
    def transform_frame(extract_key: str, run_id: str) -> str:
        df = read_stage(extract_key)
        df["value_x10"] = df["value"] * 10
        return write_stage(df, f"{run_id}/transform.parquet") # ← returns S3 key

    @task
    def load_frame(transform_key: str) -> None:
        df = read_stage(transform_key)
        df.to_sql("events_processed", conn, if_exists="append", index=False)

    extract_key   = extract_frame("{{ ds }}", "{{ run_id }}")
    transform_key = transform_frame(extract_key, "{{ run_id }}")
    load_frame(transform_key)

s3_stage_pattern()

Step-by-step explanation.

In the anti-pattern, extract_frame returns a 20MB DataFrame. Airflow tries to JSON-serialize it via a custom serializer (or fails), pushes the serialized blob to XCom, and the metadata DB either rejects the row or accepts it grudgingly. The transform_frame task pulls the same 20MB blob back, deserializes it, and repeats the cycle.
Every task boundary shuffles the full DataFrame through JSON. For a three-task DAG, that's three serializations, three deserializations, and three metadata-DB round trips of 20MB each. The scheduler's metadata DB fills up quickly; the DAG is slow and fragile.
In the S3 stage pattern, extract_frame writes the DataFrame to S3 at a key derived from the DAG run_id. The return is the S3 key — a ~50-byte string. XCom carries only the pointer.
transform_frame takes the S3 key as an argument, reads the DataFrame from S3, transforms it, writes the result to a new S3 key, returns the new key. The XCom traffic across the task boundary is 50 bytes; the actual data flows through S3 (fast, unlimited, priced at S3 rates).
The load_frame task takes the final key, reads once, writes to the warehouse. The XCom volume across the whole DAG is a few hundred bytes; the DataFrame is materialized on the worker that needs it and nowhere else.

Output.

Metric	Anti-pattern	S3 stage pattern
XCom volume per DAG run	~60MB (3 × 20MB)	~150 bytes (3 × 50 bytes)
MetadataDB row count	3 large rows	3 tiny rows
Data transport	XCom (slow)	S3 (fast, parallel)
Failure mode	MetadataDB rejects row	Rare (bounded by S3 SLA)
Cost of DataFrame doubling	2x XCom volume	S3 cost only

Rule of thumb. Any XCom payload that could plausibly exceed a few KB belongs in a stage. The run_id template variable makes stage keys naturally unique across DAG runs; cleanup can be a downstream sweep task or an S3 lifecycle policy.

Worked example — fan-out with the same XComArg

Detailed explanation. A single upstream task produces a list of items; five downstream tasks each process a different aspect of that list — one aggregates, one dedupes, one enriches, one alerts, one archives. All five need the same input. The clean pattern: pass the same XComArg as the argument to all five.

The upstream. extract_events returns a list of event dicts.
The downstreams. Five tasks, each reading the same list.
The DAG shape. One-to-five fan-out; five parallel tasks converging (or not) into a final sink.

Question. Write the DAG and show that all five downstream tasks depend on extract_events and pull the same XCom.

Input.

Downstream	Purpose
aggregate	Compute totals
dedupe	Remove duplicates
enrich	Add lookup fields
alert	Page if error count > 0
archive	Write raw to cold storage

Code.

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

@dag(dag_id="fanout_demo", start_date=datetime(2026, 6, 1), schedule="@hourly", catchup=False)
def fanout_demo():

    @task
    def extract_events(ds: str) -> list[dict]:
        return fetch_events(ds)

    @task
    def aggregate(events: list[dict]) -> dict:
        return {"total": sum(e["value"] for e in events)}

    @task
    def dedupe(events: list[dict]) -> list[dict]:
        seen, out = set(), []
        for e in events:
            if e["id"] not in seen:
                seen.add(e["id"])
                out.append(e)
        return out

    @task
    def enrich(events: list[dict]) -> list[dict]:
        return [{**e, "region": lookup_region(e["ip"])} for e in events]

    @task
    def alert(events: list[dict]) -> None:
        errs = [e for e in events if e.get("error")]
        if errs:
            page(f"{len(errs)} error events")

    @task
    def archive(events: list[dict]) -> None:
        write_cold_storage(events)

    events = extract_events("{{ ds }}")
    aggregate(events)
    dedupe(events)
    enrich(events)
    alert(events)
    archive(events)

fanout_demo()

Step-by-step explanation.

events = extract_events("{{ ds }}") registers the extract task and returns an XComArg. events is the reference; it will be resolved to the actual list at run time by whichever task pulls it.
Each of the five downstream calls — aggregate(events), dedupe(events), enrich(events), alert(events), archive(events) — takes the same XComArg. Each call establishes a dependency edge from extract_events to the specific downstream, and each records that the downstream task's events argument reads extract's return XCom.
At run time, the scheduler dispatches extract_events first. Its return is pushed once to XCom.
The scheduler then dispatches all five downstream tasks in parallel (they have no dependencies among themselves). Each worker independently pulls the same XCom row and deserializes it.
Note that the XCom row is read five times. If the list is large, this is five deserializations of the same JSON. That's fine for small lists; if the list is 5MB, five deserializations = 25MB of aggregate deserialization work. Consider the S3 stage pattern for large fan-outs.

Output.

Downstream	Depends on	XCom pulled
aggregate	extract_events	return_value
dedupe	extract_events	return_value
enrich	extract_events	return_value
alert	extract_events	return_value
archive	extract_events	return_value

Rule of thumb. Fan-out via the same XComArg is the idiomatic TaskFlow pattern. For large payloads, stage in S3 first and fan out the S3 key — the pattern is identical, just with a small string XCom instead of a large one.

Senior interview question on XCom automation

A senior interviewer might ask: "Explain XCom automation in TaskFlow. What actually happens when I return a value from a @task function? What actually happens when I pass it into another @task call? And what's the practical size limit — when does the pattern break, and what's the fix?"

Solution Using the push/pull/XComArg trio + the S3-stage escape hatch

# ============================================================
# XCom automation in TaskFlow — the three-step story
# ============================================================
#
# 1. PUSH — after a @task's function body returns, the worker
#    JSON-serializes the return value and writes a row into the
#    xcom table under key "return_value" (or per-key if
#    multiple_outputs=True was set).
#
# 2. XCOMARG — inside the @dag body, calling extract("{{ ds }}")
#    does NOT run extract. It registers a task instance and
#    returns an XComArg — a lightweight reference to the future
#    XCom entry. XComArgs are the currency of DAG wiring.
#
# 3. PULL — when a downstream @task receives an XComArg as an
#    argument, the worker starting the downstream task SELECTs
#    the referenced xcom row, deserializes it, and passes it to
#    the function body.
#
# The 48KB reality check:
#   The xcom table's value column has a practical size ceiling
#   of ~48KB (Postgres text row + JSON serialization overhead).
#   Any payload larger than a few KB should be staged in S3 with
#   only the pointer flowing through XCom.
#
# The S3 stage escape hatch — templated key + parquet body:
#   Every upstream task writes to s3://stage/{run_id}/{step}.parquet
#   and returns the key. Every downstream reads by key.
#   run_id makes keys unique per DAG run; S3 lifecycle rules
#   handle cleanup automatically.

from datetime import datetime
from airflow.decorators import dag, task
import io, boto3, pandas as pd

BUCKET = "airflow-stage"
s3 = boto3.client("s3")

def _stage_write(df: pd.DataFrame, key: str) -> str:
    buf = io.BytesIO(); df.to_parquet(buf); buf.seek(0)
    s3.put_object(Bucket=BUCKET, Key=key, Body=buf.getvalue())
    return key

def _stage_read(key: str) -> pd.DataFrame:
    return pd.read_parquet(io.BytesIO(
        s3.get_object(Bucket=BUCKET, Key=key)["Body"].read()
    ))

@dag(
    dag_id="xcom_automation_reference",
    start_date=datetime(2026, 6, 1),
    schedule="@daily",
    catchup=False,
    tags=["taskflow", "xcom"],
)
def xcom_automation_reference():

    @task
    def summarise(ds: str) -> dict:
        # Small dict → XCom directly (well under 48KB)
        return {"date": ds, "row_count": count_rows(ds)}

    @task
    def extract_large(ds: str, run_id: str) -> str:
        # Large DataFrame → S3 stage, return key only
        df = fetch_large_dataset(ds)
        return _stage_write(df, f"{run_id}/extract.parquet")

    @task
    def transform_large(key: str, run_id: str) -> str:
        df = _stage_read(key)
        df["v10"] = df["value"] * 10
        return _stage_write(df, f"{run_id}/transform.parquet")

    @task
    def load(summary: dict, transform_key: str) -> None:
        df = _stage_read(transform_key)
        warehouse.insert(df)
        print(f"loaded {len(df)} rows for {summary['date']}")

    summary        = summarise("{{ ds }}")               # small — direct XCom
    extract_key    = extract_large("{{ ds }}", "{{ run_id }}")  # large — stage
    transform_key  = transform_large(extract_key, "{{ run_id }}")
    load(summary, transform_key)                          # mix of both

xcom_automation_reference()

Step-by-step trace.

Step	Task	XCom action	Size
1	summarise	push dict `{"date":..., "row_count":...}`	~60 bytes
2	extract_large	push S3 key string	~50 bytes
3	transform_large	pull extract's S3 key; push new key	~50 bytes
4	load	pull summary dict + transform key	60 + 50 bytes

At no task boundary does XCom carry the actual DataFrame — S3 does. Every XCom row is comfortably under 1KB. The metadata DB stays healthy; the pipeline is fast because S3 GETs and PUTs are parallelisable and cheap.

Output:

Metric	Direct XCom (naive)	Direct XCom + stage (mix)
XCom rows per DAG run	4 large	4 tiny
Peak XCom row size	20MB (unstable)	60 bytes (safe)
Metadata DB pressure	high	none
Data movement	XCom + PostgreSQL	S3 (fast, parallel)
Pattern	fragile	canonical

Why this works — concept by concept:

Push/pull/XComArg trio — every TaskFlow XCom question resolves into (1) the return-becomes-XCom push, (2) the argument-becomes-XCom pull, (3) the XComArg reference that carries the dependency. Understand this trio and every downstream question follows.
multiple_outputs — the escape hatch for structured returns where different consumers want different keys. Splits the dict at push time; downstream signatures become type-safe on the individual key.
48KB reality — the practical XCom size ceiling. Not a hard schema limit, but the point past which the JSON-serialization + metadata-DB cost dominates the pipeline. Anything larger belongs in a stage.
S3 stage pattern — the canonical fix. Upstream writes to s3://bucket/{run_id}/{step}, returns the key. Downstream reads by key. run_id keeps runs isolated; S3 lifecycle rules handle cleanup.
Cost — XCom cost is proportional to payload size × task boundaries; stage cost is proportional to payload size × 2 (one write + one read per task) but on S3, not on the metadata DB. For any payload above a few KB, the stage pattern is cheaper and safer.

ETL
Topic — etl
ETL problems on data-transfer patterns

Practice →

Optimization Topic — optimization Optimization problems on pipeline throughput

Practice →

4. Mixing TaskFlow with classic Operators

`@task` DAGs and classic Operators coexist in one file — `.output` bridges Operator XComs into TaskFlow arguments, `>>` still wires explicit dependencies

The mental model in one line: classic Airflow operators expose an .output attribute that is an XComArg referencing the operator's return_value XCom, and @task function calls return XComArgs that can be passed as op_kwargs (or as templated arguments) to classic operators — the two APIs interoperate via a common XComArg currency, and the >> shift-arrow still works for any dependency that isn't already implied by argument passing. This is what makes migration incremental and what makes hybrid DAGs practical.

The four axes interviewers actually probe on interop.

When to reach for a classic operator. Anywhere you're calling an external system: S3KeySensor, KubernetesPodOperator, BigQueryInsertJobOperator, SnowflakeOperator, EmailOperator, HttpOperator. The classic operators wrap SDK boilerplate — reinventing them as @task functions loses connection pooling, retry semantics, and the ecosystem's hardened error handling.
.output on classic operators. Every classic operator exposes .output as an XComArg pointing at the operator's return_value XCom. A @task function can consume it as an argument: my_task(sensor.output).
@task return into a classic operator. A @task function's return XCom is available to classic operators via templated op_kwargs: BashOperator(..., bash_command="echo {{ ti.xcom_pull(task_ids='my_task') }}"). Or in modern Airflow: BashOperator(..., bash_command=my_task_output) where my_task_output is the XComArg.
Explicit deps with >>. For dependencies not implied by an XComArg link — "run cleanup after archive even though cleanup doesn't need archive's output" — the >> operator still works. Mix implicit and explicit deps as needed.

When classic operators still win.

Sensors. S3KeySensor, SqlSensor, HttpSensor, ExternalTaskSensor — sensors have specialised poke intervals, mode (poke vs reschedule vs deferrable), and timeout semantics that a @task wrapping a while-loop cannot match without a significant reinvention.
External-system Operators. KubernetesPodOperator, BigQueryInsertJobOperator, SnowflakeOperator, DatabricksSubmitRunOperator — these wrap SDK boilerplate that would take 100+ lines to rewrite as a @task. Use the operator; don't rebuild.
Deferrable/async operators. Newer operators run in the Airflow triggerer process without blocking a worker slot; wrapping them in a @task loses the deferral benefit and forces a worker slot to be occupied while waiting.
Custom operators with hardened error handling. Company-internal operators with domain-specific retry logic, credentials handling, or lineage integration should stay as classes; wrapping them in @task obscures that logic.

The .output bridge — classic → @task.

The signature. Every classic operator instance has an .output attribute that is an XComArg bound to (dag_id, task_id, key="return_value").
Consumption. A @task function receiving operator.output as an argument automatically depends on the operator and pulls its return_value XCom.
Multi-XCom. Classic operators can push multiple XComs (e.g. via context["ti"].xcom_push(key="custom", value=...)); .output only points at return_value. For a specific key, use XComArg(operator, key="custom") explicitly.

The reverse bridge — @task → classic operator.

Modern form. BashOperator(task_id="echo", bash_command=my_task_output, ...) where my_task_output is the XComArg returned by calling a @task function. Airflow renders the XComArg at run time by pulling the referenced XCom.
Legacy form. BashOperator(task_id="echo", bash_command="{{ ti.xcom_pull(task_ids='my_task') }}") — the same result via a Jinja template. Works but is stringly-typed.
Which to use. Prefer the XComArg form where the operator's parameter accepts one; fall back to the templated form only when the operator's parameter is documented as accepting only strings.

>> still works — when to use it.

Non-XCom dependencies. "Run cleanup after archive, even though cleanup takes no arguments from archive." archive >> cleanup states the dependency; no argument passing needed.
Trigger-rule dependencies. "Run the notify task if any upstream fails" — set the classic trigger_rule on the notify task and wire [extract, transform, load] >> notify; the trigger rule kicks in based on upstream states.
Sensor → next. sensor >> next_task establishes the ordering. If next_task also consumes sensor.output as an argument, the >> is redundant but harmless.
Explicit control-plane order. When you want the DAG to visually communicate an ordering that would otherwise be obscure, >> adds no ambiguity and no runtime cost.

Common interview probes on interop.

"How do you call a @task function's return value from a BashOperator?" — pass the XComArg as the operator's parameter, or use a Jinja template.
"When would you use a classic operator instead of a @task in a new DAG?" — sensors, external-system SDK operators, deferrable operators, custom operators with hardened logic.
"How does .output on a classic operator work?" — it's an XComArg pointing at the operator's return_value XCom; consumable by @task arguments.
"Can I use >> on @task calls?" — yes; TaskFlow doesn't replace >>, it just makes it usually unnecessary because XComArg arguments imply the dep.

Worked example — S3KeySensor → `@task` transform → BashOperator chain

Detailed explanation. The archetypal hybrid DAG. An S3KeySensor waits for an input file to appear; a @task function reads and transforms it; a BashOperator archives the result. The sensor is a specialised classic operator (deferrable, with a poke_interval); the transform is a pure-Python @task (best-in-class ergonomics); the archive is a shell command wrapped in BashOperator. All three coexist; .output and templated arguments bridge them.

Node 1. S3KeySensor polls s3://input/{{ ds }}/events.jsonl and pushes the S3 key as its return_value.
Node 2. @task transform_file(s3_key) reads the file, transforms rows, writes to s3://staged/{{ ds }}/events.parquet, returns the output key.
Node 3. BashOperator runs aws s3 cp <output_key> s3://archive/... using the @task's return.

Question. Write the DAG with all three nodes and show the XCom flow across the two boundaries.

Input.

Node	Type	XCom in	XCom out
wait_for_input	S3KeySensor	(none)	S3 key of the appeared file
transform_file	`@task`	sensor's key	new S3 key of transformed file
archive	BashOperator	transform's key (via template)	(none)

Code.

from datetime import datetime
from airflow.decorators import dag, task
from airflow.providers.amazon.aws.sensors.s3 import S3KeySensor
from airflow.operators.bash import BashOperator

@dag(
    dag_id="hybrid_sensor_task_bash",
    start_date=datetime(2026, 6, 1),
    schedule="@daily",
    catchup=False,
    tags=["taskflow", "hybrid"],
)
def hybrid_pipeline():

    wait_for_input = S3KeySensor(
        task_id="wait_for_input",
        bucket_key="s3://input/{{ ds }}/events.jsonl",
        aws_conn_id="aws_default",
        poke_interval=60,
        timeout=60 * 30,
        mode="reschedule",             # release worker slot between pokes
    )

    @task
    def transform_file(input_key: str, ds: str) -> str:
        rows = read_jsonl_from_s3(input_key)
        parquet_key = f"s3://staged/{ds}/events.parquet"
        write_parquet_to_s3(parquet_key, transform_rows(rows))
        return parquet_key

    # Classic → @task bridge: wait_for_input.output is an XComArg
    transformed_key = transform_file(wait_for_input.output, "{{ ds }}")

    # @task → BashOperator bridge: pass XComArg into templated parameter
    archive = BashOperator(
        task_id="archive",
        bash_command=(
            "aws s3 cp {{ ti.xcom_pull(task_ids='transform_file') }} "
            "s3://archive/{{ ds }}/events.parquet"
        ),
    )

    transformed_key >> archive         # explicit ordering; the template already implies it

hybrid_pipeline()

Step-by-step explanation.

S3KeySensor is a classic operator instantiated inside the @dag body. Airflow registers it with the surrounding DAG. Its mode="reschedule" means when the file isn't there yet, the sensor task marks itself for re-run in poke_interval seconds and releases its worker slot — no worker is held blocking.
When the file appears, the sensor's execute() returns the found S3 key. Airflow pushes this to XCom under key return_value. wait_for_input.output is the XComArg bound to that XCom entry.
transform_file(wait_for_input.output, "{{ ds }}") calls the @task with the XComArg as its first argument. At run time, the worker resolves wait_for_input.output by SELECTing the XCom row and passing the value to the function body as input_key.
transform_file returns the new parquet key. This return is pushed to XCom as transform_file.return_value.
BashOperator uses a Jinja template that calls xcom_pull(task_ids='transform_file') to fetch the return_value at render time. The bash_command is rendered just before the shell command executes, substituting the actual key. The transformed_key >> archive line is optional (the template already implies the dep) but makes the ordering explicit for readers.

Output.

Boundary	Bridge mechanism	XCom key
Sensor → @task	`wait_for_input.output` XComArg → function arg	return_value
@task → BashOp	`{{ ti.xcom_pull(task_ids='transform_file') }}` template	return_value
Explicit ordering	`>>` on the XComArg → operator	(visual only)

Rule of thumb. Sensors go before @task transforms because sensors are best-in-class for polling external systems. Shell-command tail steps go into BashOperator because reinventing shell wrapping as a @task is friction. The @task middle does the pure-Python work; classic operators bracket it for external-system I/O.

Worked example — `@task` return feeding a KubernetesPodOperator argument

Detailed explanation. A pattern common in ML pipelines. A @task function computes a training-job spec (image tag, resources, hyperparameters) and returns it as a dict. A KubernetesPodOperator launches the job with those parameters. The bridge is passing the @task's XComArg into the KubernetesPodOperator's arguments or env_vars parameter, or into a templated field.

Node 1. @task compute_job_spec() returns a dict with image_tag, epochs, learning_rate.
Node 2. KubernetesPodOperator launches a training pod using the spec.
Bridge. Individual dict fields threaded into the operator's arguments.

Question. Write the two-task DAG and show two bridging strategies: multiple_outputs=True on the @task (each field is its own XCom), and env-var-templated on the operator.

Input.

Task	Purpose	Bridge
compute_job_spec	Returns dict of image + hyperparameters	XCom entries
run_training	Launches k8s pod with those parameters	env_vars

Code.

from datetime import datetime
from airflow.decorators import dag, task
from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator

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

    @task(multiple_outputs=True)
    def compute_job_spec(ds: str) -> dict:
        # Pick hyperparameters based on recent metrics
        recent = metrics_store.recent()
        return {
            "image_tag":     pick_image_tag(recent),        # e.g. "ml-trainer:v2.6-2026-06-22"
            "epochs":        str(pick_epochs(recent)),      # str because env_vars are strings
            "learning_rate": str(pick_lr(recent)),
        }

    spec = compute_job_spec("{{ ds }}")
    # spec["image_tag"], spec["epochs"], spec["learning_rate"] are XComArgs

    train = KubernetesPodOperator(
        task_id="run_training",
        name="ml-trainer",
        namespace="ml",
        image=f"registry.internal/{{{{ ti.xcom_pull(task_ids='compute_job_spec', key='image_tag') }}}}",
        env_vars={
            "EPOCHS":        spec["epochs"],
            "LEARNING_RATE": spec["learning_rate"],
            "DATE":          "{{ ds }}",
        },
        cmds=["python", "/app/train.py"],
        get_logs=True,
        is_delete_operator_pod=True,
    )

ml_training()

Step-by-step explanation.

@task(multiple_outputs=True) on compute_job_spec splits the returned dict into three per-key XComs: image_tag, epochs, learning_rate. Each is individually addressable as an XComArg through the dict-style indexing on spec.
spec["epochs"] and spec["learning_rate"] are XComArgs. Passing them into the KubernetesPodOperator's env_vars dict is a supported interop point — the operator renders each XComArg to its actual value at run time.
The image parameter uses a Jinja template with an explicit xcom_pull call because image is a string parameter that doesn't accept XComArgs directly in all operator versions. The template works universally.
is_delete_operator_pod=True tells the operator to clean up the pod after completion — a specialised knob that a @task reimplementation would need to rebuild against the Kubernetes API. This is why the classic operator earns its keep here.
At run time: compute_job_spec runs first, pushes three XComs. run_training pulls the three needed values, materialises the env_vars dict, launches the k8s pod with those environment variables, waits for pod completion, streams logs.

Output.

Field	Source	Bridge
image_tag	compute_job_spec XCom	Jinja template in `image`
epochs	compute_job_spec XCom	env_vars dict via XComArg
learning_rate	compute_job_spec XCom	env_vars dict via XComArg
DATE	Jinja	template on `{{ ds }}`

Rule of thumb. Use multiple_outputs=True on the @task producer so downstream classic operators can consume individual fields as clean XComArgs (via env_vars, params, or arguments dicts). Fall back to Jinja templates when the operator's parameter is documented as string-only.

Worked example — mixing `>>` with implicit deps in a fan-in DAG

Detailed explanation. A gnarlier case: three parallel branches feed into a synthesis @task. Two branches produce XComs that the synthesis task consumes as arguments; the third branch produces no XCom but must complete before the synthesis (for example, it warms a cache that synthesis reads). The synthesis needs an explicit >> dep on the third branch even though its XCom flow only covers the first two.

Branches A, B. Produce XComs consumed by the synthesis.
Branch C. Produces no XCom (warms a Redis cache). Synthesis must wait for it.
Wiring. XComArg args for A + B (implicit dep); >> for C (explicit dep).

Question. Write the DAG and show which dependencies are implicit and which are explicit.

Input.

Branch	Produces	Consumed by
A	dict of stats	synthesis (arg)
B	list of ids	synthesis (arg)
C	(side effect on cache)	synthesis (must run first, no XCom)

Code.

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

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

    @task
    def branch_a(ds: str) -> dict:
        return compute_stats(ds)

    @task
    def branch_b(ds: str) -> list[int]:
        return fetch_ids(ds)

    @task
    def branch_c(ds: str) -> None:
        # Warms a Redis cache; no XCom to speak of
        warm_cache_for(ds)

    @task
    def synthesise(stats: dict, ids: list[int]) -> dict:
        # Reads from cache (populated by branch_c)
        enriched = enrich_from_cache(ids)
        return {"stats": stats, "count": len(enriched)}

    stats = branch_a("{{ ds }}")
    ids   = branch_b("{{ ds }}")
    warm  = branch_c("{{ ds }}")

    # Explicit dep on branch_c (no XCom link)
    result = synthesise(stats, ids)
    warm >> result

fan_in_mixed()

Step-by-step explanation.

stats = branch_a(...), ids = branch_b(...), warm = branch_c(...) register three parallel tasks. stats and ids are XComArgs; warm is also an XComArg but bound to a None return.
synthesise(stats, ids) records dependencies from branch_a and branch_b implicitly (via the XComArg arguments) and returns a new XComArg result.
warm >> result states the explicit dependency: branch_c must complete before synthesise runs. The >> between XComArgs works because XComArgs proxy to their underlying task operators.
At run time, the scheduler dispatches branch_a, branch_b, branch_c in parallel. synthesise waits for all three (two implicit deps from XComArg args + one explicit from >>) before running.
synthesise pulls stats and ids XComs at start, calls enrich_from_cache(ids) which reads from the cache warmed by branch_c, computes the result, and returns.

Output.

Downstream	Depends on	Dep type
synthesise	branch_a	implicit (XComArg arg)
synthesise	branch_b	implicit (XComArg arg)
synthesise	branch_c	explicit (`>>`)

Rule of thumb. Use XComArg arguments for dependencies that also carry data. Use >> for dependencies that are purely about ordering (side effects, cache warm-ups, log rotations). Both are first-class; use whichever the reader understands faster.

Senior interview question on the mixed-API pattern

A senior interviewer might ask: "You have a DAG with an S3 sensor, three parallel @task transforms, a KubernetesPodOperator that runs a Spark job, and a final @task that writes lineage metadata. Walk me through the wiring — which dependencies are implicit, which are explicit, which XComs flow where, and how you'd bridge each boundary."

Solution Using XComArg fan-out, `.output` bridges, and `>>` for the k8s dep

from datetime import datetime
from airflow.decorators import dag, task
from airflow.providers.amazon.aws.sensors.s3 import S3KeySensor
from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator

@dag(
    dag_id="mixed_reference",
    start_date=datetime(2026, 6, 1),
    schedule="@daily",
    catchup=False,
    tags=["taskflow", "hybrid", "reference"],
)
def mixed_reference():

    # 1. Classic sensor — wait for input file
    wait = S3KeySensor(
        task_id="wait_for_input",
        bucket_key="s3://input/{{ ds }}/events.jsonl",
        poke_interval=60,
        timeout=60 * 30,
        mode="reschedule",
    )

    # 2. Three parallel @task transforms — all consume sensor's key
    @task
    def normalise(input_key: str, ds: str) -> str:
        return write_stage(read(input_key), f"{ds}/normalised.parquet")

    @task
    def enrich(input_key: str, ds: str) -> str:
        return write_stage(enrich_rows(read(input_key)), f"{ds}/enriched.parquet")

    @task
    def anonymise(input_key: str, ds: str) -> str:
        return write_stage(anonymise_rows(read(input_key)), f"{ds}/anon.parquet")

    normalised = normalise(wait.output, "{{ ds }}")
    enriched   = enrich(wait.output,   "{{ ds }}")
    anonymised = anonymise(wait.output, "{{ ds }}")

    # 3. Classic k8s operator — runs Spark; needs all three stage keys
    spark = KubernetesPodOperator(
        task_id="spark_merge",
        name="spark-merge",
        namespace="etl",
        image="registry.internal/spark-merge:v3.2",
        env_vars={
            "NORMALISED": normalised,
            "ENRICHED":   enriched,
            "ANONYMISED": anonymised,
            "OUTPUT":     "s3://final/{{ ds }}/merged.parquet",
        },
        get_logs=True,
    )

    # 4. Final @task — writes lineage metadata
    @task
    def write_lineage(spark_ret: str, ds: str) -> None:
        lineage.record(
            date=ds,
            output_key="s3://final/{{ ds }}/merged.parquet",
            job_id=spark_ret,
        )

    write_lineage(spark.output, "{{ ds }}")

mixed_reference()

Step-by-step trace.

Boundary	Wiring	Dep type	XCom
sensor → normalise	`normalise(wait.output, ...)`	implicit	`wait.output` (S3 key)
sensor → enrich	`enrich(wait.output, ...)`	implicit	`wait.output` (same key, fan-out)
sensor → anonymise	`anonymise(wait.output, ...)`	implicit	`wait.output`
3 @tasks → spark k8s	env_vars XComArgs	implicit	three parquet keys
spark → write_lineage	`write_lineage(spark.output, ...)`	implicit	spark's return_value

Every dep is implicit — no >> required. The DAG reads top-to-bottom as a function-call graph with three parallel middle branches converging into a Spark job and a lineage tail. Interviewers love this shape because the wiring is the code; the reviewer's mental model is the code.

Output:

Node	Task type	Input XCom	Output XCom
wait_for_input	S3KeySensor	(none)	S3 input key
normalise / enrich / anonymise	@task (×3, parallel)	wait.output	3 parquet keys
spark_merge	KubernetesPodOperator	3 keys via env_vars	job return
write_lineage	@task	spark.output	(none)

Why this works — concept by concept:

XComArg is the universal currency — both classic operators (via .output) and @task calls return XComArgs. Any consumer that accepts an XComArg accepts either.
Fan-out via reuse — passing wait.output to three parallel @task calls creates three edges from the sensor. The sensor's XCom is pulled three times, but the read is cheap (small S3 key string).
Classic operator env_vars accept XComArgs — modern KubernetesPodOperator + BashOperator + BigQueryInsertJobOperator accept XComArgs in their arg dicts, no template gymnastics needed.
Sensor + Spark stay classic, transforms stay @task — each API used where it's best-in-class. Sensors + external-system operators stay classic; pure-Python transforms use @task.
Cost — one XCom per boundary; O(edges) at run time. The mixed DAG has the operational profile of a pure-classic DAG (same scheduling, same worker semantics) with the authoring ergonomics of a pure-TaskFlow DAG.

ETL
Topic — etl
ETL problems on hybrid pipeline design

Practice →

Optimization Topic — optimization Optimization problems on DAG topology

Practice →

5. Typing, retries, and testing TaskFlow DAGs

Type hints, retry policies, and `dag.test()` are the three habits that separate a hobbyist DAG from one that ships

The mental model in one line: type hints on @task functions turn the IDE into a compile-time-ish safety net (mypy-friendly, autocomplete on XComArg fields, misuse flagged at edit time), @task(retries=N, retry_delay=timedelta(minutes=M)) is the single most under-used knob for transient-failure resilience, and dag.test() gives you a pytest-native way to unit-test a TaskFlow function outside a running Airflow scheduler — combine the three and your DAGs are shippable to a demanding SRE without eye-rolls.

The four axes interviewers actually probe on production hygiene.

Type hints. @task def transform(rows: list[dict]) -> list[dict]: — the annotations flow into the IDE, mypy, and the XComArg proxy. Passing an XComArg of the wrong element type into another @task is flagged at edit time, not at run time.
Retries. @task(retries=3, retry_delay=timedelta(minutes=5)) — controls how many times Airflow retries a failed task and how long it waits between attempts. Combined with retry_exponential_backoff=True, this gets you exponential backoff with jitter that would otherwise take 30 lines of manual code.
Testing. dag.test() runs the entire DAG in-process with a minimal scheduler shim, returning a dict of task states. Wrapped in a pytest fixture, this gives you real DAG-level integration tests without a running Airflow deployment.
Senior signals. Type-hinting discipline, XCom-size awareness, retry-with-backoff, dag.test() in CI — these are the four habits senior interviewers hear as "this candidate has actually shipped an Airflow DAG."

Type hints — what they buy you.

IDE autocomplete on XComArgs. With @task def extract(...) -> list[dict]:, hovering over the XComArg returned by extract(...) shows XComArg[list[dict]]. Passing it into transform(rows: list[dict]) is checked; passing it into transform(rows: int) is flagged.
mypy compatibility. Run mypy in CI over the DAG file; type mismatches surface at PR-review time. This is the closest Airflow gets to a compile step.
Documentation. A team member reads def summarise(rows: list[dict]) -> dict and knows what the task consumes and produces without running it. Docstrings and type hints together make DAGs self-documenting.
multiple_outputs=True unlocks per-key typing. If you type-annotate the return as a TypedDict, individual key accesses on the XComArg holder inherit the correct type from the TypedDict definition.

Retries — the senior-signal knob.

The basic form. @task(retries=3) — retry up to 3 times on any exception raised in the task body. Default retry_delay is 5 minutes.
The polished form. @task(retries=5, retry_delay=timedelta(minutes=2), retry_exponential_backoff=True, max_retry_delay=timedelta(hours=1)) — five retries with exponential backoff starting at 2 minutes and capping at 1 hour.
DAG-level defaults. default_args={"retries": 3, "retry_delay": timedelta(minutes=5)} on the @dag sets the default for every task. Per-task @task(retries=...) overrides.
What to retry. Transient errors (network timeouts, temporary rate limits, connection resets). Never retry deterministic errors (bad SQL, missing config) — retries just delay the failure without fixing anything.

Testing with dag.test() — the pytest bridge.

The call. dag.test() runs the DAG in-process with a minimal scheduler; no Airflow scheduler process, no metadata DB, no workers. Returns a mapping of task_id → TaskInstance.
Fixture use. Wrap dag.test() in a pytest fixture; assert on task states, XCom values, and side effects. This is real DAG-level integration testing.
Function-level testing. For unit-testing individual @task functions, extract the wrapped function via .function on the decorated callable and call it directly. extract.function("2026-06-22") calls the underlying extract function with no Airflow context.
CI integration. pytest tests/test_dags.py runs in seconds; no Airflow install needed beyond the library import. This is the single biggest testability win over classic DAGs.

Common interview probes on production hygiene.

"How would you unit-test a @task function?" — extract via .function or use dag.test().
"What's a sane retry policy for an external HTTP call?" — 3–5 retries, exponential backoff starting at 30s, max delay 5–15 minutes.
"How do you make sure your DAG passes mypy?" — type-annotate every @task return; run mypy on the DAG file in CI.
"What's dag.test() and when do you use it?" — in-process DAG runner, used for integration testing in CI without a live Airflow.

Worked example — type hints + `multiple_outputs=True` with TypedDict

Detailed explanation. The highest-quality typing pattern for TaskFlow. A @task(multiple_outputs=True) function returns a dict; annotating the return as a TypedDict gives every downstream consumer of individual keys the correct type via the XComArg holder's indexing. The IDE sees the shape of the dict; mypy checks the field types; refactors flow through both the producer and the consumers.

The TypedDict. Defined once; used both as the return annotation and as the key-type source for XComArg indexing.
The producer. @task(multiple_outputs=True) returning the TypedDict.
The consumers. Each downstream @task takes one field, correctly typed.

Question. Write the DAG with a TypedDict-typed multiple_outputs producer and three type-checked consumers. Show a mypy-flagged mistake.

Input.

Field	Type
row_count	int
latest_ts	str (ISO 8601)
revenue	float

Code.

from datetime import datetime, timedelta
from typing import TypedDict
from airflow.decorators import dag, task

class Summary(TypedDict):
    row_count: int
    latest_ts: str
    revenue: float

@dag(
    dag_id="typed_summary",
    start_date=datetime(2026, 6, 1),
    schedule="@daily",
    catchup=False,
    default_args={"retries": 3, "retry_delay": timedelta(minutes=5)},
)
def typed_summary():

    @task(multiple_outputs=True)
    def compute_summary(ds: str) -> Summary:
        return Summary(
            row_count=count_rows(ds),
            latest_ts=max_ts(ds),
            revenue=sum_revenue(ds),
        )

    @task
    def alert_on_low_count(row_count: int) -> None:
        if row_count < 1000:
            page(f"Low row count: {row_count}")

    @task
    def log_latest(latest_ts: str) -> None:
        audit.write(latest_ts)

    @task
    def push_revenue(revenue: float) -> None:
        metrics.gauge("revenue", revenue)

    s = compute_summary("{{ ds }}")

    alert_on_low_count(s["row_count"])   # ✓ int
    log_latest(s["latest_ts"])           # ✓ str
    push_revenue(s["revenue"])           # ✓ float

    # This line would be flagged by mypy — 'row_count' is int, log_latest expects str:
    # log_latest(s["row_count"])         # ✗ mypy: Argument has incompatible type "int"; expected "str"

typed_summary()

Step-by-step explanation.

class Summary(TypedDict) defines the shape of the return. Each key is typed; the TypedDict is a shared source of truth for both the producer's return annotation and the consumers' key-access types.
@task(multiple_outputs=True) def compute_summary(...) -> Summary: tells Airflow to split the return into three XComs (row_count, latest_ts, revenue) and tells mypy that the return conforms to the TypedDict schema.
s = compute_summary("{{ ds }}") binds s to an XComArg dict-holder. s["row_count"] returns an XComArg typed as int; s["latest_ts"] returns str; s["revenue"] returns float — the types come from the TypedDict.
alert_on_low_count(s["row_count"]) type-checks: the argument type of row_count: int matches the XComArg type of s["row_count"]: XComArg[int]. mypy is happy.
The commented-out log_latest(s["row_count"]) would be flagged: log_latest expects str, but s["row_count"] is int. mypy catches the bug at PR-review time; without the TypedDict, this would blow up at run time with a type error deep in the task's audit.write call.

Output.

Access	Type inferred	Consumed as	Compatible
s["row_count"]	int	alert_on_low_count(int)	yes
s["latest_ts"]	str	log_latest(str)	yes
s["revenue"]	float	push_revenue(float)	yes
s["row_count"]	int	log_latest(str)	no (mypy flag)

Rule of thumb. Whenever a @task(multiple_outputs=True) function returns a dict, define a TypedDict and annotate the return. The compile-time-ish safety net compounds across every consumer; refactors become safe.

Worked example — retry policy with exponential backoff for an HTTP call

Detailed explanation. A task calls an external HTTP API. Transient failures (network timeouts, rate limits, 5xx responses) should be retried; deterministic failures (4xx client errors, malformed responses) should not. Configure the @task retry knobs and layer in a distinguisher between retryable and non-retryable errors.

The task. fetch_from_api(ds) calls an HTTP endpoint and returns JSON.
The transient failures. Timeouts, connection resets, 429 (rate limited), 5xx.
The deterministic failures. 400, 401, 403, 404 — don't retry.

Question. Configure the @task with retries + exponential backoff and inside the function distinguish retryable from non-retryable errors.

Input.

Parameter	Value
retries	5
retry_delay	30 s (base)
retry_exponential_backoff	True
max_retry_delay	15 min
Retryable exceptions	requests.Timeout, requests.ConnectionError, 429, 5xx
Non-retryable	4xx (except 429)

Code.

from datetime import datetime, timedelta
from airflow.decorators import dag, task
from airflow.exceptions import AirflowFailException
import requests

class TransientAPIError(Exception):
    """Retryable — network / server-side."""

@dag(
    dag_id="retry_http",
    start_date=datetime(2026, 6, 1),
    schedule="@hourly",
    catchup=False,
    default_args={
        "retries": 5,
        "retry_delay": timedelta(seconds=30),
        "retry_exponential_backoff": True,
        "max_retry_delay": timedelta(minutes=15),
    },
)
def retry_http():

    @task
    def fetch_from_api(ds: str) -> dict:
        try:
            r = requests.get(
                "https://api.internal/events",
                params={"date": ds},
                timeout=10,
            )
        except (requests.Timeout, requests.ConnectionError) as e:
            # Network transient — let Airflow retry
            raise TransientAPIError(f"network: {e}") from e

        if r.status_code == 429:
            raise TransientAPIError("rate limited")
        if 500 <= r.status_code < 600:
            raise TransientAPIError(f"server {r.status_code}")
        if 400 <= r.status_code < 500:
            # Deterministic — do NOT retry
            raise AirflowFailException(f"client error {r.status_code}: {r.text[:200]}")

        return r.json()

    fetch_from_api("{{ ds }}")

retry_http()

Step-by-step explanation.

default_args on the @dag sets the DAG-wide retry policy: 5 retries with a 30-second base delay, exponential backoff, capped at 15 minutes. Every task in the DAG inherits this unless overridden.
Inside fetch_from_api, network exceptions (Timeout, ConnectionError) are wrapped and re-raised as a TransientAPIError. Any exception raised inside a @task triggers Airflow's retry policy, so this will retry.
HTTP 429 (rate limited) and 5xx (server errors) are also treated as transient — raise TransientAPIError to trigger retry.
HTTP 4xx client errors (except 429) are deterministic — retrying won't fix a 401 unauthorised or a 400 bad request. Raise AirflowFailException — a special Airflow exception that fails the task immediately without triggering retry.
The retry schedule with exponential backoff and jitter: 30s, 60s, 120s, 240s, 480s (capped at 900s). Five attempts cover ~15 minutes total — enough to absorb a short outage or a rate-limit window without masking a real problem.

Output.

Failure	Exception raised	Retry?	Delay before next
Timeout	TransientAPIError	yes	30s, 60s, 120s, 240s, 480s
ConnectionError	TransientAPIError	yes	same
HTTP 429	TransientAPIError	yes	same
HTTP 503	TransientAPIError	yes	same
HTTP 401	AirflowFailException	no	(fail immediately)
HTTP 400	AirflowFailException	no	(fail immediately)

Rule of thumb. Distinguish retryable from non-retryable failures at the exception level; use AirflowFailException to short-circuit retries on deterministic errors. Never retry a task that will fail identically on the next attempt — you're just delaying the outage.

Worked example — unit-testing a `@task` function outside Airflow

Detailed explanation. The testability win is one of TaskFlow's underrated selling points. A @task-decorated function still has a .function attribute pointing at the original undecorated function; you can call it directly in a pytest test with no Airflow context. For DAG-level integration testing, dag.test() runs the whole DAG in-process.

Function-level test. Call extract.function("2026-06-22") and assert on the return value.
DAG-level test. Call my_dag.test(execution_date=...) and assert on task states.
Fixture pattern. pytest fixture that returns the DAG; tests exercise both levels.

Question. Write two pytest tests: one that unit-tests the transform @task in isolation, one that integration-tests the whole DAG end-to-end via dag.test().

Input.

Test	Level	Assertion
test_transform_normalises	function	Given raw dict, returns normalised dict
test_dag_end_to_end	DAG	All tasks succeed; final XCom matches expected

Code.

# ---------- The DAG under test ----------
# dags/events_etl.py
from datetime import datetime
from airflow.decorators import dag, task

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

    @task
    def extract(ds: str) -> list[dict]:
        return fake_source_for_test(ds)

    @task
    def transform(rows: list[dict]) -> list[dict]:
        out = []
        for r in rows:
            out.append({
                "id":       int(r["id"]),
                "value":    float(r["value"]),
                "ts_iso":   r["ts"].isoformat() if hasattr(r["ts"], "isoformat") else r["ts"],
            })
        return out

    @task
    def summarise(rows: list[dict]) -> dict:
        return {"count": len(rows), "total": sum(r["value"] for r in rows)}

    summarise(transform(extract("{{ ds }}")))

events_etl_dag = events_etl()


# ---------- The tests ----------
# tests/test_events_etl.py
import pytest
from datetime import datetime, timezone
from dags.events_etl import events_etl_dag

# Unit test — call the @task function directly via .function
def test_transform_normalises():
    from dags.events_etl import events_etl

    # Access the inner @task by name via the DAG
    transform = events_etl_dag.task_dict["transform"].python_callable

    raw = [
        {"id": "1", "value": "10.5", "ts": datetime(2026, 6, 22, tzinfo=timezone.utc)},
        {"id": "2", "value": "3.14", "ts": "2026-06-22T00:00:00+00:00"},
    ]
    out = transform(raw)
    assert out == [
        {"id": 1, "value": 10.5, "ts_iso": "2026-06-22T00:00:00+00:00"},
        {"id": 2, "value": 3.14, "ts_iso": "2026-06-22T00:00:00+00:00"},
    ]


# Integration test — run the whole DAG in-process via dag.test()
def test_dag_end_to_end(monkeypatch):
    # Stub the source
    def fake_source_for_test(ds: str) -> list[dict]:
        return [
            {"id": "1", "value": "10.5", "ts": "2026-06-22T00:00:00+00:00"},
            {"id": "2", "value": "3.14", "ts": "2026-06-22T00:00:00+00:00"},
        ]
    import dags.events_etl as m
    monkeypatch.setattr(m, "fake_source_for_test", fake_source_for_test)

    # Run the DAG in-process
    dag_run = events_etl_dag.test(execution_date=datetime(2026, 6, 22, tzinfo=timezone.utc))

    # Every task should succeed
    for ti in dag_run.get_task_instances():
        assert ti.state == "success", f"{ti.task_id} failed with state {ti.state}"

    # The summarise task's XCom should match
    summarise_ti = dag_run.get_task_instance("summarise")
    assert summarise_ti.xcom_pull(task_ids="summarise", key="return_value") == {
        "count": 2,
        "total": pytest.approx(13.64),
    }

Step-by-step explanation.

The unit test accesses the wrapped function via the DAG's task_dict. Every @task-registered task in the DAG has a python_callable attribute pointing at the underlying function. Calling transform(raw) runs the pure Python function with no Airflow context and no scheduler — just Python.
The unit test asserts on the exact return value. Fast (~5ms), no external dependencies, no Airflow install needed beyond the library import.
The integration test uses monkeypatch to stub the external source (fake_source_for_test) so the test is hermetic. Every real network / DB call in the DAG should have a corresponding monkeypatch or fake in the test.
events_etl_dag.test(execution_date=...) runs the whole DAG in-process. Airflow spins up a minimal in-memory metadata store, runs each task in dependency order on the local process, and populates task instances with states and XComs.
The assertions check every task's state ("success") and pull the final XCom. pytest.approx handles floating-point equality. The whole test runs in a few seconds and gives real DAG-level coverage.

Output.

Test	Level	Runtime	Coverage
test_transform_normalises	function	~5 ms	transform logic only
test_dag_end_to_end	DAG	~2 s	all tasks + XCom flow

Rule of thumb. Every @task function gets a unit test that calls it directly via the wrapper's python_callable. Every DAG gets one integration test that runs dag.test() with monkeypatched sources. Wire both into CI; catch regressions before a scheduler ever sees them.

Senior interview question on production hygiene for TaskFlow

A senior interviewer might ask: "Walk me through the four habits that keep a TaskFlow DAG production-grade: type hints, retry policy, testing, and observability. Show me the config knobs, the exception classes, the test fixtures, and the CI wiring you'd expect on any DAG that gets merged to main."

Solution Using a four-habit template DAG + pytest harness

# ============================================================
# Four-habit template — every DAG should look like this
# ============================================================
#
# Habit 1: Type hints on every @task
#   - Return annotations feed XComArg types
#   - Argument annotations feed IDE + mypy
#   - TypedDict for multiple_outputs
#
# Habit 2: Retry policy in default_args
#   - retries=3 (min), exponential backoff, max_retry_delay
#   - AirflowFailException for deterministic errors
#
# Habit 3: pytest-native testing
#   - Function-level: .python_callable on the task
#   - DAG-level: dag.test() with monkeypatched sources
#
# Habit 4: Observability
#   - task-level SLA via sla parameter
#   - Structured logs from the function body
#   - on_failure_callback for alerts
#
# CI wiring: mypy dags/ && pytest tests/

from datetime import datetime, timedelta
from typing import TypedDict
from airflow.decorators import dag, task
from airflow.exceptions import AirflowFailException
from airflow.utils.email import send_email
import requests

class Summary(TypedDict):
    date: str
    row_count: int
    revenue: float

def _on_failure(context: dict) -> None:
    """Sent on any task failure — page / email / slack."""
    ti = context["task_instance"]
    send_email(
        to=["data-eng-oncall@example.com"],
        subject=f"[Airflow] {ti.dag_id}.{ti.task_id} failed",
        html_content=(
            f"<p>Task {ti.task_id} in DAG {ti.dag_id} failed at "
            f"{context['execution_date']}.</p>"
            f"<p>Log URL: {ti.log_url}</p>"
        ),
    )

@dag(
    dag_id="production_grade_reference",
    start_date=datetime(2026, 6, 1),
    schedule="@hourly",
    catchup=False,
    default_args={
        "retries": 3,
        "retry_delay": timedelta(seconds=30),
        "retry_exponential_backoff": True,
        "max_retry_delay": timedelta(minutes=15),
        "on_failure_callback": _on_failure,
        "sla": timedelta(minutes=30),
    },
    tags=["taskflow", "production"],
)
def production_grade():

    @task
    def extract(ds: str) -> list[dict]:
        try:
            r = requests.get(f"https://api.internal/events?date={ds}", timeout=10)
        except (requests.Timeout, requests.ConnectionError) as e:
            raise RuntimeError(f"transient: {e}") from e
        if 400 <= r.status_code < 500 and r.status_code != 429:
            raise AirflowFailException(f"client {r.status_code}")
        r.raise_for_status()
        return r.json()

    @task
    def transform(rows: list[dict]) -> list[dict]:
        return [
            {"id": int(r["id"]), "value": float(r["value"]), "ds": r["ds"]}
            for r in rows
        ]

    @task(multiple_outputs=True)
    def summarise(rows: list[dict], ds: str) -> Summary:
        return Summary(
            date=ds,
            row_count=len(rows),
            revenue=sum(r["value"] for r in rows),
        )

    @task
    def load(rows: list[dict], summary_date: str, row_count: int) -> None:
        warehouse.insert(rows)
        print(f"loaded {row_count} rows for {summary_date}")

    raw     = extract("{{ ds }}")
    cleaned = transform(raw)
    summary = summarise(cleaned, "{{ ds }}")
    load(cleaned, summary["date"], summary["row_count"])

production_grade_dag = production_grade()

# ============================================================
# pytest harness — CI runs `pytest tests/` on every PR
# ============================================================

# tests/test_production_grade.py
def test_transform_types():
    transform = production_grade_dag.task_dict["transform"].python_callable
    out = transform([{"id": "1", "value": "2.5", "ds": "2026-06-22"}])
    assert out == [{"id": 1, "value": 2.5, "ds": "2026-06-22"}]

def test_end_to_end(monkeypatch):
    import requests as _req
    class FakeResp:
        status_code = 200
        def json(self): return [{"id": "1", "value": "2.5", "ds": "2026-06-22"}]
        def raise_for_status(self): pass
    monkeypatch.setattr(_req, "get", lambda *a, **k: FakeResp())
    dag_run = production_grade_dag.test(execution_date=datetime(2026, 6, 22))
    assert all(ti.state == "success" for ti in dag_run.get_task_instances())

Step-by-step trace.

Habit	Config surface	Runtime effect
1. Type hints	`Summary(TypedDict)`, `@task -> list[dict]`	mypy checks; IDE autocomplete on XComArgs
2. Retry policy	`default_args={retries, retry_delay, exp_backoff, max_retry_delay}`	3 retries, 30s→15min backoff
3. Testing	`.python_callable` + `dag.test()`	Function + DAG tests in CI
4. Observability	`on_failure_callback=_on_failure`, `sla=timedelta(minutes=30)`	Email on fail; SLA miss alerts
Bonus: FailException	`AirflowFailException` on 4xx	Deterministic errors skip retry

After the template is adopted, every merged DAG has the same production hygiene shape. Interviewers hear "type hints, retries, dag.test() in CI, on_failure_callback" and immediately mark the candidate as ship-ready.

Output:

Surface	Before hygiene	After
Type errors caught	at run time	at PR review
Transient failures	manual retry / outage	3 auto-retries
DAG regression risk	untested	pytest gates
Ops notification lead time	when someone notices	on_failure_callback

Why this works — concept by concept:

Type hints propagate — TypedDict at the producer means every downstream indexer inherits the field type. mypy + IDE autocomplete + XComArg proxies compound into compile-time-ish safety.
Retry policy at DAG level — default_args sets the policy for every task; individual tasks override only if they need something specific. Exponential backoff with a max cap absorbs transient outages without paging.
.python_callable + dag.test() — the two pytest hooks. Function-level tests are fast (ms); DAG-level tests are hermetic (monkeypatched sources); both run in CI in seconds.
on_failure_callback + sla — the observability tail. Any failure pages the on-call; any SLA miss alerts even if the task eventually succeeds. Both without a separate monitoring pipeline.
Cost — one-time template adoption; per-DAG cost is a docstring, a TypedDict, a monkeypatched test. The saved incidents pay for the template a hundred times over the DAG's lifetime.

ETL
Topic — etl
ETL problems on pipeline testing patterns

Practice →

Optimization
Topic — optimization
Optimization problems on retry and backoff

Practice →

Cheat sheet — TaskFlow recipes

The canonical 4-task TaskFlow DAG. @dag(dag_id=..., start_date=..., schedule="@hourly", catchup=False) on a factory function; four @task-decorated functions inside; a chain of function calls at the bottom of the factory body; a single factory() call at module level. The trailing factory() line is not optional — without it the module registers zero DAGs.
@task with retries + trigger rule. @task(retries=5, retry_delay=timedelta(seconds=30), retry_exponential_backoff=True, max_retry_delay=timedelta(minutes=15), trigger_rule="all_success"). Retries handle transient errors; exponential backoff absorbs outages without paging; trigger_rule controls when the task runs based on upstream states.
XComArg chaining pattern. load(transform(extract("{{ ds }}"))) is legal and idiomatic. For readability with 5+ tasks, use local variables: raw = extract("{{ ds }}"), cleaned = transform(raw), n = validate(cleaned), load(cleaned, n). Both forms compile to the same DAG.
multiple_outputs=True with TypedDict. Define class Summary(TypedDict) with typed fields; return Summary(...) from a @task(multiple_outputs=True); each downstream consumer indexes into the XComArg holder with typed access: summary["row_count"] is an XComArg[int].
Sensor → TaskFlow bridge. Classic sensors expose .output as an XComArg pointing at their return_value XCom. @task def transform(input_key: str, ds: str) -> str: receives sensor.output as the first arg; the dependency is implicit; no >> needed.
@task → BashOperator bridge. Prefer XComArg-parameter form: BashOperator(task_id=..., bash_command=my_task_output) where my_task_output is the XComArg returned by calling a @task. Fall back to Jinja {{ ti.xcom_pull(task_ids='my_task') }} for older operators.
The 48KB XCom limit. Any payload above a few KB belongs in a stage. Upstream writes to s3://stage/{{ run_id }}/{step}.parquet and returns the key; downstream reads by key. run_id isolates concurrent runs; S3 lifecycle rules handle cleanup.
dag.test() invocation for CI. my_dag.test(execution_date=datetime(...)) runs the DAG in-process with a minimal scheduler; returns a DagRun with per-task states and XComs. Wrap in a pytest fixture with monkeypatched sources for hermetic integration tests.
Unit-testing a single @task. Access via my_dag.task_dict["task_id"].python_callable; call directly with plain arguments; no Airflow context needed. Fast (~5ms) and hermetic.
Type hints for the IDE. Annotate every @task argument and return. TypedDict for multiple_outputs; list[dict] and int for the common cases. mypy catches type-mismatches at PR-review time; without hints, mismatches surface only when the DAG runs.
on_failure_callback for alerts. default_args={"on_failure_callback": my_alert_fn} on the @dag; the callback receives the task context and can email, page, or post to Slack. Set once at the DAG level; every task inherits.
AirflowFailException for deterministic errors. Raise AirflowFailException("bad config") to fail a task without triggering retries. Use for 4xx HTTP errors, malformed inputs, missing config — anything that will fail identically on retry.
Parametrised DAG factory. Wrap the @dag-decorated function in an outer plain function that returns the DAG; call the outer function once per variant at module level. Each call registers a separate DAG with its own dag_id, schedule, and closed-over config.
Mixed-DAG rule. Keep classic operators for sensors, external-system SDK operators, and deferrable operators. Use @task for pure-Python transforms and control-flow. Bridge via .output (classic → @task) and XComArg-parameter form (@task → classic).

Frequently asked questions

What is the Airflow TaskFlow API and why does it matter?

The airflow taskflow api is a decorator-based DAG authoring API introduced in Airflow 2.0 that turns ordinary Python functions into tasks (@task) and ordinary Python factory functions into DAGs (@dag). Under the hood it compiles to the same DAG object the classic PythonOperator-plus->> API produces — the scheduler, workers, XCom system, and UI are all identical — but the authoring layer is roughly half the lines of code and eliminates the stringly-typed task_id and XCom-key drift that classic DAGs suffer. It matters because as of 2026, TaskFlow is the default in new Airflow 3 DAGs, most Airflow tutorials and provider examples now use it, and interviewers expect senior candidates to explain the two-decorator model plus XCom automation plus the .output interop bridge without prompting. The mental shift is one sentence: return values become XComs automatically, arguments pull them automatically, dependencies are inferred from the function-call graph.

TaskFlow vs PythonOperator — when do I still use the classic operator?

Default to TaskFlow for any pure-Python transform work. Reach for classic operators specifically when: (a) you need a sensor (S3KeySensor, SqlSensor, ExternalTaskSensor) — sensors have specialised poke, mode (reschedule, deferrable), and timeout semantics you don't want to reimplement in a @task; (b) you're calling an external system via a hardened SDK operator (KubernetesPodOperator, BigQueryInsertJobOperator, SnowflakeOperator, DatabricksSubmitRunOperator) — the operator wraps 50–200 lines of SDK boilerplate, connection pooling, and error handling you'd otherwise rebuild; (c) you need deferrable execution — modern deferrable operators run in the Airflow triggerer without holding a worker slot, and wrapping them in a @task loses that; (d) you're calling a company-internal operator with hardened credentials and lineage integration. The pattern is @task for pure Python; classic operators for external I/O and specialised runtime behaviour. Mixed DAGs are first-class; use .output on classic operators and XComArg-parameter passing on @task calls to bridge the two APIs.

How do XComs work with TaskFlow — do I still need to call `xcom_push` and `xcom_pull`?

No, and that's the whole point. In TaskFlow, every @task function's return value is automatically pushed to XCom under key return_value, and every @task function argument that receives an XComArg is automatically pulled from XCom at task start. You never write context["ti"].xcom_push(...) or xcom_pull(task_ids=..., key=...) for the common case. Under the hood the same xcom metadata table stores the entries, but you interact with them via Python return statements and function arguments instead of stringly-typed key lookups. Two escape hatches remain: @task(multiple_outputs=True) on a function returning a dict splits each key into a separate XCom (letting downstream tasks depend on specific fields), and inside a @task body you can still call context["ti"].xcom_pull(...) for cross-task lookups the argument-passing model can't express. The 48KB practical size limit on XCom rows is unchanged — any payload above a few KB belongs in an S3 stage with the key flowing through XCom.

Can I mix classic Operators with TaskFlow `@task` functions in the same DAG?

Yes — mixed DAGs are the norm in production. Two bridge patterns cover 99% of interop. Classic → @task: every classic operator exposes an .output attribute that is an XComArg bound to the operator's return_value XCom. A @task function taking operator.output as an argument automatically depends on the operator and pulls its return XCom. @task → classic: modern classic operators accept XComArgs directly in their parameter dicts (env_vars, op_args, bash_command on new versions, etc.); for older or strict-string parameters, use a Jinja template like {{ ti.xcom_pull(task_ids='my_task') }}. Explicit >> still works for dependencies that aren't implied by argument passing — for example, "run cleanup after archive, no data flow between them." A typical hybrid DAG has 2–5 classic operators (sensor + external-system operators) and 5–15 @task functions doing pure-Python transforms; the two APIs share the DAG file, the XCom backend, and the scheduler.

How do I unit-test a TaskFlow DAG without running the whole Airflow?

Two levels of testing, both native. Function-level unit tests: every @task-decorated function exposes the wrapped callable via dag.task_dict["task_id"].python_callable. Call it directly with plain arguments in a pytest test; no Airflow context, no scheduler, no metadata DB needed. Runs in a few milliseconds. DAG-level integration tests: call my_dag.test(execution_date=datetime(...)) — Airflow spins up an in-memory metadata store, runs each task in dependency order in the current process, and returns a DagRun object with per-task states and XComs. Wrap it in a pytest fixture with monkeypatched external sources (monkeypatch.setattr(mod, "fetch_from_api", fake_fetch)) and assert on task states + XCom values. Both tests run in a pytest tests/ invocation with no live Airflow deployment; wire them into CI to catch DAG regressions before merge. This is the single biggest testability win over classic PythonOperator DAGs, where the equivalent required either spinning up a full scheduler or extracting and calling python_callable with a hand-constructed context dict.

Is TaskFlow the default in Airflow 3, and what changes if I upgrade?

Yes — Airflow 3 (released late 2025) treats TaskFlow as first-class in the UI (function-level docs from docstrings), the CLI (airflow tasks test dag_id task_id works without extra wiring), and the testing tooling (dag.test() is the recommended integration-test entry point). The API surface is stable — DAGs written on Airflow 2.7+ using @dag + @task from airflow.decorators continue to work on Airflow 3 unchanged. What changes: (a) classic operator wrappers now expose .output uniformly across all provider packages, making mixed-DAG interop friction-free; (b) the UI's task detail pane shows the wrapped function's docstring, source code, and type signature — you get autogenerated per-task docs from the Python code itself; (c) the scheduler is smarter about dispatching XComArg-chained tasks (fewer round trips between the scheduler and workers for tightly-coupled @task chains); (d) @task(multiple_outputs=True) gains better TypedDict integration. If you're on Airflow 2.7+ with TaskFlow-first DAGs, the upgrade to Airflow 3 is a low-risk minor version bump; if you're on Airflow 2.4 or earlier with mostly PythonOperator DAGs, the upgrade is a good forcing function to modernise. Either way, TaskFlow is the API interviewers ask about — knowing it well is a hard requirement for any senior data-engineering role in 2026.

Practice on PipeCode

Drill the ETL practice library → for the DAG-authoring, XCom-flow, and hybrid-pipeline problems senior interviewers love.
Rehearse on the SQL practice library → for the source and warehouse-side query patterns your TaskFlow DAGs will call.
Sharpen the tuning axis with the optimization practice library → for the retry-policy, XCom-size, and factory-pattern problems.
Stack the prerequisites against PipeCode's broader 450+ data-engineering catalogue to anchor the TaskFlow + XCom + testing intuition against real graded inputs.

Lock in TaskFlow muscle memory

Airflow docs explain the API. PipeCode drills explain the decision — when `@task` beats PythonOperator, when XCom breaks and needs an S3 stage, when `.output` bridges a sensor to a functional DAG, when `dag.test()` catches a regression that would have paged an on-call. Pipecode.ai is Leetcode for Data Engineering — pattern-first practice tuned for the production trade-offs senior data engineers actually face.

Practice ETL problems →
Practice optimization problems →

1. Why TaskFlow replaced most PythonOperators in 2026

The verbose old world — PythonOperator(task_id=..., python_callable=..., op_kwargs=...) — collapses to @task plus a return statement

Worked example — the same ETL, classic PythonOperator vs TaskFlow

Worked example — the op_kwargs leak

Worked example — mixing decorators with existing PythonOperator DAGs

Senior interview question on the TaskFlow value proposition

Solution Using opportunistic migration + a 90-day rollout plan

2. @task and @dag decorators — the two-decorator pipeline

Two decorators that convert plain Python into a DAG — @dag on the factory function, @task on each task function

Worked example — a 4-task DAG in 20 lines

Worked example — parametrised DAG factory (same shape, three schedules)

Worked example — multiple_outputs=True for structured returns

Senior interview question on decorator semantics

Solution Using the parse-time / run-time split as the mental model

3. XCom automation and XComArg chaining

Return values push XComs, arguments pull XComs, XComArg is the object that carries the reference — with a hard 48KB size limit that forces a stage for real data

Worked example — passing a dict of results between tasks

Worked example — the 48KB gotcha and the S3 stage pattern

Worked example — fan-out with the same XComArg

Senior interview question on XCom automation

Solution Using the push/pull/XComArg trio + the S3-stage escape hatch

4. Mixing TaskFlow with classic Operators

@task DAGs and classic Operators coexist in one file — .output bridges Operator XComs into TaskFlow arguments, >> still wires explicit dependencies

Worked example — S3KeySensor → @task transform → BashOperator chain

Worked example — @task return feeding a KubernetesPodOperator argument

Worked example — mixing >> with implicit deps in a fan-in DAG

Senior interview question on the mixed-API pattern

Solution Using XComArg fan-out, .output bridges, and >> for the k8s dep

5. Typing, retries, and testing TaskFlow DAGs

Type hints, retry policies, and dag.test() are the three habits that separate a hobbyist DAG from one that ships

Worked example — type hints + multiple_outputs=True with TypedDict

Worked example — retry policy with exponential backoff for an HTTP call

Worked example — unit-testing a @task function outside Airflow

Senior interview question on production hygiene for TaskFlow

Solution Using a four-habit template DAG + pytest harness

Cheat sheet — TaskFlow recipes

Frequently asked questions

What is the Airflow TaskFlow API and why does it matter?

TaskFlow vs PythonOperator — when do I still use the classic operator?

How do XComs work with TaskFlow — do I still need to call xcom_push and xcom_pull?

Can I mix classic Operators with TaskFlow @task functions in the same DAG?

How do I unit-test a TaskFlow DAG without running the whole Airflow?

Is TaskFlow the default in Airflow 3, and what changes if I upgrade?

Practice on PipeCode

Lock in TaskFlow muscle memory

The verbose old world — `PythonOperator(task_id=..., python_callable=..., op_kwargs=...)` — collapses to `@task` plus a return statement

Worked example — the `op_kwargs` leak

2. `@task` and `@dag` decorators — the two-decorator pipeline

Two decorators that convert plain Python into a DAG — `@dag` on the factory function, `@task` on each task function

Worked example — `multiple_outputs=True` for structured returns

Return values push XComs, arguments pull XComs, `XComArg` is the object that carries the reference — with a hard 48KB size limit that forces a stage for real data

`@task` DAGs and classic Operators coexist in one file — `.output` bridges Operator XComs into TaskFlow arguments, `>>` still wires explicit dependencies

Worked example — S3KeySensor → `@task` transform → BashOperator chain

Worked example — `@task` return feeding a KubernetesPodOperator argument

Worked example — mixing `>>` with implicit deps in a fan-in DAG

Solution Using XComArg fan-out, `.output` bridges, and `>>` for the k8s dep

Type hints, retry policies, and `dag.test()` are the three habits that separate a hobbyist DAG from one that ships

Worked example — type hints + `multiple_outputs=True` with TypedDict

Worked example — unit-testing a `@task` function outside Airflow

How do XComs work with TaskFlow — do I still need to call `xcom_push` and `xcom_pull`?

Can I mix classic Operators with TaskFlow `@task` functions in the same DAG?