Cannot use custom XCOM keys in task mapping expand function

[qcha41]

Apache Airflow version

2.4.0

What happened

I am trying to use a custom XCOM key in task mapping, other than the default "return_value" key. This results in a ValueError :

ValueError: cannot map over XCom with custom key [...]

What you think should happen instead

I don't understand this limitation.

How to reproduce

import datetime as dt

from airflow import DAG, XComArg
from airflow.operators.trigger_dagrun import TriggerDagRunOperator
from airflow.operators.python import PythonOperator

with DAG(
    'dag_test_v1.0.0',
    start_date = dt.datetime(2022, 10, 12),
) as dag:

    def produce_parameters(**kwargs):
        kwargs['ti'].xcom_push("my_custom_key", [{"message":1}, {"message":2}])

    t1 = PythonOperator(
        task_id='task1', 
        python_callable=produce_parameters, 
        provide_context=True
    )

    t2 = TriggerDagRunOperator.partial(
        task_id = 'task2',
        trigger_dag_id = 'dag_target_v1.0.0'
    ).expand(
        conf=XComArg(t1, 'my_custom_key')
    )

    t1 >> t2

Traceback (most recent call last):
  File "dag_test.py", line 17, in <module>
    t2 = TriggerDagRunOperator.partial(
  File "/usr/local/lib/python3.10/site-packages/airflow/models/mappedoperator.py", line 179, in expand
    return self._expand(DictOfListsExpandInput(mapped_kwargs), strict=False)
  File "/usr/local/lib/python3.10/site-packages/airflow/models/mappedoperator.py", line 196, in _expand
    ensure_xcomarg_return_value(expand_input.value)
  File "/usr/local/lib/python3.10/site-packages/airflow/models/mappedoperator.py", line 133, in ensure_xcomarg_return_value
    ensure_xcomarg_return_value(v)
  File "/usr/local/lib/python3.10/site-packages/airflow/models/mappedoperator.py", line 128, in ensure_xcomarg_return_value
    raise ValueError(f"cannot map over XCom with custom key {key!r} from {operator}")
ValueError: cannot map over XCom with custom key 'my_custom_key' from <Task(PythonOperator): python_task>

Operating System

Debian 11

Versions of Apache Airflow Providers

No response

Deployment

Virtualenv installation

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!

[uranusjr]

This is an intentional design at the current time. Is there a concrete reason why you need to use a custom key?

[potiuk]

Converted it into discussion - because this is what it is (and also while waiting gor @qcha41 answers).

[dinigo]

I have several Google Workflows executions I have some resources batch-transformed with it

    run_workflow = WorkflowsCreateExecutionOperator\
    .partial(
        task_id="run_workflow",
        location=LOCATION,
        project_id=PROJECT_ID,
        workflow_id=WORKFLOW_ID,
    )\
    .expand(
        execution={
            "argument": {
                "some_params": ["a", "b"],
                # ...
                "dependant_on_previous": XComArg(previous_task),
                "pipeline_name": PIPELINE_NAME,
            }
        },
    )

    await_workflow_complete = WorkflowExecutionSensor\
    .partial(
        task_id="await_workflow_complete",
        location=LOCATION,
        project_id=PROJECT_ID,
        workflow_id=WORKFLOW_ID,
    )\
    .expand(
        execution_id=XComArg(run_workflow, 'execution_id'),
    )

I need several sensors, one per each workflow triggered, The WorkflowsCreateExecutionOperator XCom stores the execution_id for each exec, but it uses a different key than the default one (It actually stores some more data)

Also, would this "dependant_on_previous": XComArg(previous_task) line be ok? Not sure if this is the right way. Is this dictionary automagically expanded into a list of dicts?

[dinigo]

Also, why is it the expected behavior? Couldn't find it in the docs nor think of a reason