How to read multiple parquet files from an s3 bucket folder into dataframes efficiently

Viewed 34

I'm currently reading parquet files from an s3 bucket with the structure that goes like day/country/geohash/file. There's 30 sub folders for day and multiple countries with multiple geohashs containing multiple parquet files. As you can imagine its a lot to process and my current code is painstakingly slow. I'm using pandas to store them into a dataframe and a paginator to read the s3 bucket keys for each parquet object. Any suggestions on speeding it up? I've considered using pyspark instead of pandas. Or continue using pandas and somehow incorporate threading.

class AmazonS3:
    def __init__(self, access_key_id ,secret_access_key,region):


        self.logger = logging.getLogger("S3")
        
        logging.basicConfig(level=logging.INFO)

        session = boto3.Session(profile_name="default")
        credentials = session.get_credentials()
        self.s3 = boto3.client('s3', aws_access_key_id = credentials.access_key,
                                    aws_secret_access_key= credentials.secret_key,
                                    region_name ='us-east-1')

    def get_matching_s3_objects(self,bucket, prefix):
        """
        Generate objects in an S3 bucket.
        :param bucket: Name of the S3 bucket.
        :param prefix: Only fetch objects whose key starts with
            this prefix (optional).
        """
        paginator = self.s3.get_paginator('list_objects_v2')
        pages = paginator.paginate(Bucket=bucket, Prefix=prefix)
        for page in pages:
            contents = page['Contents']
            for obj in contents:
                key = obj["Key"]
                if key.startswith(prefix):
                    yield obj
        
    def get_matching_s3_keys(self,bucket, prefix):
        """
        Generate the keys in an S3 bucket.
        :param bucket: Name of the S3 bucket.
        :param prefix: Only fetch keys that start with this prefix (optional).
        """
        for obj in self.get_matching_s3_objects(bucket=bucket, prefix=prefix):
            yield obj["Key"]
            
    def read_parquet_objects(self,bucket,prefix):
        """
        read parquet objects into one dataframe with consistent metadata. 
        If object is not parquet type then convert it.
        :param bucket: Name of the S3 bcuket.
        :Param prefix: Only fetch keys that start with this prefix
        """
        keys_generator = self.get_matching_s3_keys(bucket=bucket, prefix=prefix)
        for key in keys_generator:
            df_holder=[]
            if key.endswith('.snappy.parquet') == True:
                self.logger.info(f"reading the parquet objects into a dataframe with the following key: {key}")
                object = self.s3.get_object(Bucket=bucket, Key=key)
                bytes_ = BytesIO(object['Body'].read())
                df = pd.read_parquet(bytes_)
                df_holder.append(df) 
            else:
                    self.logger.warning(f"The following s3 key:{key} has non csv/parquet format")
        return df_holder
      


0 Answers
Related