Subsetting a BGEN file in hail

I am trying to read a bgen file and subset a list of variants to use the final subset genotype to calculate PRS score.
My hail code is in below which is working fine, but it needs to allocate a lots of memory. I am seeking help to make it more efficient so that I can run it over all chromosomes.
Currently, running it over a spark cluster I need to allocate 50 worker with 16GB of RAM to execute this task on a BGEN file of size 50GB.
Does it make sense to run this about 20 minutes?
The another issue is if I run this code for the larger chromosome with 70-80 GB of BGEN file size, my spark crashes due to stage failure.
As far I know, hail is dealing with large file such a this as a stream of data. considering that I should not face any problem with the resource allocated.

Any help os appreciate it.

###########
import hail as hl
mt = hl.import_bgen(“gs:\…\inputBGEN.chr”+chr+“.bgen”,
entry_fields=[‘GT’],
sample_file= “gs:\…\sample_fileName.sample”)
mt = hl.variant_qc(mt)

scores=hl.import_table(“variant_list.txt”, delimiter=’ ')
scores=scores.key_by(“rsid” )

mt= mt.annotate_rows(**scores[mt.rsid])

flip = hl.case().when(mt.alleleB == mt.alleles[0], True).when(mt.alleleB == mt.alleles[1], False).or_missing()

mt = mt.annotate_rows(flip=flip)
mt = mt.annotate_rows(prior=2 * hl.if_else(mt.flip, mt.variant_qc.AF[0], mt.variant_qc.AF[1]))

mt = mt.annotate_cols( prs=hl.agg.sum( mt.score * hl.coalesce( hl.if_else(mt.flip, 2 - mt.GT.n_alt_alleles(), mt.GT.n_alt_alleles()), mt.prior)))

mt.cols().export(gcPRSscoresDir+“prs_scores_chr”+chr+“_”+pg_tail+“.tsv”)

First of all, how many partitions does your bgen import as? That is, how many tasks do you see when running this pipeline for a given chromosome?

I’m not seeing anything obviously wrong with the code here. What is the error that you are seeing specifically?