Map length determined incorrectly when mapping over an output of a multi-output task

[Dev-iL]

Apache Airflow version

2.11.0

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

No response

What happened?

The mapping length in expandinput is determined incorrectly to be the number of returns from a task, instead of the number of elements in the sequence we map on (which is one of the values returned by the aforementioned task).

Note how all_lengths is always 2 regardless of the contents of value:

Depending on the relative size of the determined and actual map size, the dag either reports the error below, works correctly, or ends up not processing the later elements.

.../.venv/lib64/python3.11/site-packages/airflow/models/expandinput.py", line 197, in _expand_mapped_field
    return value[found_index]
           ~~~~~^^^^^^^^^^^^^
IndexError: list index out of range

What you think should happen instead?

all_length should correctly reflect the number of elements we're mapping on.

How to reproduce

from __future__ import annotations

import pendulum
from airflow.decorators import task, task_group
from airflow.models.dag import DAG

TYPE_CHECKING: bool = False
if TYPE_CHECKING:
    from datetime import datetime
    from typing import Any


@task.python(multiple_outputs=True)
def determine_run_params() -> dict[str, Any]:
    return {
        "targets": [f"target{n}" for n in range(1)],  # <<<<<<<<<< Modify for different results
        "date": pendulum.now(),
        # "str1": "1",  # Uncomment these for different results
        # "str2": "2",
        # "str3": "3",
    }


@task.python
def do_something(
    target: str,
    date: datetime,
) -> str:
    return f"All good for {target} on {date}!"


@task_group
def process(
    date: datetime,
    target: str,
) -> None:
    do_something(target, date)


with DAG(
    "mcve",
    start_date=pendulum.datetime(2020, 1, 1),
) as dag:
    run_params = determine_run_params()
    process.partial(date=run_params["date"]).expand(target=run_params["targets"])


if __name__ == "__main__":
    dag.test()

Workaround

Introduce an intermediate passthrough task:

@task.python
def unpack(something: list[str]) -> list[str]:
    return something

- process.partial(date=run_params["date"]).expand(target=run_params["targets"])
+ process.partial(date=run_params["date"]).expand(target=unpack(run_params["targets"])) 

Operating System

Linux

Versions of Apache Airflow Providers

No response

Deployment

Other

Deployment details

No response

Anything else?

When the number of map elements (length of targets) is equal to the number of outputs from determine_run_params, the code works correctly. Perhaps Airflow's tests related to this functionality are hitting this exact sweet spot.

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

[jroachgolf84]

Just throwing a comment in here to stay up-to-date, this is one I'd be curious to try and reproduce.

[jroachgolf84]

I'm able to reproduce this exception. The # of dynamically mapped Tasks within a TaskGroup matches the # of elements in the dictionary returned by determine_run_params. Just to validate, you're trying to dynamically map TaskGroups (i.e. spin up 10 TaskGroups, each with 1 Task) or dynamically map Tasks (i.e. 1 TaskGroup, each with 10 Tasks)?

[Dev-iL]

10 groups each with 1 task

[jroachgolf84]

Seems like the issue is parsing the MappedArgument(_input=DictOfListsExpandInput(value={'target': XComArg(<Task(_PythonDecoratedOperator): determine_run_params>, 'targets')}), _key='target'). I'm able to successfully get something to work when just passing a list to expand. However, this is where it gets a bit more interesting.

Here's my updated code (just for reproduction). This DAG actually won't parse, with the message below. From what I can tell, we cannot expand by trying to parse a key that returns a dictionary.

Broken DAG: [/usr/local/airflow/dags/issue_51109_solution.py] Traceback (most recent call last):
  File "/usr/local/lib/python3.12/site-packages/airflow/models/mappedoperator.py", line 137, in ensure_xcomarg_return_value
    ensure_xcomarg_return_value(v)
  File "/usr/local/lib/python3.12/site-packages/airflow/models/mappedoperator.py", line 132, 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 'targets' from <Task(_PythonDecoratedOperator): start>

from airflow.models import DAG
from airflow.decorators import task, task_group

from datetime import datetime

@task.python(multiple_outputs=True)
def start():
    return {"targets": ["0", "1", "2", "3", "4"]}

@task
def do_something(date, targets):
    print(f"date: {date}")
    print(f"targets: {targets}")

@task_group
def dynamic_task_group(date, targets):
    do_something(date, targets)

@task
def end():
    pass


with DAG(
    dag_id="issue_51109_solution",
    start_date=datetime(2025, 1, 1),
    schedule="@once",
    catchup=False,
) as dag:

    start_output = start()
    dynamic_task_group.partial(date="2025-01-01").expand(targets=start_output["targets"])
    do_something.partial(date="2025-01-01").expand(targets=start_output["targets"]) >> end()

[jroachgolf84]

In the meantime, you'll want to decouple your Task that's returning the date, and the one that's returning the targets. I think the change to the OSS project is going to be adding validation when parsing the DAG itself.

[Dev-iL]

(@jroachgolf84 I was under the impression that when the needs-triage tag is removed, it is replaced by the determined priority level - low/medium/..., no?)