fixup! Add dynamic task mapping into TaskSDK runtime

Author: ashb

airflow/models/skipmixin.py

@@ -97,7 +97,7 @@ def skip(
         dag_id: str,
         run_id: str,
         tasks: Iterable[DAGNode],
-        map_index: int = -1,
+        map_index: int | None = -1,
         session: Session = NEW_SESSION,
     ):
         """
@@ -126,6 +126,9 @@ def skip(
         if task_id is not None:
             from airflow.models.xcom import XCom
 
+            if map_index is None:
+                map_index = -1
+
             XCom.set(
                 key=XCOM_SKIPMIXIN_KEY,
                 value={XCOM_SKIPMIXIN_SKIPPED: task_ids_list},

airflow/serialization/serialized_objects.py

@@ -1792,7 +1792,7 @@ def serialize_task_group(cls, task_group: TaskGroup) -> dict[str, Any] | None:
         if isinstance(task_group, MappedTaskGroup):
             expand_input = task_group._expand_input
             encoded["expand_input"] = {
-                "type": type(expand_input).EXPAND_INPUT_TYPE,
+                "type": expand_input.EXPAND_INPUT_TYPE,
                 "value": cls.serialize(expand_input.value),
             }
             encoded["is_mapped"] = True

providers/src/airflow/providers/cncf/kubernetes/operators/pod.py

@@ -489,7 +489,7 @@ def _get_ti_pod_labels(context: Context | None = None, include_try_number: bool
         }
 
         map_index = ti.map_index
-        if map_index >= 0:
+        if map_index is not None and map_index >= 0:
             labels["map_index"] = str(map_index)
 
         if include_try_number:

providers/src/airflow/providers/microsoft/azure/operators/msgraph.py

@@ -241,7 +241,6 @@ def pull_xcom(self, context: Context) -> list:
                 key=self.key,
                 task_ids=self.task_id,
                 dag_id=self.dag_id,
-                map_indexes=map_index,
             )
             or []
         )

providers/standard/tests/provider_tests/standard/decorators/test_python.py

@@ -18,7 +18,7 @@
 import sys
 import typing
 from collections import namedtuple
-from datetime import date, timedelta
+from datetime import date
 from typing import Union
 
 import pytest
@@ -28,8 +28,6 @@
 from airflow.exceptions import AirflowException, XComNotFound
 from airflow.models.taskinstance import TaskInstance
 from airflow.models.taskmap import TaskMap
-from airflow.sdk import DAG, BaseOperator, TaskGroup, XComArg
-from airflow.sdk.definitions.mappedoperator import MappedOperator
 from airflow.utils import timezone
 from airflow.utils.state import State
 from airflow.utils.task_instance_session import set_current_task_instance_session
@@ -41,12 +39,17 @@
 from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
 
 if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import DAG, BaseOperator, TaskGroup, XComArg
     from airflow.sdk.definitions._internal.expandinput import DictOfListsExpandInput
     from airflow.sdk.definitions.mappedoperator import MappedOperator
     from airflow.utils.types import DagRunTriggeredByType
 else:
+    from airflow.models.baseoperator import BaseOperator
+    from airflow.models.dag import DAG  # type: ignore[assignment]
     from airflow.models.expandinput import DictOfListsExpandInput
     from airflow.models.mappedoperator import MappedOperator
+    from airflow.models.xcom_arg import XComArg
+    from airflow.utils.task_group import TaskGroup
 
 pytestmark = pytest.mark.db_test
 
@@ -794,36 +797,6 @@ def task2(arg1, arg2): ...
     assert set(unmapped.op_kwargs) == {"arg1", "arg2"}
 
 
-def test_mapped_decorator_converts_partial_kwargs(dag_maker, session):
-    with dag_maker(session=session):
-
-        @task_decorator
-        def task1(arg):
-            return ["x" * arg]
-
-        @task_decorator(retry_delay=30)
-        def task2(arg1, arg2): ...
-
-        task2.partial(arg1=1).expand(arg2=task1.expand(arg=[1, 2]))
-
-    run = dag_maker.create_dagrun()
-
-    # Expand and run task1.
-    dec = run.task_instance_scheduling_decisions(session=session)
-    assert [ti.task_id for ti in dec.schedulable_tis] == ["task1", "task1"]
-    for ti in dec.schedulable_tis:
-        ti.run(session=session)
-        assert not isinstance(ti.task, MappedOperator)
-        assert ti.task.retry_delay == timedelta(seconds=300)  # Operator default.
-
-    # Expand task2.
-    dec = run.task_instance_scheduling_decisions(session=session)
-    assert [ti.task_id for ti in dec.schedulable_tis] == ["task2", "task2"]
-    for ti in dec.schedulable_tis:
-        unmapped = ti.task.unmap((ti.get_template_context(session),))
-        assert unmapped.retry_delay == timedelta(seconds=30)
-
-
 def test_mapped_render_template_fields(dag_maker, session):
     @task_decorator
     def fn(arg1, arg2): ...

tests_common/test_utils/mock_context.py

@@ -56,14 +56,16 @@ def xcom_pull(
             default: Any = None,
             run_id: str | None = None,
         ) -> Any:
-            if map_indexes:
-                return values.get(
-                    f"{task_ids or self.task_id}_{dag_id or self.dag_id}_{key}_{map_indexes}", default
-                )
-            return values.get(f"{task_ids or self.task_id}_{dag_id or self.dag_id}_{key}", default)
+            key = f"{self.task_id}_{self.dag_id}_{key}"
+            if map_indexes is not None and (not isinstance(map_indexes, int) or map_indexes >= 0):
+                key += f"_{map_indexes}"
+            return values.get(key, default)
 
         def xcom_push(self, key: str, value: Any, session: Session = NEW_SESSION, **kwargs) -> None:
-            values[f"{self.task_id}_{self.dag_id}_{key}_{self.map_index}"] = value
+            key = f"{self.task_id}_{self.dag_id}_{key}"
+            if self.map_index is not None and self.map_index >= 0:
+                key += f"_{self.map_index}"
+            values[key] = value
 
     values["ti"] = MockedTaskInstance(task=task)