3
\$\begingroup\$

I am working with a large dataset expanding > 50 years. Each year has ~10 million lines of records with multiple variables/columns. I need to perform groupby operations by location and time. My code runs extremely slow - it takes 2-5 hours to process 1 year's data depending on the number of stations in the year. I looked at a few posts on multiprocessing, but since I have no experiences with it, I am not sure if that method applies to my problem. I'd appreciate it if someone can point out how I can make the code more efficient.

#!/usr/bin/env python
# encoding: utf-8
import numpy as np
import pandas as pd
import datetime
import argparse
from scipy.stats.mstats import hmean

def Nstat(df):
    duMW = [6,7,8,9,30,31,32,33,34,35,98]
    d = {}
    d['NMW']  = df['MW'].count()
    d['NPW']  = df['PW'].count()
    d['NDU']  = df.loc[ isd['RH']<=90,'MW'].isin(duMW).sum()
    d['NDU6']  = (df.loc[ df['RH']<=90,'MW']==6 ).sum()
    d['NDU7']  = (df.loc[ df['RH']<=90,'MW']==7 ).sum()
    d['NDU8']  = (df.loc[ df['RH']<=90,'MW']==8 ).sum()
    d['NDU9']  = (df.loc[ df['RH']<=90,'MW']==9 ).sum()
    d['NDU30'] = (df.loc[ df['RH']<=90,'MW']==30).sum()
    d['NDU31'] = (df.loc[ df['RH']<=90,'MW']==31).sum()
    d['NDU32'] = (df.loc[ df['RH']<=90,'MW']==32).sum()
    d['NDU33'] = (df.loc[ df['RH']<=90,'MW']==33).sum()
    d['NDU34'] = (df.loc[ df['RH']<=90,'MW']==34).sum()
    d['NDU35'] = (df.loc[ df['RH']<=90,'MW']==35).sum()
    d['NDU98'] = (df.loc[ df['RH']<=90,'MW']==98).sum()
    d['NDUpw'] = (df.loc[ df['RH']<=90,'PW']==3).sum()
    d['VIS_Hvg'] = hmean(df.loc[df['VIS']>0,'VIS'])
    d['Vi_Avg'] = df['Vi'].mean()

    return pd.Series(d,index=['NMW','NPW',\
'NDU','NDU6','NDU7','NDU8','NDU9','NDU30','NDU31','NDU32',\
'NDU33','NDU34','NDU35','NDU98','NDUpw','VIS_Hvg','Vi_Avg'])


if __name__ =='__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("start_year",type=int,help='4-digit start year')
    parser.add_argument("end_year",type=int,help='4-digit end year')
    args = parser.parse_args()
    years = np.arange(args.start_year,args.end_year)

    dTypes = {
        'NMW':'Int32',\
        'NPW':'Int32',\
        'NDU':'Int32',\
        'NDU6':'Int32',\
        'NDU7':'Int32',\
        'NDU8':'Int32',\
        'NDU9':'Int32',\
        'NDU30':'Int32',\
        'NDU31':'Int32',\
        'NDU32':'Int32',\
        'NDU33':'Int32',\
        'NDU34':'Int32',\
        'NDU35':'Int32',\
        'NDU98':'Int32',\
        'NDUpw':'Int32'\
    }

    for iyr,yr in enumerate(years):
        print('process year {:d} at {:s}'.format(yr,datetime.datetime.now().strftime('%m-%d %H:%M:%S')))
        isd = pd.read_hdf('isd_lite_'+str(yr)+'.h5',dtype={'STATION':'str'})

        isd['YYYYMM'] = pd.to_datetime(isd['YYYYMMDDHH'],format='%Y%m%d%H').dt.strftime('%Y%m')
        isd['VIS'] = isd['VIS']/1000.
        isd['Vi'] = isd['VIS'].apply(lambda x: 1/x if x>0 else np.nan)

        print('>> groupby and output at {:s}'.format(datetime.datetime.now().strftime('%m-%d %H:%M:%S')))

        stn_month = isd.groupby(['STATION','YYYYMM']).apply(Nstat).reset_index().astype(dTypes)
        stn_month.to_csv('stn_month_'+str(yr)+'.csv',index=False,float_format='%.3f')

The last groupby (by STATION and YYYYMM) operation is most time consuming. I have a fairly good work station (256 cores) and want to maximize the use of it.

A sample file is provided here. It takes 7 min to process this file. Not horribly long due to a small number of stations.

\$\endgroup\$
  • \$\begingroup\$ Welcome to Code Review! Please state a bit more about the goal of your program, perhaps read the FAQ on asking questions to get the most out of your question. \$\endgroup\$ – Mast Apr 20 at 20:49
  • 1
    \$\begingroup\$ Have you considered using other libraries besides or in addition to pandas, like Dask? It has parallel versions of .groupby and can be set up to use multiple cores or computers. \$\endgroup\$ – RootTwo Apr 22 at 15:50
  • \$\begingroup\$ @Juho sample file is provided. It takes ~7 minutes to process this file. drive.google.com/file/d/1hJN7dYYpcG73PSJOB8zukAR14FnfZt_e/… \$\endgroup\$ – peteron30 Apr 25 at 16:52
  • \$\begingroup\$ @Juho sorry about that. code fixed. should work now. \$\endgroup\$ – peteron30 Apr 25 at 17:19
  • \$\begingroup\$ @Juho use python script.py 1947 1948 to run for 1947 only. It will take <7 minutes on your end, since I removed some lines. The last groupby is what I am trying to optimize. \$\endgroup\$ – peteron30 Apr 25 at 17:26
3
\$\begingroup\$

Consider the following points:

  1. Pandas has a good datetime functionality; you shouldn't cast into strings and then later on group by those. It's unnatural and slow. Instead, just do:

    isd['YYYYMM'] = pd.to_datetime(isd['YYYYMMDDHH'],format='%Y%m%d%H')
    

    And then in the groupby, you can simply do

    stn_month = isd.groupby(['STATION', isd['YYYYMM'].dt.to_period('M')]) ...
    
  2. In general, using apply is usually not great for performance. First, notice that you are doing a lot of things inside Nstat that are not necessary: all the lines like d['NDU6'] = (df.loc[ df['RH']<=90,'MW']==6 ).sum() are unnecessary in a sense that you can just precompute this outside of the function. As a side note, the way that you write is unnatural to me and I would more simply do:

    df[(df['RH'] <= 90) & (df['MW'] == 6)]
    

    Second, the agg function also takes a dictionary so that you can just do:

    isd.groupby(['STATION', isd['YYYYMM'].dt.to_period('M')]).agg({'MW' : 'count', 'PW' : 'count', 'Vi': 'mean'})
    

    I hope this will get you started.

| improve this answer | |
\$\endgroup\$
1
\$\begingroup\$

It turns out that my script had an error, hence the ridiculously long runtime. After fixing the error, runtime is shortened, but the code itself is still inefficient. The real problem is in Nstat - The row-based computation is both CPU and memory-inefficient. For those interested, read this.

Thanks to @Juho, I removed Nstat and switched to agg. Runtime is reduced by more than half!

        #prescreening by RH>90%
        isd.loc[ isd.RH>90, 'MW'] = 0
        isd.loc[ isd.RH>90, 'PW'] = 0

        stn_month = isd.groupby(['STATION',isd.DATE.dt.to_period('M')]).agg(
            NMW=('MW','count'),\
            NPW=('PW','count'),\
            NDU=('MW',lambda x: x.isin(duMW).sum()),\
            NDU6=('MW',lambda x: x.eq(6).sum()),\
            NDU7=('MW',lambda x: x.eq(7).sum()),\
            NDU8=('MW',lambda x: x.eq(8).sum()),\
            NDU9=('MW',lambda x: x.eq(9).sum()),\
            NDU30=('MW',lambda x: x.eq(30).sum()),\
            NDU31=('MW',lambda x: x.eq(31).sum()),\
            NDU32=('MW',lambda x: x.eq(32).sum()),\
            NDU33=('MW',lambda x: x.eq(33).sum()),\
            NDU34=('MW',lambda x: x.eq(34).sum()),\
            NDU35=('MW',lambda x: x.eq(35).sum()),\
            NDU98=('MW',lambda x: x.eq(98).sum()),\
            NDUPW=('PW',lambda x: x.eq(3).sum()),\
            VIS=('VIS',lambda x: hmean(x[x>0])),\
            Vi=('Vi','mean'),\
            DUP=('DUP','mean')\
            ).reset_index().astype(dTypes)

stn_month.to_csv('../stat/yearly/stn_all/stn_month_{:d}.csv'.format(yr),index=False,float_format='%.3f')
| improve this answer | |
\$\endgroup\$
  • \$\begingroup\$ Great to see such an improvement! \$\endgroup\$ – Juho Apr 26 at 12:42

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.