You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "shahar1 (via GitHub)" <gi...@apache.org> on 2026/06/27 06:01:05 UTC
[PR] Fix mapped task group XCom collapsing to a bare value for a single expansion [airflow]
shahar1 opened a new pull request, #69075:
URL: https://github.com/apache/airflow/pull/69075
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:
- A mapped task group that expanded **only once** handed the raw return value to the downstream task instead of a one-element list (#69036 — e.g. the consumer received `'14'` and iterated `'1'`, `'4'` instead of receiving `['14']`).
- A mapped task group whose expansions **all returned `None`** would hand `None` instead of an empty list (the #48005 class of bug, via the task-group resolution path).
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:
- single expansion now yields `['14']` (was `'14'`),
- multiple expansions still yield the full list,
- all-`None` group expansions yield `[]`,
- the directly-mapped all-`None` case from #48005 continues to yield `[]`.
closes: #69036
related: #48005
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes — Claude Code (Opus 4.8)
Generated-by: Claude Code (Opus 4.8) following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix mapped task group XCom collapsing to a bare value for a single expansion [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #69075:
URL: https://github.com/apache/airflow/pull/69075#discussion_r3488532028
##########
task-sdk/src/airflow/sdk/definitions/xcom_arg.py:
##########
@@ -354,8 +354,11 @@ def resolve(self, context: Mapping[str, Any]) -> Any:
ti_count=ti_count,
session=None, # Not used in SDK implementation
)
- # None means "no filtering needed" -> use NOTSET to pull all values
- map_indexes = NOTSET if computed is None else computed
+ if computed is None:
+ # Aggregate the mapped task group as a list, even for a single expansion (#69036) or all-None values (#48005)
+ # Materialise eagerly (one slice request) so a task returning the value unchanged can still serialize it
+ return LazyXComSequence(xcom_arg=self, ti=ti)[:]
+ map_indexes = computed
Review Comment:
Hmm, this has the potential to kill performance is the map size is large. I wonder if it would be appropriate to return the lazy sequence directly, and handle it on XCom serialization _if_ it’s actually returned instead.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix mapped task group XCom collapsing to a bare value for a single expansion [airflow]
Posted by "shahar1 (via GitHub)" <gi...@apache.org>.
shahar1 merged PR #69075:
URL: https://github.com/apache/airflow/pull/69075
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix mapped task group XCom collapsing to a bare value for a single expansion [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #69075:
URL: https://github.com/apache/airflow/pull/69075#discussion_r3497078204
##########
task-sdk/tests/task_sdk/definitions/test_xcom_arg.py:
##########
@@ -350,3 +350,72 @@ def xcom_get(msg):
states = [run_ti(dag, "pull_one", map_index) for map_index in range(5)]
assert states == [TaskInstanceState.SUCCESS] * 5
assert agg_results == {"a", "b", "c", 1, 2}
+
+
+class TestPlainXComArgResolveMappedGroup:
+ """Resolving a task inside a mapped task group from a task outside that group.
+
+ Regression tests for #69036 and #48005: the combined return value of a
+ mapped task group must always serialize to a list (one element per
+ expansion), even when the group expanded only once or every expansion
+ returned ``None``. Previously this case was routed through ``xcom_pull``
+ pulling all map indexes, which collapsed a single value to a bare scalar
+ and an empty set of values to ``None``. ``resolve`` stays lazy and serde
+ materialises the sequence only when the value is actually returned.
+ """
+
+ @staticmethod
+ def _make_ti(*, computed):
+ ti = mock.MagicMock()
+ ti._upstream_map_indexes = None
+ ti._cached_template_context = {"expanded_ti_count": 1}
+ ti.run_id = "run-1"
+ ti.get_relevant_upstream_map_indexes.return_value = computed
+ return ti
+
+ @staticmethod
+ def _make_arg():
+ from airflow.sdk.definitions.xcom_arg import PlainXComArg
Review Comment:
I think we can move it to the top level
##########
task-sdk/tests/task_sdk/definitions/test_xcom_arg.py:
##########
@@ -350,3 +350,72 @@ def xcom_get(msg):
states = [run_ti(dag, "pull_one", map_index) for map_index in range(5)]
assert states == [TaskInstanceState.SUCCESS] * 5
assert agg_results == {"a", "b", "c", 1, 2}
+
+
+class TestPlainXComArgResolveMappedGroup:
+ """Resolving a task inside a mapped task group from a task outside that group.
+
+ Regression tests for #69036 and #48005: the combined return value of a
+ mapped task group must always serialize to a list (one element per
+ expansion), even when the group expanded only once or every expansion
+ returned ``None``. Previously this case was routed through ``xcom_pull``
+ pulling all map indexes, which collapsed a single value to a bare scalar
+ and an empty set of values to ``None``. ``resolve`` stays lazy and serde
+ materialises the sequence only when the value is actually returned.
+ """
+
+ @staticmethod
+ def _make_ti(*, computed):
+ ti = mock.MagicMock()
+ ti._upstream_map_indexes = None
+ ti._cached_template_context = {"expanded_ti_count": 1}
+ ti.run_id = "run-1"
+ ti.get_relevant_upstream_map_indexes.return_value = computed
+ return ti
+
+ @staticmethod
+ def _make_arg():
+ from airflow.sdk.definitions.xcom_arg import PlainXComArg
+
+ operator = mock.MagicMock()
+ operator.is_mapped = False
+ operator.task_id = "do_something"
+ operator.dag_id = "test_dag"
+ operator.get_closest_mapped_task_group.return_value = mock.MagicMock()
+ return PlainXComArg(operator=operator, key="test")
+
+ @pytest.mark.parametrize(
+ ("root", "expected"),
+ [
+ pytest.param(["14"], ["14"], id="single-expansion-stays-a-list"),
+ pytest.param([], [], id="all-none-expansions-give-empty-list"),
+ pytest.param(["a", "b"], ["a", "b"], id="multiple-expansions"),
+ ],
+ )
+ def test_resolve_stays_lazy_and_serializes_as_list(self, root, expected, mock_supervisor_comms):
+ from airflow.sdk.execution_time.comms import XComSequenceSliceResult
Review Comment:
same here
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix mapped task group XCom collapsing to a bare value for a single expansion [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on PR #69075:
URL: https://github.com/apache/airflow/pull/69075#issuecomment-4840405849
Worth checking whether XComArg still accepts ArgNotSet after this change. Maybe we can remove handling for that entirely?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Fix mapped task group XCom collapsing to a bare value for a single expansion [airflow]
Posted by "shahar1 (via GitHub)" <gi...@apache.org>.
shahar1 commented on PR #69075:
URL: https://github.com/apache/airflow/pull/69075#issuecomment-4841300446
Thanks TP for approval.
Clarifying - issue #48005 is out of scope of this PR.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org