From 7c0398f97fd0550504afb7823618e874085bb469 Mon Sep 17 00:00:00 2001 From: Shahar Epstein <60007259+shahar1@users.noreply.github.com> Date: Sat, 27 Jun 2026 10:04:53 +0300 Subject: [PATCH 1/6] Fix mapped task group XCom collapsing to a bare value for a single expansion 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. --- .../src/airflow/sdk/definitions/xcom_arg.py | 7 +- .../task_sdk/definitions/test_xcom_arg.py | 66 +++++++++++++++++++ 2 files changed, 71 insertions(+), 2 deletions(-) diff --git a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py index 595490832597f..880517be4f498 100644 --- a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py +++ b/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 result = ti.xcom_pull( task_ids=task_id, key=self.key, diff --git a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py index 6014d26f2208b..623c00ff8c98f 100644 --- a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py +++ b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py @@ -350,3 +350,69 @@ 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 be an eager 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``. The list must be materialised eagerly so a task that + returns the value unchanged can still push it to XCom. + """ + + @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_aggregates_mapped_group_as_eager_list(self, root, expected, mock_supervisor_comms): + from airflow.sdk.execution_time.comms import XComSequenceSliceResult + + mock_supervisor_comms.send.return_value = XComSequenceSliceResult(root=root) + + arg = self._make_arg() + ti = self._make_ti(computed=None) + + resolved = arg.resolve({"ti": ti}) + + assert resolved == expected + assert isinstance(resolved, list) + ti.xcom_pull.assert_not_called() + + def test_resolve_uses_xcom_pull_for_specific_index(self): + arg = self._make_arg() + ti = self._make_ti(computed=0) + ti.xcom_pull.return_value = "value-0" + + resolved = arg.resolve({"ti": ti}) + + assert resolved == "value-0" + ti.xcom_pull.assert_called_once() + assert ti.xcom_pull.call_args.kwargs["map_indexes"] == 0 From ec1f78199e5cd5c557f92f92c4e7af6fe4b51576 Mon Sep 17 00:00:00 2001 From: Shahar Epstein <60007259+shahar1@users.noreply.github.com> Date: Sat, 27 Jun 2026 10:05:06 +0300 Subject: [PATCH 2/6] Add newsfragment for mapped task group XCom list fix --- airflow-core/newsfragments/69075.bugfix.rst | 1 + 1 file changed, 1 insertion(+) create mode 100644 airflow-core/newsfragments/69075.bugfix.rst diff --git a/airflow-core/newsfragments/69075.bugfix.rst b/airflow-core/newsfragments/69075.bugfix.rst new file mode 100644 index 0000000000000..d992083c37820 --- /dev/null +++ b/airflow-core/newsfragments/69075.bugfix.rst @@ -0,0 +1 @@ +Fix mapped task group return value handed to a downstream task being a bare value instead of a one-element list when the group expanded only once (and ``None`` instead of an empty list when every expansion returned ``None``). From 3b856fd7628c3f3f3b315c044f30f4869f1c4b6e Mon Sep 17 00:00:00 2001 From: Shahar Epstein <60007259+shahar1@users.noreply.github.com> Date: Tue, 30 Jun 2026 00:37:14 +0300 Subject: [PATCH 3/6] Resolve mapped task group XCom lazily and materialize on serialization 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. --- .../src/airflow/sdk/definitions/xcom_arg.py | 7 ++++--- task-sdk/src/airflow/sdk/serde/__init__.py | 8 ++++++++ .../task_sdk/definitions/test_xcom_arg.py | 20 ++++++++++--------- 3 files changed, 23 insertions(+), 12 deletions(-) diff --git a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py index 880517be4f498..4e6e1ed06644a 100644 --- a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py +++ b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py @@ -355,9 +355,10 @@ def resolve(self, context: Mapping[str, Any]) -> Any: session=None, # Not used in SDK implementation ) 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)[:] + # Resolve the mapped task group as a list, even for a single expansion (#69036) + # or all-None values (#48005). Stay lazy: serde materialises the sequence only if + # the value is actually returned and pushed to XCom (see airflow.sdk.serde.serialize). + return LazyXComSequence(xcom_arg=self, ti=ti) map_indexes = computed result = ti.xcom_pull( task_ids=task_id, diff --git a/task-sdk/src/airflow/sdk/serde/__init__.py b/task-sdk/src/airflow/sdk/serde/__init__.py index 0d94ebc328238..274fcaf494342 100644 --- a/task-sdk/src/airflow/sdk/serde/__init__.py +++ b/task-sdk/src/airflow/sdk/serde/__init__.py @@ -254,6 +254,14 @@ def serialize(o: object, depth: int = 0) -> U | None: return o + # Materialise a lazily-resolved mapped XCom (e.g. mapped task group output) to a list here; + # otherwise the attrs branch below would serialize its internal fields. The single slice fetch + # happens only if the value is actually returned and serialized. + from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence + + if isinstance(o, LazyXComSequence): + return serialize(o[:], depth + 1) + # custom serializers dct = { CLASSNAME: qn, diff --git a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py index 623c00ff8c98f..6272d739e585a 100644 --- a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py +++ b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py @@ -356,12 +356,12 @@ 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 be an eager 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``. The list must be materialised eagerly so a task that - returns the value unchanged can still push it to XCom. + 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 @@ -392,8 +392,10 @@ def _make_arg(): pytest.param(["a", "b"], ["a", "b"], id="multiple-expansions"), ], ) - def test_resolve_aggregates_mapped_group_as_eager_list(self, root, expected, mock_supervisor_comms): + def test_resolve_stays_lazy_and_serializes_as_list(self, root, expected, mock_supervisor_comms): from airflow.sdk.execution_time.comms import XComSequenceSliceResult + from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence + from airflow.sdk.serde import serialize mock_supervisor_comms.send.return_value = XComSequenceSliceResult(root=root) @@ -402,9 +404,9 @@ def test_resolve_aggregates_mapped_group_as_eager_list(self, root, expected, moc resolved = arg.resolve({"ti": ti}) - assert resolved == expected - assert isinstance(resolved, list) + assert isinstance(resolved, LazyXComSequence) ti.xcom_pull.assert_not_called() + assert serialize(resolved) == expected def test_resolve_uses_xcom_pull_for_specific_index(self): arg = self._make_arg() From eba2697024fe2c25d83d19f5f4a254e6eaa7a71c Mon Sep 17 00:00:00 2001 From: Shahar Epstein <60007259+shahar1@users.noreply.github.com> Date: Tue, 30 Jun 2026 09:57:38 +0300 Subject: [PATCH 4/6] Handle mapped task group XCom via a serde serializer and drop dead ArgNotSet 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. --- .../src/airflow/sdk/definitions/xcom_arg.py | 4 +- task-sdk/src/airflow/sdk/serde/__init__.py | 8 ---- .../serde/serializers/lazy_xcom_sequence.py | 46 +++++++++++++++++++ .../task_sdk/definitions/test_xcom_arg.py | 5 +- 4 files changed, 51 insertions(+), 12 deletions(-) create mode 100644 task-sdk/src/airflow/sdk/serde/serializers/lazy_xcom_sequence.py diff --git a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py index 4e6e1ed06644a..a38e47a466bcf 100644 --- a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py +++ b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py @@ -30,7 +30,7 @@ from airflow.sdk.definitions._internal.abstractoperator import AbstractOperator from airflow.sdk.definitions._internal.mixins import DependencyMixin, ResolveMixin from airflow.sdk.definitions._internal.setup_teardown import SetupTeardownContext -from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet, is_arg_set +from airflow.sdk.definitions._internal.types import NOTSET, is_arg_set from airflow.sdk.exceptions import AirflowException, XComNotFound from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence from airflow.sdk.execution_time.xcom import BaseXCom @@ -338,7 +338,7 @@ def resolve(self, context: Mapping[str, Any]) -> Any: tg = self.operator.get_closest_mapped_task_group() if tg is None: # No mapped task group - pull from unmapped instance - map_indexes: int | range | None | ArgNotSet = None + map_indexes: int | range | None = None else: # Check for pre-computed value from server (backward compatibility) upstream_map_indexes = getattr(ti, "_upstream_map_indexes", None) diff --git a/task-sdk/src/airflow/sdk/serde/__init__.py b/task-sdk/src/airflow/sdk/serde/__init__.py index 274fcaf494342..0d94ebc328238 100644 --- a/task-sdk/src/airflow/sdk/serde/__init__.py +++ b/task-sdk/src/airflow/sdk/serde/__init__.py @@ -254,14 +254,6 @@ def serialize(o: object, depth: int = 0) -> U | None: return o - # Materialise a lazily-resolved mapped XCom (e.g. mapped task group output) to a list here; - # otherwise the attrs branch below would serialize its internal fields. The single slice fetch - # happens only if the value is actually returned and serialized. - from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence - - if isinstance(o, LazyXComSequence): - return serialize(o[:], depth + 1) - # custom serializers dct = { CLASSNAME: qn, diff --git a/task-sdk/src/airflow/sdk/serde/serializers/lazy_xcom_sequence.py b/task-sdk/src/airflow/sdk/serde/serializers/lazy_xcom_sequence.py new file mode 100644 index 0000000000000..3648cb0c1243f --- /dev/null +++ b/task-sdk/src/airflow/sdk/serde/serializers/lazy_xcom_sequence.py @@ -0,0 +1,46 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from airflow.sdk.module_loading import qualname + +if TYPE_CHECKING: + from airflow.sdk.serde import U + +__version__ = 1 + +serializers = ["airflow.sdk.execution_time.lazy_sequence.LazyXComSequence"] +deserializers = serializers + + +def serialize(o: object) -> tuple[U, str, int, bool]: + """Materialize a lazily-resolved mapped XCom to a list (single slice fetch).""" + from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence + + if isinstance(o, LazyXComSequence): + return list(o[:]), qualname(o), __version__, True + return "", "", 0, False + + +def deserialize(cls: type, version: int, data: list) -> list: + """Deserialize to a plain list; the lazy sequence cannot be rebuilt off-worker.""" + if version > __version__: + raise TypeError(f"serialized version {version} is newer than class version {__version__}") + return list(data) diff --git a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py index 6272d739e585a..bb731c6d2877f 100644 --- a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py +++ b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py @@ -395,7 +395,7 @@ def _make_arg(): def test_resolve_stays_lazy_and_serializes_as_list(self, root, expected, mock_supervisor_comms): from airflow.sdk.execution_time.comms import XComSequenceSliceResult from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence - from airflow.sdk.serde import serialize + from airflow.sdk.serde import deserialize, serialize mock_supervisor_comms.send.return_value = XComSequenceSliceResult(root=root) @@ -406,7 +406,8 @@ def test_resolve_stays_lazy_and_serializes_as_list(self, root, expected, mock_su assert isinstance(resolved, LazyXComSequence) ti.xcom_pull.assert_not_called() - assert serialize(resolved) == expected + # The lazy sequence materializes to a plain list only when actually serialized. + assert deserialize(serialize(resolved)) == expected def test_resolve_uses_xcom_pull_for_specific_index(self): arg = self._make_arg() From c468738c2757daec8fdcbd139c66b1361b17a95d Mon Sep 17 00:00:00 2001 From: Shahar Epstein <60007259+shahar1@users.noreply.github.com> Date: Tue, 30 Jun 2026 10:47:23 +0300 Subject: [PATCH 5/6] Support lazy item access in mapped task group resolution test 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. --- .../tests/task_sdk/definitions/test_mappedoperator.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py b/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py index 93c5cc19aed47..4aadf923e9c85 100644 --- a/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py +++ b/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py @@ -32,12 +32,15 @@ from airflow.sdk.definitions.mappedoperator import MappedOperator from airflow.sdk.definitions.xcom_arg import XComArg from airflow.sdk.execution_time.comms import ( + ErrorResponse, GetTICount, GetXCom, + GetXComSequenceItem, GetXComSequenceSlice, SetXCom, TICount, XComResult, + XComSequenceIndexResult, XComSequenceSliceResult, ) @@ -659,6 +662,13 @@ def mock_comms_response(msg): task_id = msg.task_id values = [v for k, v in expected_values.items() if k[0] == task_id and k[1] is not None] return XComSequenceSliceResult(root=values) + elif isinstance(msg, GetXComSequenceItem): + # The aggregated task-group value now resolves lazily, so iterating it fetches one item at a time. + task_id = msg.task_id + values = [v for k, v in expected_values.items() if k[0] == task_id and k[1] is not None] + if 0 <= msg.offset < len(values): + return XComSequenceIndexResult(root=values[msg.offset]) + return ErrorResponse() elif isinstance(msg, GetTICount): # Handle TI count queries for upstream_map_indexes computation if msg.task_ids: From 2cdce9c73d42fd267aca5f9c077687788d39706e Mon Sep 17 00:00:00 2001 From: Shahar Epstein <60007259+shahar1@users.noreply.github.com> Date: Tue, 30 Jun 2026 22:21:46 +0300 Subject: [PATCH 6/6] Move test imports for mapped task group XCom resolution to module top level --- task-sdk/tests/task_sdk/definitions/test_xcom_arg.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py index bb731c6d2877f..2d1f8b2a4fa52 100644 --- a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py +++ b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py @@ -26,8 +26,11 @@ from airflow.sdk import TaskInstanceState from airflow.sdk.bases.xcom import BaseXCom from airflow.sdk.definitions.dag import DAG +from airflow.sdk.definitions.xcom_arg import PlainXComArg from airflow.sdk.exceptions import AirflowSkipException -from airflow.sdk.execution_time.comms import GetXCom, XComResult +from airflow.sdk.execution_time.comms import GetXCom, XComResult, XComSequenceSliceResult +from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence +from airflow.sdk.serde import deserialize, serialize log = structlog.get_logger(__name__) @@ -375,8 +378,6 @@ def _make_ti(*, computed): @staticmethod def _make_arg(): - from airflow.sdk.definitions.xcom_arg import PlainXComArg - operator = mock.MagicMock() operator.is_mapped = False operator.task_id = "do_something" @@ -393,10 +394,6 @@ def _make_arg(): ], ) def test_resolve_stays_lazy_and_serializes_as_list(self, root, expected, mock_supervisor_comms): - from airflow.sdk.execution_time.comms import XComSequenceSliceResult - from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence - from airflow.sdk.serde import deserialize, serialize - mock_supervisor_comms.send.return_value = XComSequenceSliceResult(root=root) arg = self._make_arg()