Dynamic Task Mapping

Concepts covered: Dynamic task mapping

The shape the file cannot know

Every DAG the earlier lessons built had its shape fixed in the file: 3 tasks written down means 3 cells in the grid on every run, whatever the data looked like that night. A partner feed job breaks that model. Some nights 3 files land and some nights 40. The file can't name a task for a file it has never seen. Dynamic task mapping keeps a single task definition in the file and lets the run decide how many copies of it exist, from a value produced upstream in that same run.

from airflow.sdk import dag, task


@dag(schedule=None)
def load_partner_feeds():
    @task
    def list_files():
        return ["feeds/acme.csv", "feeds/globex.csv", "feeds/initech.csv"]

    @task
    def load(path):
        print(f"loading {path}")
        return path.split("/")[-1]

    @task
    def report(names):
        print(sorted(names))

    report(load.expand(path=list_files()))


load_partner_feeds()

The one new call is load.expand(path=list_files()) in place of load(path=...). At parse time the processor sees 3 tasks and the graph draws 3 boxes, one each for list_files, load and report. Nothing has fanned out yet, because the parser never calls list_files. When a run is triggered and list_files finishes, the scheduler reads its return value. It finds 3 strings and creates 3 task instances of load, numbered 0 to 2 by map_index. Each copy receives one string as its path. The grid keeps load on a single row and labels the row with the count; opening it lists the copies with their indexes and states.

report is an ordinary task that takes names. Because its argument is the output of a mapped task, it waits for all 3 copies to finish, then receives their return values as a sequence in index order: ['acme.csv', 'globex.csv', 'initech.csv']. If 1 copy fails, report never runs, the same as it would behind a single failed task. A downstream task wired with >> and no argument from load waits for every copy in the same way, because the mapped task counts as done only when its last copy is.

partial, expand, and the product of 2 lists

expand takes the arguments that vary per copy. Everything shared across copies goes in partial, which is also where a classic operator receives its task_id, since the operator is never constructed directly. Passing 2 lists to expand gives the cartesian product: 3 paths times 2 formats is 6 copies. When the values must stay paired, expand_kwargs takes a list of dicts, 1 dict per copy, each holding that copy's complete keyword set. A name given to partial cannot also be given to expand; the file fails to parse if it is.

from airflow.providers.standard.operators.bash import BashOperator
from airflow.sdk import dag, task


@dag(schedule=None)
def reshape_feeds():
    @task
    def convert(path, fmt):
        return f"{path}.{fmt}"

    @task
    def parse(path, delimiter):
        return path.split("/")[-1]

    BashOperator.partial(task_id="compress", env={"LEVEL": "9"}).expand(
        bash_command=["gzip -k feeds/acme.csv", "gzip -k feeds/globex.csv"],
    )

    convert.expand(path=["a.csv", "b.csv", "c.csv"], fmt=["parquet", "orc"])

    parse.expand_kwargs([
        {"path": "feeds/acme.csv", "delimiter": ","},
        {"path": "feeds/globex.csv", "delimiter": "|"},
    ])


reshape_feeds()

• 1,024

Knowing which copy you are

Inside a mapped task, get_current_context()["ti"].map_index returns the copy's index, which lets a copy write to a file or a partition named after itself without colliding with its siblings. An unmapped task reads -1 there. The context holds no total, so a copy that must know how many siblings it has is handed the number: the upstream task returns 1 dict per file carrying both the path and len() of the list, and expand_kwargs passes each dict to its own copy. A task group can be mapped too: decorate it with @task_group, give it 1 argument, and call .expand() on the group; every task inside is copied once per value, and the grid shows the group with the count on its row, so a 3-step pipeline per file stays readable as 1 group with a number beside it.

from airflow.sdk import get_current_context, task


@task
def load(path):
    idx = get_current_context()["ti"].map_index
    out = f"staging/part-{idx:03d}.parquet"
    print(f"{path} -> {out}")
    return out

About this lesson

This section is part of the DAG Anatomy: Advanced lesson on DataDriven, the home of data engineering. Lessons cover SQL, Python, Spark, data modeling and pipeline architecture, and each section pairs its explanation with worked examples and practice you run as you go.