Mapped `task_group` does not output list of combined results to consumer

[retsilla]

Apache Airflow version

3.0.0

If "Other Airflow 2 version" selected, which one?

No response

What happened?

When dynamically creating task groups and creating a task to consume combined results, the input received by the consuming task is None rather than a list of combined results.

What you think should happen instead?

According to the example in docs and behaviour in previous versions of Airflow the result passed to the consuming tasks should be a combined list of results from the task group instances.

How to reproduce

DAG derived from the example in the docs:

from datetime import datetime

from airflow.sdk import DAG, task, task_group

with DAG(
    dag_id="dynamic_task_group_dag",
    start_date=datetime(2025, 1, 1),
    schedule=None,
):

    @task_group
    def add_one_and_double(value: int):
        @task
        def add_one(value: int):
            return value + 1

        @task
        def double(value: int):
            return value * 2

        new_value = add_one(value)
        return double(new_value)

    @task
    def consumer(values: list[int]):
        print(f"New values are: {values}")

    results = add_one_and_double.expand(value=[1, 2, 3])
    consumer(results)  # Should receive [4, 6, 8] but gets null

Logs of consumer task showing null was received:

[2025-04-24, 13:14:50] INFO - DAG bundles loaded: dags-folder: source="airflow.dag_processing.bundles.manager.DagBundlesManager"
[2025-04-24, 13:14:50] INFO - Filling up the DagBag from /home/airflow/dags/dynamic_task_group_dag.py: source="airflow.models.dagbag.DagBag"
[2025-04-24, 13:14:50] INFO - Done. Returned value was: None: source="airflow.task.operators.airflow.providers.standard.decorators.python._PythonDecoratedOperator"
[2025-04-24, 13:14:50] INFO - New values are: None: chan="stdout": source="task"

Operating System

Debian GNU/Linux 12

Versions of Apache Airflow Providers

No response

Deployment

Docker-Compose

Deployment details

No response

Anything else?

No response

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[boring-cyborg]

Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.

[Lee-W]

I've verified this bug. This seems to be a mapped operator issue that happens to all tasks. Will continue on the investigation

[tirkarthi]

@Lee-W I was checking on this while debugging #49737 . It seems in the example above while trying to resolve consumer(results) and results from double is an xcom arg that is attached to a mapped task group. No map_indexes are passed which causes xcoms to be fetched for add_one_and_double.double with map_index as -1 from self.map_index for consumer which is a normal task inside xcom_pull function.

One possible approach would be in case of the normal task having a closes task group which is also mapped get the parse time count of map_indexes and fetch all their values which would return [4, 6, 8] . But get_parse_time_mapped_ti_count will not be possible in cases where expansion depends on another task like add_one_and_double.expand(value=producer()) . Hope I am correct in my understanding here.

diff --git a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
index ab3a5a82b2..b85a14d5a7 100644
--- a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
+++ b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
@@ -350,11 +350,26 @@ class PlainXComArg(XComArg):
             )
         else:
             # task from a task group
-            result = ti.xcom_pull(
-                task_ids=task_id,
-                key=self.key,
-                default=NOTSET,
-            )
+            from airflow.sdk.definitions.taskgroup import MappedTaskGroup
+
+            # Upstream of this normal task is a mapped task group. Hence all
+            # xcom values for the upstream map_indexes need to be collected to be passed
+            # to the normal task.
+            if (tg, MappedTaskGroup) and ti.map_index == -1:
+                map_indexes = list(range(tg.get_parse_time_mapped_ti_count()))
+
+                result = ti.xcom_pull(
+                    task_ids=task_id,
+                    key=self.key,
+                    default=NOTSET,
+                    map_indexes=map_indexes
+                )
+            else:
+                result = ti.xcom_pull(
+                    task_ids=task_id,
+                    key=self.key,
+                    default=NOTSET,
+                )
         if not isinstance(result, ArgNotSet):
             return result
         if self.key == XCOM_RETURN_KEY:

[tirkarthi]

cc: @uranusjr

[Lee-W]

@Lee-W I was checking on this while debugging #49737 . It seems in the example above while trying to resolve consumer(results) and results from double is an xcom arg that is attached to a mapped task group. No map_indexes are passed which causes xcoms to be fetched for add_one_and_double.double with map_index as -1 from self.map_index for consumer which is a normal task inside xcom_pull function.

One possible approach would be in case of the normal task having a closes task group which is also mapped get the parse time count of map_indexes and fetch all their values which would return [4, 6, 8] . But get_parse_time_mapped_ti_count will not be possible in cases where expansion depends on another task like add_one_and_double.expand(value=producer()) . Hope I am correct in my understanding here.

diff --git a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
index ab3a5a82b2..b85a14d5a7 100644
--- a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
+++ b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
@@ -350,11 +350,26 @@ class PlainXComArg(XComArg):
             )
         else:
             # task from a task group
-            result = ti.xcom_pull(
-                task_ids=task_id,
-                key=self.key,
-                default=NOTSET,
-            )
+            from airflow.sdk.definitions.taskgroup import MappedTaskGroup
+
+            # Upstream of this normal task is a mapped task group. Hence all
+            # xcom values for the upstream map_indexes need to be collected to be passed
+            # to the normal task.
+            if (tg, MappedTaskGroup) and ti.map_index == -1:
+                map_indexes = list(range(tg.get_parse_time_mapped_ti_count()))
+
+                result = ti.xcom_pull(
+                    task_ids=task_id,
+                    key=self.key,
+                    default=NOTSET,
+                    map_indexes=map_indexes
+                )
+            else:
+                result = ti.xcom_pull(
+                    task_ids=task_id,
+                    key=self.key,
+                    default=NOTSET,
+                )
         if not isinstance(result, ArgNotSet):
             return result
         if self.key == XCOM_RETURN_KEY:

Yep, it was checked by the common ancestor function in 2.10. I was thinking whether we want to do it the old way or doing something similar to what you did