Title here
Summary here
Lesson 1 - The TaskFlow API
Welcome to the TaskFlow API
Module 2 wrote every DAG with @dag and @task and called it a day — the TaskFlow API is genuinely the easier way to write most DAGs, and there was no reason to slow down and show the alternative. This lesson stops and shows it: the exact same real computation, written the classic way Airflow used before TaskFlow existed, run side by side with its TaskFlow twin, with a real diff proving they produce identical output.
By the end of this lesson, you will be able to:
- Write the same DAG with
PythonOperatorinstances and explicit>>, instead of@dag/@task - Explain what TaskFlow’s decorators are actually doing under the hood: pushing and pulling XCom values by hand, the same mechanism, just spelled out
- Prove two independently-written DAGs are equivalent with a real file diff, not just a read-through
The Same Report, Written Twice
Both DAGs read CityFlow’s real 265-zone lookup table, count how many zones fall in each borough, sanity-check the total, and write the result to a JSON file. TaskFlow first:
# gate: skip
"""CityFlow's real zone-by-borough report, written with the TaskFlow API.
Lesson 1 (Module 3) pairs this with cityflow_zone_report_classic_dag.py --
same three tasks, same real data, same output, two different ways of
wiring an Airflow DAG."""
from datetime import datetime
from airflow.decorators import dag, task
REPORT_PATH = "/opt/airflow/dags/data/work/zone_counts_taskflow.json"
@dag(
dag_id="cityflow_zone_report_taskflow",
schedule=None,
start_date=datetime(2026, 1, 1),
catchup=False,
tags=["cityflow", "module-3", "taskflow"],
)
def cityflow_zone_report_taskflow_dag():
@task
def extract_zone_counts() -> dict:
"""Read CityFlow's real 265-zone lookup table and count zones per borough."""
import pandas as pd
zones = pd.read_csv("/opt/airflow/dags/data/taxi_zone_lookup.csv")
# One real row (LocationID 265, "Outside of NYC") has no Borough value --
# label it explicitly instead of letting value_counts() silently drop it.
counts = zones["Borough"].fillna("Outside of NYC").value_counts().to_dict()
return counts
@task
def validate_zone_counts(counts: dict) -> dict:
"""Sanity check: the boroughs must add back up to all 265 real zones."""
total = sum(counts.values())
assert total == 265, f"expected 265 zones total, got {total}"
return counts
@task
def report_zone_counts(counts: dict) -> None:
import json
import os
os.makedirs(os.path.dirname(REPORT_PATH), exist_ok=True)
with open(REPORT_PATH, "w") as f:
json.dump(counts, f, sort_keys=True, indent=2)
print(f"wrote {REPORT_PATH}")
print(json.dumps(counts, sort_keys=True, indent=2))
report_zone_counts(validate_zone_counts(extract_zone_counts()))
cityflow_zone_report_taskflow_dag()Notice the real data caught something a synthetic example never would: taxi_zone_lookup.csv’s last row (LocationID 265, "Outside of NYC") has no Borough value at all. value_counts() silently drops NaN by default, which would have made the report’s boroughs sum to 264, not 265 — the assertion in validate_zone_counts exists specifically to catch a bug like that, and it did, the first time this lesson’s code ran against the real file.
Now the same report, written the classic way:
# gate: skip
"""CityFlow's real zone-by-borough report, written with classic operator
instantiation. Lesson 1 (Module 3) pairs this with
cityflow_zone_report_taskflow_dag.py -- same three tasks, same real data,
same output, wired with PythonOperator + explicit `>>` instead of the
TaskFlow API's decorators and nested calls."""
from datetime import datetime
from airflow import DAG
from airflow.operators.python import PythonOperator
REPORT_PATH = "/opt/airflow/dags/data/work/zone_counts_classic.json"
def extract_zone_counts(**context):
"""Read CityFlow's real 265-zone lookup table and count zones per borough.
No decorator does this for us here -- the return value only becomes an
XCom because PythonOperator pushes it under the key 'return_value' by
default; validate_zone_counts pulls it back explicitly below."""
import pandas as pd
zones = pd.read_csv("/opt/airflow/dags/data/taxi_zone_lookup.csv")
return zones["Borough"].fillna("Outside of NYC").value_counts().to_dict()
def validate_zone_counts(**context):
ti = context["ti"]
counts = ti.xcom_pull(task_ids="extract_zone_counts")
total = sum(counts.values())
assert total == 265, f"expected 265 zones total, got {total}"
return counts
def report_zone_counts(**context):
import json
import os
ti = context["ti"]
counts = ti.xcom_pull(task_ids="validate_zone_counts")
os.makedirs(os.path.dirname(REPORT_PATH), exist_ok=True)
with open(REPORT_PATH, "w") as f:
json.dump(counts, f, sort_keys=True, indent=2)
print(f"wrote {REPORT_PATH}")
print(json.dumps(counts, sort_keys=True, indent=2))
with DAG(
dag_id="cityflow_zone_report_classic",
schedule=None,
start_date=datetime(2026, 1, 1),
catchup=False,
tags=["cityflow", "module-3", "classic"],
) as dag:
extract = PythonOperator(
task_id="extract_zone_counts",
python_callable=extract_zone_counts,
)
validate = PythonOperator(
task_id="validate_zone_counts",
python_callable=validate_zone_counts,
)
report = PythonOperator(
task_id="report_zone_counts",
python_callable=report_zone_counts,
)
extract >> validate >> reportLine up the two side by side and TaskFlow’s whole value becomes concrete: report_zone_counts(validate_zone_counts(extract_zone_counts())) and extract >> validate >> report wire the identical dependency chain, but the classic version has to say the XCom exchange out loud — ti.xcom_pull(task_ids="extract_zone_counts") in every downstream function — where TaskFlow does the same push-then-pull silently, just by nesting function calls. Nothing about what runs changed; only how much of the plumbing you have to write yourself.
Both, Deployed and Triggered for Real
Both files are saved to .airflow-local/dags/. Same pattern as Module 2: confirm the scheduler picked up both, unpause and trigger both, poll each to a terminal state.
import subprocess
import time
import json
dag_ids = ("cityflow_zone_report_taskflow", "cityflow_zone_report_classic")
for dag_id in dag_ids:
deadline = time.time() + 320
found = False
while time.time() < deadline:
result = subprocess.run(
["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "list", "--output", "json"],
capture_output=True, text=True, check=True,
)
if any(d["dag_id"] == dag_id for d in json.loads(result.stdout)):
found = True
break
time.sleep(5)
assert found, f"{dag_id} was never picked up by the scheduler"
print(f"scheduler picked up {dag_id!r}")
run_ids = {}
for dag_id in dag_ids:
subprocess.run(
["docker", "exec", "airflow-local-airflow-scheduler-1", "airflow", "dags", "unpause", dag_id],
capture_output=True, text=True, check=True,
)
run_id = f"lesson1_verify_{int(time.time())}_{dag_id[-4:]}"
subprocess.run(
["docker", "exec", "airflow-local-airflow-scheduler-1",
"airflow", "dags", "trigger", dag_id, "--run-id", run_id],
capture_output=True, text=True, check=True,
)
run_ids[dag_id] = run_id
print(f"triggered {dag_id!r} as {run_id!r}")
time.sleep(1)
for dag_id, run_id in run_ids.items():
deadline = time.time() + 60
state = "queued"
while state not in ("success", "failed") and time.time() < deadline:
time.sleep(3)
result = subprocess.run(
["docker", "exec", "airflow-local-airflow-scheduler-1",
"airflow", "dags", "list-runs", "-d", dag_id, "--output", "json"],
capture_output=True, text=True, check=True,
)
runs = json.loads(result.stdout)
state = next(r["state"] for r in runs if r["run_id"] == run_id)
print(f"{dag_id} run {run_id!r} reached state: {state}")
assert state == "success", f"expected success, got {state}"scheduler picked up 'cityflow_zone_report_taskflow'
scheduler picked up 'cityflow_zone_report_classic'
triggered 'cityflow_zone_report_taskflow' as 'lesson1_verify_1784885113_flow'
triggered 'cityflow_zone_report_classic' as 'lesson1_verify_1784885116_ssic'
cityflow_zone_report_taskflow run 'lesson1_verify_1784885113_flow' reached state: success
cityflow_zone_report_classic run 'lesson1_verify_1784885116_ssic' reached state: successBoth DAGs’ three tasks all reached success:
dag_id | task_id | state
===============================+======================+========
cityflow_zone_report_taskflow | extract_zone_counts | success
cityflow_zone_report_taskflow | validate_zone_counts | success
cityflow_zone_report_taskflow | report_zone_counts | success
cityflow_zone_report_classic | extract_zone_counts | success
cityflow_zone_report_classic | validate_zone_counts | success
cityflow_zone_report_classic | report_zone_counts | successProof, Not a Read-Through: Diffing the Real Output
Reading both files side by side is one kind of evidence. A real diff on what each DAG actually wrote to disk is stronger — it can’t be fooled by two implementations that merely look equivalent:
import subprocess
taskflow_json = subprocess.run(
["docker", "exec", "airflow-local-airflow-scheduler-1",
"cat", "/opt/airflow/dags/data/work/zone_counts_taskflow.json"],
capture_output=True, text=True, check=True,
).stdout
classic_json = subprocess.run(
["docker", "exec", "airflow-local-airflow-scheduler-1",
"cat", "/opt/airflow/dags/data/work/zone_counts_classic.json"],
capture_output=True, text=True, check=True,
).stdout
print(taskflow_json)
assert taskflow_json == classic_json, "TaskFlow and classic output diverged!"
print("IDENTICAL: TaskFlow and classic DAGs produced byte-for-byte the same output"){
"Bronx": 43,
"Brooklyn": 61,
"EWR": 1,
"Manhattan": 69,
"Outside of NYC": 1,
"Queens": 69,
"Staten Island": 20,
"Unknown": 1
}
IDENTICAL: TaskFlow and classic DAGs produced byte-for-byte the same outputEight boroughs (including the one real “Outside of NYC” placeholder row and the one genuinely unclassified “Unknown” zone), summing to all 265 real zones — and the two DAGs, written independently with two different mechanisms, agree on every single number.
Key Takeaways
- The TaskFlow API’s
@dag/@taskand classicDAG(...)+PythonOperator+>>are two syntaxes for the same underlying engine — this lesson’s real diff proves it, not just asserts it. - TaskFlow’s nested function calls (
report_zone_counts(validate_zone_counts(extract_zone_counts()))) both wire dependencies and pass data through XCom automatically; the classic version has to callti.xcom_pull(task_ids="...")explicitly in every downstream task to get the same data. - A real bug (one zone with no
Boroughvalue) surfaced in real data the first time this DAG ran, and thevalidate_zone_countsassertion caught it in both versions identically — proof the two DAGs don’t just look the same, they behave the same, including how they fail. - When two implementations claim to be equivalent, a real file diff on their actual output is stronger evidence than reading the code and agreeing it looks right.
Exercises
Exercise 1 — Add a fourth borough-level statistic to both DAGs (e.g. the borough with the most zones) and confirm both versions still produce identical output after the change — proving the equivalence holds for a change you make yourself, not just the original code.
Exercise 2 — In the classic DAG, delete the ti.xcom_pull(task_ids="extract_zone_counts") line from validate_zone_counts and replace it with a hardcoded dict. Trigger the DAG again and use airflow tasks states-for-dag-run to confirm extract_zone_counts still runs and succeeds even though nothing downstream uses its result anymore — a task with no consumer still executes, it just becomes dead weight.
Exercise 3 — Rewrite cityflow_zone_report_classic_dag.py’s three functions to use @task (keep the PythonOperator/>> DAG structure otherwise) and confirm in the Airflow UI’s Code tab that mixing TaskFlow-decorated callables into a classic with DAG(...) block is legal — TaskFlow tasks can live inside either style of DAG definition.
Next lesson
Lesson 2 - Operators