Read a mapped upstream's XCom sequence in chunks instead of one item per request

LazyXComSequence fetched one item per GetXComSequenceItem request: iterating a
mapped upstream's results, or an iterated task reading them by index, cost one
supervisor round trip per item, which is the dominant cost for large inputs.

Reads now come in chunks. An index outside the slice held fetches chunk_size
items starting there with one GetXComSequenceSlice request, the message the
slice path already used, and consecutive reads are served from it. Sequential
access costs one request per chunk instead of one per item, both through
__getitem__/__iter__ (send) and aget/__aiter__ (asend), and at most one chunk
is held. A jump backwards fetches again from there.

The chunk size is the new [core] xcom_sequence_chunk_size option, default 32,
documented in config.yml and in the iteration docs; a LazyXComSequence can
also be given one explicitly.

A negative index no longer needs its own per-item request: it is mapped from the
end through the cached length, one GetXComCount at most. The client no longer
sends GetXComSequenceItem at all; the request handler keeps serving it.

Author: dabla

airflow-core/src/airflow/config_templates/config.yml

@@ -493,6 +493,15 @@ core:
       type: integer
       example: ~
       default: "1024"
+    xcom_sequence_chunk_size:
+      description: |
+        Number of items fetched per request when a mapped task's XCom results are read as a sequence,
+        for example by a downstream ``.iterate()`` or a ``for`` loop over them. Larger values mean fewer
+        round trips to the API server; each fetched chunk is held in memory until the reader moves past it.
+      version_added: 3.4.0
+      type: integer
+      example: ~
+      default: "32"
 
     allow_double_dot_in_ids:
       description: |

task-sdk/docs/mapped-tasks-vs-iterable-tasks.rst

@@ -319,6 +319,12 @@ concurrent approach is dramatically faster — see the
     ``asyncio.to_thread()`` or ``contextvars.copy_context().run()``, which carry the iteration's
     context over.
 
+.. note::
+
+    When the input of ``iterate()`` is the output of a mapped task, its items are fetched from the
+    API server in chunks of ``[core] xcom_sequence_chunk_size`` items (32 by default), one request
+    per chunk, with one chunk held in memory at a time.
+
 Why Iterable Tasks?
 ---------------------------
 

task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py

@@ -25,6 +25,8 @@
 import attrs
 import structlog
 
+from airflow.sdk.configuration import conf
+
 if TYPE_CHECKING:
     from airflow.sdk.definitions.xcom_arg import PlainXComArg
     from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
@@ -38,6 +40,15 @@
 
 log = structlog.get_logger(logger_name=__name__)
 
+#: Default for ``[core] xcom_sequence_chunk_size``: items a ``LazyXComSequence`` fetches per request
+#: when read by index or iterated. One slice request serves this many consecutive reads, and at
+#: most this many values are held at a time.
+XCOM_SEQUENCE_CHUNK_SIZE = 32
+
+
+def _configured_chunk_size() -> int:
+    return conf.getint("core", "xcom_sequence_chunk_size", fallback=XCOM_SEQUENCE_CHUNK_SIZE)
+
 
 @attrs.define
 class LazyXComIterator(Iterator[T]):
@@ -87,10 +98,23 @@ async def __anext__(self) -> T:
 
 @attrs.define
 class LazyXComSequence(Sequence[T]):
+    """
+    A mapped upstream's return values, read from the API server on demand.
+
+    Reads come in chunks: an index outside the slice held fetches ``chunk_size`` items starting at
+    that index with one ``GetXComSequenceSlice`` request, so iterating, or an iterated task reading
+    this by index, costs one request per chunk instead of one per item, while no more than one
+    chunk is held in memory. ``chunk_size`` defaults to ``[core] xcom_sequence_chunk_size``. A
+    negative index counts from the end through the cached length.
+    """
+
     _len: int | None = attrs.field(init=False, default=None)
     _xcom_arg: PlainXComArg = attrs.field(alias="xcom_arg")
     _ti: RuntimeTaskInstance = attrs.field(alias="ti")
     _previous_iteration: bool = attrs.field(default=False, alias="previous_iteration", kw_only=True)
+    chunk_size: int = attrs.field(kw_only=True, factory=_configured_chunk_size)
+    _chunk_start: int = attrs.field(init=False, default=0)
+    _chunk: list[T] = attrs.field(init=False, factory=list)
 
     def __repr__(self) -> str:
         if self._len is not None:
@@ -117,31 +141,50 @@ def __aiter__(self) -> AsyncIterator[T]:
         return AsyncLazyXComIterator(seq=self)
 
     async def aget(self, index: int) -> T:
-        """Async counterpart of ``self[index]``; the same ``GetXComSequenceItem`` request via ``asend``."""
-        from airflow.sdk.execution_time.comms import (
-            ErrorResponse,
-            GetXComSequenceItem,
-            XComSequenceIndexResult,
-        )
+        """Async counterpart of ``self[index]``: the same chunked reads, sent with ``asend``."""
         from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
-        from airflow.sdk.execution_time.xcom import XCom
+
+        if index < 0:
+            # range() maps a negative index from the end and raises IndexError past the front;
+            # only a negative one needs the length, so a forward read sends no count request.
+            index = range(await self.alen())[index]
+        if not self._holds(index):
+            msg = await SUPERVISOR_COMMS.asend(self._slice_request(index, index + self.chunk_size))
+            self._hold(index, self._slice_values(msg))
+        return self._chunk[index - self._chunk_start]
+
+    def _holds(self, index: int) -> bool:
+        return self._chunk_start <= index < self._chunk_start + len(self._chunk)
+
+    def _hold(self, index: int, values: list[T]) -> None:
+        """Keep the slice fetched for ``index``; an empty slice means ``index`` is past the end."""
+        if not values:
+            raise IndexError(index)
+        self._chunk_start, self._chunk = index, values
+
+    def _slice_request(self, start: int | None, stop: int | None, step: int | None = None) -> Any:
+        from airflow.sdk.execution_time.comms import GetXComSequenceSlice
 
         source = (xcom_arg := self._xcom_arg).operator
-        msg = await SUPERVISOR_COMMS.asend(
-            GetXComSequenceItem(
-                previous_iteration=self._previous_iteration,
-                key=xcom_arg.key,
-                dag_id=source.dag_id,
-                task_id=source.task_id,
-                run_id=self._ti.run_id,
-                offset=index,
-            ),
+        return GetXComSequenceSlice(
+            previous_iteration=self._previous_iteration,
+            key=xcom_arg.key,
+            dag_id=source.dag_id,
+            task_id=source.task_id,
+            run_id=self._ti.run_id,
+            start=start,
+            stop=stop,
+            step=step,
         )
-        if isinstance(msg, ErrorResponse):
-            raise IndexError(index)
-        if not isinstance(msg, XComSequenceIndexResult):
-            raise TypeError(f"Got unexpected response to GetXComSequenceItem: {msg!r}")
-        return XCom.deserialize_value(_XComWrapper(msg.root))
+
+    @staticmethod
+    def _slice_values(msg: Any) -> list[T]:
+        from airflow.sdk.execution_time.comms import XComSequenceSliceResult
+        from airflow.sdk.execution_time.xcom import XCom
+
+        if not isinstance(msg, XComSequenceSliceResult):
+            raise TypeError(f"Got unexpected response to GetXComSequenceSlice: {msg!r}")
+        return [XCom._load(_XComWrapper(value)) for value in msg.root]
 
     def __len__(self) -> int:
         if self._len is None:
@@ -197,56 +240,21 @@ def __getitem__(self, key: int) -> T: ...
     def __getitem__(self, key: slice) -> Sequence[T]: ...
 
     def __getitem__(self, key: int | slice) -> T | Sequence[T]:
-        from airflow.sdk.execution_time.comms import (
-            ErrorResponse,
-            GetXComSequenceItem,
-            GetXComSequenceSlice,
-            XComSequenceIndexResult,
-            XComSequenceSliceResult,
-        )
         from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
-        from airflow.sdk.execution_time.xcom import XCom
 
         if isinstance(key, slice):
-            start, stop, step = _coerce_slice(key)
-            source = (xcom_arg := self._xcom_arg).operator
-            msg = SUPERVISOR_COMMS.send(
-                GetXComSequenceSlice(
-                    previous_iteration=self._previous_iteration,
-                    key=xcom_arg.key,
-                    dag_id=source.dag_id,
-                    task_id=source.task_id,
-                    run_id=self._ti.run_id,
-                    start=start,
-                    stop=stop,
-                    step=step,
-                ),
-            )
-            if not isinstance(msg, XComSequenceSliceResult):
-                raise TypeError(f"Got unexpected response to GetXComSequenceSlice: {msg!r}")
-            return [XCom._load(_XComWrapper(value)) for value in msg.root]
+            return self._slice_values(SUPERVISOR_COMMS.send(self._slice_request(*_coerce_slice(key))))
 
         if not isinstance(key, int):
             if (index := getattr(key, "__index__", None)) is not None:
                 key = index()
             raise TypeError(f"Sequence indices must be integers or slices not {type(key).__name__}")
 
-        source = (xcom_arg := self._xcom_arg).operator
-        msg = SUPERVISOR_COMMS.send(
-            GetXComSequenceItem(
-                previous_iteration=self._previous_iteration,
-                key=xcom_arg.key,
-                dag_id=source.dag_id,
-                task_id=source.task_id,
-                run_id=self._ti.run_id,
-                offset=key,
-            ),
-        )
-        if isinstance(msg, ErrorResponse):
-            raise IndexError(key)
-        if not isinstance(msg, XComSequenceIndexResult):
-            raise TypeError(f"Got unexpected response to GetXComSequenceItem: {msg!r}")
-        return XCom._load(_XComWrapper(msg.root))
+        index = range(len(self))[key] if key < 0 else key
+        if not self._holds(index):
+            msg = SUPERVISOR_COMMS.send(self._slice_request(index, index + self.chunk_size))
+            self._hold(index, self._slice_values(msg))
+        return self._chunk[index - self._chunk_start]
 
 
 def _coerce_slice_index(value: Any) -> int | None:

task-sdk/tests/task_sdk/definitions/test_mappedoperator.py

@@ -32,15 +32,12 @@
 from airflow.sdk.definitions.mappedoperator import MappedOperator
 from airflow.sdk.definitions.xcom_arg import XComArg
 from airflow.sdk.execution_time.comms import (
-    ErrorResponse,
     GetTICount,
     GetXCom,
-    GetXComSequenceItem,
     GetXComSequenceSlice,
     SetXCom,
     TICount,
     XComResult,
-    XComSequenceIndexResult,
     XComSequenceSliceResult,
 )
 
@@ -695,17 +692,10 @@ def mock_comms_response(msg):
                 value = [v for k, v in expected_values.items() if k[0] == msg.task_id]
                 return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=value)
         elif isinstance(msg, GetXComSequenceSlice):
-            # Handle sequence slicing for pulling all XCom values from mapped tasks
+            # The aggregated task-group value resolves lazily and is read in chunks: honour the bounds.
             task_id = msg.task_id
             values = [v for k, v in expected_values.items() if k[0] == task_id and k[1] is not None]
-            return XComSequenceSliceResult(root=values)
-        elif isinstance(msg, GetXComSequenceItem):
-            # The aggregated task-group value now resolves lazily, so iterating it fetches one item at a time.
-            task_id = msg.task_id
-            values = [v for k, v in expected_values.items() if k[0] == task_id and k[1] is not None]
-            if 0 <= msg.offset < len(values):
-                return XComSequenceIndexResult(root=values[msg.offset])
-            return ErrorResponse()
+            return XComSequenceSliceResult(root=values[msg.start : msg.stop : msg.step])
         elif isinstance(msg, GetTICount):
             # Handle TI count queries for upstream_map_indexes computation
             if msg.task_ids: