Improve writing time for GWAS results

Hi, I’m running linear regression with about 100K samples and 90K variants. After GWAS, I will also export the results. The whole process took about 2 hours which seems to be too long. Below are my codes and may I kindly ask how can I improve it? Thanks a lot!

import hail as hl

### prepare the inputs:
y = 'Diabetes'  # name of the phenotype in the text file

### import the covariates and phenotype info:
sample_info = hl.import_table(paths=Input, types={'IID':hl.tstr, y:hl.tint32}, impute=True).key_by('IID')
sample_info = sample_info.rename({y: 'phenotype'})  # rename the regression Y to be "phenotype"

### Import PLINK bfiles: 
mt = hl.import_plink(bed='genotype/chr'+chr+'.bed',
                     quant_pheno=True, missing='-9', 

### add phenotypes to genotype matrix and keep only GWAS sample
mt = mt.annotate_cols(pheno = sample_info[mt.s])  # add phenotype and covariate info
mt = mt.filter_cols(~hl.is_nan(mt.pheno.phenotype))  # filter NA phenotypes

### variant QC
mt = hl.variant_qc(mt)  
mt = mt.filter_rows((mt.variant_qc.AF[1] >= 0.01)&
                    (mt.variant_qc.p_value_hwe >= 1e-6)&
                    (mt.variant_qc.call_rate >= 0.99))       

### run linear regressions
gwas = mt.annotate_rows(linreg = hl.agg.linreg(y = mt.pheno.phenotype, 
                                                x = [1, mt.GT.n_alt_alleles()])) 
gwas = gwas.rows().key_by()  
gwas = =, SNP = gwas.rsid, BP =, 
                A1 = gwas.alleles[1], A2 = gwas.alleles[0], EAF = gwas.variant_qc.AF[1],
                BETA = gwas.linreg.beta[1], SE = gwas.linreg.standard_error[1], 
                P = gwas.linreg.p_value[1], N = gwas.linreg.n)

### export result
# first write to Hail’s efficient and fast on-disk format, then read back in and convert to a text file
hl.read_table('chr'+chr+'').export('chr'+chr+'linear.tsv.bgz', header=True)

One possible code change is to use the linear_regression_rows method instead of hl.agg.linreg – it is purpose-built for efficient GWAS execution, and we’re still working on improving the performance of the linreg aggregator.