diff --git a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py index 36b393c165f42..037777f74e654 100644 --- a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py +++ b/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) 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 53b1139dacb41..f4065474bc3e0 100644 --- a/task-sdk/tests/task_sdk/definitions/test_xcom_arg.py +++ b/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