dabla commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r4252539275
##########
task-sdk/src/airflow/sdk/bases/xcom.py:
##########
@@ -566,3 +569,174 @@ def delete(
),
)
cls.purge(xcom_result)
+
+
+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.
+
+ Inside a loop, the result an iterated terminal pushed in the previous
iteration
+ (``loop.previous``) has ``previous_iteration`` set, so its elements are
read from that
+ iteration too and not from the current one.
+ """
+
+ def __init__(
+ self,
+ task_id: str,
+ dag_id: str,
+ run_id: str,
+ map_index: int | None = None,
+ length: int | None = None,
+ skipped: Sequence[int] = (),
+ previous_iteration: bool = False,
+ ):
+ 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)
+ self.previous_iteration = previous_iteration
+
+ def _index_of(self, position: int) -> int:
+ """Map ``position`` in the sequence to its input index, stepping over
the skipped indices."""
+ # range() maps a negative position from the end and raises IndexError
past either end.
+ index = range(len(self))[position]
+ 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(**self._item_query(key))
+
+ async def aget(self, index: int) -> Any:
+ """
+ Async counterpart of ``self[index]``: fetch one value through
``XCom.aget_one``.
+
+ Use it, or ``async for``, from code running on an event loop that has
other SDK calls in
+ flight (an iterated task consuming this iterable as its input): a
synchronous read there
+ would block the loop thread on the supervisor channel and deadlock
with them.
+ """
+ from airflow.sdk.execution_time.xcom import XCom
+
+ return await XCom.aget_one(**self._item_query(index))
+
+ def __aiter__(self) -> AsyncIterator[Any]:
+ return _AsyncXComIterator(self)
+
+ def _item_query(self, position: int) -> dict[str, Any]:
+ """
+ Build the arguments of the XCom read for the item at ``position``.
+
+ The previous-iteration flag is passed only when set, so a read outside
a loop calls the
+ XCom backend's ``get_one``/``aget_one`` with the same arguments as
before.
+ """
+ query: dict[str, Any] = {
+ "key": f"{BaseXCom.XCOM_RETURN_KEY}_{self._index_of(position)}",
+ "dag_id": self.dag_id,
+ "task_id": self.task_id,
+ "run_id": self.run_id,
+ "map_index": self.map_index,
+ }
+ if self.previous_iteration:
+ query["_previous_iteration"] = True
+ return query
+
+ def serialize(self) -> dict:
+ """Ensure the object is JSON serializable."""
+ data = {
+ "task_id": self.task_id,
+ "dag_id": self.dag_id,
+ "run_id": self.run_id,
+ "map_index": self.map_index,
+ "length": self.length,
+ "skipped": self.skipped,
+ }
+ if self.previous_iteration:
+ # Only when set, so a value pushed outside a loop keeps its
earlier form.
+ data["previous_iteration"] = True
Review Comment:
Good question: it is not supported, and it now says so instead of reading
the wrong items (d6d5844266). "Previous" is resolved against the reader's
iteration on the server, so a `loop.previous` result pushed on unchanged points
every later reader, in the next iteration or after the loop, at another
iteration's items. `XComIterable.serialize()` now refuses an iterable with
`previous_iteration` set with a `TypeError` that says to push the values
instead (`list(loop.previous)`), and the flag is no longer written into the
pushed form. serde goes through it for both an XCom push and an item's
checkpoint, so both refuse it the same way; reading through `loop.previous` is
unchanged.
I looked at the two alternatives. Materializing at push, as the
`LazyXComSequence` serializer does for a mapped terminal, would store a
previous result of any size as one XCom, which `XComIterable` is there to
avoid. Pinning the iteration, so the pointer stays lazy and right, needs an
XCom read by absolute iteration in the Execution API; that is an API change I
would rather put in a stacked PR of its own than in this one.
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]
