apache / apache/airflow

create a new index for span_status, dag_version_id in task_instance table

Open
#53,401 13 comments 4 reactions 0 assignees View on GitHub
area:core area:DAG-processing area:MetaDB area:performance area:scheduler needs-triage
Dominant language
Python
Stars
46.9k
Forks
17.8k
Avg merge
2d 10h
Merged PRs (30d)
483

Description

### Description

When the TI table is huge(~ 10M records), , `_end_spans_of_externally_ended_ops` method([Link](https://github.com/apache/airflow/blob/912ccace3532e3c05a46f6a1cf9457c0de41db6b/airflow-core/src/airflow/jobs/scheduler_job_runner.py#L1085C9-L1085C43)) takes very long to finish(in our case, about 75 seconds).

This particular sqlalchemy query was taking too long:
```
tis_should_end: list[TaskInstance] = session.scalars(
select(TaskInstance).where(TaskInstance.span_status == SpanStatus.SHOULD_END)
).all()
```

After adding the following index, performance improved significantly

```
create index idx_span_status on task_instance (id, span_status);
```

Along with above, creating an index on dag_version_id(`create index idx_dag_version_id on task_instance (dag_version_id);`) also helped to improve DAG Processor performance.
### Use case/motivation

Improve scheduler and Dag Processor performance when TI table has lot of records

### Related issues

_No response_

### Are you willing to submit a PR?

- [ ] 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)

Contributor guide

Open the contributing guide

Research direction

Start with the `_end_spans_of_externally_ended_ops` entry point in `airflow-core/src/airflow/jobs/scheduler_job_runner.py` and the shown SQLAlchemy query on the `task_instance` table. Review the project’s schema-change entry points, then add and verify indexes for `span_status` and `dag_version_id`; done means both scheduler and DAG Processor queries have the intended indexes.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, sqlalchemy
Domain
databases, performance
Issue type
Feature
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.