Optimisation: Allow prefetch of LazyXComSequence in next method of LazyXComIterator

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

@@ -19,12 +19,16 @@
 
 import collections
 import itertools
+from collections import deque
 from collections.abc import Iterator, Sequence
 from typing import TYPE_CHECKING, Any, Literal, TypeVar, overload
 
 import attrs
 import structlog
 
+from airflow.configuration import conf
+from airflow.sdk.definitions._internal.logging_mixin import LoggingMixin
+
 if TYPE_CHECKING:
     from airflow.sdk.definitions.xcom_arg import PlainXComArg
     from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
@@ -35,25 +39,40 @@
 # ``XCom.deserialize_value``. We don't want to wrap the API values in a nested
 # {"value": value} dict since it wastes bandwidth.
 _XComWrapper = collections.namedtuple("_XComWrapper", "value")
-
 log = structlog.get_logger(logger_name=__name__)
 
 
+def _deque_factory() -> deque:
+    return deque(maxlen=conf.getint("core", "parallelism"))
+
+
 @attrs.define
-class LazyXComIterator(Iterator[T]):
+class LazyXComIterator(LoggingMixin, Iterator[T]):
     seq: LazyXComSequence[T]
     index: int = 0
     dir: Literal[1, -1] = 1
+    _buffer: deque[T] = attrs.field(factory=_deque_factory, init=False)
+
+    @property
+    def prefetch_size(self) -> int:
+        return self._buffer.maxlen or conf.getint("core", "parallelism")
 
     def __next__(self) -> T:
         if self.index < 0:
             # When iterating backwards, avoid extra HTTP request
             raise StopIteration()
-        try:
-            val = self.seq[self.index]
-        except IndexError:
-            raise StopIteration from None
+
+        # If the buffer is empty, fetch the next chunk
+        if not self._buffer:
+            chunk = list(self.seq[self.index : self.index + self.prefetch_size])
+            if not chunk:
+                raise StopIteration()
+            self._buffer.extend(chunk)
+            self.log.debug("Buffered %s XCom's", len(self._buffer))
+
+        val = self._buffer.popleft()
         self.index += self.dir
+        self.log.debug("Popped buffered XCom for index %s: %s", self.index, val)
         return val
 
     def __iter__(self) -> Iterator[T]:

task-sdk/tests/task_sdk/execution_time/test_lazy_sequence.py

@@ -22,6 +22,7 @@
 import pytest
 
 import airflow
+from airflow.configuration import conf
 from airflow.sdk.bases.xcom import BaseXCom
 from airflow.sdk.exceptions import ErrorType
 from airflow.sdk.execution_time.comms import (
@@ -74,35 +75,75 @@ def test_len(mock_supervisor_comms, lazy_sequence):
 
 
 def test_iter(mock_supervisor_comms, lazy_sequence):
+    stop = conf.getint("core", "parallelism")
     it = iter(lazy_sequence)
 
     mock_supervisor_comms.send.side_effect = [
-        XComSequenceIndexResult(root="f"),
-        ErrorResponse(error=ErrorType.XCOM_NOT_FOUND, detail={"oops": "sorry!"}),
+        XComSequenceSliceResult(root=["f"]),
     ]
     assert list(it) == ["f"]
-    mock_supervisor_comms.send.assert_has_calls(
-        [
-            call(
-                msg=GetXComSequenceItem(
-                    key=BaseXCom.XCOM_RETURN_KEY,
-                    dag_id="dag",
-                    task_id="task",
-                    run_id="run",
-                    offset=0,
-                ),
-            ),
-            call(
-                msg=GetXComSequenceItem(
-                    key=BaseXCom.XCOM_RETURN_KEY,
-                    dag_id="dag",
-                    task_id="task",
-                    run_id="run",
-                    offset=1,
-                ),
-            ),
-        ]
-    )
+
+    assert mock_supervisor_comms.send.call_args_list == [
+        call(
+            GetXComSequenceSlice(
+                key=BaseXCom.XCOM_RETURN_KEY,
+                dag_id="dag",
+                task_id="task",
+                run_id="run",
+                start=0,
+                stop=stop,
+                step=None,
+            )
+        ),
+        call(
+            GetXComSequenceSlice(
+                key=BaseXCom.XCOM_RETURN_KEY,
+                dag_id="dag",
+                task_id="task",
+                run_id="run",
+                start=1,
+                stop=stop + 1,
+                step=None,
+            )
+        ),
+    ]
+
+
+def test_iter_when_xcom_not_found(mock_supervisor_comms, lazy_sequence):
+    stop = conf.getint("core", "parallelism")
+    it = iter(lazy_sequence)
+
+    mock_supervisor_comms.send.side_effect = [
+        XComSequenceSliceResult(root=["f"]),
+        ErrorResponse(error=ErrorType.XCOM_NOT_FOUND, detail={"oops": "sorry!"}),
+    ]
+    with pytest.raises(TypeError):
+        list(it)
+
+    assert mock_supervisor_comms.send.call_args_list == [
+        call(
+            GetXComSequenceSlice(
+                key=BaseXCom.XCOM_RETURN_KEY,
+                dag_id="dag",
+                task_id="task",
+                run_id="run",
+                start=0,
+                stop=stop,
+                step=None,
+            )
+        ),
+        call(
+            GetXComSequenceSlice(
+                key=BaseXCom.XCOM_RETURN_KEY,
+                dag_id="dag",
+                task_id="task",
+                run_id="run",
+                start=1,
+                stop=stop + 1,
+                step=None,
+            )
+        ),
+    ]
 
 
 def test_getitem_index(mock_supervisor_comms, lazy_sequence):