This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 37f6259b708 Fix mapped task group XCom collapsing to a bare value for
a single expansion (#69075)
37f6259b708 is described below
commit 37f6259b708ce9b357f78300aebf60f280a25242
Author: Shahar Epstein <[email protected]>
AuthorDate: Thu Jul 2 00:31:06 2026 +0300
Fix mapped task group XCom collapsing to a bare value for a single
expansion (#69075)
---
airflow-core/newsfragments/69075.bugfix.rst | 1 +
task-sdk/src/airflow/sdk/definitions/xcom_arg.py | 12 ++--
.../sdk/serde/serializers/lazy_xcom_sequence.py | 46 +++++++++++++++
.../task_sdk/definitions/test_mappedoperator.py | 10 ++++
.../tests/task_sdk/definitions/test_xcom_arg.py | 68 +++++++++++++++++++++-
5 files changed, 132 insertions(+), 5 deletions(-)
diff --git a/airflow-core/newsfragments/69075.bugfix.rst
b/airflow-core/newsfragments/69075.bugfix.rst
new file mode 100644
index 00000000000..d992083c378
--- /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``).
diff --git a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
index 59549083259..a38e47a466b 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 import TriggerRule
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 @@ class PlainXComArg(XComArg):
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)
@@ -354,8 +354,12 @@ class PlainXComArg(XComArg):
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:
+ # 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,
key=self.key,
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 00000000000..3648cb0c124
--- /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_mappedoperator.py
b/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py
index 93c5cc19aed..4aadf923e9c 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.dag import DAG
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
test_operator_mapped_task_group_receives_value(create_runtime_ti, mock_super
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:
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 6014d26f220..2d1f8b2a4fa 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 @@ import structlog
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__)
@@ -350,3 +353,66 @@ def test_xcom_concat(run_ti, mock_supervisor_comms):
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():
+ 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):
+ 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 isinstance(resolved, LazyXComSequence)
+ ti.xcom_pull.assert_not_called()
+ # 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()
+ 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
