dabla commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r4252638451
##########
task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py:
##########
@@ -184,6 +316,16 @@ def iter_references(self) -> Iterable[tuple[Operator,
str]]:
if isinstance(x, XComArg):
yield from x.iter_references()
+ async def aresolve(self, context: Mapping[str, Any]) -> Resolved:
+ sources = {key: await Source.from_argument(value, context) for key,
value in self.value.items()}
+ lengths = {key: await source.alen() for key, source in sources.items()}
+
+ async def aget(index: int) -> Mapping[str, Any]:
+ positions = index_for_each_field(index, lengths)
+ return {key: await source.aget(positions[key]) for key, source in
sources.items()}
Review Comment:
Right, done in aa4abe38de, but deliberately not as a cache of every
position: a lazy source is usually an upstream's `XComIterable` or
`LazyXComSequence`, whose elements can be arbitrarily large, and keeping them
for the run would materialize the upstream in the worker's memory, which those
classes exist to avoid.
What the cross product needs is smaller. The last argument varies fastest
and iterations start in position order, so every argument but the last is asked
for the same element over a run of consecutive positions. A lazy source of a
cross product now keeps its last read (`SharedReadsSource`, one value per
source, the one the running iterations already use) and shares concurrent reads
of a position through `asyncio.shield`; a failed read is forgotten so the next
iteration retries it. Counted with 16 iterations at a time and a 1 ms read:
| `iterate(...)` | iterations | reads before | reads now |
|---|---|---|---|
| `a=<lazy 1000>, b=<lazy 50>` | 50,000 | 100,000 | 51,000 |
| `a=<lazy 1000>, b=[50 literals]` | 50,000 | 50,000 | 1,000 |
| `a=<lazy 100>, b=<lazy 10>, c=[5 literals]` | 5,000 | 10,000 | 1,100 |
The last argument is read as it is, since it uses each element for one
iteration only; sparing its reads too would mean holding all of it. Since the
order of the arguments now decides the cost, the iterable tasks page has a note
saying to put the longest upstream first and literal lists last.
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]
