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``). diff --git a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py index 595490832597f..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) @@ -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, 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_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: 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..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__) @@ -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