Add map_index filter option to GetTICount and GetTaskStates

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

@@ -593,6 +593,7 @@ def get_previous_successful_dagrun(
 def get_task_instance_count(
     dag_id: str,
     session: SessionDep,
+    map_index: Annotated[int | None, Query()] = None,
     task_ids: Annotated[list[str] | None, Query()] = None,
     task_group_id: Annotated[str | None, Query()] = None,
     logical_dates: Annotated[list[UtcDateTime] | None, Query()] = None,
@@ -605,6 +606,9 @@ def get_task_instance_count(
     if task_ids:
         query = query.where(TI.task_id.in_(task_ids))
 
+    if map_index is not None:
+        query = query.where(TI.map_index == map_index)
+
     if logical_dates:
         query = query.where(TI.logical_date.in_(logical_dates))
 
@@ -615,7 +619,12 @@ def get_task_instance_count(
         group_tasks = _get_group_tasks(dag_id, task_group_id, session, logical_dates, run_ids)
 
         # Get unique (task_id, map_index) pairs
+
         task_map_pairs = [(ti.task_id, ti.map_index) for ti in group_tasks]
+
+        if map_index is not None:
+            task_map_pairs = [(ti.task_id, ti.map_index) for ti in group_tasks if ti.map_index == map_index]
+
         if not task_map_pairs:
             # If no task group tasks found, default to checking the task group ID itself
             # This matches the behavior in _get_external_task_group_task_ids
@@ -643,6 +652,7 @@ def get_task_instance_count(
 def get_task_instance_states(
     dag_id: str,
     session: SessionDep,
+    map_index: Annotated[int | None, Query()] = None,
     task_ids: Annotated[list[str] | None, Query()] = None,
     task_group_id: Annotated[str | None, Query()] = None,
     logical_dates: Annotated[list[UtcDateTime] | None, Query()] = None,
@@ -664,12 +674,21 @@ def get_task_instance_states(
 
     results = session.scalars(query).all()
 
-    [run_id_task_state_map[task.run_id].update({task.task_id: task.state}) for task in results]
-
     if task_group_id:
         group_tasks = _get_group_tasks(dag_id, task_group_id, session, logical_dates, run_ids)
 
-        [run_id_task_state_map[task.run_id].update({task.task_id: task.state}) for task in group_tasks]
+        results = results + group_tasks if task_ids else group_tasks
+
+    if map_index is not None:
+        results = [task for task in results if task.map_index == map_index]
+    [
+        run_id_task_state_map[task.run_id].update(
+            {task.task_id: task.state}
+            if task.map_index < 0
+            else {f"{task.task_id}_{task.map_index}": task.state}
+        )
+        for task in results
+    ]
 
     return TaskStatesResponse(task_states=run_id_task_state_map)
 

airflow-core/src/airflow/jobs/triggerer_job_runner.py

@@ -430,6 +430,7 @@ def _handle_request(self, msg: ToTriggerSupervisor, log: FilteringBoundLogger) -
         elif isinstance(msg, GetTICount):
             resp = self.client.task_instances.get_count(
                 dag_id=msg.dag_id,
+                map_index=msg.map_index,
                 task_ids=msg.task_ids,
                 task_group_id=msg.task_group_id,
                 logical_dates=msg.logical_dates,
@@ -440,6 +441,7 @@ def _handle_request(self, msg: ToTriggerSupervisor, log: FilteringBoundLogger) -
         elif isinstance(msg, GetTaskStates):
             run_id_task_state_map = self.client.task_instances.get_task_states(
                 dag_id=msg.dag_id,
+                map_index=msg.map_index,
                 task_ids=msg.task_ids,
                 task_group_id=msg.task_group_id,
                 logical_dates=msg.logical_dates,