dabla commented on code in PR #62922: URL: https://github.com/apache/airflow/pull/62922#discussion_r4221064298
########## task-sdk/docs/mapped-tasks-vs-iterable-tasks.rst: ########## @@ -0,0 +1,474 @@ + .. Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + .. http://www.apache.org/licenses/LICENSE-2.0 + + .. Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. + +.. _sdk-mapped-tasks-vs-iterable-tasks: + +Mapped tasks vs iterable tasks +============================== + +.. versionadded:: 3.4.0 + +Airflow provides two complementary ways to process collections of data: + +- **Mapped tasks** distribute work **across multiple workers**. + Each item becomes a separate Task Instance that can run on a different worker, + giving you horizontal scalability and per-item observability. + +- **Iterable Tasks (IT)** improves concurrency **within a single task**. + All items are processed inside one Task Instance on one worker, eliminating + scheduling overhead and — when combined with async operators — enabling true + I/O multiplexing through a shared event loop. + +In short: **mapping spreads load across workers; IT speeds up work within one worker.** + +While both approaches allow you to apply an operation over a collection, +they differ significantly in execution model, scheduler impact, and observability. +This page explains the trade-offs and when to use each. + +Real-World Motivation +--------------------- + +Consider a workflow that downloads ~17,000 XML files from an SFTP server and loads +them into a data warehouse. Community benchmarks compare a mapped operator with a loop +written by hand inside a single ``@task``, the pattern IT runs for you (none of these rows +uses ``iterate()`` itself): + +.. list-table:: + :header-rows: 1 + + * - Approach + - Execution Time + * - Mapped ``SFTPOperator`` + - 3 h 25 m + * - Sync ``@task`` with ``SFTPHook`` (sequential loop) + - 1 h 21 m + * - Async ``@task`` with ``SFTPHookAsync`` (concurrent loop) + - 8 m 29 s + * - Async ``@task`` with ``SFTPHookAsync`` and connection pooling + - 3 m 32 s + +The ~60× improvement comes from eliminating per-item scheduling overhead and +sharing a single event loop for concurrent I/O. This is the kind of workload +IT is built for: many small, I/O-bound operations processed within one task. + +Mapped tasks +------------ + +Task mapping allows you to expand a single task definition into multiple +Task Instances (TIs). + +For more details, see :ref:`task mapping <sdk-dynamic-task-mapping>`. + +Key characteristics: + +- Each item in the iterable creates a separate Task Instance. +- The scheduler is responsible for creating and managing all mapped tasks. +- Tasks can run in parallel across multiple worker slots. +- Fine-grained retry, logging, and observability per item. +- Well suited for workloads where each item should be independently scheduled and tracked. + +The following example fetches Pokémon data from a REST API. Each Pokémon becomes +a separate Task Instance, individually scheduled, retried, and visible in the UI: + +.. code-block:: python + + from datetime import datetime + + from airflow.providers.http.operators.http import HttpOperator + from airflow.sdk import DAG, task + + with DAG(dag_id="dtm-http-pokemon-example", start_date=datetime(2026, 1, 1)): + list_pokemon_task = HttpOperator( + task_id="list_pokemon", + http_conn_id="pokeapi", + method="GET", + endpoint="api/v2/pokemon?limit=100", + response_filter=lambda response: [ + pokemon["url"].replace("https://pokeapi.co/";, "") for pokemon in response.json()["results"] + ], + log_response=False, + ) + + get_pokemon_task = HttpOperator.partial( + task_id="get_pokemon", + http_conn_id="pokeapi", + method="GET", + ).expand(endpoint=list_pokemon_task.output) + + list_pokemon_task >> get_pokemon_task + + +With 100 Pokémon the scheduler creates 100 Task Instances, each occupying +a worker slot. This is fine for small lists, but for thousands of items the +scheduler and database overhead becomes significant. + +Iterable Tasks (IT) +---------------------------- + +Iterable Tasks allows you to iterate over an iterable (typically an XCom result) +*within a single Task Instance*, applying an operator multiple times without creating +separate Task Instances. + +This means that iteration happens inside the task execution itself rather than at the +scheduler level. + +Key characteristics: + +- A single Task Instance processes all items in the iterable. +- No task expansion; the scheduler manages only one task. +- Lower scheduler overhead compared to mapping. +- Iterations share the same execution context (e.g., memory, event loop). +- Particularly well suited for async operators and high-throughput workloads. + +The same Pokémon fetching problem can be solved with IT. Here, a single Task +Instance processes all Pokémon concurrently using the sync +:class:`~airflow.providers.http.operators.http.HttpOperator`: + +.. code-block:: python + + from datetime import datetime + + from airflow.providers.http.operators.http import HttpOperator + from airflow.sdk import DAG, task + + with DAG(dag_id="it-http-pokemon-example", start_date=datetime(2026, 1, 1)): + list_pokemon_task = HttpOperator( + task_id="list_pokemon", + http_conn_id="pokeapi", + method="GET", + endpoint="api/v2/pokemon?limit=100", + response_filter=lambda response: [ + pokemon["url"].replace("https://pokeapi.co/";, "") for pokemon in response.json()["results"] + ], + log_response=False, + ) + + get_pokemon_task = HttpOperator.partial( + task_id="get_pokemon", + http_conn_id="pokeapi", + method="GET", + ).iterate(endpoint=list_pokemon_task.output) + + list_pokemon_task >> get_pokemon_task + + +The scheduler only manages a single task. With a sync operator, iterations run in +a pool of up to ``task_concurrency`` threads, so blocking I/O such as these HTTP +requests overlaps up to that many at a time. Async operators scale further: a +coroutine waiting on I/O costs far less than a thread, so ``task_concurrency`` can +be set much higher. CPU-bound Python code speeds up with neither, since the +iterations share one process and its GIL. + +To **multiplex** many I/O-bound operations on one event loop, use an async task with +:class:`~airflow.providers.http.hooks.http.HttpAsyncHook`: + +.. code-block:: python + + from datetime import datetime + + from airflow.providers.http.hooks.http import HttpAsyncHook, HttpHook + from airflow.sdk import dag, task + + + @dag( + dag_id="it-async-http-pokemon-example", + start_date=datetime(2026, 1, 1), + ) + def it_async_http_pokemon_example(): + @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.iterate( + url=list_pokemon(), + ) + + + it_async_http_pokemon_example() + + +When ``iterate()`` is used with an async task, all iterations share the same +event loop, enabling true multiplexing of I/O-bound operations without any +manual concurrency management by the DAG author. For a handful of items the +difference is negligible, but for hundreds or thousands of items the +concurrent approach is dramatically faster — see the +:ref:`benchmarks above <sdk-mapped-tasks-vs-iterable-tasks>`. + +.. note:: + + ``multiple_outputs`` is ignored by ``iterate()``. Each iteration's return value is pushed + whole as ``return_value_<index>`` and the task's own return value is the lazy sequence over + them, so a ``dict`` return annotation on the task does not fan its keys out into separate + XComs the way it does with ``expand()``. Every key an iteration pushes or stores carries its + index the same way: ``ti.xcom_push("foo", v)`` in iteration 2 lands under ``foo_2``, and so + does ``task_state_store.set("foo", v)``, so iterations never overwrite each other's values. + This holds in the iteration's own thread or coroutine. A ``threading.Thread`` the task starts, + or ``loop.run_in_executor()``, does not inherit it: ``get_current_context()`` there returns + the task's own context, whose ``ti`` and ``task_state_store`` add no index, so the iterations' + keys overwrite each other. Use ``context["ti"]`` passed to the task, or start helpers with + ``asyncio.to_thread()`` or ``contextvars.copy_context().run()``, which carry the iteration's + context over. + +.. warning:: + + Inputs are shared between iterations. All iterations run in one process, so a value handed + to several of them is the same object in each: every value passed through ``partial()``, and + with ``iterate(a=..., b=...)`` every element of ``a`` and of ``b``, which the cross product + combines more than once. With ``expand()`` each task instance runs in its own process and + gets its own copy. Treat inputs as read-only, or copy what the task changes in place. + +Why Iterable Tasks? +--------------------------- + +IT is designed to address limitations of task mapping in specific scenarios: + +- **Scheduler scalability**: + Mapping creates one Task Instance per item, which can put pressure on the scheduler + for very large datasets. IT avoids this by keeping execution within a single task. + +- **Async multiplexing**: + With Python-native async support in Airflow 3.2, IT allows multiple + operations to share the same event loop within a single Task Instance. + This enables efficient multiplexing of I/O-bound workloads. + +- **Lower overhead**: + No need to serialize, schedule, and track thousands of Task Instances. + +- **Triggerer and deferrable-operator bottleneck**: + Deferrable operators delegate async work to triggerers, which store yielded + events directly in the Airflow metadata database. Unlike workers, triggerers + cannot leverage a custom XCom backend to offload large payloads. This makes + triggerers a bottleneck for sustained high-load async execution or workloads + that return large results. Mapping deferrable operators + amplifies the problem further. IT sidesteps triggerers entirely — iterations + execute on workers, which scale more effectively and support custom XCom + backends. Review Comment: Agreed, that was overselling it. 28e125fb11 adds a paragraph under the triggerer bullet and extends the "XCom backend" row: each item's result is also checkpointed in the `task_state_store` table for the store's retention unless a `[state_store] state_store_backend` is configured, in which case only a reference is stored, so large results need both a custom XCom backend and a state store backend. The class docstring says the same where it describes the checkpoints. Drafted-by: Claude Fable 5.1; reviewed by @dabla before posting ########## task-sdk/src/airflow/sdk/definitions/iterableoperator.py: ########## @@ -0,0 +1,1050 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import asyncio +import hashlib +import json +import os +import threading +import warnings +from collections.abc import AsyncIterable, AsyncIterator, Iterable, Mapping, Sequence +from functools import partial +from typing import TYPE_CHECKING, Any + +try: + # Python 3.11+ + BaseExceptionGroup +except NameError: + from exceptiongroup import BaseExceptionGroup + +from airflow.sdk import BaseXCom, TaskInstanceState, TriggerRule +from airflow.sdk.bases.operator import BaseAsyncOperator, BaseOperator, event_loop +from airflow.sdk.bases.skipmixin import SkipMixin +from airflow.sdk.bases.xcom import XComIterable +from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetAliasEvent, AssetUniqueKey +from airflow.sdk.definitions.retry_policy import RetryAction +from airflow.sdk.definitions.xcom_arg import XComArg +from airflow.sdk.exceptions import ( + AirflowFailException, + AirflowRescheduleException, + AirflowSensorTimeout, + AirflowSkipException, + AirflowTaskTerminated, + AirflowTaskTimeout, + DagRunTriggerException, + DownstreamTasksSkipped, + TaskDeferred, +) +from airflow.sdk.execution_time.comms import DeadlockImminentError +from airflow.sdk.execution_time.context import OutletEventAccessors, context_update_for_unmapped +from airflow.sdk.execution_time.executor import AsyncAwareExecutor +from airflow.sdk.execution_time.task_runner import ( + IndexedTaskInstance, + IndexedTaskRunner, + IndexedTaskState, + _push_xcom_if_needed, +) +from airflow.sdk.serde import serialize + +if TYPE_CHECKING: + import jinja2 + + from airflow.sdk.definitions._internal.expandinput import ExpandInput, Resolved + from airflow.sdk.definitions.context import Context + from airflow.sdk.definitions.mappedoperator import MappedOperator + from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance + from airflow.sdk.types import OutletEventAccessorsProtocol + + +# The trigger rules under which one skipped upstream task instance skips a task, whatever the other +# upstream task instances did (see TriggerRuleDep). A skipped iteration has the same effect on a +# downstream task with one of these rules as a skipped mapped task instance would. +SKIPPED_WITH_A_SKIPPED_UPSTREAM = frozenset( + {TriggerRule.ALL_SUCCESS, TriggerRule.NONE_SKIPPED, TriggerRule.ALL_DONE_MIN_ONE_SUCCESS} +) + + +# Raised by an item, these fail the task without a retry, as the runner does for a task that raises +# them itself (see _run_task_and_map_outcome and _handle_handler_failure). +FAIL_WITHOUT_RETRY = (AirflowFailException, AirflowSensorTimeout, AirflowTaskTerminated) + +# How strongly a retry policy decision speaks for the task when several items failed: one item the +# policy says must not be retried fails the task, one it says to retry makes it retry on its terms. +_DECISION_WEIGHT = {RetryAction.FAIL: 2, RetryAction.RETRY: 1, RetryAction.DEFAULT: 0} + + +def refuse_operators_that_skip_downstream(operator: MappedOperator) -> None: + """ + Refuse to iterate an operator that can skip downstream tasks. + + An iteration has no downstream tasks of its own, so ``ShortCircuitOperator``, the branch + operators and any other ``SkipMixin`` would skip nothing and let every downstream task run. + Checked on the class: ``MappedOperator._can_skip_downstream`` is only derived from ``SkipMixin`` + on the classic path, while the ``@task`` path copies a class default that is ``False`` even for + ``@task.short_circuit`` and ``@task.branch``. + """ + if issubclass(operator.operator_class, SkipMixin): + raise TypeError( + f"{operator.operator_name} can skip downstream tasks and cannot be iterated: an iteration " + f"of {operator.task_id!r} has no downstream tasks of its own, so it would skip nothing and " + "every downstream task would run. Use .expand() for it instead." + ) + + +def _unprefixed_task_id(operator: MappedOperator) -> str: + """ + Return the wrapped operator's task id without its task group's prefix. + + ``partial()`` already gave the wrapped operator the prefixed id, and ``BaseOperator.__init__`` + prefixes the id it gets once more, since an IterableOperator is not built from a mapped + operator. Handing it the bare id keeps the two equal. The same rule as ``label``, which cannot + be used here because it returns the display name when there is one. + """ + task_group = operator.task_group + if task_group and task_group.node_id and task_group.prefix_group_id: + return operator.task_id[len(task_group.node_id) + 1 :] + return operator.task_id + + +def _fingerprint(mapped_kwargs: Mapping[str, Any]) -> str | None: + """ + Digest one sub-task's input, stored on its checkpoint to tell whether the checkpoint still applies. + + A retry may run on another input than the attempt that wrote the checkpoints: the upstream was + cleared together with this task and produced other items. An index then no longer means the + same work, and replaying its result would hand downstream a value computed from the old item. + An input serde cannot serialize has no digest, and its checkpoint is honoured by index alone. + """ + try: + serialized = json.dumps(serialize(mapped_kwargs), sort_keys=True) + except (TypeError, ValueError, AttributeError, RecursionError): + return None + return hashlib.sha256(serialized.encode()).hexdigest() + + +def _partial_inputs_from_upstream( + partial_kwargs: Mapping[str, Any], unmapped_task: BaseOperator +) -> dict[str, Any]: + """ + Collect the rendered values of the partial kwargs an upstream task provides. + + They belong in the fingerprint next to the iterated kwargs: clearing the upstream together with + this task can change them while the items stay the same, and a checkpoint written with the old + value must not be replayed. Only XComArg values count, read back from the unmapped operator + once rendered, at the top level or inside a mapping such as a ``@task``'s ``op_kwargs``. Other + templated values are left out on purpose: one like ``{{ ti.try_number }}`` changes with every + attempt and would make every checkpoint look stale. + """ + inputs: dict[str, Any] = {} + for key, value in partial_kwargs.items(): + if isinstance(value, XComArg): + inputs[key] = getattr(unmapped_task, key, None) + elif isinstance(value, Mapping): + rendered = getattr(unmapped_task, key, None) + for name, nested in value.items(): + if isinstance(nested, XComArg): + inputs[f"{key}.{name}"] = rendered.get(name) if isinstance(rendered, Mapping) else None + return inputs + + +def _serialize_outlet_events(accessors: OutletEventAccessors) -> list[dict[str, Any]]: + """ + Snapshot the outlet asset events one sub-task recorded into a JSON-safe list. + + Persisted on the sub-task's checkpoint so a later attempt can replay them via + ``_replay_outlet_events`` when the sub-task is skipped because it already succeeded. + """ + events: list[dict[str, Any]] = [] + for _asset_or_alias, accessor in accessors.items(): + if isinstance(accessor.key, AssetUniqueKey): + events.append( + { + "kind": "asset", + "name": accessor.key.name, + "uri": accessor.key.uri, + "extra": accessor.extra, + "partition_keys": sorted(accessor.partition_keys), + } + ) + for alias_event in accessor.asset_alias_events: + events.append( + { + "kind": "asset_alias", + "source_alias_name": alias_event.source_alias_name, + "dest_asset_key": { + "name": alias_event.dest_asset_key.name, + "uri": alias_event.dest_asset_key.uri, + }, + "dest_asset_extra": alias_event.dest_asset_extra, + "extra": alias_event.extra, + } + ) + return events + + +def _merge_outlet_events(target: OutletEventAccessorsProtocol, source: OutletEventAccessors) -> None: + """ + Merge every outlet asset event recorded in ``source`` into ``target``. + + Used both to fold a sub-task's isolated accessor into the IterableOperator's shared + ``context["outlet_events"]`` right after it succeeds, and to replay a checkpointed + snapshot (via ``_replay_outlet_events``) for a sub-task skipped on retry. + """ + for asset_or_alias, accessor in source.items(): + target_accessor = target[asset_or_alias] + target_accessor.extra.update(accessor.extra) Review Comment: Right, the task instance can only send one event per asset, so this stays as it is and the docs say it now: c50026d6c9 adds an "Asset events" row to the comparison table (one event per mapped task instance vs one per asset for the whole task instance, `extra` from the last item to finish, partitions and alias events accumulated, a per-file `Metadata` pattern ported from `.expand()` produces one event) and a note on `_merge_outlet_events`. Drafted-by: Claude Fable 5.1; reviewed by @dabla before posting -- 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]
