Give every mapped expansion a region and store its index as region_index

airflow-core/docs/core-concepts/task-and-asset-state-store.rst

@@ -38,7 +38,7 @@ Task and Asset state store provide two key/value stores to persist data like a j
      - Default lifetime
      - Primary use case
    * - **Task state store**
-     - A single task Instance (``dag_id`` + ``run_id`` + ``task_id`` + ``region_id`` + ``map_index``)
+     - A single task Instance (``dag_id`` + ``run_id`` + ``task_id`` + ``region_id`` + ``region_index``)
      - Configurable retention; cleared on task success when ``clear_on_success = True``
      - Survive retries, track in-flight jobs, checkpoint progress within a run, resume progress from checkpoint set by a past run
    * - **Asset state store**

airflow-core/docs/migrations-ref.rst

@@ -39,7 +39,10 @@ Here's the list of all the Database Migrations that are executed via when you ru
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | Revision ID             | Revises ID       | Airflow Version   | Description                                                  |
 +=========================+==================+===================+==============================================================+
-| ``54a27b6f9d01`` (head) | ``e7c2a91bd540`` | ``3.4.0``         | Add dynamic region storage.                                  |
+| ``7f8c9a2d410e`` (head) | ``54a27b6f9d01`` | ``3.4.0``         | Rename the stored map index to region_index and add the      |
+|                         |                  |                   | dynamic region slot key.                                     |
++-------------------------+------------------+-------------------+--------------------------------------------------------------+
+| ``54a27b6f9d01``        | ``e7c2a91bd540`` | ``3.4.0``         | Add dynamic region storage.                                  |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | ``e7c2a91bd540``        | ``90e4d18ccadf`` | ``3.4.0``         | Unify task attempt ownership without rewriting legacy XCom   |
 |                         |                  |                   | data.                                                        |

airflow-core/src/airflow/api_fastapi/common/parameters/misc.py

@@ -245,7 +245,7 @@ def _optional_boolean(value: bool | None) -> bool | None:
     FilterParam[int | None],
     Depends(
         filter_param_factory(
-            TaskInstance.map_index,
+            TaskInstance.region_index,
             int | None,
             filter_name="map_index",
         )

airflow-core/src/airflow/api_fastapi/common/parameters/task_instance.py

@@ -237,7 +237,11 @@ def _transform_ti_states(states: list[str] | None) -> list[TaskInstanceState | N
     FilterParam[list[int]],
     Depends(
         filter_param_factory(
-            TaskInstance.map_index, list[int], FilterOptionEnum.ANY_EQUAL, default_factory=list
+            TaskInstance.region_index,
+            list[int],
+            FilterOptionEnum.ANY_EQUAL,
+            filter_name="map_index",
+            default_factory=list,
         )
     ),
 ]

airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml

@@ -3083,6 +3083,12 @@ paths:
               schema:
                 $ref: '#/components/schemas/HTTPExceptionResponse'
           description: Not Found
+        '409':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Conflict
         '422':
           description: Validation Error
           content:

airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py

@@ -863,7 +863,7 @@ def trigger_dag_run(
     summary="Experimental: Wait for a dag run to complete, and return task results if requested.",
     description="🚧 This is an experimental endpoint and may change or be removed without notice.Successful response are streamed as newline-delimited JSON (NDJSON). Each line is a JSON object representing the Dag run state.",
     responses={
-        **create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND]),
+        **create_openapi_http_exception_doc([status.HTTP_404_NOT_FOUND, status.HTTP_409_CONFLICT]),
         status.HTTP_200_OK: {
             "description": "Successful Response",
             "content": {
@@ -886,6 +886,7 @@ def trigger_dag_run(
 def wait_dag_run_until_finished(
     dag_id: str,
     dag_run_id: str,
+    dag_bag: DagBagDep,
     session: SessionDep,
     user: GetUserDep,
     interval: Annotated[float, Query(gt=0.0, description="Seconds to wait between dag run state checks")],
@@ -917,7 +918,8 @@ def wait_dag_run_until_finished(
                 "User is not authorized to read XCom data for this Dag",
             )
         result_task_ids = []  # Explicitly not returning any XCom results.
-    if not session.scalar(select(1).where(DagRun.dag_id == dag_id, DagRun.run_id == dag_run_id)):
+    dag_run = session.scalar(select(DagRun).where(DagRun.dag_id == dag_id, DagRun.run_id == dag_run_id))
+    if dag_run is None:
         raise HTTPException(
             status.HTTP_404_NOT_FOUND,
             f"The DagRun with dag_id: `{dag_id}` and run_id: `{dag_run_id}` was not found",
@@ -928,6 +930,7 @@ def wait_dag_run_until_finished(
         interval=interval,
         result_task_ids=result_task_ids,
     )
+    waiter.validate_result_selection(dag_run, dag_bag, session=session)
     return StreamingResponse(waiter.wait())
 
 

airflow-core/src/airflow/api_fastapi/core_api/routes/public/task_state_store.py

@@ -42,6 +42,8 @@
 from airflow.state.metastore import _get_db_backend
 
 if TYPE_CHECKING:
+    from uuid import UUID
+
     from sqlalchemy.orm import Session
 
 task_state_store_router = AirflowRouter(
@@ -56,16 +58,17 @@ def _require_task_instance(
     task_id: str,
     map_index: int | None,
     session: Session,
-) -> None:
-    """Raise 404 unless the task instance exists. ``map_index=None`` matches any map index."""
-    statement = select(TI.task_id).where(
+) -> UUID:
+    """Return the region of the addressed task instance, raising 404 when it does not exist."""
+    statement = select(TI.region_id).where(
         TI.dag_id == dag_id,
         TI.run_id == dag_run_id,
         TI.task_id == task_id,
     )
     if map_index is not None:
-        statement = statement.where(TI.map_index == map_index)
-    if session.scalar(statement.limit(1)) is None:
+        statement = statement.where(TI.region_index == map_index)
+    region_id = session.scalar(statement.order_by(TI.region_id).limit(1))
+    if region_id is None:
         addressed_by = "all_map_indices=True" if map_index is None else f"map_index={map_index}"
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND,
@@ -74,6 +77,7 @@ def _require_task_instance(
                 f"task_id={task_id!r}, {addressed_by}"
             ),
         )
+    return region_id
 
 
 def _resolve_scope(
@@ -91,8 +95,10 @@ def _resolve_scope(
 
 def _validate_scope(scope: TaskScopeDep, session: SessionDep) -> TaskScope:
     """Resolve the scope, 404ing when the task instance it addresses does not exist."""
-    _require_task_instance(scope.dag_id, scope.run_id, scope.task_id, scope.map_index, session)
-    return scope
+    region_id = _require_task_instance(scope.dag_id, scope.run_id, scope.task_id, scope.map_index, session)
+    return TaskScope(
+        scope.dag_id, scope.run_id, scope.task_id, region_index=scope.region_index, region_id=region_id
+    )
 
 
 ValidatedTaskScopeDep = Annotated[TaskScope, Depends(_validate_scope)]
@@ -109,10 +115,12 @@ def _validate_clear_scope(
     ``all_map_indices`` addresses the task across every index, so it is validated against any
     instance -- an expanded mapped task has no ``map_index=-1`` instance to check.
     """
-    _require_task_instance(
+    region_id = _require_task_instance(
         scope.dag_id, scope.run_id, scope.task_id, None if all_map_indices else scope.map_index, session
     )
-    return scope
+    return TaskScope(
+        scope.dag_id, scope.run_id, scope.task_id, region_index=scope.region_index, region_id=region_id
+    )
 
 
 ValidatedClearTaskScopeDep = Annotated[TaskScope, Depends(_validate_clear_scope)]
@@ -161,7 +169,7 @@ def list_task_state_store(
             TaskStateStoreModel.dag_id == scope.dag_id,
             TaskStateStoreModel.run_id == scope.run_id,
             TaskStateStoreModel.task_id == scope.task_id,
-            TaskStateStoreModel.map_index == scope.map_index,
+            TaskStateStoreModel.region_index == scope.region_index,
         )
         .order_by(TaskStateStoreModel.key.asc())
     )
@@ -204,7 +212,7 @@ def get_task_state_store(
             TaskStateStoreModel.dag_id == scope.dag_id,
             TaskStateStoreModel.run_id == scope.run_id,
             TaskStateStoreModel.task_id == scope.task_id,
-            TaskStateStoreModel.map_index == scope.map_index,
+            TaskStateStoreModel.region_index == scope.region_index,
             TaskStateStoreModel.key == key,
         )
     ).one_or_none()
@@ -256,7 +264,7 @@ def patch_task_state_store(
             TaskStateStoreModel.dag_id == scope.dag_id,
             TaskStateStoreModel.run_id == scope.run_id,
             TaskStateStoreModel.task_id == scope.task_id,
-            TaskStateStoreModel.map_index == scope.map_index,
+            TaskStateStoreModel.region_index == scope.region_index,
             TaskStateStoreModel.key == key,
         )
     ).one_or_none()

airflow-core/src/airflow/api_fastapi/core_api/routes/ui/gantt.py

@@ -73,7 +73,7 @@ def get_gantt_data(
     ).where(
         TaskInstance.dag_id == dag_id,
         TaskInstance.run_id == run_id,
-        TaskInstance.map_index == -1,
+        TaskInstance.region_index == -1,
         or_(TaskInstance.state != TaskInstanceState.UP_FOR_RETRY, TaskInstance.state.is_(None)),
     )
 

airflow-core/src/airflow/api_fastapi/core_api/services/public/dag_run.py

@@ -20,7 +20,6 @@
 import asyncio
 import itertools
 import json
-import operator
 from typing import TYPE_CHECKING, Any
 
 import attrs
@@ -61,13 +60,15 @@
 from airflow.listeners.listener import get_listener_manager
 from airflow.models.dagrun import DagRun, clear_partition_runs
 from airflow.models.renderedtifields import load_legacy_rendered_fields
+from airflow.models.task_coordinates import enclosing_loop, public_map_index_expression
 from airflow.models.taskinstance import TaskInstance
 from airflow.models.xcom import XCOM_RETURN_KEY, XComModel, xcom_entity
 from airflow.utils.session import create_session_async
 from airflow.utils.state import State, TaskInstanceState
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator, Iterator
+    from collections.abc import AsyncGenerator
+    from uuid import UUID
 
     from airflow.serialization.definitions.dag import SerializedDAG
 
@@ -256,44 +257,90 @@ class DagRunWaiter:
     interval: float
     result_task_ids: list[str] | None
 
+    def validate_result_selection(self, dag_run: DagRun, dag_bag: DagBagDep, *, session: Session) -> None:
+        if self.result_task_ids == []:
+            return
+        query = select(TaskInstance.task_id, TaskInstance.dag_version_id).where(
+            TaskInstance.dag_id == self.dag_id, TaskInstance.run_id == self.run_id
+        )
+        if self.result_task_ids is not None:
+            query = query.where(TaskInstance.task_id.in_(self.result_task_ids))
+        versions: dict[UUID | None, set[str]] = {}
+        for task_id, version in session.execute(query.distinct()):
+            versions.setdefault(version or dag_run.created_dag_version_id, set()).add(task_id)
+        current_task_ids = {task_id for task_ids in versions.values() for task_id in task_ids}
+        created_dag = (
+            dag_bag.get_dag(dag_run.created_dag_version_id, session=session)
+            if dag_run.created_dag_version_id is not None
+            else None
+        )
+        if created_dag is not None:
+            requested = (
+                set(self.result_task_ids) if self.result_task_ids is not None else set(created_dag.task_ids)
+            )
+            versions.setdefault(dag_run.created_dag_version_id, set()).update(requested - current_task_ids)
+        for version, task_ids in versions.items():
+            if version is None:
+                continue
+            dag = (
+                created_dag
+                if version == dag_run.created_dag_version_id
+                else dag_bag.get_dag(version, session=session)
+            )
+            if dag is None:
+                continue
+            for task_id in task_ids:
+                if not dag.has_task(task_id):
+                    continue
+                task = dag.get_task(task_id)
+                if self.result_task_ids is None and not task.returns_dag_result:
+                    continue
+                if enclosing_loop(task) is not None:
+                    raise HTTPException(
+                        status.HTTP_409_CONFLICT,
+                        f"Cannot collect loop task {task_id!r} as a DagRun result; "
+                        "publish the result through a task outside the loop.",
+                    )
+
     async def _get_dag_run(self) -> DagRun:
         async with create_session_async() as session:
             return await session.scalar(select(DagRun).filter_by(dag_id=self.dag_id, run_id=self.run_id))
 
     async def _serialize_xcoms(self) -> dict[str, Any]:
+        xcom_query = XComModel.get_many(
+            run_id=self.run_id,
+            key=XCOM_RETURN_KEY,
+            dag_ids=self.dag_id,
+            task_ids=self.result_task_ids,
+            region_id=None,
+        )
+        entity = xcom_entity(xcom_query)
         if self.result_task_ids is None:  # Return dag-author-specified results.
-            xcom_read = XComModel.get_many(
-                run_id=self.run_id,
-                key=XCOM_RETURN_KEY,
-                dag_ids=self.dag_id,
-            )
-            entity = xcom_entity(xcom_read)
-            xcom_query = xcom_read.where(entity.dag_result.is_(True))
-        else:  # Explicitly API user-specified results.
-            xcom_read = XComModel.get_many(
-                run_id=self.run_id,
-                key=XCOM_RETURN_KEY,
-                task_ids=self.result_task_ids,
-                dag_ids=self.dag_id,
-            )
-            xcom_query = xcom_read
-            entity = xcom_entity(xcom_read)
+            xcom_query = xcom_query.where(entity.dag_result.is_(True))
         # XComModel.get_many() orders XCom by timestamp. Reset this to make
         # mapped task results stable since execution order is not guaranteed.
-        xcom_query = xcom_query.order_by(None).order_by(entity.task_id, entity.map_index)
+        xcom_query = (
+            xcom_query.add_columns(public_map_index_expression(entity))
+            .order_by(None)
+            .order_by(entity.task_id, entity.region_index)
+        )
         async with create_session_async() as session:
-            xcom_results = (await session.scalars(xcom_query)).all()
-
-        def _group_xcoms(g: Iterator[XComModel | tuple[XComModel]]) -> Any:
-            entries = [row[0] if isinstance(row, tuple) else row for row in g]
-            if len(entries) == 1 and entries[0].map_index < 0:  # Unpack non-mapped task xcom.
-                return entries[0].value
-            return [entry.value for entry in entries]  # Task is mapped; return all xcoms in a list.
-
-        return {
-            task_id: _group_xcoms(g)
-            for task_id, g in itertools.groupby(xcom_results, key=operator.attrgetter("task_id"))
-        }
+            xcom_results = (await session.execute(xcom_query)).all()
+
+        results = {}
+        for task_id, group in itertools.groupby(xcom_results, key=lambda row: row[0].task_id):
+            entries = list(group)
+            if len({entry.region_id for entry, _ in entries}) > 1 or (
+                len(entries) > 1 and any(map_index < 0 for _, map_index in entries)
+            ):
+                raise ValueError(
+                    f"Ambiguous current producers for task {task_id!r} in DAG run {self.run_id!r}"
+                )
+            if len(entries) == 1 and entries[0][1] < 0:
+                results[task_id] = entries[0][0].value
+            else:
+                results[task_id] = [entry.value for entry, _ in entries]
+        return results
 
     async def _serialize_response(self, dag_run: DagRun) -> str:
         resp = {"state": dag_run.state}

airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_state_store.py

@@ -60,7 +60,7 @@ def _fetch_ti_writer_fields(token: TIToken, session: SessionDep) -> _TIWriterFie
             TaskInstance.dag_id,
             TaskInstance.run_id,
             TaskInstance.task_id,
-            TaskInstance.map_index,
+            TaskInstance.region_index.label("map_index"),
         ).where(TaskInstance.id == token.id)
     ).one_or_none()
     if row is None: