AIP-104: Task Iteration and Dynamic Task Batching

.github/CODEOWNERS

@@ -162,6 +162,9 @@ Dockerfile.ci @potiuk @ashb @gopidesupavan @amoghrajesh @jscheffl @bugraoz93 @ja
 # AIP-72 - Task SDK
 # Python SDK
 /task-sdk/ @ashb @kaxil @amoghrajesh
+/task-sdk/**/batchedoperator.py @dabla
+/task-sdk/**/executor.py @dabla
+/task-sdk/**/iterableoperator.py @dabla
 
 # AIP-108 - Coordinators
 /task-sdk/src/airflow/sdk/coordinators/ @jason810496 @uranusjr

airflow-core/src/airflow/serialization/definitions/mappedoperator.py

@@ -345,6 +345,10 @@ def on_failure_fail_dagrun(self) -> bool:
     def on_failure_fail_dagrun(self, v) -> None:
         self.partial_kwargs["on_failure_fail_dagrun"] = bool(v)
 
+    @property
+    def batch_size(self) -> int:
+        return self.partial_kwargs.get("batch_size", 0)
+
     @classmethod
     def get_serialized_fields(cls):
         """Fields to extract from JSON-Serialized DAG."""
@@ -470,15 +474,20 @@ def get_parse_time_mapped_ti_count(self) -> int:
         current_count = self._get_specified_expand_input().get_parse_time_mapped_ti_count()
 
         def _get_parent_count() -> int:
-            if (group := self.get_closest_mapped_task_group()) is None:
-                raise NotMapped()
-            return group.get_parse_time_mapped_ti_count()
+            try:
+                if (group := self.get_closest_mapped_task_group()) is None:
+                    raise NotMapped()
+                return group.get_parse_time_mapped_ti_count()
+            except NotMapped:
+                return 1
+
+        parent_count = _get_parent_count()
+        mapped_ti_count = parent_count * current_count
 
-        try:
-            parent_count = _get_parent_count()
-        except NotMapped:
-            return current_count
-        return parent_count * current_count
+        if self.batch_size > 0:
+            if self.batch_size < mapped_ti_count:
+                return self.batch_size
+        return mapped_ti_count
 
     def iter_mapped_dependencies(self) -> Iterator[Operator]:
         """Upstream dependencies that provide XComs used by this task for task mapping."""
@@ -544,9 +553,14 @@ def _(task: SerializedMappedOperator | TaskSDKMappedOperator, run_id: str, *, se
 
     group = task.get_closest_mapped_task_group()
     if group is None:
-        return current_count
-    parent_count = get_mapped_ti_count(group, run_id, session=session)
-    return parent_count * current_count
+        mapped_ti_count = current_count
+    else:
+        parent_count = get_mapped_ti_count(group, run_id, session=session)
+        mapped_ti_count = parent_count * current_count
+
+    if task.batch_size > 0:
+        return min(task.batch_size, mapped_ti_count)
+    return mapped_ti_count
 
 
 @get_mapped_ti_count.register

airflow-core/tests/unit/models/test_mappedoperator.py

@@ -1828,3 +1828,26 @@ def test_mapped_operator_retry_delay_explicit(dag_maker):
 
     # Should return the explicitly set value
     assert mapped_deser.retry_delay == custom_retry_delay
+
+
+@pytest.mark.parametrize(
+    ("batch_size", "items", "expected"),
+    [
+        pytest.param(5, [1, 2, 3], 3, id="batch_size-larger-than-items-clamped-to-items"),
+        pytest.param(2, [1, 2, 3], 2, id="batch_size-smaller-than-items-returned-as-is"),
+        pytest.param(3, [1, 2, 3], 3, id="batch_size-equal-to-items"),
+    ],
+)
+def test_get_mapped_ti_count_clamps_to_items_count_when_batch_size_exceeds_it(
+    dag_maker, session, batch_size, items, expected
+):
+    from airflow.serialization.definitions.mappedoperator import get_mapped_ti_count
+
+    with dag_maker(dag_id=f"test_batch_size_clamp_{batch_size}", session=session, serialized=True) as dag:
+        MockOperator.partial(task_id="task").batch(size=batch_size).iterate(arg1=items)
+
+    dr = dag_maker.create_dagrun()
+    task = dag.task_dict["task"]
+
+    result = get_mapped_ti_count(task, dr.run_id, session=session)
+    assert result == expected

airflow-core/tests/unit/models/test_taskinstance.py

@@ -46,6 +46,7 @@
 from airflow.exceptions import (
     AirflowException,
     AirflowSkipException,
+    NotMapped,
 )
 from airflow.models.asset import (
     AssetActive,
@@ -3487,6 +3488,74 @@ def show(a, b):
             dag_maker.run_ti(ti.task_id, map_index=ti.map_index, dag_run=dag_run, session=session)
         assert outputs == [(2, 5), (2, 10), (4, 5), (4, 10), (8, 5), (8, 10)]
 
+    def test_iterate_literal_cross_product(self, dag_maker, session):
+        """Test an iterated task with literal cross product args properly."""
+        outputs = []
+
+        with dag_maker(dag_id="product_same_types", session=session, serialized=True) as dag:
+
+            @dag.task
+            def show(a, b):
+                outputs.append((a, b))
+
+            show.iterate(a=[2, 4, 8], b=[5, 10])
+
+        dag_run = dag_maker.create_dagrun()
+
+        show_task = dag.get_task("show")
+        with pytest.raises(NotMapped):
+            show_task.get_parse_time_mapped_ti_count()
+        with pytest.raises(NotMapped):
+            TaskMap.expand_mapped_task(show_task, dag_run.run_id, session=session)
+
+        tis = session.scalars(
+            select(TaskInstance)
+            .where(
+                TaskInstance.dag_id == dag.dag_id,
+                TaskInstance.task_id == "show",
+                TaskInstance.run_id == dag_run.run_id,
+            )
+            .order_by(TaskInstance.map_index)
+        ).all()
+        for ti in tis:
+            ti.refresh_from_task(show_task)
+            dag_maker.run_ti(ti.task_id, map_index=ti.map_index, dag_run=dag_run, session=session)
+        assert outputs == [(2, 5), (2, 10), (4, 5), (4, 10), (8, 5), (8, 10)]
+
+    def test_iterate_literal_cross_product_in_batch(self, dag_maker, session):
+        """Test a batched iterated task with literal cross product args properly."""
+        outputs = []
+
+        with dag_maker(dag_id="product_same_types", session=session, serialized=True) as dag:
+
+            @dag.task
+            def show(a, b):
+                outputs.append((a, b))
+
+            show.batch(size=2).iterate(a=[2, 4, 8], b=[5, 10])
+
+        dag_run = dag_maker.create_dagrun()
+
+        show_task = dag.get_task("show")
+        assert show_task.get_parse_time_mapped_ti_count() == 2
+        mapped_tis, max_map_index = TaskMap.expand_mapped_task(show_task, dag_run.run_id, session=session)
+        assert len(mapped_tis) == 0  # Expanded at parse!
+        assert max_map_index == 1
+
+        tis = session.scalars(
+            select(TaskInstance)
+            .where(
+                TaskInstance.dag_id == dag.dag_id,
+                TaskInstance.task_id == "show",
+                TaskInstance.run_id == dag_run.run_id,
+            )
+            .order_by(TaskInstance.map_index)
+        ).all()
+        for ti in tis:
+            ti.refresh_from_task(show_task)
+            dag_maker.run_ti(ti.task_id, map_index=ti.map_index, dag_run=dag_run, session=session)
+        assert outputs == [(2, 5), (4, 5), (8, 5), (2, 10), (4, 10), (8, 10)]
+
     def test_map_in_group(self, tmp_path: pathlib.Path, dag_maker, session):
         out = tmp_path.joinpath("out")
         out.touch()

airflow-core/tests/unit/serialization/test_dag_serialization.py

@@ -2994,6 +2994,7 @@ def test_operator_expand_serde():
         "template_ext": [".sh", ".bash"],
         "template_fields_renderers": {"bash_command": "bash", "env": "json"},
         "ui_color": "#f0ede4",
+        "_register_with_dag": True,
         "_disallow_kwargs_override": False,
         "_expand_input_attr": "expand_input",
     }
@@ -3037,6 +3038,7 @@ def test_operator_expand_xcomarg_serde():
         },
         "task_id": "task_2",
         "template_fields": ["arg1", "arg2"],
+        "_register_with_dag": True,
         "_disallow_kwargs_override": False,
         "_expand_input_attr": "expand_input",
     }
@@ -3092,6 +3094,7 @@ def test_operator_expand_kwargs_literal_serde(strict):
         },
         "task_id": "task_2",
         "template_fields": ["arg1", "arg2"],
+        "_register_with_dag": True,
         "_disallow_kwargs_override": strict,
         "_expand_input_attr": "expand_input",
     }
@@ -3139,6 +3142,7 @@ def test_operator_expand_kwargs_xcomarg_serde(strict):
         },
         "task_id": "task_2",
         "template_fields": ["arg1", "arg2"],
+        "_register_with_dag": True,
         "_disallow_kwargs_override": strict,
         "_expand_input_attr": "expand_input",
     }
@@ -3247,6 +3251,7 @@ def x(arg1, arg2, arg3):
             "op_args": "py",
             "op_kwargs": "py",
         },
+        "_register_with_dag": True,
         "_disallow_kwargs_override": False,
         "_expand_input_attr": "op_kwargs_expand_input",
         "python_callable_name": "test_taskflow_expand_serde.<locals>.x",
@@ -3340,6 +3345,7 @@ def x(arg1, arg2, arg3):
             "op_args": "py",
             "op_kwargs": "py",
         },
+        "_register_with_dag": True,
         "_disallow_kwargs_override": strict,
         "_expand_input_attr": "op_kwargs_expand_input",
     }
@@ -3451,6 +3457,7 @@ def operator_extra_links(self):
         "partial_kwargs": {
             "retry_delay": {"__type": "timedelta", "__var": 300.0},
         },
+        "_register_with_dag": True,
         "_disallow_kwargs_override": False,
         "_expand_input_attr": "expand_input",
         "_operator_extra_links": {"airflow": "_link_AirflowLink2"},

devel-common/src/tests_common/test_utils/mock_context.py

@@ -20,14 +20,29 @@
 from typing import TYPE_CHECKING, Any
 from unittest import mock
 
+from airflow.models import DagRun
+from airflow.utils.types import DagRunType
+
 from tests_common.test_utils.compat import Context
 from tests_common.test_utils.taskinstance import create_task_instance
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_1_PLUS
+
+if AIRFLOW_V_3_1_PLUS:
+    from airflow.sdk import timezone
+else:
+    from airflow.utils import timezone  # type: ignore[attr-defined,no-redef]
 
 if TYPE_CHECKING:
     from sqlalchemy.orm import Session
 
 
-def mock_context(task) -> Context:
+def generate_run_id() -> str:
+    if AIRFLOW_V_3_0_PLUS:
+        return DagRun.generate_run_id(run_type=DagRunType.MANUAL, run_after=timezone.utcnow())
+    return DagRun.generate_run_id(run_type=DagRunType.MANUAL, execution_date=timezone.utcnow())  # type: ignore[call-arg]
+
+
+def mock_context(task, run_id: str | None = None) -> Context:
     from airflow.models import TaskInstance
     from airflow.utils.session import NEW_SESSION
 
@@ -64,6 +79,6 @@ def xcom_push(self, key: str, value: Any, session: Session = NEW_SESSION, **kwar
             values[key] = value
 
     values["ti"] = create_task_instance(task, dag_version_id=mock.MagicMock(), ti_type=MockedTaskInstance)
+    values["run_id"] = generate_run_id() if run_id is None else run_id
 
-    # See https://github.com/python/mypy/issues/8890 - mypy does not support passing typed dict to TypedDict
     return Context(values)  # type: ignore[misc]

docs/spelling_wordlist.txt

@@ -1266,6 +1266,7 @@ PodManager
 PodSpec
 podSpec
 podspec
+Pokémon
 polars
 poller
 polyfill

task-sdk/.pre-commit-config.yaml

@@ -39,7 +39,9 @@ repos:
           ^src/airflow/sdk/definitions/asset/__init__\.py$|
           ^src/airflow/sdk/definitions/asset/decorators\.py$|
           ^src/airflow/sdk/definitions/taskgroup\.py$|
+          ^src/airflow/sdk/definitions/iterableoperator\.py$|
           ^src/airflow/sdk/definitions/mappedoperator\.py$|
+          ^src/airflow/sdk/definitions/batchedoperator\.py$|
           ^src/airflow/sdk/definitions/deadline\.py$|
           ^src/airflow/sdk/definitions/dag\.py$|
           ^src/airflow/sdk/definitions/_internal/types\.py$|

task-sdk/docs/deferred-vs-async-operators.rst

@@ -205,7 +205,7 @@ concurrently using ``asyncio.gather`` while limiting concurrency with a semaphor
 
 .. note::
 
-   The upcoming *Dynamic Task Iteration* feature will simplify patterns like this.
+   The new :ref:`Dynamic Task Iteration <sdk-dynamic-task-mapping-vs-iteration>`. feature will simplify patterns like this.
    Instead of manually managing concurrency with constructs such as
    ``asyncio.gather`` and ``asyncio.Semaphore``, authors will be able to iterate
    over asynchronous results directly in downstream tasks while still benefiting