dabla opened a new pull request, #73683:
URL: https://github.com/apache/airflow/pull/73683
<!-- 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 Batching** (called *Task Spreading* in
[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 `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 batching work is the last commit, *"Add task batching on top of
iterable tasks"*.
### 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 batch 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 batched tasks, the docs section *Combining DTM and
IT (Batched Task Mapping)*, and tests for all of the above.
## Example
### Task Iteration with Task Batching
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_batched_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_batched_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 Batching 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.** Whether `.batch(size=N)` should be called something else,
since `size` is the number of task instances rather than an
`itertools.batched`-style chunk length. AIP-104 calls the feature *Task
Spreading*; a vote is running on the
[devlist](https://lists.apache.org/thread/x5bd26zfqkkhp7sx7nb0crjlsctf2z5j).
The rename, if any, lands here.
🤖 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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
