Fix upstream map index resolution after placeholder expansion

[SameerMesiah97]

Description

Fix upstream map index resolution for dynamically mapped task groups when placeholder task instances are replaced during expansion.

After expansion, the upstream placeholder (map_index = -1) is replaced by the first expanded task instance (map_index = 0). Downstream task instances that are still unexpanded may continue to reference the placeholder, causing incorrect upstream dependency resolution.

This change introduces a small helper (_should_use_post_expansion_placeholder) that is invoked during upstream map index resolution to correctly handle this transition, while preserving existing behavior for already-expanded downstream tasks.

Rationale

Placeholder task instances are a temporary pre-expansion representation. Once expansion occurs, they are no longer valid references, but downstream resolution logic may still encounter them. Handling this transition explicitly prevents downstream tasks from treating completed upstream work as unresolved.

This could cause implicit dependencies to be skipped during DAG evaluation under certain trigger rules (such as NONE_FAILED_MIN_ONE_SUCCESS), particularly when multiple parallel task streams within the MappedTaskGroup converge on a single downstream task outside the group. As a result, DAG runs could be incorrectly marked as failed or skipped despite valid upstream execution.

Please refer to issue #59289 for more context. This PR was opened in response to that. The author of the issue reported the bug in Airflow 3.0.6 but I can confirm that the same issue is present in Airflow 3.1.5 (as well as the main branch at the time this PR was last updated).

Tests

Added a test covering the post-expansion state where:

• The upstream placeholder has been expanded.
• The downstream task instance is either unexpanded or expanded.

The test asserts that get_relevant_upstream_map_indexes returns the correct upstream map index in both cases.

Closes: #59289

[SameerMesiah97]

@ashb@XD-DENG

This is just a soft reminder for review. It fixes a genuine bug (please refer to issue #59289) that subtly breaks DAGs under certain mapped task group setups. I believe it should be addressed before the next release.

[potiuk]

Also @uranusjr can help with review - there ias related WIP he eworks on now I think

[SameerMesiah97]

@uranusjr

Since, according to @potiuk your recent work is adjacent to what is being addressed in this PR, I am tagging you as a soft reminder for review and/or closing #59289 if you feel that your work has addressed it.

[SameerMesiah97]

I have checked the main branch using a variety of DAG shapes following the same theme as the one in the bug report and the issue is still present. So it appears any recent work in this specific area (that is separate from this PR) which has been merged into main, has not resolved the bug.

As of this comment, I cannot see any open PRs that could resolve the issue (anyone reading this comment is free to correct me).

This PR has been rebased and is ready for review.

[SameerMesiah97]

I have checked main again using the same test cases and the issue is still present so the bug being addressed by this PR has not been resolved. I am unable to locate any open PRs that I would expect to fix this issue either. The bug silently renders certain DAG shapes unusable with Dynamic Task Mapping so I believe this should not be deprioritized.

This PR has been open for approximately 2 months. It has now been rebased again and is ready for review. I have run the aforementioned test cases with my changes and I can confirm that they give the expected result.

@eladkal@ephraimbuddy@uranusjr

[Nataneljpwd]

Hello @SameerMesiah97, I think there have been a lot of changes since you created the pr, the code you have won't work now as all parameters but the upstream, session and ti_count have been removed from the get_relevant_upstream_map_indexes method

Could you check if the issue still persists?

[SameerMesiah97]

Hello @SameerMesiah97, I think there have been a lot of changes since you created the pr, the code you have won't work now as all parameters but the upstream, session and ti_count have been removed from the get_relevant_upstream_map_indexes method

Could you check if the issue still persists?

I checked 3 weeks ago but I will check again today.

[SameerMesiah97]

Hello @SameerMesiah97, I think there have been a lot of changes since you created the pr, the code you have won't work now as all parameters but the upstream, session and ti_count have been removed from the get_relevant_upstream_map_indexes method

Could you check if the issue still persists?

Just checked again right now. The bug is still present in the main branch using the exact DAG shapes I used in my previous tests and my fix resolves it for all of these cases. I have rebased and force-pushed again.

[SameerMesiah97]

@uranusjr

Your feedback has been addressed.

[Quantum0uasar]

Reviewed the diff carefully. Great fix for a subtle race condition in dynamic task group expansion. A few observations:

On _should_use_post_expansion_placeholder:
The function issues a DB query on every invocation inside what can be a hot path (get_relevant_upstream_map_indexes is called for each upstream dependency during scheduler evaluation). For DAGs with many mapped downstream TIs this could add meaningful scheduler latency. Consider:

1. Caching the result keyed by (run_id, task_id) within the scope of a single scheduling loop — the expansion state won't change mid-evaluation.
2. Alternatively, checking ti.state or ti.expanded_map_indexes fields that are likely already loaded in the session identity map before falling back to the DB query.

On the condition logic:

return (
-1intask_to_map_indexes[task.task_id]
and-1notintask_to_map_indexes[relative.task_id]
and0intask_to_map_indexes[relative.task_id]
)

Is the third condition (0 in task_to_map_indexes[relative.task_id]) strictly necessary? If -1 is not in relative's map indexes and expansion has happened, 0 should always be present. The condition is correct but the third check is implied by not -1 in relative. Worth a comment clarifying the intent.

On test coverage:
The regression test covers NONE_FAILED_MIN_ONE_SUCCESS — would be good to add a case for ALL_SUCCESS trigger rule too, since the upstream resolution path is shared and the skip logic may differ.

Overall the approach is sound and the fix is minimal. Appreciate the thorough rebasing history on this one.

[Quantum0uasar]

Thanks for tackling this edge case — placeholder-to-expanded index transition bugs are notoriously hard to reproduce in dynamic task mapping. I reviewed the logic and have a few observations:

On the core helper

The guard condition:

return (
-1intask_to_map_indexes[task.task_id] and-1notintask_to_map_indexes[relative.task_id] and0intask_to_map_indexes[relative.task_id]
)

This reads cleanly, but I’d flag one subtle case: what if the upstream expands to more than one instance (e.g., map_index in {0, 1, 2, ...})? The function returns a boolean for translating -1 → 0, but for a fan-out case where downstream map_index=-1 should resolve to all expanded instances, not just 0, is that handled elsewhere? Worth a comment in the code or a targeted test for fan-out upstreams.

On DB query cost

The reviewer note about unnecessary DB queries is important. If task_to_map_indexes is built by querying per-task rather than via a single bulk query, this will fire N queries for N upstream tasks in complex DAGs. Consider pre-fetching once:

all_indexes=session.execute(
select(TI.task_id, TI.map_index)
.where(TI.dag_id==dag_id, TI.run_id==run_id)
).all()
task_to_map_indexes=defaultdict(set)
fortask_id, map_indexinall_indexes:
task_to_map_indexes[task_id].add(map_index)

This avoids the query amplification problem especially for wide DAGs.

Test coverage gap

The existing regression tests cover:

• Placeholder already expanded (upstream map_index=-1 replaced by 0)
• Downstream in both unexpanded and expanded states

One gap: a test where the upstream has not yet been expanded (placeholder still active). The fix should be a no-op in that case, but it’s easy to accidentally break this when changing the condition.

Also agree with the reviewer’s request for an ALL_SUCCESS trigger rule test — trigger rule handling is another layer where index resolution can silently produce wrong results.

Minor

• The helper name resolve_placeholder_map_index is descriptive; consider adding a one-line docstring explaining the -1 → 0 contract so future maintainers don’t have to re-derive it.
• Is there a linked issue number for this fix? Would help close the feedback loop (e.g., Closes #59289 in the PR description if that’s the upstream issue).

Overall the fix is targeted and correct for the primary case. The main risk area is the fan-out upstream scenario and DB query cost at scale. Happy to look at the full diff in taskinstance.py if you want a more detailed line-level review.

[SameerMesiah97]

@Quantum0uasar

Thanks for the review. I have added another phase in the existing test to cover the initial state where the upstream placeholder has not yet been expanded.

A few clarifications on the other points:

1. I do agree that this DB query is in a scheduler-sensitive codepath and will therefore run frequently. However, it is a very narrow query scoped to a single DAG run, a single upstream/downstream task pair, and map indexes -1 and 0. Given the very small result set involved here, I do not currently think introducing caching/state-management complexity would be worth it without evidence that this lookup is a measurable bottleneck.

2. I agree the helper would likely work fine without the third condition since it is implied by the second condition. However, this is a fairly subtle placeholder-transition edge case, so I think the explicitness is beneficial for readability and future maintenance unless there is a stronger rationale for simplifying it.

3. I believe the trigger rule discussion is orthogonal to this PR. This fix is focused on resolving the correct upstream mapped task instance during placeholder transition handling, not determining state resolution semantics for the task instances themselves. Trigger rule evaluation lives within trigger_rule_dep.py, which is a bit outisde the scope of this PR.

4. In this implementation, we are intentionally only considering the upstream placeholder (-1) and its immediate post-expansion successor (0). The goal is specifically to handle the transition where the upstream placeholder no longer exists but the downstream task is still referencing placeholder state. Matching against arbitrary expanded indexes (>0) would not reflect the semantics of the downstream placeholder task instance correctly.

5. I think there may be a misunderstanding around the caching example provided. The snippet appears to broaden the query scope to all mapped task instances for the DAG run rather than cache previously resolved state, which I suspect would increase the DB query cost rather than reduce it. If I am misunderstanding the intent there, please clarify.

6. I think you may have been looking at an earlier commit. The helper is now called _should_use_post_expansion_placeholder, and I believe the current docstring describes the intended contract clearly. The issue has been linked in the PR description.

[SameerMesiah97]

@uranusjr

I would appreciate another look.

[SameerMesiah97]

Requesting review for this.

[dabla]

@SameerMesiah97 have you tested this fix locally with Dynamic Task Expansion, any logs to prove that it the change fixes the issue? Maybe a stupid question, but that could help approve the PR as for me code is looking good, but of course I haven't tested it so I cannot confirm this will actually fix the issue. Nice work BTW.

[SameerMesiah97]

@SameerMesiah97 have you tested this fix locally with Dynamic Task Expansion, any logs to prove that it the change fixes the issue? Maybe a stupid question, but that could help approve the PR as for me code is looking good, but of course I haven't tested it so I cannot confirm this will actually fix the issue. Nice work BTW.

If by Dynamic Task Expansion you meant Dynamic Task Mapping, then yes. I have done several tests using a variety of DAG shapes at mutliple points in time whilst working on this PR. If you scroll up, you will see me reporting the results after each test run. I will do another test as soon as I am able to and I will provide all the details as well.

[SameerMesiah97]

@SameerMesiah97 have you tested this fix locally with Dynamic Task Expansion, any logs to prove that it the change fixes the issue? Maybe a stupid question, but that could help approve the PR as for me code is looking good, but of course I haven't tested it so I cannot confirm this will actually fix the issue. Nice work BTW.

@dabla

I have ran the following DAG (closely mirrors the one used in issue #59289) using the code in main:

from datetime import datetime
from airflow import DAG
from airflow.operators.empty import EmptyOperator
from airflow.decorators import dag, task, task_group
from airflow.utils.trigger_rule import TriggerRule
@task
def get_the_list():
return ["one", "two", "three"]
@task
def process_something(x: str):
output = f"{x}"
print(output)
return output
@task_group
def the_task_group(x: str):
start = EmptyOperator(task_id="start")
something = process_something(x)
end = EmptyOperator(task_id="end")
start >> something >> end
@dag(
default_args={
"start_date": datetime(year=2025, month=12, day=1),
"trigger_rule": TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS,
}
)
def dtm_bug_reproduce_dag_original_case():
end = EmptyOperator(task_id="end")
the_list = get_the_list()
mapped_group = (
the_task_group
.expand(x=the_list)
)
the_list >> mapped_group >> end
dag = dtm_bug_reproduce_dag_original_case()

This is the UI screenshot for the DAG run:

The same DAG was run again using the branch which has my fix (rebased on top of the latest version of main):

You can observe that the terminal tasks of both the DAG and the mapped group are no longer being skipped. I must admit that unlike in previous test runs, the DAG is being marked as successful. And amongst the other DAG shapes I tested, the bug is no longer present in main for a few of them. But it exists for at least this specific DAG I have presented here and mutliple others besides those few that I have tested.

[dabla]

Thanks for the nice work and patience @SameerMesiah97 !