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 92a4e3ebea7 Fix XComArg.map() failing when the upstream XCom is a dict
(#75022)
92a4e3ebea7 is described below
commit 92a4e3ebea7b788f479388788b0cc2a49cbe900b
Author: Justin Pakzad <[email protected]>
AuthorDate: Sat Nov 7 02:03:00 2026 -0500
Fix XComArg.map() failing when the upstream XCom is a dict (#75022)
---
.../dynamic-task-mapping.rst | 1 +
task-sdk/src/airflow/sdk/definitions/xcom_arg.py | 5 ++++-
.../tests/task_sdk/definitions/test_xcom_arg.py | 22 ++++++++++++++++++++++
3 files changed, 27 insertions(+), 1 deletion(-)
diff --git
a/airflow-core/docs/authoring-and-scheduling/dynamic-task-mapping.rst
b/airflow-core/docs/authoring-and-scheduling/dynamic-task-mapping.rst
index 972f7440011..02d7038a0b6 100644
--- a/airflow-core/docs/authoring-and-scheduling/dynamic-task-mapping.rst
+++ b/airflow-core/docs/authoring-and-scheduling/dynamic-task-mapping.rst
@@ -534,6 +534,7 @@ There are a couple of things to note:
#. The callable argument of :func:`map()` (``create_copy_kwargs`` in the
example) **must not** be a task, but a plain Python function. The
transformation is as a part of the "pre-processing" of the downstream task
(i.e. ``copy_files``), not a standalone task in the Dag.
#. The callable always take exactly one positional argument. This function is
called for each item in the iterable used for task-mapping, similar to how
Python's built-in :func:`map()` works.
+#. If the upstream value is a dict, the callable receives each item as a
``(key, value)`` tuple.
#. Since the callable is executed as a part of the downstream task, you can
use any existing techniques to write the task function. To mark a component as
skipped, for example, you should raise ``AirflowSkipException``. Note that
returning ``None`` **does not** work here.
Combining upstream data (aka "zipping")
diff --git a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
index 36b393c165f..fe608eea454 100644
--- a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
+++ b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
@@ -400,7 +400,10 @@ class _MapResult(Sequence):
callables: MapCallables
def __getitem__(self, index: Any) -> Any:
- value = self.value[index]
+ if isinstance(self.value, dict):
+ value = list(self.value.items())[index]
+ else:
+ value = self.value[index]
for f in self.callables:
value = f(value)
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 53b1139dacb..40bd01158c3 100644
--- a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py
+++ b/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py
@@ -295,6 +295,28 @@ def test_xcom_map_raise_to_skip(run_ti,
mock_supervisor_comms):
assert result == ["a", "b"]
+def test_xcom_map_dict(run_ti, mock_supervisor_comms):
+ result = set()
+ with DAG("test") as dag:
+
+ @dag.task
+ def push():
+ return {"a": 1, "b": 2}
+
+ @dag.task
+ def pull(value):
+ result.add(value)
+
+ pull.expand(value=push().map(lambda kv: f"{kv[0]}={kv[1]}"))
+
+ mock_supervisor_comms.send.return_value =
XComResult(key=BaseXCom.XCOM_RETURN_KEY, value={"a": 1, "b": 2})
+
+ for map_index in range(2):
+ assert run_ti(dag, "pull", map_index) == TaskInstanceState.SUCCESS
+
+ assert result == {"a=1", "b=2"}
+
+
def test_xcom_concat(run_ti, mock_supervisor_comms):
from airflow.sdk.definitions.xcom_arg import _ConcatResult
