shard project_gt_stats incr query (#1184)

* shard incr query

* bugs

* fix math

* fix query

* fix

* ruff

Author: bpblanken

v03_pipeline/lib/misc/clickhouse.py

@@ -12,6 +12,7 @@
 from v03_pipeline.lib.core import DatasetType, ReferenceGenome
 from v03_pipeline.lib.core.environment import Env
 from v03_pipeline.lib.logger import get_logger
+from v03_pipeline.lib.misc.math import split_ranges
 from v03_pipeline.lib.misc.retry import retry
 from v03_pipeline.lib.paths import (
     new_entries_parquet_path,
@@ -793,12 +794,27 @@ def rebuild_gt_stats(
         table_name_builder.dst_prefix,
         table_name_builder.staging_dst_prefix,
     )
-    logged_query(
+    # NB: encountered OOMs with large projects, necessitating sharding the insertion query.
+    max_key = logged_query(
         f"""
-        INSERT INTO {table_name_builder.staging_dst_table(ClickHouseTable.PROJECT_GT_STATS)}
-        {select_statement}
+        SELECT max(key) FROM {table_name_builder.dst_table(ClickHouseTable.GT_STATS)}
         """,
-    )
+    )[0][0]
+    for range_start, range_end in split_ranges(max_key):
+        logged_query(
+            f"""
+            INSERT INTO {
+                table_name_builder.staging_dst_table(ClickHouseTable.PROJECT_GT_STATS)
+            }
+            {
+                select_statement.replace(
+                    'GROUP BY project_guid',
+                    'WHERE key >= %(range_start)s AND key <= %(range_end)s GROUP BY project_guid',
+                )
+            }
+            """,
+            {'range_start': range_start, 'range_end': range_end},
+        )
     finalize_refresh_flow(table_name_builder, project_guids)
 
 

v03_pipeline/lib/misc/math.py

@@ -1,3 +1,19 @@
+import math
+from collections.abc import Generator
+
+
+def split_ranges(max_value: int, n: int = 5) -> Generator[tuple[int, int], None, None]:
+    if max_value < n:
+        yield (0, max_value)
+        return
+    step = math.ceil(max_value / n)
+    start, end = 0, step
+    while start < max_value:
+        yield (start, min(end, max_value))
+        start += step
+        end += step
+
+
 def constrain(
     number: int | float,
     lower_bound: int | float,