Reading S3 files from a manifest and processing them in parallel using Pandas

Viewed 2643

I have about 50k to read from S3 using a manifest file. I have to read contents of every single (JSON) file into a dataframe and process the files (normalize them as database tables). Right now I have a working code that take about 15hours to process the 50k files. I have to run this as a daily job. Is there any way to parallel process larger amount to files or any better way to speed up process?

Updating the question with the code

import json 
import pandas as pd 
import os
import gzip
import boto3
from datetime import datetime,timezone,timedelta


session = boto3.session.Session()
s3 = session.resource('s3')
client = session.client('s3')

#read the S3 inventory report, get the keys of files that are modified on sysdate-1   
dt=(datetime.now(timezone.utc) + timedelta(days=-1)).strftime('%Y-%m-%d') 
dtz=dt+'T00-00Z'
print('reading inventory report for', dtz)
inventory_bucket = 'xxx'
manifest_key='s3-bucket'+dtz+'/manifest.json'
manifest = json.load(s3.Object(inventory_bucket, manifest_key).get()['Body'])
df=pd.DataFrame()
for obj in manifest['files']:
        gzip_obj = s3.Object(bucket_name=inventory_bucket, key=obj['key'])
        print('csv obj:', gzip_obj)
        buffer = gzip.open(gzip_obj.get()["Body"], mode='rt')
        reader = pd.read_csv(buffer)
        reader.columns=['id','key','modified_date']
        print('converting csv obj to dataframe')
        df=df.append(reader[(reader.modified_date>dt)])
source_keys=list(df['key'])
s3_bucket_source='yyy'

#download the files to a tmp folder
local='/tmp/'
print("downloading from S3")
for k in source_keys:
  k_path=k.replace('/', '-')
  dest_pathname = os.path.join(local, k_path)
  if not os.path.exists(os.path.dirname(dest_pathname)):
      os.makedirs(os.path.dirname(dest_pathname))
      client.download_file(s3_bucket_source, k, dest_pathname)

#read the latest jsons from tmp folder
path_to_json =  os.path.expanduser('/tmp/') 
json_files = [pos_json for pos_json in os.listdir(path_to_json) if pos_json.endswith('.json')]

#loop through the jsons and normalize each file contents into df_table1, df_table2, df_table3
for index, js in enumerate(json_files):
    with open(os.path.join(path_to_json, js)) as json_file:
        print('processing file:', js)
        d=json.loads(json_file.read())
        v=d['key1'][0]['key2']
        if isinstance(v, list):
                for i, v2 in enumerate(v): 
                    df_table1, df_table2, df_table3 = normalizeJSON(d,i,v2['id']) 
                    #normalize is the custom function to split the nested json into relational tables 
        else:
            print('invalid json')

I use the S3 inventory report to get the list of most recent modified files from the manifest, download the files to a tmp location and read them one by one to do what I need to do

1 Answers

You use the multiprocessing module to download JSON files in parallel. You code contains 3 for blocks. You can do each one of them in parallel. An example of how to do this for the first for follows:

The first for:

df=pd.DataFrame()
for obj in manifest['files']:
        gzip_obj = s3.Object(bucket_name=inventory_bucket, key=obj['key'])
        print('csv obj:', gzip_obj)
        buffer = gzip.open(gzip_obj.get()["Body"], mode='rt')
        reader = pd.read_csv(buffer)
        reader.columns=['id','key','modified_date']
        print('converting csv obj to dataframe')
        df=df.append(reader[(reader.modified_date>dt)])

becomes:

def get_reader(obj):
    gzip_obj = s3.Object(bucket_name=inventory_bucket, key=obj['key'])
    print('csv obj:', gzip_obj)
    buffer = gzip.open(gzip_obj.get()["Body"], mode='rt')
    reader = pd.read_csv(buffer)
    reader.columns=['id','key','modified_date']
    print('converting csv obj to dataframe')
    return reader[(reader.modified_date>dt)]

num_of_workers = 4
df = pd.DataFrame()
with multiprocessing.Pool(num_of_workers) as p:
    results = p.map(get_reader, manifest['files'])

for result in results:
    df = df.append(result)

You can do the same for the other for blocks

Related