branch operators inside expand/expand_kwargs-mapped TaskGroups in 3.2.x — non-selected branch siblings run instead of being skipped
- Dominant language
- Python
- Stars
- 46.9k
- Forks
- 17.8k
- Avg merge
- 2d 9h
- Merged PRs (30d)
- 472
Description
### Under which category would you file this issue?
Airflow Core
### Apache Airflow version
3.2.1
### What happened and how to reproduce it?
In the following minimal example, run_optional_step runs for each sibling (item_type `full`, and item_type `partial`), however `partial` is intended to be skipped::
```
"""Minimal example: @task.branch inside a dynamically mapped task group.
This example replicates the structure of a real-world DAG, where:
- One outer task group (process_items) fetches items and expands
an inner task group (handle_item) via expand_kwargs.
- Inside handle_item, a @task.branch decides whether to run
run_optional_step based on item_type.
Expected behavior
------------------
item_type == "full" → branch returns absolute task ID → run_optional_step should run
item_type == "partial" → branch returns None → run_optional_step should be skipped
Actual behavior
-------------------------
run_optional_step runs for ALL items.
"""
from __future__ import annotations
import pendulum
from airflow.sdk import DAG, task, task_group
from airflow.providers.standard.operators.empty import EmptyOperator
with DAG(
dag_id="branch_in_mapped_taskgroup_example",
schedule=None,
start_date=pendulum.datetime(2024, 1, 1, tz="UTC"),
catchup=False,
) as dag:
# ------------------------------------------------------------------ #
# Inner task group: handles one mapped item. #
# The branch should skip or run an optional step per mapped instance. #
# ------------------------------------------------------------------ #
@task_group(group_id="handle_item")
def tg_handle_item(item_id: str, item_type: str):
"""Process one item. Only 'full' items need the extra step."""
@task(task_id="process")
def process(item_id: str, item_type: str) -> str:
print(f"Processing {item_id=} {item_type=}")
return item_type
result = process(item_id, item_type)
@task.branch(task_id="select_optional_step")
def select_optional_step(item_type: str):
if item_type == "full":
return ["process_items.handle_item.run_optional_step"]
return None
branch = select_optional_step(result)
run_optional_step = EmptyOperator(task_id="run_optional_step")
branch >> run_optional_step
# ------------------------------------------------------------------ #
# Outer task group: loads items and expands the inner group once per #
# item. There is only one outer group instance in the DAG. #
# ------------------------------------------------------------------ #
@task_group(group_id="process_items")
def tg_process_items():
@task(task_id="get_items")
def get_items() -> list[dict]:
return [
{"item_id": "item_1", "item_type": "full"}, # optional step should run
{"item_id": "item_2", "item_type": "partial"}, # optional step should be skipped
]
tg_handle_item.expand_kwargs(get_items())
tg_process_items()
```
### What you think should happen instead?
item_type == "full" → branch returns absolute task ID → run_optional_step should run
item_type == "partial" → branch returns None → run_optional_step should be skipped
This is what happens up through 3.1.7.
### Operating System
_No response_
### Deployment
None
### Apache Airflow Provider(s)
_No response_
### Versions of Apache Airflow Providers
_No response_
### Official Helm Chart version
Not Applicable
### Kubernetes Version
_No response_
### Helm Chart configuration
_No response_
### Docker Image customizations
_No response_
### Anything else?
This was originally discussed in https://github.com/apache/airflow/issues/65745, however, it seems like that issue is either not-reproducible or describing another issue.
[Further context from @martijn-exads](https://github.com/apache/airflow/issues/65745#issuecomment-4497401392):
We think we've traced it to https://github.com/apache/airflow/pull/62287, specifically this line in NotPreviouslySkippedDep:
xcom_map_index = ti.map_index if parent.is_mapped else -1
For a branch operator inside a mapped @task_group, the operator itself has no .expand() so parent.is_mapped is False, but its TIs run with map_index >= 0 and SkipMixin writes XCOM_SKIPMIXIN_KEY at the parent TI's per-map map_index. The dep then queries map_indexes=-1, finds nothing, returns "not skipped", and the non-selected sibling proceeds.
### Are you willing to submit PR?
- [x] Yes I am willing to submit a PR!
### Code of Conduct
- [x] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)
Contributor guide
Assessment
This issue has not been assessed yet.