AIP-104: Task Spreading

[dabla]

Review only this change: the PR's own commit · diff against #62922's branch. Everything else in this PR's commit list and file list is #62922, which this PR is stacked on.

---

Was generative AI tooling used to co-author this PR?

• [ x ] Yes (please specify the tool below)

Claude Code (Fable 5.1).

Description

This PR adds Task Spreading from AIP-104 on top of Iterable Tasks from #62922: .batch(size=N).iterate(...) spreads one iteration over exactly N task instances through Dynamic Task Mapping, and each task instance iterates over its round-robin share with Iterable Tasks. It is the second half of the original #62922, split out so that PR stays focused on .iterate().

Stacked on #62922. Until that PR merges, the diff shown here includes it. The spreading work is the last commit, "Add task batching on top of iterable tasks". The method is still called .batch() in the code; see Open points for the rename to .spread(). To review only the spreading layer, use the compare view of this branch against #62922's branch: dabla/airflow@feature/dynamic-task-iteration...feature/task-spreading (GitHub's native stacked PRs cannot be used here, since they require all branches in the same repository and the branches live in a fork). Once #62922 merges, this branch is rebased onto main and the diff here collapses to that layer.

What this PR adds

• .batch(size=...) on OperatorPartial and on @task returns a BatchedOperator / DecoratedBatchedOperator (airflow.sdk.definitions.batchedoperator). Their iterate() / iterate_kwargs() reuse the partial's own input validation and _expand(), and build a MappedIterableOperator instead of an IterableOperator.
• MappedIterableOperator: a MappedOperator whose task instances each run an IterableOperator over their share of the input. BatchedExpandInput routes item i to task instance i % size, so exactly size instances are created regardless of how many items the (possibly unbounded or paginated) input yields.
• Runtime size: size may be an XComArg, the return value of a plain, non-mapped task. The scheduler never reads the XCom value: the worker pushes the integer as the mapped_length of that push and the scheduler counts instances from that column (SerializedMappedOperator.resolve_batch_size), raising NotFullyPopulated until the upstream has run. The value must be an integer of at least 2 and at most core.max_map_length; .map()/.zip() results, pushed keys and mapped upstreams are rejected at parse time.
• Scheduler-side counting (get_parse_time_mapped_ti_count / get_mapped_ti_count) for spread tasks, the docs section Combining DTM and IT (Batched Task Mapping), and tests for all of the above.

Example

Task Iteration with Task Spreading

This example performs the same work as the Task Iteration example in #62922, but spreads the workload over two concurrent task instances. Each task instance processes roughly half of the Pokémon URLs using Task Iteration.

from airflow.sdk import dag, task
from airflow.providers.http.hooks.http import HttpHook, HttpAsyncHook

from pendulum import datetime

@dag(
    start_date=datetime(2025, 1, 1),
    schedule=None,
    catchup=False,
)
def pokemon_spread_iteration():
    @task
    def list_pokemon() -> list[str]:
        response = HttpHook(
            http_conn_id="pokeapi",
            method="GET",
        ).run(
            endpoint="api/v2/pokemon?limit=100",
        )

        return [
            pokemon["url"].replace("https://pokeapi.co/", "")
            for pokemon in response.json()["results"]
        ]

    @task(
        retries=3,
        task_concurrency=2,
        show_return_value_in_logs=False,
    )
    async def get_pokemon(url: str):
        async with HttpAsyncHook(
            http_conn_id="pokeapi",
            method="GET",
        ).session() as session:
            response = await session.run(endpoint=url)
            return await response.json()

    get_pokemon.batch(size=2).iterate(
        url=list_pokemon(),
    )

pokemon_spread_iteration()

Pokemon.Partitioned.Task.Iteration.mp4

Comparison

Pattern	Task Instances	Work Per Task
get_pokemon.expand(url=urls)	100	1 Pokémon
get_pokemon.iterate(url=urls)	1	100 Pokémon
get_pokemon.batch(size=2).iterate(url=urls)	2	~50 Pokémon each

Task Spreading keeps the reduced TaskInstance overhead of Task Iteration while allowing controlled parallelism across workers.

Note: batch(size=2) creates two task instances, each processing its assigned items with Task Iteration. size is the number of task instances, not a chunk length: items are distributed round-robin, unlike itertools.batched, because the instance count must be fixed before a possibly unbounded input is consumed.

Open points

• Naming. The method is still .batch(size=N) in the code. The devlist vote leans towards .spread(), which describes it better: exactly N task instances with the items dealt round-robin, not itertools.batched-style chunks. The rename lands in this PR.
• Chunking is a different feature. Handing each mapped task instance a list of items from a known input (for example 100 S3 keys as 10 lists of 10, for an operator that accepts a list) is a transformation of the upstream value and belongs on XComArg next to .map() and .zip(), e.g. .expand(keys=files.batch(10)). It needs no iteration and is not part of this PR.

🤖 Generated with Claude Code

[dabla]

Rebased on the reworked #62922 (see the comment there: .iterate() now resolves its input by index like .expand()). Still a single commit on top of it, f3ed32f.

What the rework changed here: BatchedExpandInput.aresolve() is a range(map_index, length, size) over the delegate instead of walking the whole input and skipping; the striding tests moved to aresolve(); and the round-robin note in the docs no longer argues from a length that is "unknowable until drained", since the count is fixed by the scheduler before the task runs, which is the reason that actually holds.

[dabla]

Restacked on the rebased #62922; still one commit, 5a96cbf.

[dabla]

Restacked on #62922 after XComIterable.flatten() moved out to its own PR; still one commit, 488139d.