rajat315315 opened a new issue, #70018:
URL: https://github.com/apache/airflow/issues/70018
### Description
### Description
In `airflow-core`, the
[ConcurrencyMap.load](file:///home/rajat/airflow/airflow-core/src/airflow/jobs/scheduler_job_runner.py#L239-L247)
method executes a query to count active task instances grouped by Dag, run,
and task:
```python
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](file:///home/rajat/airflow/airflow-core/src/airflow/models/taskinstance.py):
```python
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?
- [x] Yes I am willing to submit a PR!
### Code of Conduct
- [x] I agree to follow this project's [Code of
Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)
--
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]
