You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "dabla (via GitHub)" <gi...@apache.org> on 2026/09/24 19:38:22 UTC
[PR] AIP-104: Task Spreading [airflow]
dabla opened a new pull request, #73688:
URL: https://github.com/apache/airflow/pull/73688
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull request.
Write a good git commit message following this guide: http://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the documentation.
Feel free to ping (in general) for the review if you do not see reaction for a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR, please
change below checkbox to `[X]` followed by the name of the tool, uncomment the "Generated-by".
-->
- [ x ] Yes (please specify the tool below)
Claude Code (Fable 5.1).
<!--
Generated-by: [Tool Name] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
## Description
This PR adds **Task Spreading** from [AIP-104](https://cwiki.apache.org/confluence/pages/viewpage.action?pageId=421954527) 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()`.
### 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.
```python
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()
```
https://github.com/user-attachments/assets/87edba93-cd31-4711-be96-c49a2b4cf6d1
### 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](https://lists.apache.org/thread/x5bd26zfqkkhp7sx7nb0crjlsctf2z5j) 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](https://claude.com/claude-code)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-104: Task Spreading [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on PR #73688:
URL: https://github.com/apache/airflow/pull/73688#issuecomment-5858139754
Restacked on #62922 after `XComIterable.flatten()` moved out to its own PR; still one commit, 488139d1ea.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-104: Task Spreading [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on PR #73688:
URL: https://github.com/apache/airflow/pull/73688#issuecomment-5855905720
Restacked on the rebased #62922; still one commit, 5a96cbf2d9.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] AIP-104: Task Spreading [airflow]
Posted by "dabla (via GitHub)" <gi...@apache.org>.
dabla commented on PR #73688:
URL: https://github.com/apache/airflow/pull/73688#issuecomment-5855751497
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, f3ed32fe4d.
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org