Airflow Tutorial: Orchestrating Data Pipelines

A script that loads data is a pipeline only while someone remembers to run it. Orchestration is what turns scripts into a dependable system: running tasks on a schedule, in the right order, retrying failures, recording what happened, and making it easy to rerun a past day. Apache Airflow is the most widely used open-source orchestrator. Its documentation describes it as a platform that lets you build and run workflows, where a workflow is a DAG, a directed acyclic graph, of tasks (Airflow docs: architecture).

The pieces

• DAG files: Python files that define workflows. Airflow reads them; it does not execute your data logic while reading.
• DAG processor: parses the files and stores the result in the metadata database.
• Scheduler: decides which runs and tasks are due and submits them.
• Workers: run the tasks. In small installations the scheduler runs them itself.
• Metadata database: the record of every DAG run and task attempt, with its state.
• API server and web interface: where you watch runs, read logs and trigger reruns.

Tasks come in three kinds: operators, ready-made tasks such as "run this SQL"; sensors, which wait for something external; and TaskFlow functions, ordinary Python functions marked with a decorator.

A complete DAG

This DAG processes one day of New York taxi trips (NYC TLC: trip record data) in three tasks. It is written with the TaskFlow API and ends with a block that runs it directly, twice, for 15 January 2024:

import sqlite3
from datetime import datetime, timedelta, timezone

import pandas as pd
from airflow.sdk import dag, task

SOURCE = "yellow_tripdata_2024-01.parquet"       # downloaded once from the TLC trip record page
DB = "taxi.db"


@dag(schedule="@daily", start_date=datetime(2024, 1, 1, tzinfo=timezone.utc), catchup=False,
     default_args={"retries": 2, "retry_delay": timedelta(minutes=5)}, tags=["taxi"])
def taxi_daily():
    @task
    def extract(ds=None):
        """Pull one day's trips from the source into a staging file."""
        trips = pd.read_parquet(SOURCE, columns=["tpep_pickup_datetime", "PULocationID", "fare_amount"])
        day = trips[trips["tpep_pickup_datetime"].dt.strftime("%Y-%m-%d") == ds]
        path = f"staging_{ds}.parquet"
        day.to_parquet(path)
        return {"day": ds, "path": path, "rows": len(day)}

    @task
    def validate(batch):
        """Fail the run if the day is empty; drop rows that break the fare rule."""
        if batch["rows"] == 0:
            raise ValueError(f"no trips found for {batch['day']}")
        day = pd.read_parquet(batch["path"])
        clean = day[day["fare_amount"] >= 0]
        clean.to_parquet(batch["path"])
        return {**batch, "rejected": len(day) - len(clean), "rows": len(clean)}

    @task
    def load(batch):
        """Replace the day's rows in the reporting table, so reruns never duplicate."""
        day = pd.read_parquet(batch["path"])
        zones = day.groupby("PULocationID").agg(trips=("fare_amount", "size"), fares=("fare_amount", "sum")).reset_index()
        con = sqlite3.connect(DB)
        with con:
            con.execute("CREATE TABLE IF NOT EXISTS daily_zone (day TEXT, zone INTEGER, trips INTEGER, fares REAL)")
            con.execute("DELETE FROM daily_zone WHERE day = ?", (batch["day"],))
            con.executemany("INSERT INTO daily_zone VALUES (?, ?, ?, ?)",
                            [(batch["day"], int(z), int(t), float(f)) for z, t, f in zones.itertuples(index=False)])
        con.close()
        return batch["rows"]

    load(validate(extract()))


pipeline = taxi_daily()

if __name__ == "__main__":
    for attempt in (1, 2):                                    # the second run shows the load is idempotent
        run = pipeline.test(logical_date=datetime(2024, 1, 15, tzinfo=timezone.utc))
        states = {ti.task_id: ti.state for ti in run.get_task_instances()}
        print(f"run {attempt}: {run.state} | " + ", ".join(f"{t}={states[t]}" for t in ("extract", "validate", "load")))
    con = sqlite3.connect(DB)
    print("rows in daily_zone:", con.execute("SELECT day, COUNT(*), SUM(trips) FROM daily_zone GROUP BY day").fetchall())

run 1: success | extract=success, validate=success, load=success
run 2: success | extract=success, validate=success, load=success
rows in daily_zone: [('2024-01-15', 216, 76103)]

Reading the DAG

The decorator defines the workflow. The @dag settings say it runs daily, when its schedule starts, and that every task gets two retries five minutes apart. catchup=False stops Airflow creating a run for every day since the start date the moment the DAG is switched on.

Tasks are functions, and calls are dependencies. load(validate(extract())) does not run anything when the file is parsed. It tells Airflow that validate needs extract's result and load needs validate's, which is the graph. With classic operators the same thing is written with the >> operator.

The logical date makes runs repeatable. extract receives ds, the date the run is for, and selects that day's trips. The run for 15 January processes 15 January whenever it executes, today, tomorrow or in a backfill next year. Tasks that use the wall clock instead cannot be rerun correctly.

Return values travel as XComs. Each task returns a small dictionary, the day, a file path and counts, which Airflow stores and hands to the next task. The trips themselves go to a staging file, not through XCom: the documentation is firm that XComs are for small amounts of data, not large values such as dataframes (Airflow docs: XComs).

Validation can fail the run. If a day has no trips, validate raises an error; Airflow marks the task failed, retries it twice, and never starts load. A failed check stops bad data from being published.

Running it

dag.test() executes a real DAG run in a single Python process, with no scheduler or executor, which is the quickest way to develop and debug (Airflow docs: debugging). It needs an initialised metadata database, from airflow db migrate, and the file must be in the configured DAGs folder; a DAG defined elsewhere fails with a "not serialized" error, as it did the first time this example was run from outside that folder.

The output shows both runs succeeding and the reporting table holding 216 zone rows and 76,103 trips for the day, once. The second run did not double the data, because load deletes the day's rows before inserting them inside one transaction. That is essential in Airflow, where reruns are routine: retries, manual clears and backfills all execute tasks again. Idempotency in data pipelines explains the patterns.

In production the same file is deployed to the DAGs folder of a running Airflow, and the scheduler creates a run each day. Managed services, such as Cloud Composer, Amazon MWAA and Astronomer, operate the components for you.

Habits for dependable DAGs

• Keep DAG files light. Code at the top level runs every time the file is parsed, so put work inside tasks.
• Make every task idempotent and date-driven, using the logical date, not the current time.
• Pass references, not data. Store data in files or tables and pass paths or keys through XCom.
• Set retries and alerts in default arguments, and timeouts on tasks that can hang.
• Let Airflow orchestrate, not compute. Heavy processing belongs in Spark, the warehouse or dbt, triggered by Airflow; the dbt tutorial covers one such tool.
• Test DAGs with dag.test() and in continuous integration before deploying.

The logic inside the tasks follows the pattern from building your first data pipeline, and Airflow vs Dagster vs Prefect compares the alternatives.

Apache Airflow, workflow scheduling, DAG creation, pipeline monitoring and automation are part of the data pipelines and orchestration module of the Data Engineering course.

Sources

• Architecture overview (Apache Airflow documentation) — Apache Airflow
• Debugging Airflow Dags (Apache Airflow documentation) — Apache Airflow
• XComs (Apache Airflow documentation) — Apache Airflow
• TLC Trip Record Data (NYC Taxi and Limousine Commission) — NYC Taxi and Limousine Commission