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[bot]]

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

[tagoad]

Ran into this today. I have a task that gets a list of data and sorts it into two buckets. Each bucket has its own downstream task, it would be a lot cleaner to have the downstream task map off its specific bucket that is saved under a custom xcom key instead of needing to add additional filtering before mapping out.
I'm curious the functional reasons why to limit mapping to only return_value since we can specify custom keys and the XComArg can accept custom keys. Why does expand need to care what key it is?

[leodrivera]

I came across this today too. I have the same question as the others. Why is there this restriction regarding custom keys?

[pykenny]

From how a task instance decides what kind of downstream mapped tasks to be dynamically created, the decision to force using "return_value" in the XCom is based on the fact that a task instance uses value returned from the Python callable within the operator, instead of full XCom content of the task instance, to decide what should be mapped to the downstream. For most decorated tasks I think it might not be the problem, since what's returned from the function is exactly how the whole XCom output should look like, but for traditional task objects base on BaseOperator and decorated tasks that call TaskInstnace.xcom_push() internally, extra values can be added to XCom output during task execution and these key-value pairs won't be available to the mapping process.

That said, adding extra logic to pull full XCom output from the backbone database again may help with this use case, but that depends on when XCom output of a task is "determined/stablized", and whether it's reasonable/feasible to do that from current task execution workflow design.