Repository: apache/airflow

Issue:
XComs from another task group fail to populate dynamic task mapping metadata
### Apache Airflow version

2.3.3

### What happened

When a task returns a mappable Xcom within a task group, the dynamic task mapping feature (via `.expand`) causes the Airflow Scheduler to infinitely loop with a runtime error:

```
Traceback (most recent call last):

  File "/home/airflow/.local/bin/airflow", line 8, in <module>

    sys.exit(main())

  File "/home/airflow/.local/lib/python3.9/site-packages/airflow/__main__.py", line 38, in main

    args.func(args)

  File "/home/airflow/.local/lib/python3.9/site-packages/airflow/cli/cli_parser.py", line 51, in command

    return func(*args, **kwargs)

  File "/home/airflow/.local/lib/python3.9/site-packages/airflow/utils/cli.py", line 99, in wrapper

    return f(*args, **kwargs)

  File "/home/airflow/.local/lib/python3.9/site-packages/airflow/cli/commands/scheduler_command.py", line 75, in scheduler

    _run_scheduler_job(args=args)

  File "/home/airflow/.local/lib/python3.9/site-packages/airflow/cli/commands/scheduler_command.py", line 46, in _run_scheduler_job

    job.run()

  File "/home/airflow/.local/lib/python3.9/site-packages/airflow/jobs/base_job.py", line 244, in run

    self._execute()

  File "/home/airflow/.local/lib/python3.9/site-packages/airflow/jobs/scheduler_job.py", line 751, in _execute

    self._run_scheduler_loop()

  File "/home/airflow/.local/lib/python3.9/site-packages/airflow/jobs/scheduler_job.py", line 839, in _run_scheduler_loop

    num_queued_tis = self._do_scheduling(session)

  File "/home/airflow/.local/lib/python3.9/site-packages/airflow/jobs/scheduler_job.py", line 921, in _do_scheduling

    callback_to_run = self._schedule_dag_run(dag_run, session)

  File "/home/airflow/.local/lib/python3.9/site-packages/airflow/jobs/scheduler_job.py", line 1163, in _schedule_dag_run

    schedulable_tis, callback_to_run = dag_run.update_state(session=session, execute_callbacks=False)

  File "/home/airflow/.local/lib/python3.9/site-packages/airflow/utils/session.py", line 68, in wrapper

    return func(*args, **kwargs)

  File "/home/airflow/.local/lib/python3.9/site-packages/airflow/models/dagrun.py", line 524, in update_state

    info = self.task_instance_scheduling_decisions(session)

  File "/home/airflow/.local/lib/python3.9/site-packages/airflow/utils/session.py", line 68, in wrapper

    return func(*args, **kwargs)

  File "/home/airflow/.local/lib/python3.9/site-packages/airflow/models/dagrun.py", line 654, in task_instance_scheduling_decisions

    schedulable_tis, changed_tis, expansion_happened = self._get_ready_tis(

  File "/home/airflow/.local/lib/python3.9/site-packages/airflow/models/dagrun.py", line 710, in _get_ready_tis

    expanded_tis, _ = schedulable.task.expand_mapped_task(self.run_id, session=session)

  File "/home/airflow/.local/lib/python3.9/site-packages/airflow/models/mappedoperator.py", line 614, in expand_mapped_task

    operator.mul, self._resolve_map_lengths(run_id, session=session).values()

  File "/home/airflow/.local/lib/python3.9/site-packages/airflow/models/mappedoperator.py", line 600, in _resolve_map_lengths

    raise RuntimeError(f"Failed to populate all mapping metadata; missing: {keys}")

RuntimeError: Failed to populate all mapping metadata; missing: 'x'
```

### What you think should happen instead

Xcoms from different task groups should be mappable within other group scopes.

### How to reproduce

```
from airflow import DAG
from airflow.decorators import task
from airflow.utils.task_group import TaskGroup

import pendulum

@task
def enumerate(x):
    return [i for i in range(x)]

@task 
def addOne(x):
    return x+1

with DAG(
    dag_id="TaskGroupMappingBug",
    schedule_interval=None,
    start_date=pendulum.now().subtract(days=1),
) as dag:
    
    with TaskGroup(group_id="enumerateNine"):
        y = enumerate(9)
        
    with TaskGroup(group_id="add"):
        # airflow scheduler throws error here so this is never reached
        z = addOne.expand(x=y)
    

```

### Operating System

linux/amd64 via Docker (apache/airflow:2.3.3-python3.9)

### Versions of Apache Airflow Providers

_No response_

### Deployment

Docker-Compose

### Deployment details

docker-compose version 1.29.2, build 5becea4c

Docker Engine v20.10.14


### Anything else

_No response_

### Are you willing to submit PR?

- [ ] 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)



Discussion:
Thanks for opening your first issue here! Be sure to follow the issue template!


Resolving patch:
diff --git a/airflow/models/abstractoperator.py b/airflow/models/abstractoperator.py
--- a/airflow/models/abstractoperator.py
+++ b/airflow/models/abstractoperator.py
@@ -26,6 +26,7 @@
     Dict,
     FrozenSet,
     Iterable,
+    Iterator,
     List,
     Optional,
     Sequence,
@@ -54,6 +55,7 @@
 
     from airflow.models.baseoperator import BaseOperator, BaseOperatorLink
     from airflow.models.dag import DAG
+    from airflow.models.mappedoperator import MappedOperator
     from airflow.models.operator import Operator
     from airflow.models.taskinstance import TaskInstance
 
@@ -231,6 +233,70 @@ def get_flat_relatives(self, upstream: bool = False) -> Collection["Operator"]:
             return set()
         return [dag.task_dict[task_id] for task_id in self.get_flat_relative_ids(upstream)]
 
+    def _iter_all_mapped_downstreams(self) -> Iterator["MappedOperator"]:
+        """Return mapped nodes that are direct dependencies of the current task.
+
+        For now, this walks the entire DAG to find mapped nodes that has this
+        current task as an upstream. We cannot use ``downstream_list`` since it
+        only contains operators, not task groups. In the future, we should
+        provide a way to record an DAG node's all downstream nodes instead.
+
+        Note that this does not guarantee the returned tasks actually use the
+        current task for task mapping, but only checks those task are mapped
+        operators, and are downstreams of the current task.
+
+        To get a list of tasks that uses the current task for task mapping, use
+        :meth:`iter_mapped_dependants` instead.
+        """
+        from airflow.models.mappedoperator import MappedOperator
+        from airflow.utils.task_group import TaskGroup
+
+        def _walk_group(group: TaskGroup) -> Iterable[Tuple[str, DAGNode]]:
+            """Recursively walk children in a task group.
+
+            This yields all direct children (including both tasks and task
+            groups), and all children of any task groups.
+            """
+            for key, child in group.children.items():
+                yield key, child
+                if isinstance(child, TaskGroup):
+                    yield from _walk_group(child)
+
+        dag = self.get_dag()
+        if not dag:
+            raise RuntimeError("Cannot check for mapped dependants when not attached to a DAG")
+        for key, child in _walk_group(dag.task_group):
+            if key == self.node_id:
+                continue
+            if not isinstance(child, MappedOperator):
+                continue
+            if self.node_id in child.upstream_task_ids:
+                yield child
+
+    def iter_mapped_dependants(self) -> Iterator["MappedOperator"]:
+        """Return mapped nodes that depend on the current task the expansion.
+
+        For now, this walks the entire DAG to find mapped nodes that has this
+        current task as an upstream. We cannot use ``downstream_list`` since it
+        only contains operators, not task groups. In the future, we should
+        provide a way to record an DAG node's all downstream nodes instead.
+        """
+        return (
+            downstream
+            for downstream in self._iter_all_mapped_downstreams()
+            if any(p.node_id == self.node_id for p in downstream.iter_mapped_dependencies())
+        )
+
+    def unmap(self, resolve: Union[None, Dict[str, Any], Tuple[Context, "Session"]]) -> "BaseOperator":
+        """Get the "normal" operator from current abstract operator.
+
+        MappedOperator uses this to unmap itself based on the map index. A non-
+        mapped operator (i.e. BaseOperator subclass) simply returns itself.
+
+        :meta private:
+        """
+        raise NotImplementedError()
+
     @property
     def priority_weight_total(self) -> int:
         """
@@ -291,16 +357,6 @@ def global_operator_extra_link_dict(self) -> Dict[str, Any]:
     def extra_links(self) -> List[str]:
         return list(set(self.operator_extra_link_dict).union(self.global_operator_extra_link_dict))
 
-    def unmap(self, resolve: Union[None, Dict[str, Any], Tuple[Context, "Session"]]) -> "BaseOperator":
-        """Get the "normal" operator from current abstract operator.
-
-        MappedOperator uses this to unmap itself based on the map index. A non-
-        mapped operator (i.e. BaseOperator subclass) simply returns itself.
-
-        :meta private:
-        """
-        raise NotImplementedError()
-
     def get_extra_links(self, ti: "TaskInstance", link_name: str) -> Optional[str]:
         """For an operator, gets the URLs that the ``extra_links`` entry points to.
 
diff --git a/airflow/models/taskmixin.py b/airflow/models/taskmixin.py
--- a/airflow/models/taskmixin.py
+++ b/airflow/models/taskmixin.py
@@ -17,7 +17,7 @@
 
 import warnings
 from abc import ABCMeta, abstractmethod
-from typing import TYPE_CHECKING, Any, Iterable, Iterator, List, Optional, Sequence, Set, Tuple, Union
+from typing import TYPE_CHECKING, Any, Iterable, List, Optional, Sequence, Set, Tuple, Union
 
 import pendulum
 
@@ -28,7 +28,6 @@
     from logging import Logger
 
     from airflow.models.dag import DAG
-    from airflow.models.mappedoperator import MappedOperator
     from airflow.utils.edgemodifier import EdgeModifier
     from airflow.utils.task_group import TaskGroup
 
@@ -273,57 +272,3 @@ def get_direct_relatives(self, upstream: bool = False) -> Iterable["DAGNode"]:
     def serialize_for_task_group(self) -> Tuple[DagAttributeTypes, Any]:
         """This is used by SerializedTaskGroup to serialize a task group's content."""
         raise NotImplementedError()
-
-    def _iter_all_mapped_downstreams(self) -> Iterator["MappedOperator"]:
-        """Return mapped nodes that are direct dependencies of the current task.
-
-        For now, this walks the entire DAG to find mapped nodes that has this
-        current task as an upstream. We cannot use ``downstream_list`` since it
-        only contains operators, not task groups. In the future, we should
-        provide a way to record an DAG node's all downstream nodes instead.
-
-        Note that this does not guarantee the returned tasks actually use the
-        current task for task mapping, but only checks those task are mapped
-        operators, and are downstreams of the current task.
-
-        To get a list of tasks that uses the current task for task mapping, use
-        :meth:`iter_mapped_dependants` instead.
-        """
-        from airflow.models.mappedoperator import MappedOperator
-        from airflow.utils.task_group import TaskGroup
-
-        def _walk_group(group: TaskGroup) -> Iterable[Tuple[str, DAGNode]]:
-            """Recursively walk children in a task group.
-
-            This yields all direct children (including both tasks and task
-            groups), and all children of any task groups.
-            """
-            for key, child in group.children.items():
-                yield key, child
-                if isinstance(child, TaskGroup):
-                    yield from _walk_group(child)
-
-        tg = self.task_group
-        if not tg:
-            raise RuntimeError("Cannot check for mapped dependants when not attached to a DAG")
-        for key, child in _walk_group(tg):
-            if key == self.node_id:
-                continue
-            if not isinstance(child, MappedOperator):
-                continue
-            if self.node_id in child.upstream_task_ids:
-                yield child
-
-    def iter_mapped_dependants(self) -> Iterator["MappedOperator"]:
-        """Return mapped nodes that depend on the current task the expansion.
-
-        For now, this walks the entire DAG to find mapped nodes that has this
-        current task as an upstream. We cannot use ``downstream_list`` since it
-        only contains operators, not task groups. In the future, we should
-        provide a way to record an DAG node's all downstream nodes instead.
-        """
-        return (
-            downstream
-            for downstream in self._iter_all_mapped_downstreams()
-            if any(p.node_id == self.node_id for p in downstream.iter_mapped_dependencies())
-        )


Test patch:
diff --git a/tests/models/test_baseoperator.py b/tests/models/test_baseoperator.py
--- a/tests/models/test_baseoperator.py
+++ b/tests/models/test_baseoperator.py
@@ -833,3 +833,24 @@ def _do_render():
         assert expected_log in caplog.text
     if not_expected_log:
         assert not_expected_log not in caplog.text
+
+
+def test_find_mapped_dependants_in_another_group(dag_maker):
+    from airflow.utils.task_group import TaskGroup
+
+    @task_decorator
+    def gen(x):
+        return list(range(x))
+
+    @task_decorator
+    def add(x, y):
+        return x + y
+
+    with dag_maker():
+        with TaskGroup(group_id="g1"):
+            gen_result = gen(3)
+        with TaskGroup(group_id="g2"):
+            add_result = add.partial(y=1).expand(x=gen_result)
+
+    dependants = list(gen_result.operator.iter_mapped_dependants())
+    assert dependants == [add_result.operator]