dabla commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r4027374446
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -891,6 +906,85 @@ def mark_success_url(self) -> str:
return self.log_url
+@dataclass
+class IndexedTaskState:
+ status: TaskInstanceState
+ try_number: int
+ result: Any | None = None
+ # Outlet asset events the sub-task recorded on a previous successful
attempt. A sub-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 IterableOperator._run_task replay it instead of silently losing
those events.
+ outlet_events: list[dict[str, Any]] | None = None
+
+ def serialize(self) -> dict[str, Any]:
+ data: dict[str, Any] = {"status": self.status.value, "try_number":
self.try_number}
+ if self.result is not None:
+ data["result"] = self.result
+ if self.outlet_events:
+ data["outlet_events"] = self.outlet_events
+ return data
+
+ @classmethod
+ def deserialize(cls, raw: Any) -> IndexedTaskState | None:
+ if not isinstance(raw, Mapping):
+ return None
+ return cls(
+ status=TaskInstanceState(raw["status"]),
+ try_number=raw["try_number"],
+ result=raw.get("result"),
+ outlet_events=raw.get("outlet_events"),
+ )
+
+
+class IndexedTaskInstance(RuntimeTaskInstance):
+ """Indexed task instance to run a mapped operator."""
+
+ index: int
+
+ def __init__(self, /, **data: Any):
+ super().__init__(**data)
+
+ if self.index is None or self.index < 0:
+ raise ValueError("IndexedTaskInstance requires index >= 0")
+
+ def xcom_push(
+ self,
+ key: str,
+ value: Any,
+ ):
+ super().xcom_push(key=f"{key}_{self.index}", value=value)
+
+ async def axcom_push(
+ self,
+ key: str,
+ value: Any,
+ ):
+ await super().axcom_push(key=f"{key}_{self.index}", value=value)
+
+ async def aget_state(self) -> IndexedTaskState | None:
+ return IndexedTaskState.deserialize(await
self.task_state_store.aget(self.xcom_key))
+
+ async def aset_state(self, state: IndexedTaskState) -> None:
+ await self.task_state_store.aset(self.xcom_key, state.serialize())
+
+ @property
+ def is_async(self) -> bool:
+ return self.task.is_async
+
+ @property
+ def next_try_number(self) -> int:
+ return self.try_number + 1
+
+ @property
+ def xcom_key(self) -> str:
+ return f"{self.task_id}_{self.index}"
Review Comment:
You're right that the scope already covers it, and each batched parent has
its own scope too, so the task id guarded nothing. Changed in 242134872a: keys
are now `_iterable_<index>` and `_iterable_completed`. I kept a fixed namespace
rather than a bare `str(index)` so the operator's entries cannot collide with
keys user code stores from inside a sub-task, where names like `completed` are
plausible. The property is also renamed from `xcom_key` to `state_key`, since
that is what it is.
---
Drafted-by: Claude Fable 5.1; 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]
