Decide min_one_success trigger rules per index in mapped task groups

airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py

@@ -195,20 +195,18 @@ def _get_relevant_upstream_map_indexes(upstream_id: str) -> int | range | None:
 
             # Only the not-yet-expanded summary ti (map_index < 0) needs the broad
             # "depend on every upstream ti" behavior, so a rule that can be satisfied
-            # by a subset of upstreams (ONE_SUCCESS / ONE_FAILED / ONE_DONE /
-            # NONE_FAILED_MIN_ONE_SUCCESS) does not skip it before the mapped task
-            # group has expanded (see #34023, #39801). Once the ti is expanded, each
-            # instance must depend on the upstream instance(s) that share its map
-            # index, otherwise a single upstream failure would wrongly trigger every
-            # expanded instance (see #50210). The task may be nested in plain task
-            # groups inside the mapped one (see #39801), so check the closest mapped
-            # ancestor rather than only the immediate parent group.
+            # by a subset of upstreams (ONE_SUCCESS / ONE_FAILED / ONE_DONE) does not
+            # skip it before the mapped task group has expanded (see #34023). Once the
+            # ti is expanded, each instance must depend on the upstream instance(s) that
+            # share its map index, otherwise a single upstream failure would wrongly
+            # trigger every expanded instance (see #50210). The task may be nested in
+            # plain task groups inside the mapped one (see #39801), so check the closest
+            # mapped ancestor rather than only the immediate parent group.
             if map_index < 0 and task.get_closest_mapped_task_group() is not None:
                 is_fast_triggered = task.trigger_rule in (
                     TR.ONE_SUCCESS,
                     TR.ONE_FAILED,
                     TR.ONE_DONE,
-                    TR.NONE_FAILED_MIN_ONE_SUCCESS,
                 )
                 if is_fast_triggered and upstream_id not in set(
                     _iter_expansion_dependencies(task_group=task.task_group)
@@ -422,6 +420,7 @@ def _evaluate_direct_relatives() -> Iterator[TIDepStatus]:
             success_setup = upstream_states.success_setup
             skipped_setup = upstream_states.skipped_setup
 
+            task_id_counts: list[tuple[str, int]] | None = None
             # Optimization: Don't need to hit the database if all upstreams are
             # "simple" tasks (no task or task group mapping involved).
             if not any(t.get_needs_expansion() for t in upstream_tasks.values()):
@@ -431,7 +430,6 @@ def _evaluate_direct_relatives() -> Iterator[TIDepStatus]:
                 # Regional selections share counts only for identical producer UUID sets;
                 # legacy runs can share by task IDs. Expansion invalidates both forms.
                 cache_key: tuple[str, str, frozenset[str | UUID]] | None = None
-                task_id_counts: list[tuple[str, int]] | None = None
                 if task.get_closest_mapped_task_group() is None:
                     scope: frozenset[str | UUID] = (
                         frozenset(
@@ -459,6 +457,20 @@ def _evaluate_direct_relatives() -> Iterator[TIDepStatus]:
                 upstream = sum(count for _, count in task_id_counts)
                 upstream_setup = sum(c for t, c in task_id_counts if upstream_tasks[t].is_setup)
 
+            # Once an in-group upstream has expanded, the summary ti of a task in a mapped task
+            # group no longer sees it and would judge every index by the outside upstreams alone.
+            # For these two rules the expanded tis decide instead. The cost: an outside failure seen
+            # after the expansion gives one upstream_failed ti per index instead of one in total,
+            # which is why other rules still decide here.
+            if (
+                map_index < 0
+                and trigger_rule in (TR.ALL_DONE_MIN_ONE_SUCCESS, TR.NONE_FAILED_MIN_ONE_SUCCESS)
+                and task_id_counts is not None
+                and len(task_id_counts) < len(upstream_tasks)
+                and task.get_closest_mapped_task_group() is not None
+            ):
+                return
+
             upstream_done = done >= upstream
 
             changed = False

airflow-core/tests/unit/models/test_mappedoperator.py

@@ -1956,15 +1956,16 @@ def group(value):
         assert states[task_id] == {0: "success", 1: "skipped"}
 
 
-def test_none_failed_min_one_success_trigger_rule_expands_in_mapped_task_group(dag_maker):
-    """Regression test for #39801.
+@pytest.mark.parametrize(
+    "trigger_rule", [TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS, TriggerRule.ALL_DONE_MIN_ONE_SUCCESS]
+)
+def test_min_one_success_trigger_rule_expands_in_mapped_task_group(dag_maker, trigger_rule):
+    """A ``*_MIN_ONE_SUCCESS`` task in a mapped task group runs once its in-group upstream succeeds.
 
-    A task with the ``NONE_FAILED_MIN_ONE_SUCCESS`` trigger rule inside a dynamically
-    expanded task group must expand and run once its upstream (also inside the group)
-    succeeds, instead of being skipped prematurely at its unexpanded (map_index -1)
-    summary ti before the task group has expanded.
+    It must not be skipped or marked upstream_failed prematurely at its unexpanded
+    (map_index -1) summary ti.
     """
-    with dag_maker(dag_id="test_none_failed_min_one_success_in_mapped_task_group") as dag:
+    with dag_maker(dag_id=f"test_{trigger_rule.value}_in_mapped_task_group") as dag:
 
         @task
         def init():
@@ -1976,7 +1977,7 @@ def tg(message):
             def imsleepy(message):
                 return message
 
-            @task(trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS)
+            @task(trigger_rule=trigger_rule)
             def imawake(message):
                 pass
 
@@ -1994,14 +1995,17 @@ def imawake(message):
     assert states["tg.imawake"] == {0: "success", 1: "success", 2: "success"}
 
 
-def test_none_failed_min_one_success_trigger_rule_expands_in_nested_task_group(dag_maker):
+@pytest.mark.parametrize(
+    "trigger_rule", [TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS, TriggerRule.ALL_DONE_MIN_ONE_SUCCESS]
+)
+def test_min_one_success_trigger_rule_expands_in_nested_task_group(dag_maker, trigger_rule):
     """Regression test for #39801, nested-group shape.
 
-    Same as above, but the ``NONE_FAILED_MIN_ONE_SUCCESS`` task lives in a plain task
+    Same as above, but the ``*_MIN_ONE_SUCCESS`` task lives in a plain task
     group nested inside the mapped one (the exact shape of the issue's reproducer), so
     its immediate ``task_group`` is not mapped — only an ancestor is.
     """
-    with dag_maker(dag_id="test_none_failed_min_one_success_in_nested_task_group") as dag:
+    with dag_maker(dag_id=f"test_{trigger_rule.value}_in_nested_task_group") as dag:
 
         @task
         def init():
@@ -2015,7 +2019,7 @@ def imsleepy(message):
 
             @task_group(group_id="inner")
             def inner(message):
-                @task(trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS)
+                @task(trigger_rule=trigger_rule)
                 def imawake(message):
                     pass
 
@@ -2035,6 +2039,70 @@ def imawake(message):
     assert states["tg.inner.imawake"] == {0: "success", 1: "success", 2: "success"}
 
 
+@pytest.mark.parametrize(
+    ("trigger_rule", "outside_outcome", "inner_outcome", "expected_state"),
+    [
+        (TriggerRule.ALL_DONE_MIN_ONE_SUCCESS, None, "skip", "skipped"),
+        (TriggerRule.ALL_DONE_MIN_ONE_SUCCESS, None, "fail", "upstream_failed"),
+        (TriggerRule.ALL_DONE_MIN_ONE_SUCCESS, "fail", "success", "success"),
+        (TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS, None, "fail", "upstream_failed"),
+        (TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS, "skip", "fail", "upstream_failed"),
+    ],
+)
+def test_min_one_success_decides_per_index_in_mapped_task_group(
+    dag_maker, trigger_rule, outside_outcome, inner_outcome, expected_state
+):
+    """A ``*_MIN_ONE_SUCCESS`` task in a mapped group is evaluated per expanded index.
+
+    The middle index of its in-group upstream ends with ``inner_outcome``; an optional
+    upstream outside the group ends with ``outside_outcome``.
+    """
+    with dag_maker(
+        dag_id=f"test_{trigger_rule.value}_outside_{outside_outcome}_inner_{inner_outcome}"
+    ) as dag:
+
+        @task
+        def init():
+            return ["seize", "the", "day"]
+
+        @task
+        def outside():
+            if outside_outcome == "skip":
+                raise AirflowSkipException
+            raise ValueError("boom")
+
+        outside_task = outside() if outside_outcome else None
+
+        @task_group(group_id="tg")
+        def tg(message):
+            @task
+            def imsleepy(message):
+                if message == "the":
+                    if inner_outcome == "skip":
+                        raise AirflowSkipException
+                    if inner_outcome == "fail":
+                        raise ValueError("boom")
+                return message
+
+            @task(trigger_rule=trigger_rule)
+            def imawake(message):
+                pass
+
+            awake = imawake(imsleepy(message))
+            if outside_task is not None:
+                outside_task >> awake
+
+        tg.expand(message=init())
+
+    dr = dag.test()
+
+    states: dict[str, dict[int, str | None]] = defaultdict(dict)
+    for ti in dr.get_task_instances():
+        states[ti.task_id][ti.map_index] = ti.state
+
+    assert states["tg.imawake"] == {0: "success", 1: expected_state, 2: "success"}
+
+
 def test_mapped_operator_retry_delay_default(dag_maker):
     """
     Test that MappedOperator.retry_delay returns default value when not explicitly set.