Repository: apache/airflow

Issue:
Dynamic task mapping creates too many mapped instances when task pushed non-default XCom
### Apache Airflow version

2.3.0 (latest released)

### What happened

Excess tasks are created when using dynamic task mapping with KubernetesPodOperator, but only in certain cases which I do not understand. I have a simple working example of this where the flow is:
- One task that returns a list XCom (list of lists, since I'm partial-ing to `KubernetesPodOperator`'s `arguments`) of length 3. This looks like `[["a"], ["b"], ["c"]]`
- A `partial` from this, which is expanded on the above's result. Each resulting task has an XCom of a single element list that looks like `["d"]`. We expect the `expand` to result in 3 tasks, which it does. So far so good. Why doesn't the issue occur at this stage? No clue.
- A `partial` from the above. We expect 3 tasks in this final stage, but get 9. 3 succeed and 6 fail consistently. This 3x rule scales to as many tasks as you define in step 2 (e.g. 2 tasks in step 2 -> 6 tasks in step 3, where 4 fail)

![image](https://user-images.githubusercontent.com/71299310/169179360-d1ddfe49-f20e-4f27-909f-4dd101386a5a.png)

If I recreate this using the TaskFlow API with `PythonOperator`s, I get the expected result of 1 task -> 3 tasks -> 3 tasks

![image](https://user-images.githubusercontent.com/71299310/169179409-d2f71f3e-6e8c-42e0-8120-1ecccff439c0.png)

Futhermore, if I attempt to look at the `Rendered Template` of the failed tasks in the `KubernetesPodOperator` implementation (first image), I consistently get `Error rendering template` and all the fields are `None`. The succeeded tasks look normal.

![image](https://user-images.githubusercontent.com/71299310/169181052-8d182722-197b-44a5-b145-3a983e259036.png)

Since the `Rendered Template` view fails to load, I can't confirm what is actually getting provided to these failing tasks' `argument` parameter. If there's a way I can query the meta database to see this, I'd be glad to if given instruction. 

### What you think should happen instead

I think this has to do with how XComs are specially handled with the `KubernetesPodOperator`. If we look at the XComs tab of the upstream task (`some-task-2` in the above images), we see that the return value specifies `pod_name` and `pod_namespace` along with `return_value`.

![image](https://user-images.githubusercontent.com/71299310/169179724-984682d0-c2fc-4097-9527-fa3cbf3ad93f.png)

Whereas in the `t2` task of the TaskFlow version, it only contains `return_value`. 

![image](https://user-images.githubusercontent.com/71299310/169179983-b22347de-eef0-4a9a-ae85-6b5e5d2bfa42.png)

I haven't dug through the code to verify, but I have a strong feeling these extra values `pod_name` and `pod_namespace` are being used to generate the `OperatorPartial`/`MappedOperator` as well when they shouldn't be.

### How to reproduce

Run this DAG in a k8s context:

```
from datetime import datetime
from airflow import XComArg
from airflow.models import DAG
from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator


def make_operator(
    **kwargs
):
    return KubernetesPodOperator(
        **{
            'get_logs': True,
            'in_cluster': True,
            'is_delete_operator_pod': True,
            'namespace': 'default',
            'startup_timeout_seconds': 600,
            **kwargs,
        }
    )


def make_partial_operator(
    **kwargs
):
    return KubernetesPodOperator.partial(
        **{
            'get_logs': True,
            'in_cluster': True,
            'is_delete_operator_pod': True,
            'namespace': 'default',
            'startup_timeout_seconds': 600,
            **kwargs,
        }
    )


with DAG(dag_id='test-pod-xcoms',
         schedule_interval=None,
         start_date=datetime(2020, 1, 1),
         max_active_tasks=20) as dag:

    op1 = make_operator(
        cmds=['python3', '-c' 'import json;f=open("/airflow/xcom/return.json", "w");f.write(json.dumps([["a"], ["b"], ["c"]]))'],
        image='python:3.9-alpine',
        name='airflow-private-image-pod-1',
        task_id='some-task-1',
        do_xcom_push=True
    )

    op2 = make_partial_operator(
        cmds=['python3', '-c' 'import json;f=open("/airflow/xcom/return.json", "w");f.write(json.dumps(["d"]))'],
        image='python:3.9-alpine',
        name='airflow-private-image-pod-2',
        task_id='some-task-2',
        do_xcom_push=True
    )

    op3 = make_partial_operator(
        cmds=['echo', 'helloworld'],
        image='alpine:latest',
        name='airflow-private-image-pod-3',
        task_id='some-task-3',
    )

    op3.expand(arguments=XComArg(op2.expand(arguments=XComArg(op1))))
```

For the TaskFlow version of this that works, run this DAG (doesn't have to be k8s context):

```
from datetime import datetime

from airflow.decorators import task
from airflow.models import DAG, Variable


@task
def t1():
    return [[1], [2], [3]]


@task
def t2(val):
    return val


@task
def t3(val):
    print(val)


with DAG(dag_id='test-mapping',
         schedule_interval=None,
         start_date=datetime(2020, 1, 1)) as dag:

    t3.partial().expand(val=t2.partial().expand(val=t1()))
```

### Operating System

MacOS 11.6.5

### Versions of Apache Airflow Providers

Relevant:

```
apache-airflow-providers-cncf-kubernetes==4.0.1
```

Full:

```
apache-airflow-providers-amazon==3.3.0
apache-airflow-providers-celery==2.1.4
apache-airflow-providers-cncf-kubernetes==4.0.1
apache-airflow-providers-docker==2.6.0
apache-airflow-providers-elasticsearch==3.0.3
apache-airflow-providers-ftp==2.1.2
apache-airflow-providers-google==6.8.0
apache-airflow-providers-grpc==2.0.4
apache-airflow-providers-hashicorp==2.2.0
apache-airflow-providers-http==2.1.2
apache-airflow-providers-imap==2.2.3
apache-airflow-providers-microsoft-azure==3.8.0
apache-airflow-providers-mysql==2.2.3
apache-airflow-providers-odbc==2.0.4
apache-airflow-providers-postgres==4.1.0
apache-airflow-providers-redis==2.0.4
apache-airflow-providers-sendgrid==2.0.4
apache-airflow-providers-sftp==2.6.0
apache-airflow-providers-slack==4.2.3
apache-airflow-providers-sqlite==2.1.3
apache-airflow-providers-ssh==2.4.3
```

### Deployment

Official Apache Airflow Helm Chart

### Deployment details

Docker (Docker Desktop)

- Server Version: 20.10.13
- API Version: 1.41
- Builder: 2

Kubernetes (Docker Desktop)

- Env: docker-desktop
- Context: docker-desktop
- Cluster Name: docker-desktop
- Namespace: default
- Container Runtime: docker
- Version: v1.22.5

Helm:

- version.BuildInfo{Version:"v3.6.3", GitCommit:"d506314abfb5d21419df8c7e7e68012379db2354", GitTreeState:"dirty", GoVersion:"go1.16.5"}

### Anything else

_No response_

### 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)



Discussion:
I seem to recall there’s a bug on expanding a wrong number of tasks, but can’t find it now. This _may_ have been fixed in main.

Resolving patch:
diff --git a/airflow/models/expandinput.py b/airflow/models/expandinput.py
--- a/airflow/models/expandinput.py
+++ b/airflow/models/expandinput.py
@@ -96,7 +96,7 @@ def _get_map_lengths(self, run_id: str, *, session: Session) -> dict[str, int]:
         they will not be present in the dict.
         """
         from airflow.models.taskmap import TaskMap
-        from airflow.models.xcom import XCom
+        from airflow.models.xcom import XCOM_RETURN_KEY, XCom
         from airflow.models.xcom_arg import XComArg
 
         # Populate literal mapped arguments first.
@@ -143,6 +143,7 @@ def _get_map_lengths(self, run_id: str, *, session: Session) -> dict[str, int]:
             .filter(
                 XCom.dag_id == dag_id,
                 XCom.run_id == run_id,
+                XCom.key == XCOM_RETURN_KEY,
                 XCom.task_id.in_(mapped_dep_keys),
                 XCom.map_index >= 0,
             )


Test patch:
diff --git a/tests/models/test_mappedoperator.py b/tests/models/test_mappedoperator.py
--- a/tests/models/test_mappedoperator.py
+++ b/tests/models/test_mappedoperator.py
@@ -99,6 +99,43 @@ def test_map_xcom_arg():
     assert task1.downstream_list == [mapped]
 
 
+def test_map_xcom_arg_multiple_upstream_xcoms(dag_maker, session):
+    """Test that the correct number of downstream tasks are generated when mapping with an XComArg"""
+
+    class PushExtraXComOperator(BaseOperator):
+        """Push an extra XCom value along with the default return value."""
+
+        def __init__(self, return_value, **kwargs):
+            super().__init__(**kwargs)
+            self.return_value = return_value
+
+        def execute(self, context):
+            context['task_instance'].xcom_push(key='extra_key', value="extra_value")
+            return self.return_value
+
+    with dag_maker("test-dag", session=session, start_date=DEFAULT_DATE) as dag:
+        upstream_return = [1, 2, 3]
+        task1 = PushExtraXComOperator(return_value=upstream_return, task_id="task_1")
+        task2 = PushExtraXComOperator.partial(task_id='task_2').expand(return_value=XComArg(task1))
+        task3 = PushExtraXComOperator.partial(task_id='task_3').expand(return_value=XComArg(task2))
+
+    dr = dag_maker.create_dagrun()
+    ti_1 = dr.get_task_instance("task_1", session)
+    ti_1.run()
+
+    ti_2s, _ = task2.expand_mapped_task(dr.run_id, session=session)
+    for ti in ti_2s:
+        ti.refresh_from_task(dag.get_task("task_2"))
+        ti.run()
+
+    ti_3s, _ = task3.expand_mapped_task(dr.run_id, session=session)
+    for ti in ti_3s:
+        ti.refresh_from_task(dag.get_task("task_3"))
+        ti.run()
+
+    assert len(ti_3s) == len(ti_2s) == len(upstream_return)
+
+
 def test_partial_on_instance() -> None:
     """`.partial` on an instance should fail -- it's only designed to be called on classes"""
     with pytest.raises(TypeError):