Try a checkpoint before the join

Author: bpblanken

v03_pipeline/lib/misc/io.py

@@ -200,21 +200,24 @@ def import_pedigree(pedigree_path: str) -> hl.Table:
     )
 
 
-def write(
-    t: hl.Table | hl.MatrixTable,
-    destination_path: str,
-) -> hl.Table | hl.MatrixTable:
+def checkpoint(t: hl.Table | hl.MatrixTable) -> tuple[hl.Table | hl.MatrixTable, str]:
     suffix = 'mt' if isinstance(t, hl.MatrixTable) else 'ht'
     read_fn = hl.read_matrix_table if isinstance(t, hl.MatrixTable) else hl.read_table
     checkpoint_path = os.path.join(
         Env.HAIL_TMPDIR,
         f'{uuid.uuid4()}.{suffix}',
     )
-    # not using checkpoint to read/write here because the checkpoint codec is different, leading to a different on disk size.
     t.write(checkpoint_path)
-    t = read_fn(checkpoint_path)
+    return read_fn(checkpoint_path), checkpoint_path
+
+
+def write(
+    t: hl.Table | hl.MatrixTable,
+    destination_path: str,
+) -> hl.Table | hl.MatrixTable:
+    t, path = checkpoint(t)
     t = t.repartition(
-        compute_hail_n_partitions(file_size_bytes(checkpoint_path)),
+        compute_hail_n_partitions(file_size_bytes(path)),
         shuffle=False,
     )
     return t.write(destination_path, overwrite=True)

v03_pipeline/lib/reference_data/clinvar.py

@@ -183,7 +183,7 @@ def download_and_import_clinvar_submission_summary() -> hl.Table:
             os.path.basename(tmp_file.name),
         )
         safely_move_to_gcs(tmp_file.name, gcs_tmp_file_name)
-        return hl.import_table(
+        ht = hl.import_table(
             gcs_tmp_file_name,
             force=True,
             filter='^(#[^:]*:|^##).*$',  # removes all comments except for the header line
@@ -193,5 +193,6 @@ def download_and_import_clinvar_submission_summary() -> hl.Table:
                 'ReportedPhenotypeInfo': hl.tstr,
             },
             missing='-',
-            min_partitions=MIN_HT_PARTITIONS,
         )
+        # NB: min_partitions fails with force=True but appears overrideable
+        return ht.reparition(MIN_HT_PARTITIONS)

v03_pipeline/lib/reference_data/dataset_table_operations.py

@@ -4,6 +4,7 @@
 import hail as hl
 import pytz
 
+from v03_pipeline.lib.misc.io import checkpoint
 from v03_pipeline.lib.misc.nested_field import parse_nested_field
 from v03_pipeline.lib.model import (
     DatasetType,
@@ -36,6 +37,7 @@ def update_or_create_joined_ht(
 
         # Join the new one!
         dataset_ht = get_dataset_ht(dataset, reference_genome)
+        dataset_ht, _ = checkpoint(dataset_ht)
         joined_ht = joined_ht.join(dataset_ht, 'outer')
         joined_ht = annotate_dataset_globals(joined_ht, dataset, dataset_ht)
 
@@ -213,6 +215,7 @@ def join_hts(
     )
     for dataset in reference_dataset_collection.datasets(dataset_type):
         dataset_ht = get_dataset_ht(dataset, reference_genome)
+        dataset_ht, _ = checkpoint(dataset_ht)
         joined_ht = joined_ht.join(dataset_ht, 'outer')
         joined_ht = annotate_dataset_globals(joined_ht, dataset, dataset_ht)
     return joined_ht