Dynamic Task Mapping and TaskGroups: Fan Out, Stay Tidy

The bookshop reports orders per region, and the regions aren’t fixed. Some days there are four files, some days seven, and the day they open in a new country there’s an eighth. You can’t hardcode a task per file, because you don’t know the count when you write the DAG — you only know it when today’s run inspects the bucket. The foundations series introduced .expand() and showed the shape of a fan-out. This chapter is about the parts that bite once you actually depend on mapping in production: binding the constants you don’t want to map, telling a Cartesian product from a zip before it multiplies your task count by fifty, capping the stampede against a warehouse, and collecting the results back without silently coercing a lazy sequence into a list that eats your worker’s memory.

.partial(): the constants that must not fan out

Here’s the mistake everyone makes on their second mapped task. The first one maps a single argument and reads beautifully. The second one has a target table, a batch size, a retry count — arguments that are the same for every mapped instance — and there’s a temptation to jam them into the list too. Don’t. That’s what .partial() is for. .partial() binds the arguments that stay constant across every mapped instance; .expand() supplies the ones that vary. You cannot write a realistic mapped task without both, because real tasks almost always have a mix of the two.

from airflow.sdk import dag, task


@dag(schedule="@daily")
def bookshop_regional_load():
    @task
    def list_region_files(ds: str) -> list[str]:
        # `ds` is injected from the run context. Do NOT write "{{ ds }}" in a
        # task body — Jinja doesn't render there; you'd get literal braces.
        return BookshopStorageHook().list_keys(prefix=f"orders/{ds}/")

    @task
    def load_region(key: str, table: str, batch_rows: int) -> int:
        return BookshopWarehouseHook().copy_into(
            key, table=table, batch_rows=batch_rows
        )

    keys = list_region_files()
    counts = load_region.partial(table="stg_orders", batch_rows=5_000).expand(key=keys)


bookshop_regional_load()

table and batch_rows are bound once, for the whole mapped node. Only key fans out. If eight keys come back, you get eight load_region instances, every one of them writing to stg_orders in batches of five thousand, differing only in which file they read. Leave table out of .partial() and Airflow will complain that load_region is missing a required argument, because .expand() only provides the keys you mapped. The rule is simple and worth internalising: every argument the callable takes must be accounted for — mapped by .expand()/.expand_kwargs(), or pinned by .partial(). There’s no third option, and that’s a feature. It forces you to decide, per argument, whether it varies.

.partial() also carries the task-level knobs you’d normally pass to @task: retries, retry_delay, pool, queue, max_active_tis_per_dag, execution_timeout. Those apply to the mapped node as a whole and to each instance under it, which is exactly where we’re heading once the fan-out gets wide.

One expand, many args: the Cartesian trap

.expand() accepts more than one keyword. The moment you pass two mapped arguments, Airflow builds the cross product — every combination of the two lists. This is the single most misread behaviour in the whole feature, and it’s misread in the expensive direction.

# Cartesian product: 3 regions × 2 formats = 6 mapped instances
load_region.partial(table="stg_orders").expand(
    key=["us.csv", "eu.csv", "apac.csv"],
    fmt=["gzip", "raw"],
)

Three regions and two formats produce six instances: (us, gzip), (us, raw), (eu, gzip), (eu, raw), and so on. That’s genuinely what you want when the axes are independent — reprocess every region in every format. It is emphatically not what you want when the two lists are meant to line up positionally, one key paired with its own matching table. Do that with the cross product and three regions each with their own table become nine instances, six of which write the wrong region into the wrong table. The failure isn’t an error; it’s a quietly wrong result, which is worse.

When you want the lists zipped — the first key with the first table, the second with the second, no combinatorics — you have two tools. The first is .expand_kwargs(), which takes a list of dicts and creates exactly one instance per dict:

# Zipped: one instance per dict, key paired with its own table
load_region.partial(batch_rows=5_000).expand_kwargs([
    {"key": "orders/us.csv", "table": "stg_orders_us"},
    {"key": "orders/eu.csv", "table": "stg_orders_eu"},
    {"key": "orders/apac.csv", "table": "stg_orders_apac"},
])

Three dicts, three instances, each argument taken straight from its dict. No products, no accidental multiplication. expand_kwargs is the right call whenever the per-item difference is more than a single value and the values are pre-paired. And crucially, its argument can itself be an XComArg — the output of an upstream task that assembles those dicts at runtime — so you’re not limited to literals:

@task
def plan_region_jobs(ds: str) -> list[dict]:
    keys = BookshopStorageHook().list_keys(prefix=f"orders/{ds}/")
    return [{"key": k, "table": f"stg_orders_{region_of(k)}"} for k in keys]

load_region.partial(batch_rows=5_000).expand_kwargs(plan_region_jobs())

Now the fan-out width and the per-instance argument pairing are both decided during the run. That’s the pattern to reach for when each region needs a genuinely different set of parameters, not just a different key.

Transform before you expand: .map(), .zip(), and dicts

The second way to zip lives on XComArg itself, and it’s lighter than a whole extra task. Any task output is an XComArg, and XComArg carries lazy transforms you can chain before expanding — they don’t create their own task instances, they reshape the sequence in place.

.zip() pairs two (or more) upstream outputs into tuples, the way Python’s zip does:

keys = list_region_files()          # ["us.csv", "eu.csv", ...]
tables = derive_tables(keys)        # ["stg_orders_us", "stg_orders_eu", ...]

@task
def load_pair(pair: tuple[str, str]) -> int:
    key, table = pair
    return BookshopWarehouseHook().copy_into(key, table=table)

load_pair.expand(pair=keys.zip(tables))

By default .zip() stops at the shortest input, exactly like the builtin; pass default= for zip_longest behaviour when the lists can differ in length. There’s a sibling, .concat(), that chains sequences end to end rather than pairing them — keys_a.concat(keys_b) gives you one combined list to map over, without materialising either in a task.

.map() applies a plain Python function to each element before it reaches .expand(). The callable must be an ordinary function — not a @task — takes one argument, and runs inline as part of resolving the mapping:

def to_job(key: str) -> dict:
    return {"key": key, "table": f"stg_orders_{key.split('/')[1]}"}

load_region.partial(batch_rows=5_000).expand_kwargs(
    list_region_files().map(to_job)
)

One thing that trips people up: inside a .map() function you skip an item by raising AirflowSkipException, not by returning None. Return None and you get a mapped instance that receives None and probably explodes; raise the skip and that index simply drops out of the fan-out. It’s the map-side equivalent of filtering.

And then there’s expanding over a dict directly. If the value you hand .expand() is a dict rather than a list, Airflow maps over its items — one mapped instance per key-value pair, and each instance receives a (key, value) tuple:

@task
def region_thresholds() -> dict[str, int]:
    return {"us": 10_000, "eu": 8_000, "apac": 3_000}

@task
def check_region(item: tuple[str, int]) -> None:
    region, min_rows = item
    BookshopWarehouseHook().assert_min_rows(region, min_rows)

check_region.expand(item=region_thresholds())

Three keys, three instances, each unpacking its own (region, min_rows) pair. The map index in the UI even shows the dict key, so check_region["us"] is legible at a glance instead of an anonymous [0]. Reach for dict-expansion when the identity of each item matters and you’d otherwise be threading a separate lookup table alongside a positional list.

Naming the instances: map_index_template

Dict-expansion gets legible labels for free, but list-expansion doesn’t — map over eight file keys and the grid shows load_region[0] through load_region[7], which tells you nothing when instance [5] fails at 3 a.m. and you need to know which region it was. map_index_template fixes that. It’s a Jinja template, rendered after the task body runs, that overrides the label shown in the UI and referenced in the logs. Because it renders post-execution, it can read anything you computed during the run — you stash the value on the task context and name the template after it:

from airflow.sdk import task, get_current_context


@task(map_index_template="{{ region }}")
def load_region(key: str, table: str) -> int:
    context = get_current_context()
    context["region"] = key.split("/")[1]   # e.g. "eu"
    return BookshopWarehouseHook().copy_into(key, table=table)

Now the grid shows load_region[eu] beside load_region[us], and the instance that failed names the region that failed with it. Since the template resolves against the live context, it isn’t limited to inputs — you can label an instance by the row count it loaded, the source filename it resolved, whatever makes triage a glance instead of a per-index log dive. On a wide, mixed fan-out this small annotation pays for itself the first time something breaks.

The labels aren’t only cosmetic. A mapped task instance is still individually addressable, so once you’ve identified the one bad region you can clear and re-run just that instance from the grid — right-click, clear, and only load_region[eu] re-queues while the other seven stay done. A named index turns that from guesswork into a single deliberate action, which matters most on exactly the wide fan-outs where you’d otherwise be re-running the whole mapped node to fix one file.

Bounding the fan-out: pools, max_active_tis_per_dag, max_map_length

Dynamic does not mean unbounded, and the bookshop warehouse will remind you of that the first time a mapping over two thousand micro-batches opens two thousand connections at once. There are three separate limits, and they do genuinely different jobs.

[core] max_map_length (default 1024) is the guardrail on how many instances a single .expand() may create. If an upstream task returns a list longer than this, the mapped task fails outright rather than spawning a runaway graph. It’s a blast-radius cap: a source file with a corrupt row count that suddenly yields fifty thousand keys should break loudly, not quietly try to schedule fifty thousand tasks. Raise it deliberately in airflow.cfg (or AIRFLOW__CORE__MAX_MAP_LENGTH) only when you’ve decided a genuinely wider fan-out is safe.

max_active_tis_per_dag caps concurrency, not count — how many mapped instances of a task run at the same time, within one DAG. You can still fan out to 800 instances; with max_active_tis_per_dag=8 on the task, only eight run concurrently and the rest queue. This is the knob for “I have hundreds of files but the warehouse tolerates eight parallel COPYs.”

load_region.partial(table="stg_orders", max_active_tis_per_dag=8).expand(key=keys)

And for the case where the constraint is shared across DAGs — a warehouse that permits, say, twelve concurrent loaders total no matter which pipeline issues them — that’s a pool. Create a bookshop_warehouse pool with twelve slots in the UI, assign it, and every mapped instance draws a slot before it runs. When the pool is full, instances wait. This is the correct tool for protecting a shared resource from a stampede, because it counts across every DAG and task that names the pool, not just within one mapped node.

load_region.partial(table="stg_orders", pool="bookshop_warehouse").expand(key=keys)

The three compose: max_map_length says this fan-out may not exceed N instances at all; max_active_tis_per_dag says no more than K of this task’s instances run together; the pool says no more than M things touch the warehouse across the whole deployment. On a wide, warehouse-bound mapping you’ll often set all three.

Collecting the results: the lazy sequence, and partial failures

The reduce step — the task that consumes a mapped output — is where a subtle performance and correctness detail lives. When a downstream task takes the collected output of a mapped task, Airflow does not hand it a plain list. It hands it a lazy proxy (you’ll see it logged as LazySelectSequence([N items])) that pulls each mapped return value from the metadata database only as you iterate. That laziness is deliberate: a mapping over a thousand instances shouldn’t load a thousand XComs into worker memory just because the reduce task started.

@task
def consolidate(counts) -> None:  # counts is a lazy sequence, not a list
    total = 0
    for c in counts:              # streamed from the DB one at a time
        total += c
    BookshopWarehouseHook().refresh_orders_mart(expected_rows=total)

consolidate(load_region.partial(table="stg_orders").expand(key=keys))

Iterate it, sum it, feed it to a generator — fine, it streams. But call list(counts), len(counts), sort it, or index into it, and you force the whole thing to materialise; Airflow will even log a “coercing lazy proxy to list, which may degrade performance” warning to tell you it happened. For a reduce, iterating is almost always enough, so prefer sum(counts) and for c in counts over anything that needs the sequence all at once.

The other thing the reduce has to survive is partial failure. If one of eight mapped instances fails, the collected sequence does not include a slot for it — the failed index is absent, not None. By default the downstream task uses the all_success trigger rule, so a single failed mapped instance means consolidate doesn’t run at all, which is usually the safe default: don’t refresh the mart from seven regions when the eighth is missing. When partial results are acceptable — refresh from whatever loaded, flag the gap — set the trigger rule explicitly:

from airflow.sdk import TriggerRule

@task(trigger_rule=TriggerRule.ALL_DONE)
def consolidate(counts) -> None:
    loaded = list(counts)  # only the successful instances' returns
    BookshopWarehouseHook().refresh_orders_mart(
        expected_rows=sum(loaded), partial=True
    )

ALL_DONE runs the reduce once every mapped instance has finished in some terminal state, and the lazy sequence contains only the successful returns. Skipped instances — the ones a .map() dropped with AirflowSkipException — are likewise simply not present. Deciding between all_success and all_done here is a real product choice, not a default to accept blindly: it’s the difference between “one bad region blocks the mart” and “one bad region silently shrinks it.”

Mapping a whole TaskGroup

Sometimes the unit that needs to fan out isn’t a task — it’s a mini-pipeline. Each region needs validating, then loading, then row-counting, three steps in sequence, and you want that trio repeated per region. @task_group.expand() maps the entire group: every mapped instance gets its own copy of the group’s internal wiring.

from airflow.sdk import task_group


@task_group
def stage_region(key: str):
    @task
    def validate(k: str) -> str:
        BookshopStorageHook().assert_schema(k)
        return k

    @task
    def load(k: str) -> int:
        return BookshopWarehouseHook().copy_into(k, table="stg_orders")

    return load(validate(key))


counts = stage_region.expand(key=list_region_files())
consolidate(counts)

Airflow builds this as a mapped task group: one collapsed node in the graph that, expanded, shows a validate → load chain per region, each chain isolated from the others. The group’s return value collects across all instances into the usual lazy sequence, so the reduce downstream works exactly as before.

Constants pin to a mapped group exactly the way they pin to a mapped task — the .partial()/.expand() split lifts wholesale from one task up to a whole mini-pipeline. If every region’s trio writes to the same staging table and shares a batch size, don’t smuggle those through key; bind them once on the group:

@task_group
def stage_region(key: str, table: str, batch_rows: int):
    @task
    def validate(k: str) -> str:
        BookshopStorageHook().assert_schema(k)
        return k

    @task
    def load(k: str, table: str, batch_rows: int) -> int:
        return BookshopWarehouseHook().copy_into(k, table=table, batch_rows=batch_rows)

    return load(validate(key), table, batch_rows)


counts = stage_region.partial(table="stg_orders", batch_rows=5_000).expand(
    key=list_region_files()
)

table and batch_rows are supplied once to the group and threaded to whichever inner tasks need them; only key fans the group out. Every rule from single-task mapping still applies — the same max_active_tis_per_dag and pool you’d set on a task can go on the group’s .partial() to cap how many chains run at once.

There’s one hard limit to know before you design around this: you cannot nest task mapping inside a mapped task group. A task inside a group that’s already being expanded cannot itself call .expand() — Airflow forbids the two-dimensional fan-out because the instance-index bookkeeping doesn’t compose. So stage_region above can map as a group, but load inside it must take a single key, not expand over a list of its own. If you find yourself wanting a mapped task within a mapped group, that’s the signal to flatten: precompute the full cross product in an upstream task and expand once over the combined list, rather than nesting.

TaskGroups as structure: deps, defaults, labels

Fan-out aside, TaskGroups earn their keep as pure organisation. A TaskGroup is UI and namespacing only — it collapses related tasks into one box in the graph and prefixes their task IDs (stage_orders.load, stage_orders.validate), so the same group can appear in several DAGs or nest inside another group without ID collisions. It creates no runtime isolation: the tasks inside still run as ordinary, independent instances on whatever executor, with no shared process, no transaction boundary, nothing. If you’re expecting a group to give you atomicity or a sandbox, you’re expecting SubDAG behaviour that TaskGroups deliberately don’t provide.

What groups do give you at the structural level is clean dependency wiring. You can set dependencies between whole groups — ingest_group >> transform_group makes every task in transform_group wait on every task in ingest_group — which reads far better than hand-wiring the boundary tasks. You can pass default_args to a group so every task inside inherits the same retries, retry_delay, or pool without repeating yourself. And Label lets you annotate the edge between tasks or groups, which is invaluable when a branch leaves a group down a “processed” versus “skipped” path:

from airflow.sdk import Label

ingest_group >> Label("clean batch") >> transform_group
ingest_group >> Label("quarantine") >> triage_group

Those labels render right on the edges in the graph, turning an otherwise cryptic fork into something a teammate reads without opening a single task.

Retiring the last SubDAG

If you’re modernising an older codebase, mapped groups are also the endgame for the last SubDagOperator you’re carrying. SubDAGs are removed in Airflow 3.x, and they were never mourned: they ran a whole child DAG inside a task, brought their own scheduler and their own deadlock modes, and quietly starved pools. The migration recipe is mechanical. A SubDAG factory that returned a DAG becomes a @task_group function that returns its final task’s output. The SubDagOperator(subdag=make_subdag(...)) call site becomes a plain call to that group function, wired into the parent with >> like any other node. The child DAG’s dag_id-prefixed task IDs become the group’s automatic group_id. namespacing, so your logs and metrics stay legible. If the SubDAG existed specifically to fan out over data — the common case — the group becomes a mapped group via .expand(), and you’ve replaced a fragile nested-scheduler construct with a first-class one.

Concretely, here’s the before and after. The 2.x construct was a factory that returned a DAG, handed to a SubDagOperator at the call site:

# BEFORE — Airflow 2.x, removed in 3.x
from airflow.operators.subdag import SubDagOperator
from airflow.providers.standard.operators.python import PythonOperator


def make_stage_subdag(parent_id: str, child_id: str, args: dict) -> DAG:
    subdag = DAG(dag_id=f"{parent_id}.{child_id}", default_args=args, schedule=None)
    with subdag:
        validate = PythonOperator(task_id="validate", python_callable=_validate)
        load = PythonOperator(task_id="load", python_callable=_load)
        validate >> load
    return subdag


with DAG("bookshop_regional_load", default_args=default_args, schedule="@daily"):
    stage = SubDagOperator(
        task_id="stage_region",
        subdag=make_stage_subdag("bookshop_regional_load", "stage_region", default_args),
    )
    list_files >> stage >> consolidate

The same pipeline as a task group is a plain function that returns its final task’s output, called inline and wired with >> like any other node:

# AFTER — Airflow 3.x
@task_group
def stage_region(key: str) -> int:
    return load(validate(key))


list_files >> stage_region(one_key) >> consolidate

The dag_id-prefixed child IDs (bookshop_regional_load.stage_region.load) collapse to the group’s automatic stage_region.load namespacing, so logs and metrics keep the same shape; the child scheduler and its deadlock modes simply evaporate. And when the SubDAG was there to run a region at a time, the call site becomes stage_region.expand(key=list_region_files()) and the migration and the fan-out land in the same edit. Anything a tutorial still writes with SubDagOperator is describing an Airflow that no longer exists.

Final thoughts

Dynamic mapping and TaskGroups pull in opposite directions, and that’s exactly why they belong together. Mapping multiplies — .partial() pins what’s constant, .expand() cross-multiplies, .expand_kwargs() and .zip() keep things paired, and the whole thing turns one authored task into as many runtime instances as today’s data demands. Left unchecked, that’s also what makes a graph explode and a warehouse buckle, which is why the limits — max_map_length, max_active_tis_per_dag, and a pool — aren’t optional trivia but part of authoring a mapped task honestly. TaskGroups fold the picture back down into something a human reads at a glance, and mapped groups let a whole mini-pipeline fan out as cleanly as a single task. Get the pairing right, cap the stampede, collect the lazy sequence without coercing it, and the bookshop pipeline scales with the data while the diagram of it stays the same size.

Next: Assets: Scheduling on Data, Not the Clock