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 the
XCOM_SEQUENCE_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.

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.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Author: dabla

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

@@ -38,6 +38,10 @@
 
 log = structlog.get_logger(logger_name=__name__)
 
+#: 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
+
 
 @attrs.define
 class LazyXComIterator(Iterator[T]):
@@ -87,9 +91,21 @@ 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 the ``XCOM_SEQUENCE_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. 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")
+    _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:
@@ -116,30 +132,47 @@ 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:
+            index = _normalize_index(index, await self.alen())
+        if not self._holds(index):
+            msg = await SUPERVISOR_COMMS.asend(self._slice_request(index, index + XCOM_SEQUENCE_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(
-                key=xcom_arg.key,
-                dag_id=source.dag_id,
-                task_id=source.task_id,
-                run_id=self._ti.run_id,
-                offset=index,
-            ),
+        return GetXComSequenceSlice(
+            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.deserialize_value(_XComWrapper(value)) for value in msg.root]
 
     def __len__(self) -> int:
         if self._len is None:
@@ -193,54 +226,30 @@ 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(
-                    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.deserialize_value(_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(
-                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.deserialize_value(_XComWrapper(msg.root))
+        index = _normalize_index(key, len(self)) if key < 0 else key
+        if not self._holds(index):
+            msg = SUPERVISOR_COMMS.send(self._slice_request(index, index + XCOM_SEQUENCE_CHUNK_SIZE))
+            self._hold(index, self._slice_values(msg))
+        return self._chunk[index - self._chunk_start]
+
+
+def _normalize_index(index: int, length: int) -> int:
+    """Map a negative index onto its position from the start, or raise ``IndexError`` past the front."""
+    if index < 0:
+        index += length
+    if index < 0:
+        raise IndexError(index)
+    return index
 
 
 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,
 )
 
@@ -679,17 +676,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: