[batch] make batches query go brrrrrrr (#14629)

This PR optimises the V2 batches query along with its associated subqueries. Listing batches for the CI user now executes in 1s were previously it had taken ~14s. Resolves #14599 --------- Co-authored-by: Chris Llanwarne <[email protected]>
hail-is · Jul 31, 2024 · 21ad23a · 21ad23a
1 parent 50310c9
commit 21ad23a
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 38 deletions.
diff --git a/batch/batch/front_end/query/query.py b/batch/batch/front_end/query/query.py
@@ -455,10 +455,14 @@ def __init__(self, term: str):
         self.term = term
 
     def query(self) -> Tuple[str, List[str]]:
-        sql = """
-((job_groups.batch_id, job_groups.job_group_id) IN
- (SELECT batch_id, job_group_id FROM job_group_attributes
-  WHERE `key` = %s OR `value` = %s))
+        sql = """\
+EXISTS (
+  SELECT NULL
+  FROM job_group_attributes AS attrs
+  WHERE attrs.batch_id = job_groups.batch_id
+    AND attrs.job_group_id = job_groups.job_group_id
+    AND (attrs.`key` = %s OR attrs.`value` = %s)
+)
 """
         return (sql, [self.term, self.term])
 
@@ -478,10 +482,14 @@ def __init__(self, term: str):
         self.term = term
 
     def query(self) -> Tuple[str, List[str]]:
-        sql = """
-((job_groups.batch_id, job_groups.job_group_id) IN
- (SELECT batch_id, job_group_id FROM job_group_attributes
-  WHERE `key` LIKE %s OR `value` LIKE %s))
+        sql = """\
+EXISTS (
+  SELECT NULL
+  FROM job_group_attributes AS attrs
+  WHERE attrs.batch_id = job_groups.batch_id
+    AND attrs.job_group_id = job_groups.job_group_id
+    AND (attrs.`key` LIKE %s OR attrs.`value` LIKE %s)
+)
 """
         escaped_term = f'%{self.term}%'
         return (sql, [escaped_term, escaped_term])
@@ -505,10 +513,15 @@ def query(self) -> Tuple[str, List[str]]:
         value = self.value
         if isinstance(self.operator, PartialMatchOperator):
             value = f'%{value}%'
-        sql = f"""
-((job_groups.batch_id, job_groups.job_group_id) IN
- (SELECT batch_id, job_group_id FROM job_group_attributes
-  WHERE `key` = %s AND `value` {op} %s))
+        sql = f"""\
+EXISTS (
+  SELECT NULL
+  FROM job_group_attributes AS attrs
+  WHERE attrs.batch_id = job_groups.batch_id
+    AND attrs.job_group_id = job_groups.job_group_id
+    AND attrs.`key` = %s
+    AND attrs.`value` {op} %s
+)
 """
         return (sql, [self.key, value])
 

diff --git a/batch/batch/front_end/query/query_v2.py b/batch/batch/front_end/query/query_v2.py
@@ -125,40 +125,46 @@ def parse_list_batches_query_v2(user: str, q: str, last_batch_id: Optional[int])
         where_conditions.append(f'({cond})')
         where_args += args
 
-    sql = f"""
-SELECT batches.*,
-  cancelled_t.cancelled IS NOT NULL AS cancelled,
+    sql = f"""\
+SELECT
+  batches.*,
+  cancelled_t.cancelled <=> 1 AS cancelled,
   job_groups_n_jobs_in_complete_states.n_completed,
   job_groups_n_jobs_in_complete_states.n_succeeded,
   job_groups_n_jobs_in_complete_states.n_failed,
   job_groups_n_jobs_in_complete_states.n_cancelled,
-  cost_t.cost, cost_t.cost_breakdown
-FROM job_groups
-LEFT JOIN batches ON batches.id = job_groups.batch_id
-LEFT JOIN billing_projects ON batches.billing_project = billing_projects.name
-LEFT JOIN job_groups_n_jobs_in_complete_states ON job_groups.batch_id = job_groups_n_jobs_in_complete_states.id AND job_groups.job_group_id = job_groups_n_jobs_in_complete_states.job_group_id
-LEFT JOIN LATERAL (
-  SELECT 1 AS cancelled
-  FROM job_group_self_and_ancestors
-  INNER JOIN job_groups_cancelled
-    ON job_group_self_and_ancestors.batch_id = job_groups_cancelled.id AND
-      job_group_self_and_ancestors.ancestor_id = job_groups_cancelled.job_group_id
-  WHERE job_groups.batch_id = job_group_self_and_ancestors.batch_id AND
-    job_groups.job_group_id = job_group_self_and_ancestors.job_group_id
-) AS cancelled_t ON TRUE
-STRAIGHT_JOIN billing_project_users ON batches.billing_project = billing_project_users.billing_project
-LEFT JOIN LATERAL (
-  SELECT COALESCE(SUM(`usage` * rate), 0) AS cost, JSON_OBJECTAGG(resources.resource, COALESCE(`usage` * rate, 0)) AS cost_breakdown
-  FROM (
-    SELECT resource_id, CAST(COALESCE(SUM(`usage`), 0) AS SIGNED) AS `usage`
+  cost_t.cost,
+  cost_t.cost_breakdown
+FROM batches
+INNER JOIN job_groups
+  ON batches.id = job_groups.batch_id
+INNER JOIN billing_projects
+  ON batches.billing_project = billing_projects.name
+INNER JOIN billing_project_users
+  ON batches.billing_project = billing_project_users.billing_project
+INNER JOIN job_groups_n_jobs_in_complete_states
+  ON batches.id = job_groups_n_jobs_in_complete_states.id
+  AND job_groups.job_group_id = job_groups_n_jobs_in_complete_states.job_group_id
+LEFT JOIN (SELECT *, 1 AS cancelled FROM job_groups_cancelled) AS cancelled_t
+  ON batches.id = cancelled_t.id
+  AND job_groups.job_group_id = cancelled_t.job_group_id
+INNER JOIN LATERAL (
+  WITH resource_costs AS (
+    SELECT
+      resource_id,
+      CAST(COALESCE(SUM(`usage`), 0) AS SIGNED) AS `usage`
     FROM aggregated_job_group_resources_v3
-    WHERE job_groups.batch_id = aggregated_job_group_resources_v3.batch_id AND job_groups.job_group_id = aggregated_job_group_resources_v3.job_group_id
+    WHERE batch_id = batches.id
     GROUP BY resource_id
-  ) AS usage_t
-  LEFT JOIN resources ON usage_t.resource_id = resources.resource_id
+  )
+  SELECT
+    COALESCE(SUM(`usage` * rate), 0) AS cost,
+    JSON_OBJECTAGG(resource, COALESCE(`usage` * rate, 0)) AS cost_breakdown
+  FROM resource_costs
+  INNER JOIN resources USING (resource_id)
 ) AS cost_t ON TRUE
 WHERE {' AND '.join(where_conditions)}
-ORDER BY job_groups.batch_id DESC
+ORDER BY batches.id DESC
 LIMIT 51;
 """