AIP-104: Dynamic Task Iteration and Dynamic Task Partitioning

.github/CODEOWNERS

@@ -161,6 +161,9 @@ Dockerfile.ci @potiuk @ashb @gopidesupavan @amoghrajesh @jscheffl @bugraoz93 @ja
 # AIP-72 - Task SDK
 # Python SDK
 /task-sdk/ @ashb @kaxil @amoghrajesh
+/task-sdk/**/executor.py @dabla
+/task-sdk/**/iterableoperator.py @dabla
+/task-sdk/**/partitionedoperator.py @dabla
 
 # Golang SDK
 /go-sdk/ @ashb @amoghrajesh

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

@@ -515,6 +515,11 @@ def _(task: SerializedBaseOperator | TaskSDKBaseOperator, run_id: str, *, sessio
 def _(task: SerializedMappedOperator | TaskSDKMappedOperator, run_id: str, *, session: Session) -> int:
     from airflow.serialization.serialized_objects import BaseSerialization, _ExpandInputRef
 
+    partition_size = task.partial_kwargs.get("partition_size")
+
+    if partition_size is not None:
+        return partition_size
+
     exp_input = task._get_specified_expand_input()
     # TODO (GH-52141): 'task' here should be scheduler-bound and returns scheduler expand input.
     if not hasattr(exp_input, "get_total_map_length"):

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

@@ -2698,6 +2698,7 @@ def test_operator_expand_serde():
         "template_ext": [".sh", ".bash"],
         "template_fields_renderers": {"bash_command": "bash", "env": "json"},
         "ui_color": "#f0ede4",
+        "_apply_upstream_relationship": True,
         "_disallow_kwargs_override": False,
         "_expand_input_attr": "expand_input",
     }
@@ -2741,6 +2742,7 @@ def test_operator_expand_xcomarg_serde():
         },
         "task_id": "task_2",
         "template_fields": ["arg1", "arg2"],
+        "_apply_upstream_relationship": True,
         "_disallow_kwargs_override": False,
         "_expand_input_attr": "expand_input",
     }
@@ -2796,6 +2798,7 @@ def test_operator_expand_kwargs_literal_serde(strict):
         },
         "task_id": "task_2",
         "template_fields": ["arg1", "arg2"],
+        "_apply_upstream_relationship": True,
         "_disallow_kwargs_override": strict,
         "_expand_input_attr": "expand_input",
     }
@@ -2843,6 +2846,7 @@ def test_operator_expand_kwargs_xcomarg_serde(strict):
         },
         "task_id": "task_2",
         "template_fields": ["arg1", "arg2"],
+        "_apply_upstream_relationship": True,
         "_disallow_kwargs_override": strict,
         "_expand_input_attr": "expand_input",
     }
@@ -2951,6 +2955,7 @@ def x(arg1, arg2, arg3):
             "op_args": "py",
             "op_kwargs": "py",
         },
+        "_apply_upstream_relationship": True,
         "_disallow_kwargs_override": False,
         "_expand_input_attr": "op_kwargs_expand_input",
         "python_callable_name": "test_taskflow_expand_serde.<locals>.x",
@@ -3044,6 +3049,7 @@ def x(arg1, arg2, arg3):
             "op_args": "py",
             "op_kwargs": "py",
         },
+        "_apply_upstream_relationship": True,
         "_disallow_kwargs_override": strict,
         "_expand_input_attr": "op_kwargs_expand_input",
     }
@@ -3155,6 +3161,7 @@ def operator_extra_links(self):
         "partial_kwargs": {
             "retry_delay": {"__type": "timedelta", "__var": 300.0},
         },
+        "_apply_upstream_relationship": 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]

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/partitionedoperator\.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

@@ -196,7 +196,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