ColtenOuO opened a new pull request, #72270: URL: https://github.com/apache/airflow/pull/72270
### Sumarry `GET /execution/task-instances/states` (`get_task_instance_states` in `airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py`) is called by the worker/Task SDK whenever a task needs the states of sibling task instances (e.g. trigger-rule evaluation, `wait_for_downstream`, cross-task state checks). When the request is scoped with `task_group_id`, the handler was doing unnecessary DB work on every call: 1. It first ran a query selecting every `TaskInstance` row for the Dag (optionally narrowed by `task_ids`, `logical_dates`, `run_ids`, `map_index`), loading full ORM entities. 2. It then ran a second query to resolve the task group's members and fetch their `TaskInstance` rows. 3. If the caller had not also passed `task_ids`, the result of the first query was thrown away entirely and only the second query's rows were used — meaning step 1 was pure waste on every task-group-scoped call. For Dags with many tasks and/or many Dag runs, that first discarded query could return and hydrate a large number of rows for nothing, adding an avoidable round trip and CPU/memory cost to a request that sits on the task execution hot path (workers call this while running tasks, so the extra latency is paid per task, not once). This PR removes the wasted round trip: - The task group's task IDs are now resolved from the already-parsed Dag definition (`task_group.iter_tasks()`), which requires no extra `TaskInstance` query — only the existing Dag-version lookup that was already being performed. - Those IDs are unioned with any explicit `task_ids` from the request *before* building the query, so there is exactly one `TaskInstance` query per request, filtered with `task_id.in_(...)` — preserving the exact original semantics (task_ids alone, task_group_id alone, or the union of both). - That single query now selects only the four columns actually used (`run_id`, `task_id`, `map_index`, `state`) instead of loading full `TaskInstance` ORM objects, cutting the amount of data hydrated per call further. - The pre-existing `_get_group_tasks` helper (still used by the sibling `/count` endpoint, which needs actual `map_index` values from persisted rows for a different query shape) is left untouched aside from factoring out the shared "resolve the task group or 404" logic into `_get_task_group`. ### Change - `get_task_instance_states`: replaced the "query everything, then maybe discard it" flow with a single `TaskInstance` query. The task group's task IDs are resolved up front (from `task_group.iter_tasks()`) and unioned with any request `task_ids` into one `task_id.in_(...)` filter before the query runs, instead of running one query filtered by `task_ids` and a second, separate query for the group, then conditionally throwing the first result set away. - That single query now does `select(TI.run_id, TI.task_id, TI.map_index, TI.state)` instead of `select(TI)`, so it loads four scalar columns per row instead of hydrating full `TaskInstance` ORM objects. - Added `_get_task_group` (Dag/task-group lookup + 404 handling) factored out of the pre-existing `_get_group_tasks`, and a new `_get_group_task_ids` that returns just the group's task IDs without touching the `TaskInstance` table. `_get_group_tasks` itself is unchanged in behavior — it's still used as-is by the `/count` endpoint, which needs real persisted `(task_id, map_index)` pairs for a different query shape. - Added `test_get_task_states_with_task_group_id_query_count` (query-count regression test). ### Impact For any `task_group_id`-scoped call to this endpoint, DB round trips drop from 2 (one of which was always discarded) down to 1, and that query is narrower (4 scalar columns instead of full-row ORM hydration). This directly reduces the RTT the worker/Task SDK pays on this call, and removes the redundant row hydration/network cost on the API server and DB — proportional to Dag size and Dag-run count, so the win scales with the DAGs that would have hit this path hardest. ### Testing - Added `test_get_task_states_with_task_group_id_query_count`, a query-count regression test using `assert_queries_count`, which fails against the old code (3 queries) and passes against the fix (2 queries: one Dag-version lookup, one TaskInstance query). - All existing tests in `TestGetTaskStates` and `TestGetCount` (the `/count` endpoint that shares the untouched helper) continue to pass, confirming the task_ids/task_group_id union semantics are unchanged, including the mixed `task_ids` + `task_group_id` case. --- ##### Was generative AI tooling used to co-author this PR? - [X] Yes — Claude Code (Sonnet 5) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
