Skip to content

Commit 41d897f

Browse files
authored
fix: LatestOnlyOperator not working if direct upstream of Dynamic Task Map (#62287)
* fix: use parent's map_index for XCom lookup in NotPreviouslySkippedDep When an unmapped SkipMixin parent (e.g., LatestOnlyOperator) is upstream of a dynamically mapped task, the child task must look up XCom using map_index=-1 (the parent's XCom storage index) rather than the child's own map_index. This fix ensures mapped downstream TIs are correctly skipped when appropriate. Fixes #62118 * fix: apply ruff formatting and import sorting - Split long multi-argument function call across lines - Fix import order: move XComModel import before providers * ci: trigger CI re-run (transient docs build cache failure)
1 parent d349918 commit 41d897f

2 files changed

Lines changed: 62 additions & 2 deletions

File tree

‎airflow-core/src/airflow/ti_deps/deps/not_previously_skipped_dep.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,8 +58,15 @@ def _get_dep_statuses(self, ti, session, dep_context):
5858
# This can happen if the parent task has not yet run.
5959
continue
6060

61+
# Use the parent's map context to look up the XCom. An unmapped parent
62+
# (e.g. LatestOnlyOperator) writes XCom with map_index=-1, so we must
63+
# query with -1 instead of the child's map_index.
64+
xcom_map_index = ti.map_index if parent.is_mapped else -1
6165
prev_result = ti.xcom_pull(
62-
task_ids=parent.task_id, key=XCOM_SKIPMIXIN_KEY, session=session, map_indexes=ti.map_index
66+
task_ids=parent.task_id,
67+
key=XCOM_SKIPMIXIN_KEY,
68+
session=session,
69+
map_indexes=xcom_map_index,
6370
)
6471

6572
if prev_result is None:

‎airflow-core/tests/unit/ti_deps/deps/test_not_previously_skipped_dep.py‎

Lines changed: 54 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,10 +22,15 @@
2222
from sqlalchemy import delete
2323

2424
from airflow.models import DagRun, TaskInstance
25+
from airflow.models.xcom import XComModel
2526
from airflow.providers.standard.operators.empty import EmptyOperator
2627
from airflow.providers.standard.operators.python import BranchPythonOperator
2728
from airflow.ti_deps.dep_context import DepContext
28-
from airflow.ti_deps.deps.not_previously_skipped_dep import NotPreviouslySkippedDep
29+
from airflow.ti_deps.deps.not_previously_skipped_dep import (
30+
XCOM_SKIPMIXIN_FOLLOWED,
31+
XCOM_SKIPMIXIN_KEY,
32+
NotPreviouslySkippedDep,
33+
)
2934
from airflow.utils.state import State
3035
from airflow.utils.types import DagRunType
3136

@@ -161,3 +166,51 @@ def test_parent_not_executed(session, dag_maker):
161166
assert len(list(dep.get_dep_statuses(ti2, session, DepContext()))) == 0
162167
assert dep.is_met(ti2, session)
163168
assert ti2.state == State.NONE
169+
170+
171+
def test_unmapped_parent_skip_mapped_downstream(session, dag_maker):
172+
"""
173+
When an unmapped SkipMixin parent writes XCom with map_index=-1,
174+
mapped downstream TIs (map_index >= 0) should still be skipped
175+
by NotPreviouslySkippedDep.
176+
177+
Regression test for https://github.com/apache/airflow/issues/62118
178+
"""
179+
start_date = pendulum.datetime(2020, 1, 1)
180+
with dag_maker(
181+
"test_unmapped_skip_mapped_dag",
182+
schedule=None,
183+
start_date=start_date,
184+
session=session,
185+
):
186+
op1 = BranchPythonOperator(task_id="op1", python_callable=lambda: "op3")
187+
op2 = EmptyOperator(task_id="op2")
188+
op3 = EmptyOperator(task_id="op3")
189+
op1 >> [op2, op3]
190+
191+
dr = dag_maker.create_dagrun(run_type=DagRunType.MANUAL, state=State.RUNNING)
192+
tis = {ti.task_id: ti for ti in dr.task_instances}
193+
194+
# Simulate the unmapped branch operator having run: set it to SUCCESS
195+
# and store XCom with map_index=-1 (as SkipMixin does for unmapped tasks).
196+
tis["op1"].state = State.SUCCESS
197+
session.merge(tis["op1"])
198+
XComModel.set(
199+
key=XCOM_SKIPMIXIN_KEY,
200+
value={XCOM_SKIPMIXIN_FOLLOWED: ["op3"]},
201+
dag_id=dr.dag_id,
202+
task_id="op1",
203+
run_id=dr.run_id,
204+
map_index=-1,
205+
session=session,
206+
)
207+
208+
# Simulate a mapped downstream TI by changing map_index to 0.
209+
tis["op2"].map_index = 0
210+
session.merge(tis["op2"])
211+
session.flush()
212+
213+
dep = NotPreviouslySkippedDep()
214+
assert len(list(dep.get_dep_statuses(tis["op2"], session, DepContext()))) == 1
215+
assert not dep.is_met(tis["op2"], session)
216+
assert tis["op2"].state == State.SKIPPED

0 commit comments

Comments
 (0)