Pandas Groupby operation on a large dataset

Question

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.

Answer 1

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.

Answer 2

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')