Fix XComArg.map() failing when the upstream XCom is a dict

task-sdk/src/airflow/sdk/definitions/xcom_arg.py

@@ -400,7 +400,15 @@ class _MapResult(Sequence):
     callables: MapCallables
 
     def __getitem__(self, index: Any) -> Any:
-        value = self.value[index]
+        if isinstance(self.value, dict):
+            length = len(self.value)
+            if index < 0:
+                index += length
+            if not 0 <= index < length:
+                raise IndexError(index)
+            value = next(itertools.islice(self.value.items(), index, None))
+        else:
+            value = self.value[index]
 
         for f in self.callables:
             value = f(value)

task-sdk/tests/task_sdk/definitions/test_xcom_arg.py

@@ -27,7 +27,7 @@
 from airflow.sdk.bases.xcom import BaseXCom
 from airflow.sdk.definitions._internal.loop import create_loop
 from airflow.sdk.definitions.dag import DAG
-from airflow.sdk.definitions.xcom_arg import PlainXComArg
+from airflow.sdk.definitions.xcom_arg import PlainXComArg, _MapResult
 from airflow.sdk.exceptions import AirflowSkipException
 from airflow.sdk.execution_time.comms import (
     GetXCom,
@@ -295,6 +295,45 @@ def skip_c(v):
     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"}
+
+
+@pytest.mark.parametrize(
+    ("index", "expected"),
+    [
+        pytest.param(0, ("a", 1), id="positive"),
+        pytest.param(-1, ("b", 2), id="negative"),
+    ],
+)
+def test_map_result_dict_getitem(index, expected):
+    assert _MapResult({"a": 1, "b": 2}, [lambda kv: kv])[index] == expected
+
+
+@pytest.mark.parametrize("index", [2, -3])
+def test_map_result_dict_getitem_out_of_range(index):
+    with pytest.raises(IndexError):
+        _MapResult({"a": 1, "b": 2}, [lambda kv: kv])[index]
+
+
 def test_xcom_concat(run_ti, mock_supervisor_comms):
     from airflow.sdk.definitions.xcom_arg import _ConcatResult