dabla commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r4342062343
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -955,6 +990,633 @@ def mark_success_url(self) -> str:
return self.log_url
+class CheckpointDeserializationError(Exception):
+ """An indexed task's checkpoint holds a value serde refuses to restore,
such as a class off the allow-list."""
+
+
+@dataclass
+class IndexedTaskState:
+ status: TaskInstanceState
+ result: Any | None = None
+ # Outlet asset events the indexed task recorded on a previous successful
attempt. Outlet events
+ # only reach the server on the parent's success payload
(_handle_current_task_success), so an
+ # attempt that fails never registers what its succeeded indexed tasks
emitted. An indexed task skipped
+ # on retry (because it already succeeded) never re-executes, so it never
re-emits into the
+ # fresh OutletEventAccessors created for the new attempt; persisting a
snapshot here lets
+ # IndexedTaskLifecycle replay it, which is the only way those events
survive, and it
+ # cannot double-emit because the failed attempt sent nothing.
+ outlet_events: list[dict[str, Any]] | None = None
+ # Digest of the input the indexed task ran with. An index only means the
same work while its input
+ # is the same, so a checkpoint is honoured on a later attempt only when
this still matches:
+ # after the upstream was cleared and produced other values, the indexed
task runs again.
+ fingerprint: str | None = None
+ # The attempt that wrote the checkpoint. After a manual clear only the
checkpoints written since
+ # are resumed from (see Checkpoints), which this tells apart from those
left by the run before.
+ try_number: int = 0
+ # The other XComs the indexed task pushed, by their key before the index
suffix. The runner deletes
+ # every XCom before an attempt, so an indexed task skipped on retry has
them pushed again from here,
+ # next to its result, as a mapped task instance that succeeded keeps its
own.
+ xcoms: dict[str, Any] | None = None
+
+ def record_outlet_events(self, accessors: OutletEventAccessors) -> None:
+ """
+ Keep a JSON-safe snapshot of the outlet asset events the indexed task
recorded.
+
+ Persisted with the checkpoint so a later attempt can
:meth:`replay_outlet_events` when
+ the indexed task is skipped because it already succeeded.
+ """
+ events: list[dict[str, Any]] = []
+ for _asset_or_alias, accessor in accessors.items():
+ if isinstance(accessor.key, AssetUniqueKey):
+ events.append(
+ {
+ "kind": "asset",
+ "name": accessor.key.name,
+ "uri": accessor.key.uri,
+ "extra": accessor.extra,
+ "partition_keys": sorted(accessor.partition_keys),
+ }
+ )
+ for alias_event in accessor.asset_alias_events:
+ events.append(
+ {
+ "kind": "asset_alias",
+ "source_alias_name": alias_event.source_alias_name,
+ "dest_asset_key": {
+ "name": alias_event.dest_asset_key.name,
+ "uri": alias_event.dest_asset_key.uri,
+ },
+ "dest_asset_extra": alias_event.dest_asset_extra,
+ "extra": alias_event.extra,
+ }
+ )
+ if events:
+ self.outlet_events = events
+
+ def replay_outlet_events(self, target: OutletEventAccessorsProtocol) ->
None:
+ """
+ Re-emit into ``target`` the events the indexed task recorded on a
previous attempt.
+
+ An indexed task skipped on retry (because it already succeeded) never
re-executes, so it
+ never re-emits into the fresh ``OutletEventAccessors`` created for the
new attempt. The
+ failed attempt sent nothing to the server either (outlet events travel
only on the success
+ payload), so replaying cannot emit an event twice; without it the
events would be lost.
+ They land in the parent's accessors as
:meth:`IndexedTaskRunner.merge_outlet_events_into`
+ lands the live ones: one event per asset per task instance.
+ """
+ for event in self.outlet_events or ():
+ if event["kind"] == "asset":
+ accessor = target[Asset(name=event["name"], uri=event["uri"])]
+ accessor.extra.update(event["extra"])
+ if event["partition_keys"]:
+ accessor.add_partitions(event["partition_keys"])
+ else:
+
target[AssetAlias(name=event["source_alias_name"])].asset_alias_events.append(
+ AssetAliasEvent(
+ source_alias_name=event["source_alias_name"],
+
dest_asset_key=AssetUniqueKey(**event["dest_asset_key"]),
+ dest_asset_extra=event["dest_asset_extra"],
+ extra=event["extra"],
+ )
+ )
+
+ @staticmethod
+ def build_key(index: int) -> str:
+ # The task state store is already scoped to the parent task instance
(dag, run, task and
+ # map index), so the key carries no identity, only a namespace that
keeps the operator's own
+ # entries apart from anything user code stores from inside an indexed
task.
+ return f"_iterable_{index}"
+
+ def serialize(self) -> dict[str, Any]:
+ # The checkpoint travels to the supervisor as a JsonValue, which
rejects anything that is not
+ # plain JSON (tuples, datetimes, models, ...). Serde turns those into
JSON-compatible
+ # structures and restores them on read, exactly as XCom does with the
same result.
+ data: dict[str, Any] = {"status": self.status.value}
+ if self.result is not None:
+ data["result"] = serde_serialize(self.result)
+ if self.outlet_events:
+ data["outlet_events"] = self.outlet_events
+ if self.fingerprint:
+ data["fingerprint"] = self.fingerprint
+ if self.try_number:
+ data["try_number"] = self.try_number
+ if self.xcoms:
+ data["xcoms"] = {key: serde_serialize(value) for key, value in
self.xcoms.items()}
+ return data
+
+ @classmethod
+ def deserialize(cls, raw: Any) -> IndexedTaskState | None:
+ """
+ Restore a checkpoint, or None when there is none.
+
+ A result or XCom value serde fails to restore raises
:class:`CheckpointDeserializationError`:
+ a class off the allow-list (``ImportError``), one it has no
deserializer for or whose
+ constructor refuses the data (``TypeError``), a pydantic model that
changed since the
+ checkpoint was written (``ValidationError``, a ``ValueError``), or
whatever else a class
+ raises. Refused once, it is refused on every retry. ``raw`` is already
read, so nothing
+ transient can fail here, and every failure is taken as a refusal.
+ """
+ if not isinstance(raw, Mapping):
+ return None
+ try:
+ result = serde_deserialize(raw.get("result"))
+ xcoms = (
+ {key: serde_deserialize(value) for key, value in
raw["xcoms"].items()}
+ if raw.get("xcoms")
+ else None
+ )
+ except Exception as e:
+ raise CheckpointDeserializationError(str(e)) from e
+ return cls(
+ status=TaskInstanceState(raw["status"]),
+ result=result,
+ outlet_events=raw.get("outlet_events"),
+ fingerprint=raw.get("fingerprint"),
+ try_number=raw.get("try_number", 0),
+ xcoms=xcoms,
+ )
+
+
+class IndexedTaskInstance(RuntimeTaskInstance):
+ """
+ Indexed task instance to run a mapped operator.
+
+ It shares the parent task instance's identity, so what an iteration pushes
or stores lands in
+ the parent's scope, suffixed with the index so that iterations never
overwrite each other:
+ XComs through :meth:`xcom_push`, task state through
:attr:`task_state_store`. Reading back
+ follows the same rule: :meth:`xcom_pull` of the iteration's own XComs adds
the index, a pull
+ from another task does not, as the store's accessor adds it to every key
of its own. The
+ operator's own checkpoints go to the parent's store unsuffixed, under their
+ ``_iterable_<index>`` keys.
+ """
+
+ index: int
+ parent_task_state_store: TaskStateStoreAccessor
+ input_fingerprint: str | None = None
+ # The XComs the indexed task pushed other than its return value, by their
key before the index
+ # suffix: kept in memory while it runs and written once, with its SUCCESS
checkpoint.
+ pushed_xcoms: dict[str, Any] = Field(default_factory=dict)
+ # The params processed for this index's operator when they differ from the
iterated task's:
+ # the index iterates ``params``, or the operator set its own. None when
the parent's apply.
+ own_params: dict[str, Any] | None = None
+
+ @classmethod
+ def create_indexed_task(
+ cls, *, context: Context, index: int, operator: BaseOperator,
input_fingerprint: str | None = None
+ ) -> IndexedTaskInstance:
+ """
+ Create the runtime instance for one index of an iterated task from the
parent's context.
+
+ The instance shares the parent task instance's identity (id, run, map
index, try number and
+ retry budget), so XComs and task state land in the parent's scope, and
carries the unmapped
+ operator for that index. The budget is the parent's ``max_tries``
rather than the operator's
+ ``retries``: a manual clear raises it, and it is the parent Airflow
retries. The parent's
+ state store comes from the context, the same accessor the operator's
+ checkpoints use. ``model_construct`` skips Pydantic validation on
purpose: one instance is built per
+ indexed task, and the parent was validated already, so only the index
needs checking here.
+ """
+ if index < 0:
+ raise ValueError(f"IndexedTaskInstance requires index >= 0, got
{index}")
+ parent = context["ti"]
+ return cls.model_construct(
+ id=parent.id,
+ parent_task_state_store=context["task_state_store"],
+ # The parent's: an iteration's XComs and state belong to the task
instance that runs it.
+ task_id=parent.task_id,
+ dag_id=operator.dag_id,
+ run_id=parent.run_id,
+ map_index=parent.map_index,
+ region_id=getattr(parent, "region_id", None),
+ region_index=getattr(parent, "region_index", None),
+ index=index,
+ input_fingerprint=input_fingerprint,
+ max_tries=parent.max_tries,
+ start_date=parent.start_date,
+ state=TaskInstanceState.SCHEDULED.value,
+ is_mapped=True,
+ task=operator,
+ try_number=parent.try_number,
+ # The dag run, logical date and the rest the server sent for the
task instance that
+ # runs this iteration: its template context, get_previous_ti() and
the like read them.
+ _ti_context_from_server=getattr(parent, "_ti_context_from_server",
None),
Review Comment:
Fixed in
[`288b248fc1`](https://github.com/apache/airflow/pull/62922/commits/288b248fc13d5a36f878ce6f39b4f91dad3dc789).
Rather than adding the five fields to the hand-picked list, which would miss
the next field added to the task instance the same way, `create_indexed_task`
now takes every field of `RuntimeTaskInstance` the parent holds and sets only
what an iteration has of its own (`task`, `dag_id`, `state`, `is_mapped`,
`end_date`, `rendered_map_index`, the index and its fingerprint). The test
walks `RuntimeTaskInstance.model_fields`, so a field added later is checked
without editing it. Private attributes stay explicit: only
`_ti_context_from_server` is shared, a cached template context must not be.
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -955,6 +990,633 @@ def mark_success_url(self) -> str:
return self.log_url
+class CheckpointDeserializationError(Exception):
+ """An indexed task's checkpoint holds a value serde refuses to restore,
such as a class off the allow-list."""
+
+
+@dataclass
+class IndexedTaskState:
+ status: TaskInstanceState
+ result: Any | None = None
+ # Outlet asset events the indexed task recorded on a previous successful
attempt. Outlet events
+ # only reach the server on the parent's success payload
(_handle_current_task_success), so an
+ # attempt that fails never registers what its succeeded indexed tasks
emitted. An indexed task skipped
+ # on retry (because it already succeeded) never re-executes, so it never
re-emits into the
+ # fresh OutletEventAccessors created for the new attempt; persisting a
snapshot here lets
+ # IndexedTaskLifecycle replay it, which is the only way those events
survive, and it
+ # cannot double-emit because the failed attempt sent nothing.
+ outlet_events: list[dict[str, Any]] | None = None
+ # Digest of the input the indexed task ran with. An index only means the
same work while its input
+ # is the same, so a checkpoint is honoured on a later attempt only when
this still matches:
+ # after the upstream was cleared and produced other values, the indexed
task runs again.
+ fingerprint: str | None = None
+ # The attempt that wrote the checkpoint. After a manual clear only the
checkpoints written since
+ # are resumed from (see Checkpoints), which this tells apart from those
left by the run before.
+ try_number: int = 0
+ # The other XComs the indexed task pushed, by their key before the index
suffix. The runner deletes
+ # every XCom before an attempt, so an indexed task skipped on retry has
them pushed again from here,
+ # next to its result, as a mapped task instance that succeeded keeps its
own.
+ xcoms: dict[str, Any] | None = None
+
+ def record_outlet_events(self, accessors: OutletEventAccessors) -> None:
+ """
+ Keep a JSON-safe snapshot of the outlet asset events the indexed task
recorded.
+
+ Persisted with the checkpoint so a later attempt can
:meth:`replay_outlet_events` when
+ the indexed task is skipped because it already succeeded.
+ """
+ events: list[dict[str, Any]] = []
+ for _asset_or_alias, accessor in accessors.items():
+ if isinstance(accessor.key, AssetUniqueKey):
+ events.append(
+ {
+ "kind": "asset",
+ "name": accessor.key.name,
+ "uri": accessor.key.uri,
+ "extra": accessor.extra,
+ "partition_keys": sorted(accessor.partition_keys),
+ }
+ )
+ for alias_event in accessor.asset_alias_events:
+ events.append(
+ {
+ "kind": "asset_alias",
+ "source_alias_name": alias_event.source_alias_name,
+ "dest_asset_key": {
+ "name": alias_event.dest_asset_key.name,
+ "uri": alias_event.dest_asset_key.uri,
+ },
+ "dest_asset_extra": alias_event.dest_asset_extra,
+ "extra": alias_event.extra,
+ }
+ )
+ if events:
+ self.outlet_events = events
+
+ def replay_outlet_events(self, target: OutletEventAccessorsProtocol) ->
None:
+ """
+ Re-emit into ``target`` the events the indexed task recorded on a
previous attempt.
+
+ An indexed task skipped on retry (because it already succeeded) never
re-executes, so it
+ never re-emits into the fresh ``OutletEventAccessors`` created for the
new attempt. The
+ failed attempt sent nothing to the server either (outlet events travel
only on the success
+ payload), so replaying cannot emit an event twice; without it the
events would be lost.
+ They land in the parent's accessors as
:meth:`IndexedTaskRunner.merge_outlet_events_into`
+ lands the live ones: one event per asset per task instance.
+ """
+ for event in self.outlet_events or ():
+ if event["kind"] == "asset":
+ accessor = target[Asset(name=event["name"], uri=event["uri"])]
+ accessor.extra.update(event["extra"])
+ if event["partition_keys"]:
+ accessor.add_partitions(event["partition_keys"])
+ else:
+
target[AssetAlias(name=event["source_alias_name"])].asset_alias_events.append(
+ AssetAliasEvent(
+ source_alias_name=event["source_alias_name"],
+
dest_asset_key=AssetUniqueKey(**event["dest_asset_key"]),
+ dest_asset_extra=event["dest_asset_extra"],
+ extra=event["extra"],
+ )
+ )
+
+ @staticmethod
+ def build_key(index: int) -> str:
+ # The task state store is already scoped to the parent task instance
(dag, run, task and
+ # map index), so the key carries no identity, only a namespace that
keeps the operator's own
+ # entries apart from anything user code stores from inside an indexed
task.
+ return f"_iterable_{index}"
+
+ def serialize(self) -> dict[str, Any]:
+ # The checkpoint travels to the supervisor as a JsonValue, which
rejects anything that is not
+ # plain JSON (tuples, datetimes, models, ...). Serde turns those into
JSON-compatible
+ # structures and restores them on read, exactly as XCom does with the
same result.
+ data: dict[str, Any] = {"status": self.status.value}
+ if self.result is not None:
+ data["result"] = serde_serialize(self.result)
+ if self.outlet_events:
+ data["outlet_events"] = self.outlet_events
+ if self.fingerprint:
+ data["fingerprint"] = self.fingerprint
+ if self.try_number:
+ data["try_number"] = self.try_number
+ if self.xcoms:
+ data["xcoms"] = {key: serde_serialize(value) for key, value in
self.xcoms.items()}
+ return data
+
+ @classmethod
+ def deserialize(cls, raw: Any) -> IndexedTaskState | None:
+ """
+ Restore a checkpoint, or None when there is none.
+
+ A result or XCom value serde fails to restore raises
:class:`CheckpointDeserializationError`:
+ a class off the allow-list (``ImportError``), one it has no
deserializer for or whose
+ constructor refuses the data (``TypeError``), a pydantic model that
changed since the
+ checkpoint was written (``ValidationError``, a ``ValueError``), or
whatever else a class
+ raises. Refused once, it is refused on every retry. ``raw`` is already
read, so nothing
+ transient can fail here, and every failure is taken as a refusal.
+ """
+ if not isinstance(raw, Mapping):
+ return None
+ try:
+ result = serde_deserialize(raw.get("result"))
+ xcoms = (
+ {key: serde_deserialize(value) for key, value in
raw["xcoms"].items()}
+ if raw.get("xcoms")
+ else None
+ )
+ except Exception as e:
+ raise CheckpointDeserializationError(str(e)) from e
+ return cls(
+ status=TaskInstanceState(raw["status"]),
+ result=result,
+ outlet_events=raw.get("outlet_events"),
+ fingerprint=raw.get("fingerprint"),
+ try_number=raw.get("try_number", 0),
+ xcoms=xcoms,
+ )
+
+
+class IndexedTaskInstance(RuntimeTaskInstance):
+ """
+ Indexed task instance to run a mapped operator.
+
+ It shares the parent task instance's identity, so what an iteration pushes
or stores lands in
+ the parent's scope, suffixed with the index so that iterations never
overwrite each other:
+ XComs through :meth:`xcom_push`, task state through
:attr:`task_state_store`. Reading back
+ follows the same rule: :meth:`xcom_pull` of the iteration's own XComs adds
the index, a pull
+ from another task does not, as the store's accessor adds it to every key
of its own. The
+ operator's own checkpoints go to the parent's store unsuffixed, under their
+ ``_iterable_<index>`` keys.
+ """
+
+ index: int
+ parent_task_state_store: TaskStateStoreAccessor
+ input_fingerprint: str | None = None
+ # The XComs the indexed task pushed other than its return value, by their
key before the index
+ # suffix: kept in memory while it runs and written once, with its SUCCESS
checkpoint.
+ pushed_xcoms: dict[str, Any] = Field(default_factory=dict)
+ # The params processed for this index's operator when they differ from the
iterated task's:
+ # the index iterates ``params``, or the operator set its own. None when
the parent's apply.
+ own_params: dict[str, Any] | None = None
+
+ @classmethod
+ def create_indexed_task(
+ cls, *, context: Context, index: int, operator: BaseOperator,
input_fingerprint: str | None = None
+ ) -> IndexedTaskInstance:
+ """
+ Create the runtime instance for one index of an iterated task from the
parent's context.
+
+ The instance shares the parent task instance's identity (id, run, map
index, try number and
+ retry budget), so XComs and task state land in the parent's scope, and
carries the unmapped
+ operator for that index. The budget is the parent's ``max_tries``
rather than the operator's
+ ``retries``: a manual clear raises it, and it is the parent Airflow
retries. The parent's
+ state store comes from the context, the same accessor the operator's
+ checkpoints use. ``model_construct`` skips Pydantic validation on
purpose: one instance is built per
+ indexed task, and the parent was validated already, so only the index
needs checking here.
+ """
+ if index < 0:
+ raise ValueError(f"IndexedTaskInstance requires index >= 0, got
{index}")
+ parent = context["ti"]
+ return cls.model_construct(
+ id=parent.id,
+ parent_task_state_store=context["task_state_store"],
+ # The parent's: an iteration's XComs and state belong to the task
instance that runs it.
+ task_id=parent.task_id,
+ dag_id=operator.dag_id,
+ run_id=parent.run_id,
+ map_index=parent.map_index,
+ region_id=getattr(parent, "region_id", None),
+ region_index=getattr(parent, "region_index", None),
+ index=index,
+ input_fingerprint=input_fingerprint,
+ max_tries=parent.max_tries,
+ start_date=parent.start_date,
+ state=TaskInstanceState.SCHEDULED.value,
+ is_mapped=True,
+ task=operator,
+ try_number=parent.try_number,
+ # The dag run, logical date and the rest the server sent for the
task instance that
+ # runs this iteration: its template context, get_previous_ti() and
the like read them.
+ _ti_context_from_server=getattr(parent, "_ti_context_from_server",
None),
+ )
+
+ def xcom_push(
+ self,
+ key: str,
+ value: Any,
+ ):
+ super().xcom_push(key=f"{key}_{self.index}", value=value)
+ self._record_push(key, value)
+
+ async def axcom_push(
+ self,
+ key: str,
+ value: Any,
+ ):
+ await super().axcom_push(key=f"{key}_{self.index}", value=value)
+ self._record_push(key, value)
+
+ def _record_push(self, key: str, value: Any) -> None:
+ # The return value has its own slot on the checkpoint and is published
by the operator.
+ if key != BaseXCom.XCOM_RETURN_KEY:
Review Comment:
Fixed in
[`26fae1a35f`](https://github.com/apache/airflow/pull/62922/commits/26fae1a35ff52bd7300a3e5650f973e19ae817d9):
the indexed task instance keeps a `return_value` it pushed itself, and the
SUCCESS checkpoint takes it when execute returned nothing to push, so the retry
after a sibling failed replays it. `do_xcom_push` governs only the returned
value, so the explicit push is kept either way; the test runs with both.
##########
task-sdk/src/airflow/sdk/execution_time/lazy_sequence.py:
##########
@@ -88,6 +113,36 @@ def __hash__(self):
def __iter__(self) -> Iterator[T]:
return LazyXComIterator(seq=self)
+ 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,
+ )
+ from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+ from airflow.sdk.execution_time.xcom import XCom
+
+ 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,
+ ),
+ )
+ if isinstance(msg, ErrorResponse):
+ raise IndexError(index)
+ if not isinstance(msg, XComSequenceIndexResult):
+ raise TypeError(f"Got unexpected response to GetXComSequenceItem:
{msg!r}")
+ return XCom._load(_XComWrapper(msg.root),
previous_iteration=self._previous_iteration)
Review Comment:
Done in
[`8108c83716`](https://github.com/apache/airflow/pull/62922/commits/8108c83716d42bf4a67d09b501cf8c5f8012b3a3):
`aget` loads through `XCom._load` in `asyncio.to_thread`, as `aget_one` does,
with a test that the load runs off the loop's thread.
##########
task-sdk/src/airflow/sdk/bases/xcom.py:
##########
@@ -805,6 +821,193 @@ 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 indexed task finishes (see ``IndexedTaskLifecycle._push_result``),
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.
+
+ An iterable read from the previous iteration (``loop.previous`` inside a
loop) is built with
+ ``previous_iteration`` set (see :meth:`deserialize`), so its elements are
read from that
+ iteration too and not from the current one. "Previous" is relative to
whoever reads, so
+ such an iterable cannot be pushed on: :meth:`serialize` refuses it (see
there).
+ """
+
+ 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, and every call in flight on it, for its round
trip to the supervisor.
+ """
+ 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]:
Review Comment:
Fixed in
[`f351e4936c`](https://github.com/apache/airflow/pull/62922/commits/f351e4936ca4432ffb9ed26af8d61f802f6bcfa5),
as you suggested: the pointer now records the region and pass of the task
instance that pushed it (`region_id`, `region_index`), and its item reads name
them. The execution API XCom read takes an explicit `region_id`/`region_index`
(version 2027-02-28; the resolver already supported an explicit scope for the
core API), and `GetXCom` carries them (supervisor schema entry, Go and TS
models regenerated).
That makes the pointer absolute instead of relative to its reader, so the
reader-side fix-up goes: the `previous_iteration` context variable and
`serialize()` refusing a pointer read through `loop.previous` are removed. A
reader outside the loop, one in a later pass and one in the same pass all read
the items of the pass that produced them, and a pointer to a task outside any
loop names no scope, which closes the reverse gap you found.
The end-to-end test is
`test_iterate_in_a_loop_is_read_from_outside_the_loop_and_by_the_next_pass` in
`test_taskinstance.py`: `dag.test()` on a loop whose iterated task gets a
different input on each pass, read through `loop.previous` inside the loop and
by a task outside it. Before the change it failed with the 400 you traced.
--
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]
