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

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``).

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)
@@ -354,8 +354,12 @@ 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:
+                    # 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,

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)

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:

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__)
 
@@ -350,3 +353,66 @@ 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():
+        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