change: Changed run tag filtering query for better performance. (dags…

…ter-io#23050)

A followup to dagster-io#22833 adding an optimized query for tag filtering.

## Summary & Motivation
While running jobs on frequent schedules we have noticed that as the
amount of runs grows some ui operations become very slow. Looking at AWS
monitoring we see that one query in particular seems to be very slow:

![Screenshot_20240703_151316](https://github.com/dagster-io/dagster/assets/4254771/23040062-030b-4161-a43a-8667cbef4d56)

By analyzing the query we have noticed that it is related to run tag
filtering. Here is an example of such a query with filled parameters:
```sql
EXPLAIN (ANALYSE, BUFFERS) SELECT runs.id,
       runs.run_body,
       runs.status,
       runs.create_timestamp,
       runs.update_timestamp,
       runs.start_time,
       runs.end_time
FROM runs
WHERE runs.run_id IN (SELECT run_tags.run_id
                      FROM run_tags
                      WHERE
                          run_tags.key = 'dagster/schedule_name' AND run_tags.value = 'quick_partitioned_job_schedule'
                         OR run_tags.key = '.dagster/repository' AND
                            run_tags.value = '__repository__@example-code'
                      GROUP BY run_tags.run_id
                      HAVING count(DISTINCT run_tags.key) = 2)
ORDER BY runs.id DESC
LIMIT 1;
```

I believe there are multiple instances discussing this:
* dagster-io#18269
* dagster-io#19003

Looking at the query plan:
https://explain.dalibo.com/plan/2c1bga585e8ca45f

![image](https://github.com/dagster-io/dagster/assets/4254771/0d8fc125-ac67-4a1b-8f37-7ff69cbfa81f)

We notice that the subquery scans a lot of rows (which is correct as we
have a lot of runs with same tags), but afterwards, the filter on runs
is very slow and filters away a lot of rows. A lot of work is done to
retrieve only one row with highest matching run id which feels like it
can be much more efficient.

To improve performance of these type of queries I would like to propose
two changes:

1. Replace the subquery by multiple joins. I would expect that this
would make a much flatter execution plan and thus a potential for
earlier filtering.

The example query would result in something like this:
```sql
EXPLAIN (ANALYSE, BUFFERS)
SELECT runs.id,
       runs.run_body,
       runs.status,
       runs.create_timestamp,
       runs.update_timestamp,
       runs.start_time,
       runs.end_time
FROM runs
JOIN public.run_tags r on runs.run_id = r.run_id AND r.key = 'dagster/schedule_name' AND r.value = 'quick_partitioned_job_schedule'
JOIN public.run_tags r2 on runs.run_id = r2.run_id AND r2.key = '.dagster/repository' AND r2.value = '__repository__@example-code'
ORDER BY runs.id DESC
LIMIT 1;
```

2. As mentioned in one of the referenced threads, add an index on run_id
for run tags. This would make joins in (1) much faster.

```sql
CREATE UNIQUE INDEX run_tags_run_idx ON public.run_tags USING btree (run_id, id);
```

The changes in the PR implement both changes in Dagster.

## How I Tested These Changes

I have tested these change by first running tests to make sure they
don't break dagster. Then I have set up a local benchmark to test the
changes. I have populated the dagster instace with 5.4 million runs and
10.7 million related run tags.

Afterwards I have applied the proposed changes and measured their
performance. Each query was run five times and the performance of the
fifth run was used. This is, to make sure all the data was in shared
buffers to make the comparison fair.

Results
| Experiment | Query Plan Analysis | Runtime | 
| ----------------- | ----------------------------- | ------------ | 
| Baseline | [10.7 run tags, unoptimized -
explain.dalibo.com](https://explain.dalibo.com/plan/2c1bga585e8ca45f) |
15.579s |
| Query Optimization | [10.7 run tags, optimized query, no index -
explain.dalibo.com](https://explain.dalibo.com/plan/agf3fabc77b5ce58) |
7.560s |
| Query Optimization + Custom Index | [10.7 run tags, optimized query,
with index -
explain.dalibo.com](https://explain.dalibo.com/plan/dgf924f51g585ab5) |
0.076s |

Interestingly "Query Optimization" alone results in a more complex but
faster query plan. The "Query Optimization + Custom Index" results in a
desired much simpler query plan that doesn't do as many reads.

Overall the changes improve the performance almost 200x. The addition of
the index shouldn't provide much overhead.

I have also found that if run_tags would use `runs.id` as foreign key
and not `runs.run_id` the query performance would be much faster 0.018s.
But this change would be too large and possibly break things.

---------

Signed-off-by: Egor Dmitriev <[email protected]>

python_modules/dagster/dagster/_core/storage/runs/sql_run_storage.py

@@ -253,6 +253,12 @@ def _add_cursor_limit_to_query(
 
  return query
 
+ def _add_filters_to_table(self, table: db.Table, filters: RunsFilter) -> db.Table:
+ if filters.tags:
+ table = self._apply_tags_table_filters(table, filters.tags)
+
+ return table
+
  def _add_filters_to_query(self, query: SqlAlchemyQuery, filters: RunsFilter) -> SqlAlchemyQuery:
  check.inst_param(filters, "filters", RunsFilter)
 
@@ -290,9 +296,6 @@ def _add_filters_to_query(self, query: SqlAlchemyQuery, filters: RunsFilter) ->
  RunsTable.c.create_timestamp < filters.created_before.replace(tzinfo=None)
  )
 
- if filters.tags:
- query = self._apply_tags_table_filters(query, filters.tags)
-
  return query
 
  def _runs_query(
@@ -315,51 +318,32 @@ def _runs_query(
  if columns is None:
  columns = ["run_body", "status"]
 
- table = RunsTable
+ table = self._add_filters_to_table(RunsTable, filters)
  base_query = db_select([getattr(RunsTable.c, column) for column in columns]).select_from(
  table
  )
  base_query = self._add_filters_to_query(base_query, filters)
  return self._add_cursor_limit_to_query(base_query, cursor, limit, order_by, ascending)
 
  def _apply_tags_table_filters(
- self, query: SqlAlchemyQuery, tags: Mapping[str, Union[str, Sequence[str]]]
+ self, table: db.Table, tags: Mapping[str, Union[str, Sequence[str]]]
  ) -> SqlAlchemyQuery:
  """Efficient query pattern for filtering by multiple tags."""
- expected_count = len(tags)
- if expected_count == 1:
- key, value = next(iter(tags.items()))
- # since run tags should be much larger than runs, select where exists
- # should be more efficient than joining
- subquery = db.exists().where(
- (RunsTable.c.run_id == RunTagsTable.c.run_id)
- & (RunTagsTable.c.key == key)
- & (
- (RunTagsTable.c.value == value)
+ for i, (key, value) in enumerate(tags.items()):
+ run_tags_alias = db.alias(RunTagsTable, f"run_tags_filter{i}")
+
+ table = table.join(
+ run_tags_alias,
+ db.and_(
+ RunsTable.c.run_id == run_tags_alias.c.run_id,
+ run_tags_alias.c.key == key,
+ (run_tags_alias.c.value == value)
  if isinstance(value, str)
- else RunTagsTable.c.value.in_(value)
- )
- )
- query = query.where(subquery)
- elif expected_count > 1:
- # efficient query for filtering by multiple tags. first find all run_ids that match
- # all tags, then select from runs table where run_id in that set
- subquery = db_select([RunTagsTable.c.run_id])
- expressions = []
- for key, value in tags.items():
- expression = RunTagsTable.c.key == key
- if isinstance(value, str):
- expression &= RunTagsTable.c.value == value
- else:
- expression &= RunTagsTable.c.value.in_(value)
- expressions.append(expression)
- subquery = subquery.where(db.or_(*expressions))
- subquery = subquery.group_by(RunTagsTable.c.run_id)
- subquery = subquery.having(
- db.func.count(db.distinct(RunTagsTable.c.key)) == expected_count
+ else run_tags_alias.c.value.in_(value),
+ ),
  )
- query = query.where(RunsTable.c.run_id.in_(subquery))
- return query
+
+ return table
 
  def get_runs(
  self,