uranusjr commented on code in PR #67672:
URL: https://github.com/apache/airflow/pull/67672#discussion_r3457477936
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1520,16 +1520,26 @@ def _expand_mapped_task_if_needed(ti: TI) ->
Iterable[TI] | None:
if new_tis is not None:
additional_tis.extend(new_tis)
expansion_happened = True
+ # Expansion changes a mapped task's instance count, which
invalidates the
+ # trigger-rule upstream-count memo on this DepContext (a
downstream evaluated
+ # later in this same pass must see the post-expansion
count).
+ dep_context.upstream_task_id_counts.clear()
if new_tis is None and schedulable.state in SCHEDULEABLE_STATES:
# It's enough to revise map index once per task id,
# checking the map index for each mapped task significantly
slows down scheduling
if schedulable.task.task_id not in revised_map_index_task_ids:
- ready_tis.extend(
+ revised_tis = list(
self._revise_map_indexes_if_mapped(
schedulable.task,
dag_version_id=schedulable.dag_version_id, session=session
)
)
+ ready_tis.extend(revised_tis)
revised_map_index_task_ids.add(schedulable.task.task_id)
+ if revised_tis:
Review Comment:
Not sure if the size of `revised_tis` would be significant. We could do
```python
prev_ti_size = len(ready_tis)
ready_tis.extend(self._revise_map_indexes_if_mapped(...)
if len(ready_tis) > prev_ti_size:
...
```
to avoid building a list.
--
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]
