Skip to content

Fix mapped task group XCom collapsing to a bare value for a single expansion#69075

Merged
shahar1 merged 6 commits into
apache:mainfrom
shahar1:fix-mapped-task-group-single-xcom-list
Jul 1, 2026
Merged

Fix mapped task group XCom collapsing to a bare value for a single expansion#69075
shahar1 merged 6 commits into
apache:mainfrom
shahar1:fix-mapped-task-group-single-xcom-list

Conversation

@shahar1

@shahar1 shahar1 commented Jun 27, 2026

Copy link
Copy Markdown
Contributor

closes: #69036

Human Summary

Currently, when map_indexes is left unset, a single mapped result collapses to its bare value instead of staying one-element list.
It started to happen after #60803 was merged, when the upstream map-index computation was relocated to the task-sdk.
This PR fixes it by returning an eagerly-materialized LazyXComSequence object.

AI Summary

Click here A task consuming the output of a mapped task group from **outside** that group must always receive a list with one element per expansion. When the task-start upstream map-index computation moved into the Task SDK (#60803, released in 3.2.0), an unmapped downstream resolving a task inside a mapped task group started falling through to pulling all map indexes via `xcom_pull`, which collapses a single value to a bare scalar and an empty set of values to `None`.

As a result:

This worked correctly in 3.1.5: the pre-3.2 server returned list(range(mapped_ti_count)) for this case, guaranteeing a list. The fix resolves this case through a LazyXComSequence, mirroring how a directly mapped task already resolves, so the result is always a list regardless of how many times the group expanded or whether every value was None.

Verified end to end with airflow dags test on the reproduction DAGs from both issues:


Was generative AI tooling used to co-author this PR?
  • Yes — Claude Code (Opus 4.8)

Generated-by: Claude Code (Opus 4.8) following the guidelines

@shahar1
shahar1 marked this pull request as draft June 27, 2026 06:54
@shahar1
shahar1 force-pushed the fix-mapped-task-group-single-xcom-list branch from cdf92e6 to 2b8c984 Compare June 27, 2026 07:06
@shahar1
shahar1 marked this pull request as ready for review June 27, 2026 07:08
@shahar1
shahar1 force-pushed the fix-mapped-task-group-single-xcom-list branch from 2b8c984 to a28c3f0 Compare June 27, 2026 11:32
Comment thread task-sdk/src/airflow/sdk/definitions/xcom_arg.py Outdated
shahar1 added 3 commits June 30, 2026 00:25
…pansion

A task consuming the output of a mapped task group from outside that group
must always receive a list with one element per expansion. Since the task-start
upstream map-index computation moved into the Task SDK, an unmapped downstream
resolving a task inside a mapped task group fell through to pulling all map
indexes via xcom_pull, which collapses a single value to a bare scalar and an
empty set of values to None. A group that expanded only once therefore handed
the raw value to the downstream task instead of a one-element list, and a group
whose expansions all returned None handed None instead of an empty list. The
pre-3.2 server returned list(range(mapped_ti_count)) for this case, so the
behaviour regressed.

Resolve this case by eagerly materialising the group's per-expansion XComs into
a list (a single slice request), so the result is always a list regardless of
how many times the group expanded or whether every value was None, and stays
serializable when a task returns it unchanged.
Eagerly materializing the whole mapped task group result at resolve time
fetches every expansion even when the downstream task never iterates it,
which scales poorly for large maps. Defer the fetch to XCom serialization
so it is only paid when the value is actually returned and pushed.
@shahar1
shahar1 force-pushed the fix-mapped-task-group-single-xcom-list branch from a28c3f0 to 3b856fd Compare June 29, 2026 21:38
Comment thread task-sdk/src/airflow/sdk/serde/__init__.py Outdated
@uranusjr

Copy link
Copy Markdown
Member

Worth checking whether XComArg still accepts ArgNotSet after this change. Maybe we can remove handling for that entirely?

shahar1 added 2 commits June 30, 2026 09:57
…gNotSet path

Route the lazily-resolved mapped task group value through a registered
serde serializer instead of a special case in serialize(), matching how
every other custom type is handled, so any returned lazy XCom sequence
materializes to a list when pushed. The resolve path no longer produces
an unset map-index sentinel, so the corresponding ArgNotSet handling is
removed.
The aggregated task group value now resolves to a lazy sequence, so the
test's supervisor-comms stub must answer single-item fetches, not only
slice fetches, when the resolved value is iterated.
Comment thread task-sdk/tests/task_sdk/definitions/test_xcom_arg.py Outdated
Comment thread task-sdk/tests/task_sdk/definitions/test_xcom_arg.py Outdated
@shahar1

shahar1 commented Jun 30, 2026

Copy link
Copy Markdown
Contributor Author

Thanks TP for approval.
Clarifying - issue #48005 is out of scope of this PR.

@shahar1
shahar1 merged commit 37f6259 into apache:main Jul 1, 2026
102 checks passed
@shahar1
shahar1 deleted the fix-mapped-task-group-single-xcom-list branch July 1, 2026 21:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Dynamic Mapped Task Group does not pack single mapped tasks in list

3 participants