Try more partitions in the repartition

Author: bpblanken

v03_pipeline/lib/misc/io.py

@@ -271,26 +271,33 @@ def remap_pedigree_hash(remap_path: str, pedigree_path: str) -> hl.Int32Expressi
     return hl.int32(int(sha256.hexdigest()[:8], 16))
 
 
-def checkpoint(t: hl.Table | hl.MatrixTable) -> tuple[hl.Table | hl.MatrixTable, str]:
+def checkpoint(
+    t: hl.Table | hl.MatrixTable,
+    repartition_factor: int = 1,
+) -> 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_TMP_DIR,
         f'{uuid.uuid4()}.{suffix}',
     )
-    t.write(checkpoint_path)
+    t.write(checkpoint_path, repartition_factor=repartition_factor)
     return read_fn(checkpoint_path), checkpoint_path
 
 
 def write(
     t: hl.Table | hl.MatrixTable,
     destination_path: str,
     repartition: bool = True,
+    # May be used to increase the number of partitions beyond
+    # the optimally computed number.  A higher number will
+    # shard the table into more partitions.
+    repartition_factor: int = 1,
 ) -> hl.Table | hl.MatrixTable:
     t, path = checkpoint(t)
     if repartition:
         t = t.repartition(
-            compute_hail_n_partitions(file_size_bytes(path)),
+            (compute_hail_n_partitions(file_size_bytes(path)) * repartition_factor),
             shuffle=False,
         )
     return t.write(destination_path, overwrite=True)

v03_pipeline/lib/reference_datasets/splice_ai.py

@@ -13,7 +13,7 @@ def get_ht(
     # of file descriptors on dataproc :/
     hl._set_flags(use_new_shuffle=None, no_whole_stage_codegen='1')  # noqa: SLF001
     ht = vcf_to_ht(paths, reference_genome)
-    ht, _ = checkpoint(ht)
+    ht, _ = checkpoint(ht, repartition_factor=2)
 
     # SpliceAI INFO field description from the VCF header: SpliceAIv1.3 variant annotation. These include
     # delta scores (DS) and delta positions (DP) for acceptor gain (AG), acceptor loss (AL), donor gain (DG), and