Merge pull request #3931 from broadinstitute/parallelize-query-partit…

…ions maximize num partitions for available cpus
broadinstitute · Feb 29, 2024 · 225a5ec · 225a5ec
2 parents 911d5b0 + eb3e4e3
commit 225a5ec
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/hail_search/queries/base.py b/hail_search/queries/base.py
@@ -217,7 +217,7 @@ def __init__(self, sample_data, sort=XPOS, sort_metadata=None, num_results=100,
         self._has_secondary_annotations = False
         self._is_multi_data_type_comp_het = False
         self.max_unaffected_samples = None
-        self._load_table_kwargs = {}
+        self._load_table_kwargs = {'_n_partitions': (os.cpu_count() or 2)-1}
         self.entry_samples_by_family_guid = {}
 
         if sample_data:
@@ -1048,6 +1048,7 @@ def gene_counts(self):
     def lookup_variant(self, variant_id, sample_data=None):
         self._parse_intervals(intervals=None, variant_ids=[variant_id], variant_keys=[variant_id])
         ht = self._read_table('annotations.ht', drop_globals=['paths', 'versions'])
+        self._load_table_kwargs['_n_partitions'] = 1
         ht = ht.filter(hl.is_defined(ht[XPOS]))
 
         annotation_fields = self.annotation_fields(include_genotype_overrides=False)

diff --git a/hail_search/queries/snv_indel.py b/hail_search/queries/snv_indel.py
@@ -75,7 +75,7 @@ class SnvIndelHailTableQuery(MitoHailTableQuery):
 
     def _prefilter_entries_table(self, ht, *args, num_projects=1, **kwargs):
         ht = super()._prefilter_entries_table(ht, *args, **kwargs)
-        if num_projects > 1 or not self._load_table_kwargs:
+        if num_projects > 1 or not self._load_table_kwargs.get('_filter_intervals'):
             af_ht = self._get_loaded_filter_ht(
                 GNOMAD_GENOMES_FIELD, 'high_af_variants.ht', self._get_gnomad_af_prefilter, **kwargs)
             if af_ht: