Improve ConcurrencyMap query performance by adding composite index on TaskInstance state

[rajat315315]

Description

Description

In airflow-core, the ConcurrencyMap.load method executes a query to count active task instances grouped by Dag, run, and task:

query = session.execute(
    select(TI.dag_id, TI.task_id, TI.run_id, TI.state, func.count("*"))
    .where(TI.state.in_(ACTIVE_STATES))
    .group_by(TI.dag_id, TI.task_id, TI.run_id, TI.state)
)

Currently, the database has to perform a lookup on the single-column index ti_state and then fetch pages from the table heap (random I/O) to retrieve the values of dag_id, task_id, and run_id. Furthermore, because the index does not cover the group-by columns, the query planner has to build a temporary B-Tree/filesort table to group the results.

Adding a composite index starting with state and including dag_id, task_id, and run_id resolves these performance bottlenecks by enabling a Covering Index Scan (Index Only Scan) and leveraging the pre-sorted index keys to skip the temporary sorting/grouping steps.

Current Database Indexes on task_instance

• ti_dag_state (dag_id, state)
• ti_dag_run (dag_id, run_id)
• ti_state (state) -- Single-column index used by the query
• ti_state_lkp (dag_id, task_id, run_id, state)
• ti_pool (pool, state, priority_weight)

None of the existing composite indexes start with state, meaning they cannot cover the query filter on state.

Benchmarking Results

Benchmarks were executed on a dataset containing 100,000 finished task instances and 500 active task instances:

Metric	Before Index	After Index
Query Plan	SEARCH ... USING INDEX ti_state + USE TEMP B-TREE FOR GROUP BY	SEARCH ... USING COVERING INDEX
Average Time	5.13 ms	4.63 ms
Improvement	-	~10% reduction in execution time

Note: On production databases (PostgreSQL/MySQL), avoiding 500 heap page reads via an Index Only Scan will result in significantly higher relative performance wins, database CPU savings, and reduced lock contention under concurrent scheduler stress.

Proposed Solution

Add the following composite index to TaskInstance.__table_args__ in taskinstance.py:

Index("ti_state_active_composite", state, dag_id, task_id, run_id)

Create a corresponding Alembic migration to apply this index.

Use case/motivation

No response

Related issues

No response

Are you willing to submit a PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[boring-cyborg]

Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.