dabla commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r4248810310
##########
task-sdk/src/airflow/sdk/bases/xcom.py:
##########
@@ -566,3 +567,166 @@ def delete(
),
)
cls.purge(xcom_result)
+
+
+def _normalize_index(index: int, length: int) -> int:
+ """Map a sequence index, negative ones included, onto a position in ``[0,
length)``."""
+ if index < 0:
+ index += length
+ if not (0 <= index < length):
+ raise IndexError(index)
+ return index
+
+
+class XComIterable(Sequence):
+ """
+ An iterable that lazily fetches XCom values one by one instead of loading
all at once.
+
+ This is a read-only :class:`collections.abc.Sequence` over the
``return_value_<index>`` XComs an
+ iterated task pushed, one per index: the values are written by the
producing task's runner as
+ each sub-task finishes (see ``IterableOperator.axcom_push``), and the
iterable only ever reads
+ them. Nothing on this class mutates the underlying XComs.
+
+ Indexing follows the usual sequence rules, negative indices included:
``result[-1]`` is the last
+ value. Iterations that were skipped pushed nothing and are left out, as
the XComs of skipped
+ mapped task instances are: ``length`` counts every input item, ``skipped``
lists the indices
+ that were skipped, and positions in the sequence run over the others only.
An iteration that
+ returned ``None`` pushed nothing either but keeps its position: reading it
gives ``None``,
+ where ``.expand()`` would not count such a mapped task instance at all.
+
+ Every element is a remote fetch, so random access costs one XCom read per
element, and so do
+ iterating and slicing: N elements are N requests. Reading several in one
request needs an
+ Execution API endpoint that takes a list of keys, which does not exist yet.
+ """
+
+ def __init__(
+ self,
+ task_id: str,
+ dag_id: str,
+ run_id: str,
+ map_index: int | None = None,
+ length: int | None = None,
+ skipped: Sequence[int] = (),
+ ):
+ self.task_id = task_id
+ self.dag_id = dag_id
+ self.run_id = run_id
+ self.map_index = map_index
+ self.length = length or 0
+ self.skipped: list[int] = sorted(skipped)
+
+ def _index_of(self, position: int) -> int:
+ """Map ``position`` in the sequence to its input index, stepping over
the skipped indices."""
+ index = _normalize_index(position, len(self))
+ for skipped_index in self.skipped:
+ if skipped_index > index:
+ break
+ index += 1
+ return index
+
+ def __iter__(self) -> Iterator[Any]:
+ return _XComIterator(self)
+
+ def __len__(self) -> int:
+ return self.length - len(self.skipped)
+
+ async def alen(self) -> int:
+ """Async twin of ``len(self)``, for readers on the event loop that
take a length before each read."""
+ return len(self)
+
+ def __getitem__(self, key: int | slice) -> Any | Sequence[Any]:
+ """Allow direct indexing so this works like a sequence."""
+ from airflow.sdk.execution_time.xcom import XCom
+
+ if isinstance(key, slice):
+ # TODO: This issues one XCom.get_one call per element — N
round-trips for a full slice.
+ # XComIterable stores results under distinct keys (return_value_0,
return_value_1, …)
+ # with the same map_index, so the existing GetXComSequenceSlice
endpoint (which ranges
+ # over map_index for a single key) cannot be reused. A new POST
endpoint that accepts
+ # a list of keys and returns values in a single query is needed;
once that lands, replace
+ # this loop with a single batched fetch.
+ start, stop, step = key.indices(len(self))
+ return [self[i] for i in range(start, stop, step)]
+
+ return XCom.get_one(
Review Comment:
Right, reproduced through `loop.previous`: the outer read went out with the
flag, the item reads without it. `XComIterable` now carries
`previous_iteration`, set by `LoopContextAccessor._result` when it reads an
iterable for the previous iteration, and passes it on in its item reads, sync
and async; only when set, so reads outside a loop call the XCom backend with
the same arguments as before. `aget_one` takes the same private
`_previous_iteration` as `get_one` for the async read. The new test in
`test_loop.py` checks the outer read and both item reads for the current and
the previous iteration (4c9ce91d92). A local read benchmark (old vs new, 3000
reads) showed no measurable difference.
Drafted-by: Claude Opus 5.5; 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]
